source: inundation/pyvolution/shallow_water.py @ 2494

Last change on this file since 2494 was 2494, checked in by ole, 18 years ago

Introduced duration keyword in evolve and updated dependencies
  • Property svn:keywords set to LastChangedDate LastChangedRevision LastChangedBy HeadURL Id
File size: 55.2 KB
Line 
1"""Class Domain -
22D triangular domains for finite-volume computations of
3the shallow water wave equation.
4
5This module contains a specialisation of class Domain from module domain.py
6consisting of methods specific to the Shallow Water Wave Equation
7
8
9U_t + E_x + G_y = S
10
11where
12
13U = [w, uh, vh]
14E = [uh, u^2h + gh^2/2, uvh]
15G = [vh, uvh, v^2h + gh^2/2]
16S represents source terms forcing the system
17(e.g. gravity, friction, wind stress, ...)
18
19and _t, _x, _y denote the derivative with respect to t, x and y respectively.
20
21The quantities are
22
23symbol    variable name    explanation
24x         x                horizontal distance from origin [m]
25y         y                vertical distance from origin [m]
26z         elevation        elevation of bed on which flow is modelled [m]
27h         height           water height above z [m]
28w         stage            absolute water level, w = z+h [m]
29u                          speed in the x direction [m/s]
30v                          speed in the y direction [m/s]
31uh        xmomentum        momentum in the x direction [m^2/s]
32vh        ymomentum        momentum in the y direction [m^2/s]
33
34eta                        mannings friction coefficient [to appear]
35nu                         wind stress coefficient [to appear]
36
37The conserved quantities are w, uh, vh
38
39
40For details see e.g.
41Christopher Zoppou and Stephen Roberts,
42Catastrophic Collapse of Water Supply Reservoirs in Urban Areas,
43Journal of Hydraulic Engineering, vol. 127, No. 7 July 1999
44
45
46Ole Nielsen, Stephen Roberts, Duncan Gray, Christopher Zoppou
47Geoscience Australia, 2004
48"""
49
50#Subversion keywords:
51#
52#$LastChangedDate: 2006-03-08 03:26:58 +0000 (Wed, 08 Mar 2006) $
53#$LastChangedRevision: 2494 $
54#$LastChangedBy: ole $
55
56
57from domain import *
58from region import *#
59
60Generic_Domain = Domain #Rename
61
62#Shallow water domain
63class Domain(Generic_Domain):
64
65    def __init__(self, coordinates, vertices, boundary = None,
66                 tagged_elements = None, geo_reference = None,
67                 use_inscribed_circle=False):
68
69        conserved_quantities = ['stage', 'xmomentum', 'ymomentum']
70        other_quantities = ['elevation', 'friction']
71        Generic_Domain.__init__(self, coordinates, vertices, boundary,
72                                conserved_quantities, other_quantities,
73                                tagged_elements, geo_reference, use_inscribed_circle)
74
75        from config import minimum_allowed_height, g
76        self.minimum_allowed_height = minimum_allowed_height
77        self.g = g
78
79        self.forcing_terms.append(gravity)
80        self.forcing_terms.append(manning_friction)
81
82        #Realtime visualisation
83        self.visualiser = None
84        self.visualise  = False
85        self.visualise_color_stage = False
86        self.visualise_stage_range = 1.0
87        self.visualise_timer = True
88        self.visualise_range_z = None
89
90        #Stored output
91        self.store = True
92        self.format = 'sww'
93        self.smooth = True
94
95        #Reduction operation for get_vertex_values
96        from util import mean
97        self.reduction = mean
98        #self.reduction = min  #Looks better near steep slopes
99
100        self.quantities_to_be_stored = ['stage','xmomentum','ymomentum']
101
102
103    def set_quantities_to_be_stored(self, q):
104        """Specify which quantities will be stored in the sww file.
105
106        q must be either:
107          - the name of a quantity
108          - a list of quantity names
109          - None
110
111        In the two first cases, the named quantities will be stored at each yieldstep
112        (This is in addition to the quantities elevation and friction) 
113        If q is None, storage will be switched off altogether.
114        """
115
116
117        if q is None:
118            self.quantities_to_be_stored = []           
119            self.store = False
120            return
121
122        if isinstance(q, basestring):
123            q = [q] # Turn argument into a list
124
125        #Check correcness   
126        for quantity_name in q:
127            msg = 'Quantity %s is not a valid conserved quantity' %quantity_name
128            assert quantity_name in self.conserved_quantities, msg
129       
130        self.quantities_to_be_stored = q
131       
132
133    def initialise_visualiser(self,scale_z=1.0,rect=None):
134        #Realtime visualisation
135        if self.visualiser is None:
136            from realtime_visualisation_new import Visualiser
137            self.visualiser = Visualiser(self,scale_z,rect)
138            self.visualiser.setup['elevation']=True
139            self.visualiser.updating['stage']=True
140        self.visualise = True
141        if self.visualise_color_stage == True:
142            self.visualiser.coloring['stage'] = True
143            self.visualiser.qcolor['stage'] = (0.0, 0.0, 0.8)
144        print 'initialise visualiser'
145        print self.visualiser.setup
146        print self.visualiser.updating
147
148    def check_integrity(self):
149        Generic_Domain.check_integrity(self)
150
151        #Check that we are solving the shallow water wave equation
152
153        msg = 'First conserved quantity must be "stage"'
154        assert self.conserved_quantities[0] == 'stage', msg
155        msg = 'Second conserved quantity must be "xmomentum"'
156        assert self.conserved_quantities[1] == 'xmomentum', msg
157        msg = 'Third conserved quantity must be "ymomentum"'
158        assert self.conserved_quantities[2] == 'ymomentum', msg
159
160    def extrapolate_second_order_sw(self):
161        #Call correct module function
162        #(either from this module or C-extension)
163        extrapolate_second_order_sw(self)
164
165    def compute_fluxes(self):
166        #Call correct module function
167        #(either from this module or C-extension)
168        compute_fluxes(self)
169
170    def distribute_to_vertices_and_edges(self):
171        #Call correct module function
172        #(either from this module or C-extension)
173        distribute_to_vertices_and_edges(self)
174
175
176    #FIXME: Under construction
177#     def set_defaults(self):
178#         """Set default values for uninitialised quantities.
179#         This is specific to the shallow water wave equation
180#         Defaults for 'elevation', 'friction', 'xmomentum' and 'ymomentum'
181#         are 0.0. Default for 'stage' is whatever the value of 'elevation'.
182#         """
183
184#         for name in self.other_quantities + self.conserved_quantities:
185#             print name
186#             print self.quantities.keys()
187#             if not self.quantities.has_key(name):
188#                 if name == 'stage':
189
190#                     if self.quantities.has_key('elevation'):
191#                         z = self.quantities['elevation'].vertex_values
192#                         self.set_quantity(name, z)
193#                     else:
194#                         self.set_quantity(name, 0.0)
195#                 else:
196#                     self.set_quantity(name, 0.0)
197
198
199
200#         #Lift negative heights up
201#         #z = self.quantities['elevation'].vertex_values
202#         #w = self.quantities['stage'].vertex_values
203
204#         #h = w-z
205
206#         #for k in range(h.shape[0]):
207#         #    for i in range(3):
208#         #        if h[k, i] < 0.0:
209#         #            w[k, i] = z[k, i]
210
211
212#         #self.quantities['stage'].interpolate()
213
214
215    def evolve(self,
216               yieldstep = None,
217               finaltime = None,
218               duration = None,               
219               skip_initial_step = False):
220        """Specialisation of basic evolve method from parent class
221        """
222
223        #Call check integrity here rather than from user scripts
224        #self.check_integrity()
225
226        msg = 'Parameter beta_h must be in the interval [0, 1['
227        assert 0 <= self.beta_h < 1.0, msg
228        msg = 'Parameter beta_w must be in the interval [0, 1['
229        assert 0 <= self.beta_w < 1.0, msg
230
231
232        #Initial update of vertex and edge values before any storage
233        #and or visualisation
234        self.distribute_to_vertices_and_edges()
235
236        #Initialise real time viz if requested
237        if self.visualise is True and self.time == 0.0:
238            if self.visualiser is None:
239                self.initialise_visualiser()
240
241            self.visualiser.update_timer()
242            self.visualiser.setup_all()
243
244        #Store model data, e.g. for visualisation
245        if self.store is True and self.time == 0.0:
246            self.initialise_storage()
247            #print 'Storing results in ' + self.writer.filename
248        else:
249            pass
250            #print 'Results will not be stored.'
251            #print 'To store results set domain.store = True'
252            #FIXME: Diagnostic output should be controlled by
253            # a 'verbose' flag living in domain (or in a parent class)
254
255        #Call basic machinery from parent class
256        for t in Generic_Domain.evolve(self,
257                                       yieldstep=yieldstep,
258                                       finaltime=finaltime,
259                                       duration=duration,
260                                       skip_initial_step=skip_initial_step):
261            #Real time viz
262            if self.visualise is True:
263                self.visualiser.update_all()
264                self.visualiser.update_timer()
265
266
267            #Store model data, e.g. for subsequent visualisation
268            if self.store is True:
269                self.store_timestep(self.quantities_to_be_stored)
270
271            #FIXME: Could maybe be taken from specified list
272            #of 'store every step' quantities
273
274            #Pass control on to outer loop for more specific actions
275            yield(t)
276
277    def initialise_storage(self):
278        """Create and initialise self.writer object for storing data.
279        Also, save x,y and bed elevation
280        """
281
282        import data_manager
283
284        #Initialise writer
285        self.writer = data_manager.get_dataobject(self, mode = 'w')
286
287        #Store vertices and connectivity
288        self.writer.store_connectivity()
289
290
291    def store_timestep(self, name):
292        """Store named quantity and time.
293
294        Precondition:
295           self.write has been initialised
296        """
297        self.writer.store_timestep(name)
298
299
300#=============== End of Shallow Water Domain ===============================
301
302
303
304#Rotation of momentum vector
305def rotate(q, normal, direction = 1):
306    """Rotate the momentum component q (q[1], q[2])
307    from x,y coordinates to coordinates based on normal vector.
308
309    If direction is negative the rotation is inverted.
310
311    Input vector is preserved
312
313    This function is specific to the shallow water wave equation
314    """
315
316    from Numeric import zeros, Float
317
318    assert len(q) == 3,\
319           'Vector of conserved quantities must have length 3'\
320           'for 2D shallow water equation'
321
322    try:
323        l = len(normal)
324    except:
325        raise 'Normal vector must be an Numeric array'
326
327    assert l == 2, 'Normal vector must have 2 components'
328
329
330    n1 = normal[0]
331    n2 = normal[1]
332
333    r = zeros(len(q), Float) #Rotated quantities
334    r[0] = q[0]              #First quantity, height, is not rotated
335
336    if direction == -1:
337        n2 = -n2
338
339
340    r[1] =  n1*q[1] + n2*q[2]
341    r[2] = -n2*q[1] + n1*q[2]
342
343    return r
344
345
346
347####################################
348# Flux computation
349def flux_function(normal, ql, qr, zl, zr):
350    """Compute fluxes between volumes for the shallow water wave equation
351    cast in terms of w = h+z using the 'central scheme' as described in
352
353    Kurganov, Noelle, Petrova. 'Semidiscrete Central-Upwind Schemes For
354    Hyperbolic Conservation Laws and Hamilton-Jacobi Equations'.
355    Siam J. Sci. Comput. Vol. 23, No. 3, pp. 707-740.
356
357    The implemented formula is given in equation (3.15) on page 714
358
359    Conserved quantities w, uh, vh are stored as elements 0, 1 and 2
360    in the numerical vectors ql an qr.
361
362    Bed elevations zl and zr.
363    """
364
365    from config import g, epsilon
366    from math import sqrt
367    from Numeric import array
368
369    #Align momentums with x-axis
370    q_left  = rotate(ql, normal, direction = 1)
371    q_right = rotate(qr, normal, direction = 1)
372
373    z = (zl+zr)/2 #Take average of field values
374
375    w_left  = q_left[0]   #w=h+z
376    h_left  = w_left-z
377    uh_left = q_left[1]
378
379    if h_left < epsilon:
380        u_left = 0.0  #Could have been negative
381        h_left = 0.0
382    else:
383        u_left  = uh_left/h_left
384
385
386    w_right  = q_right[0]  #w=h+z
387    h_right  = w_right-z
388    uh_right = q_right[1]
389
390
391    if h_right < epsilon:
392        u_right = 0.0  #Could have been negative
393        h_right = 0.0
394    else:
395        u_right  = uh_right/h_right
396
397    vh_left  = q_left[2]
398    vh_right = q_right[2]
399
400    soundspeed_left  = sqrt(g*h_left)
401    soundspeed_right = sqrt(g*h_right)
402
403    #Maximal wave speed
404    s_max = max(u_left+soundspeed_left, u_right+soundspeed_right, 0)
405
406    #Minimal wave speed
407    s_min = min(u_left-soundspeed_left, u_right-soundspeed_right, 0)
408
409    #Flux computation
410
411    #FIXME(Ole): Why is it again that we don't
412    #use uh_left and uh_right directly in the first entries?
413    flux_left  = array([u_left*h_left,
414                        u_left*uh_left + 0.5*g*h_left**2,
415                        u_left*vh_left])
416    flux_right = array([u_right*h_right,
417                        u_right*uh_right + 0.5*g*h_right**2,
418                        u_right*vh_right])
419
420    denom = s_max-s_min
421    if denom == 0.0:
422        edgeflux = array([0.0, 0.0, 0.0])
423        max_speed = 0.0
424    else:
425        edgeflux = (s_max*flux_left - s_min*flux_right)/denom
426        edgeflux += s_max*s_min*(q_right-q_left)/denom
427
428        edgeflux = rotate(edgeflux, normal, direction=-1)
429        max_speed = max(abs(s_max), abs(s_min))
430
431    return edgeflux, max_speed
432
433
434def compute_fluxes(domain):
435    """Compute all fluxes and the timestep suitable for all volumes
436    in domain.
437
438    Compute total flux for each conserved quantity using "flux_function"
439
440    Fluxes across each edge are scaled by edgelengths and summed up
441    Resulting flux is then scaled by area and stored in
442    explicit_update for each of the three conserved quantities
443    stage, xmomentum and ymomentum
444
445    The maximal allowable speed computed by the flux_function for each volume
446    is converted to a timestep that must not be exceeded. The minimum of
447    those is computed as the next overall timestep.
448
449    Post conditions:
450      domain.explicit_update is reset to computed flux values
451      domain.timestep is set to the largest step satisfying all volumes.
452    """
453
454    import sys
455    from Numeric import zeros, Float
456
457    N = domain.number_of_elements
458
459    #Shortcuts
460    Stage = domain.quantities['stage']
461    Xmom = domain.quantities['xmomentum']
462    Ymom = domain.quantities['ymomentum']
463    Bed = domain.quantities['elevation']
464
465    #Arrays
466    stage = Stage.edge_values
467    xmom =  Xmom.edge_values
468    ymom =  Ymom.edge_values
469    bed =   Bed.edge_values
470
471    stage_bdry = Stage.boundary_values
472    xmom_bdry =  Xmom.boundary_values
473    ymom_bdry =  Ymom.boundary_values
474
475    flux = zeros(3, Float) #Work array for summing up fluxes
476
477
478    #Loop
479    timestep = float(sys.maxint)
480    for k in range(N):
481
482        flux[:] = 0.  #Reset work array
483        for i in range(3):
484            #Quantities inside volume facing neighbour i
485            ql = [stage[k, i], xmom[k, i], ymom[k, i]]
486            zl = bed[k, i]
487
488            #Quantities at neighbour on nearest face
489            n = domain.neighbours[k,i]
490            if n < 0:
491                m = -n-1 #Convert negative flag to index
492                qr = [stage_bdry[m], xmom_bdry[m], ymom_bdry[m]]
493                zr = zl #Extend bed elevation to boundary
494            else:
495                m = domain.neighbour_edges[k,i]
496                qr = [stage[n, m], xmom[n, m], ymom[n, m]]
497                zr = bed[n, m]
498
499
500            #Outward pointing normal vector
501            normal = domain.normals[k, 2*i:2*i+2]
502
503            #Flux computation using provided function
504            edgeflux, max_speed = flux_function(normal, ql, qr, zl, zr)
505            flux -= edgeflux * domain.edgelengths[k,i]
506
507            #Update optimal_timestep
508            try:
509                timestep = min(timestep, 0.5*domain.radii[k]/max_speed)
510            except ZeroDivisionError:
511                pass
512
513        #Normalise by area and store for when all conserved
514        #quantities get updated
515        flux /= domain.areas[k]
516        Stage.explicit_update[k] = flux[0]
517        Xmom.explicit_update[k] = flux[1]
518        Ymom.explicit_update[k] = flux[2]
519
520
521    domain.timestep = timestep
522
523#MH090605 The following method belongs to the shallow_water domain class
524#see comments in the corresponding method in shallow_water_ext.c
525def extrapolate_second_order_sw_c(domain):
526    """Wrapper calling C version of extrapolate_second_order_sw
527    """
528    import sys
529    from Numeric import zeros, Float
530
531    N = domain.number_of_elements
532
533    #Shortcuts
534    Stage = domain.quantities['stage']
535    Xmom = domain.quantities['xmomentum']
536    Ymom = domain.quantities['ymomentum']
537    from shallow_water_ext import extrapolate_second_order_sw
538    extrapolate_second_order_sw(domain,domain.surrogate_neighbours,
539                                domain.number_of_boundaries,
540                                domain.centroid_coordinates,
541                                Stage.centroid_values,
542                                Xmom.centroid_values,
543                                Ymom.centroid_values,
544                                domain.vertex_coordinates,
545                                Stage.vertex_values,
546                                Xmom.vertex_values,
547                                Ymom.vertex_values)
548
549def compute_fluxes_c(domain):
550    """Wrapper calling C version of compute fluxes
551    """
552
553    import sys
554    from Numeric import zeros, Float
555
556    N = domain.number_of_elements
557
558    #Shortcuts
559    Stage = domain.quantities['stage']
560    Xmom = domain.quantities['xmomentum']
561    Ymom = domain.quantities['ymomentum']
562    Bed = domain.quantities['elevation']
563
564    timestep = float(sys.maxint)
565    from shallow_water_ext import compute_fluxes
566    domain.timestep = compute_fluxes(timestep, domain.epsilon, domain.g,
567                                     domain.neighbours,
568                                     domain.neighbour_edges,
569                                     domain.normals,
570                                     domain.edgelengths,
571                                     domain.radii,
572                                     domain.areas,
573                                     Stage.edge_values,
574                                     Xmom.edge_values,
575                                     Ymom.edge_values,
576                                     Bed.edge_values,
577                                     Stage.boundary_values,
578                                     Xmom.boundary_values,
579                                     Ymom.boundary_values,
580                                     Stage.explicit_update,
581                                     Xmom.explicit_update,
582                                     Ymom.explicit_update,
583                                     domain.already_computed_flux)
584
585
586####################################
587# Module functions for gradient limiting (distribute_to_vertices_and_edges)
588
589def distribute_to_vertices_and_edges(domain):
590    """Distribution from centroids to vertices specific to the
591    shallow water wave
592    equation.
593
594    It will ensure that h (w-z) is always non-negative even in the
595    presence of steep bed-slopes by taking a weighted average between shallow
596    and deep cases.
597
598    In addition, all conserved quantities get distributed as per either a
599    constant (order==1) or a piecewise linear function (order==2).
600
601    FIXME: more explanation about removal of artificial variability etc
602
603    Precondition:
604      All quantities defined at centroids and bed elevation defined at
605      vertices.
606
607    Postcondition
608      Conserved quantities defined at vertices
609
610    """
611
612    from config import optimised_gradient_limiter
613
614    #Remove very thin layers of water
615    protect_against_infinitesimal_and_negative_heights(domain)
616
617    #Extrapolate all conserved quantities
618    if optimised_gradient_limiter:
619        #MH090605 if second order,
620        #perform the extrapolation and limiting on
621        #all of the conserved quantitie
622
623        if (domain.order == 1):
624            for name in domain.conserved_quantities:
625                Q = domain.quantities[name]
626                Q.extrapolate_first_order()
627        elif domain.order == 2:
628            domain.extrapolate_second_order_sw()
629        else:
630            raise 'Unknown order'
631    else:
632        #old code:
633        for name in domain.conserved_quantities:
634            Q = domain.quantities[name]
635            if domain.order == 1:
636                Q.extrapolate_first_order()
637            elif domain.order == 2:
638                Q.extrapolate_second_order()
639                Q.limit()
640            else:
641                raise 'Unknown order'
642
643
644    #Take bed elevation into account when water heights are small
645    balance_deep_and_shallow(domain)
646
647    #Compute edge values by interpolation
648    for name in domain.conserved_quantities:
649        Q = domain.quantities[name]
650        Q.interpolate_from_vertices_to_edges()
651
652
653def protect_against_infinitesimal_and_negative_heights(domain):
654    """Protect against infinitesimal heights and associated high velocities
655    """
656
657    #Shortcuts
658    wc = domain.quantities['stage'].centroid_values
659    zc = domain.quantities['elevation'].centroid_values
660    xmomc = domain.quantities['xmomentum'].centroid_values
661    ymomc = domain.quantities['ymomentum'].centroid_values
662    hc = wc - zc  #Water depths at centroids
663
664    #Update
665    for k in range(domain.number_of_elements):
666
667        if hc[k] < domain.minimum_allowed_height:
668            #Control stage
669            if hc[k] < domain.epsilon:
670                wc[k] = zc[k] # Contain 'lost mass' error
671
672            #Control momentum
673            xmomc[k] = ymomc[k] = 0.0
674
675
676def protect_against_infinitesimal_and_negative_heights_c(domain):
677    """Protect against infinitesimal heights and associated high velocities
678    """
679
680    #Shortcuts
681    wc = domain.quantities['stage'].centroid_values
682    zc = domain.quantities['elevation'].centroid_values
683    xmomc = domain.quantities['xmomentum'].centroid_values
684    ymomc = domain.quantities['ymomentum'].centroid_values
685
686    from shallow_water_ext import protect
687
688    protect(domain.minimum_allowed_height, domain.epsilon,
689            wc, zc, xmomc, ymomc)
690
691
692
693def h_limiter(domain):
694    """Limit slopes for each volume to eliminate artificial variance
695    introduced by e.g. second order extrapolator
696
697    limit on h = w-z
698
699    This limiter depends on two quantities (w,z) so it resides within
700    this module rather than within quantity.py
701    """
702
703    from Numeric import zeros, Float
704
705    N = domain.number_of_elements
706    beta_h = domain.beta_h
707
708    #Shortcuts
709    wc = domain.quantities['stage'].centroid_values
710    zc = domain.quantities['elevation'].centroid_values
711    hc = wc - zc
712
713    wv = domain.quantities['stage'].vertex_values
714    zv = domain.quantities['elevation'].vertex_values
715    hv = wv-zv
716
717    hvbar = zeros(hv.shape, Float) #h-limited values
718
719    #Find min and max of this and neighbour's centroid values
720    hmax = zeros(hc.shape, Float)
721    hmin = zeros(hc.shape, Float)
722
723    for k in range(N):
724        hmax[k] = hmin[k] = hc[k]
725        for i in range(3):
726            n = domain.neighbours[k,i]
727            if n >= 0:
728                hn = hc[n] #Neighbour's centroid value
729
730                hmin[k] = min(hmin[k], hn)
731                hmax[k] = max(hmax[k], hn)
732
733
734    #Diffences between centroids and maxima/minima
735    dhmax = hmax - hc
736    dhmin = hmin - hc
737
738    #Deltas between vertex and centroid values
739    dh = zeros(hv.shape, Float)
740    for i in range(3):
741        dh[:,i] = hv[:,i] - hc
742
743    #Phi limiter
744    for k in range(N):
745
746        #Find the gradient limiter (phi) across vertices
747        phi = 1.0
748        for i in range(3):
749            r = 1.0
750            if (dh[k,i] > 0): r = dhmax[k]/dh[k,i]
751            if (dh[k,i] < 0): r = dhmin[k]/dh[k,i]
752
753            phi = min( min(r*beta_h, 1), phi )
754
755        #Then update using phi limiter
756        for i in range(3):
757            hvbar[k,i] = hc[k] + phi*dh[k,i]
758
759    return hvbar
760
761
762
763def h_limiter_c(domain):
764    """Limit slopes for each volume to eliminate artificial variance
765    introduced by e.g. second order extrapolator
766
767    limit on h = w-z
768
769    This limiter depends on two quantities (w,z) so it resides within
770    this module rather than within quantity.py
771
772    Wrapper for c-extension
773    """
774
775    from Numeric import zeros, Float
776
777    N = domain.number_of_elements
778    beta_h = domain.beta_h
779
780    #Shortcuts
781    wc = domain.quantities['stage'].centroid_values
782    zc = domain.quantities['elevation'].centroid_values
783    hc = wc - zc
784
785    wv = domain.quantities['stage'].vertex_values
786    zv = domain.quantities['elevation'].vertex_values
787    hv = wv - zv
788
789    #Call C-extension
790    from shallow_water_ext import h_limiter_sw as h_limiter
791    hvbar = h_limiter(domain, hc, hv)
792
793    return hvbar
794
795
796def balance_deep_and_shallow(domain):
797    """Compute linear combination between stage as computed by
798    gradient-limiters limiting using w, and stage computed by
799    gradient-limiters limiting using h (h-limiter).
800    The former takes precedence when heights are large compared to the
801    bed slope while the latter takes precedence when heights are
802    relatively small.  Anything in between is computed as a balanced
803    linear combination in order to avoid numerical disturbances which
804    would otherwise appear as a result of hard switching between
805    modes.
806
807    The h-limiter is always applied irrespective of the order.
808    """
809
810    #Shortcuts
811    wc = domain.quantities['stage'].centroid_values
812    zc = domain.quantities['elevation'].centroid_values
813    hc = wc - zc
814
815    wv = domain.quantities['stage'].vertex_values
816    zv = domain.quantities['elevation'].vertex_values
817    hv = wv-zv
818
819    #Limit h
820    hvbar = h_limiter(domain)
821
822    for k in range(domain.number_of_elements):
823        #Compute maximal variation in bed elevation
824        #  This quantitiy is
825        #    dz = max_i abs(z_i - z_c)
826        #  and it is independent of dimension
827        #  In the 1d case zc = (z0+z1)/2
828        #  In the 2d case zc = (z0+z1+z2)/3
829
830        dz = max(abs(zv[k,0]-zc[k]),
831                 abs(zv[k,1]-zc[k]),
832                 abs(zv[k,2]-zc[k]))
833
834
835        hmin = min( hv[k,:] )
836
837        #Create alpha in [0,1], where alpha==0 means using the h-limited
838        #stage and alpha==1 means using the w-limited stage as
839        #computed by the gradient limiter (both 1st or 2nd order)
840
841        #If hmin > dz/2 then alpha = 1 and the bed will have no effect
842        #If hmin < 0 then alpha = 0 reverting to constant height above bed.
843
844        if dz > 0.0:
845            alpha = max( min( 2*hmin/dz, 1.0), 0.0 )
846        else:
847            #Flat bed
848            alpha = 1.0
849
850        #Let
851        #
852        #  wvi be the w-limited stage (wvi = zvi + hvi)
853        #  wvi- be the h-limited state (wvi- = zvi + hvi-)
854        #
855        #
856        #where i=0,1,2 denotes the vertex ids
857        #
858        #Weighted balance between w-limited and h-limited stage is
859        #
860        #  wvi := (1-alpha)*(zvi+hvi-) + alpha*(zvi+hvi)
861        #
862        #It follows that the updated wvi is
863        #  wvi := zvi + (1-alpha)*hvi- + alpha*hvi
864        #
865        # Momentum is balanced between constant and limited
866
867
868        #for i in range(3):
869        #    wv[k,i] = zv[k,i] + hvbar[k,i]
870
871        #return
872
873        if alpha < 1:
874
875            for i in range(3):
876                wv[k,i] = zv[k,i] + (1-alpha)*hvbar[k,i] + alpha*hv[k,i]
877
878            #Momentums at centroids
879            xmomc = domain.quantities['xmomentum'].centroid_values
880            ymomc = domain.quantities['ymomentum'].centroid_values
881
882            #Momentums at vertices
883            xmomv = domain.quantities['xmomentum'].vertex_values
884            ymomv = domain.quantities['ymomentum'].vertex_values
885
886            # Update momentum as a linear combination of
887            # xmomc and ymomc (shallow) and momentum
888            # from extrapolator xmomv and ymomv (deep).
889            xmomv[k,:] = (1-alpha)*xmomc[k] + alpha*xmomv[k,:]
890            ymomv[k,:] = (1-alpha)*ymomc[k] + alpha*ymomv[k,:]
891
892
893def balance_deep_and_shallow_c(domain):
894    """Wrapper for C implementation
895    """
896
897    #Shortcuts
898    wc = domain.quantities['stage'].centroid_values
899    zc = domain.quantities['elevation'].centroid_values
900    hc = wc - zc
901
902    wv = domain.quantities['stage'].vertex_values
903    zv = domain.quantities['elevation'].vertex_values
904    hv = wv - zv
905
906    #Momentums at centroids
907    xmomc = domain.quantities['xmomentum'].centroid_values
908    ymomc = domain.quantities['ymomentum'].centroid_values
909
910    #Momentums at vertices
911    xmomv = domain.quantities['xmomentum'].vertex_values
912    ymomv = domain.quantities['ymomentum'].vertex_values
913
914    #Limit h
915    hvbar = h_limiter(domain)
916
917    #This is how one would make a first order h_limited value
918    #as in the old balancer (pre 17 Feb 2005):
919    #from Numeric import zeros, Float
920    #hvbar = zeros( (len(hc), 3), Float)
921    #for i in range(3):
922    #    hvbar[:,i] = hc[:]
923
924    from shallow_water_ext import balance_deep_and_shallow
925    balance_deep_and_shallow(wc, zc, hc, wv, zv, hv, hvbar,
926                             xmomc, ymomc, xmomv, ymomv)
927
928
929
930
931###############################################
932#Boundaries - specific to the shallow water wave equation
933class Reflective_boundary(Boundary):
934    """Reflective boundary returns same conserved quantities as
935    those present in its neighbour volume but reflected.
936
937    This class is specific to the shallow water equation as it
938    works with the momentum quantities assumed to be the second
939    and third conserved quantities.
940    """
941
942    def __init__(self, domain = None):
943        Boundary.__init__(self)
944
945        if domain is None:
946            msg = 'Domain must be specified for reflective boundary'
947            raise msg
948
949        #Handy shorthands
950        self.stage   = domain.quantities['stage'].edge_values
951        self.xmom    = domain.quantities['xmomentum'].edge_values
952        self.ymom    = domain.quantities['ymomentum'].edge_values
953        self.normals = domain.normals
954
955        from Numeric import zeros, Float
956        self.conserved_quantities = zeros(3, Float)
957
958    def __repr__(self):
959        return 'Reflective_boundary'
960
961
962    def evaluate(self, vol_id, edge_id):
963        """Reflective boundaries reverses the outward momentum
964        of the volume they serve.
965        """
966
967        q = self.conserved_quantities
968        q[0] = self.stage[vol_id, edge_id]
969        q[1] = self.xmom[vol_id, edge_id]
970        q[2] = self.ymom[vol_id, edge_id]
971
972        normal = self.normals[vol_id, 2*edge_id:2*edge_id+2]
973
974
975        r = rotate(q, normal, direction = 1)
976        r[1] = -r[1]
977        q = rotate(r, normal, direction = -1)
978
979        return q
980
981
982
983class Transmissive_Momentum_Set_Stage_boundary(Boundary):
984    """Returns same momentum conserved quantities as
985    those present in its neighbour volume. Sets stage
986
987    Underlying domain must be specified when boundary is instantiated
988    """
989
990    def __init__(self, domain = None, function=None):
991        Boundary.__init__(self)
992
993        if domain is None:
994            msg = 'Domain must be specified for this type boundary'
995            raise msg
996
997        if function is None:
998            msg = 'Function must be specified for this type boundary'
999            raise msg
1000
1001        self.domain   = domain
1002        self.function = function
1003
1004    def __repr__(self):
1005        return 'Transmissive_Momentum_Set_Stage_boundary(%s)' %self.domain
1006
1007    def evaluate(self, vol_id, edge_id):
1008        """Transmissive Momentum Set Stage boundaries return the edge momentum
1009        values of the volume they serve.
1010        """
1011
1012        q = self.domain.get_conserved_quantities(vol_id, edge = edge_id)
1013        value = self.function(self.domain.time)
1014        q[0] = value[0]
1015        return q
1016
1017
1018        #FIXME: Consider this (taken from File_boundary) to allow
1019        #spatial variation
1020        #if vol_id is not None and edge_id is not None:
1021        #    i = self.boundary_indices[ vol_id, edge_id ]
1022        #    return self.F(t, point_id = i)
1023        #else:
1024        #    return self.F(t)
1025
1026
1027
1028class Dirichlet_Discharge_boundary(Boundary):
1029    """
1030    Sets stage (stage0)
1031    Sets momentum (wh0) in the inward normal direction.
1032
1033    Underlying domain must be specified when boundary is instantiated
1034    """
1035
1036    def __init__(self, domain = None, stage0=None, wh0=None):
1037        Boundary.__init__(self)
1038
1039        if domain is None:
1040            msg = 'Domain must be specified for this type boundary'
1041            raise msg
1042
1043        if stage0 is None:
1044            raise 'set stage'
1045
1046        if wh0 is None:
1047            wh0 = 0.0
1048
1049        self.domain   = domain
1050        self.stage0  = stage0
1051        self.wh0 = wh0
1052
1053    def __repr__(self):
1054        return 'Dirichlet_Discharge_boundary(%s)' %self.domain
1055
1056    def evaluate(self, vol_id, edge_id):
1057        """Set discharge in the (inward) normal direction
1058        """
1059
1060        normal = self.domain.get_normal(vol_id,edge_id)
1061        q = [self.stage0, -self.wh0*normal[0], -self.wh0*normal[1]]
1062        return q
1063
1064
1065        #FIXME: Consider this (taken from File_boundary) to allow
1066        #spatial variation
1067        #if vol_id is not None and edge_id is not None:
1068        #    i = self.boundary_indices[ vol_id, edge_id ]
1069        #    return self.F(t, point_id = i)
1070        #else:
1071        #    return self.F(t)
1072
1073
1074
1075#class Spatio_temporal_boundary(Boundary):
1076#    """The spatio-temporal boundary, reads values for the conserved
1077#    quantities from an sww NetCDF file, and returns interpolated values
1078#    at the midpoints of each associated boundaty segment.
1079#    Time dependency is interpolated linearly as in util.File_function.#
1080#
1081#    Example:
1082#    Bf = Spatio_temporal_boundary('source_file.sww', domain)
1083#
1084#    """
1085Spatio_temporal_boundary = File_boundary
1086
1087
1088
1089
1090#########################
1091#Standard forcing terms:
1092#
1093def gravity(domain):
1094    """Apply gravitational pull in the presence of bed slope
1095    """
1096
1097    from util import gradient
1098    from Numeric import zeros, Float, array, sum
1099
1100    xmom = domain.quantities['xmomentum'].explicit_update
1101    ymom = domain.quantities['ymomentum'].explicit_update
1102
1103    Stage = domain.quantities['stage']
1104    Elevation = domain.quantities['elevation']
1105    h = Stage.edge_values - Elevation.edge_values
1106    v = Elevation.vertex_values
1107
1108    x = domain.get_vertex_coordinates()
1109    g = domain.g
1110
1111    for k in range(domain.number_of_elements):
1112        avg_h = sum( h[k,:] )/3
1113
1114        #Compute bed slope
1115        x0, y0, x1, y1, x2, y2 = x[k,:]
1116        z0, z1, z2 = v[k,:]
1117
1118        zx, zy = gradient(x0, y0, x1, y1, x2, y2, z0, z1, z2)
1119
1120        #Update momentum
1121        xmom[k] += -g*zx*avg_h
1122        ymom[k] += -g*zy*avg_h
1123
1124
1125def gravity_c(domain):
1126    """Wrapper calling C version
1127    """
1128
1129    xmom = domain.quantities['xmomentum'].explicit_update
1130    ymom = domain.quantities['ymomentum'].explicit_update
1131
1132    Stage = domain.quantities['stage']
1133    Elevation = domain.quantities['elevation']
1134    h = Stage.edge_values - Elevation.edge_values
1135    v = Elevation.vertex_values
1136
1137    x = domain.get_vertex_coordinates()
1138    g = domain.g
1139
1140
1141    from shallow_water_ext import gravity
1142    gravity(g, h, v, x, xmom, ymom)
1143
1144
1145def manning_friction(domain):
1146    """Apply (Manning) friction to water momentum
1147    """
1148
1149    from math import sqrt
1150
1151    w = domain.quantities['stage'].centroid_values
1152    z = domain.quantities['elevation'].centroid_values
1153    h = w-z
1154
1155    uh = domain.quantities['xmomentum'].centroid_values
1156    vh = domain.quantities['ymomentum'].centroid_values
1157    eta = domain.quantities['friction'].centroid_values
1158
1159    xmom_update = domain.quantities['xmomentum'].semi_implicit_update
1160    ymom_update = domain.quantities['ymomentum'].semi_implicit_update
1161
1162    N = domain.number_of_elements
1163    eps = domain.minimum_allowed_height
1164    g = domain.g
1165
1166    for k in range(N):
1167        if eta[k] >= eps:
1168            if h[k] >= eps:
1169                S = -g * eta[k]**2 * sqrt((uh[k]**2 + vh[k]**2))
1170                S /= h[k]**(7.0/3)
1171
1172                #Update momentum
1173                xmom_update[k] += S*uh[k]
1174                ymom_update[k] += S*vh[k]
1175
1176
1177def manning_friction_c(domain):
1178    """Wrapper for c version
1179    """
1180
1181
1182    xmom = domain.quantities['xmomentum']
1183    ymom = domain.quantities['ymomentum']
1184
1185    w = domain.quantities['stage'].centroid_values
1186    z = domain.quantities['elevation'].centroid_values
1187
1188    uh = xmom.centroid_values
1189    vh = ymom.centroid_values
1190    eta = domain.quantities['friction'].centroid_values
1191
1192    xmom_update = xmom.semi_implicit_update
1193    ymom_update = ymom.semi_implicit_update
1194
1195    N = domain.number_of_elements
1196    eps = domain.minimum_allowed_height
1197    g = domain.g
1198
1199    from shallow_water_ext import manning_friction
1200    manning_friction(g, eps, w, z, uh, vh, eta, xmom_update, ymom_update)
1201
1202
1203def linear_friction(domain):
1204    """Apply linear friction to water momentum
1205
1206    Assumes quantity: 'linear_friction' to be present
1207    """
1208
1209    from math import sqrt
1210
1211    w = domain.quantities['stage'].centroid_values
1212    z = domain.quantities['elevation'].centroid_values
1213    h = w-z
1214
1215    uh = domain.quantities['xmomentum'].centroid_values
1216    vh = domain.quantities['ymomentum'].centroid_values
1217    tau = domain.quantities['linear_friction'].centroid_values
1218
1219    xmom_update = domain.quantities['xmomentum'].semi_implicit_update
1220    ymom_update = domain.quantities['ymomentum'].semi_implicit_update
1221
1222    N = domain.number_of_elements
1223    eps = domain.minimum_allowed_height
1224    g = domain.g #Not necessary? Why was this added?
1225
1226    for k in range(N):
1227        if tau[k] >= eps:
1228            if h[k] >= eps:
1229                S = -tau[k]/h[k]
1230
1231                #Update momentum
1232                xmom_update[k] += S*uh[k]
1233                ymom_update[k] += S*vh[k]
1234
1235
1236
1237def check_forcefield(f):
1238    """Check that f is either
1239    1: a callable object f(t,x,y), where x and y are vectors
1240       and that it returns an array or a list of same length
1241       as x and y
1242    2: a scalar
1243    """
1244
1245    from Numeric import ones, Float, array
1246
1247
1248    if callable(f):
1249        N = 3
1250        x = ones(3, Float)
1251        y = ones(3, Float)
1252        try:
1253            q = f(1.0, x=x, y=y)
1254        except Exception, e:
1255            msg = 'Function %s could not be executed:\n%s' %(f, e)
1256            #FIXME: Reconsider this semantics
1257            raise msg
1258
1259        try:
1260            q = array(q).astype(Float)
1261        except:
1262            msg = 'Return value from vector function %s could ' %f
1263            msg += 'not be converted into a Numeric array of floats.\n'
1264            msg += 'Specified function should return either list or array.'
1265            raise msg
1266
1267        #Is this really what we want?
1268        msg = 'Return vector from function %s ' %f
1269        msg += 'must have same lenght as input vectors'
1270        assert len(q) == N, msg
1271
1272    else:
1273        try:
1274            f = float(f)
1275        except:
1276            msg = 'Force field %s must be either a scalar' %f
1277            msg += ' or a vector function'
1278            raise msg
1279    return f
1280
1281
1282class Wind_stress:
1283    """Apply wind stress to water momentum in terms of
1284    wind speed [m/s] and wind direction [degrees]
1285    """
1286
1287    def __init__(self, *args, **kwargs):
1288        """Initialise windfield from wind speed s [m/s]
1289        and wind direction phi [degrees]
1290
1291        Inputs v and phi can be either scalars or Python functions, e.g.
1292
1293        W = Wind_stress(10, 178)
1294
1295        #FIXME - 'normal' degrees are assumed for now, i.e. the
1296        vector (1,0) has zero degrees.
1297        We may need to convert from 'compass' degrees later on and also
1298        map from True north to grid north.
1299
1300        Arguments can also be Python functions of t,x,y as in
1301
1302        def speed(t,x,y):
1303            ...
1304            return s
1305
1306        def angle(t,x,y):
1307            ...
1308            return phi
1309
1310        where x and y are vectors.
1311
1312        and then pass the functions in
1313
1314        W = Wind_stress(speed, angle)
1315
1316        The instantiated object W can be appended to the list of
1317        forcing_terms as in
1318
1319        Alternatively, one vector valued function for (speed, angle)
1320        can be applied, providing both quantities simultaneously.
1321        As in
1322        W = Wind_stress(F), where returns (speed, angle) for each t.
1323
1324        domain.forcing_terms.append(W)
1325        """
1326
1327        from config import rho_a, rho_w, eta_w
1328        from Numeric import array, Float
1329
1330        if len(args) == 2:
1331            s = args[0]
1332            phi = args[1]
1333        elif len(args) == 1:
1334            #Assume vector function returning (s, phi)(t,x,y)
1335            vector_function = args[0]
1336            s = lambda t,x,y: vector_function(t,x=x,y=y)[0]
1337            phi = lambda t,x,y: vector_function(t,x=x,y=y)[1]
1338        else:
1339           #Assume info is in 2 keyword arguments
1340
1341           if len(kwargs) == 2:
1342               s = kwargs['s']
1343               phi = kwargs['phi']
1344           else:
1345               raise 'Assumes two keyword arguments: s=..., phi=....'
1346
1347        self.speed = check_forcefield(s)
1348        self.phi = check_forcefield(phi)
1349
1350        self.const = eta_w*rho_a/rho_w
1351
1352
1353    def __call__(self, domain):
1354        """Evaluate windfield based on values found in domain
1355        """
1356
1357        from math import pi, cos, sin, sqrt
1358        from Numeric import Float, ones, ArrayType
1359
1360        xmom_update = domain.quantities['xmomentum'].explicit_update
1361        ymom_update = domain.quantities['ymomentum'].explicit_update
1362
1363        N = domain.number_of_elements
1364        t = domain.time
1365
1366        if callable(self.speed):
1367            xc = domain.get_centroid_coordinates()
1368            s_vec = self.speed(t, xc[:,0], xc[:,1])
1369        else:
1370            #Assume s is a scalar
1371
1372            try:
1373                s_vec = self.speed * ones(N, Float)
1374            except:
1375                msg = 'Speed must be either callable or a scalar: %s' %self.s
1376                raise msg
1377
1378
1379        if callable(self.phi):
1380            xc = domain.get_centroid_coordinates()
1381            phi_vec = self.phi(t, xc[:,0], xc[:,1])
1382        else:
1383            #Assume phi is a scalar
1384
1385            try:
1386                phi_vec = self.phi * ones(N, Float)
1387            except:
1388                msg = 'Angle must be either callable or a scalar: %s' %self.phi
1389                raise msg
1390
1391        assign_windfield_values(xmom_update, ymom_update,
1392                                s_vec, phi_vec, self.const)
1393
1394
1395def assign_windfield_values(xmom_update, ymom_update,
1396                            s_vec, phi_vec, const):
1397    """Python version of assigning wind field to update vectors.
1398    A c version also exists (for speed)
1399    """
1400    from math import pi, cos, sin, sqrt
1401
1402    N = len(s_vec)
1403    for k in range(N):
1404        s = s_vec[k]
1405        phi = phi_vec[k]
1406
1407        #Convert to radians
1408        phi = phi*pi/180
1409
1410        #Compute velocity vector (u, v)
1411        u = s*cos(phi)
1412        v = s*sin(phi)
1413
1414        #Compute wind stress
1415        S = const * sqrt(u**2 + v**2)
1416        xmom_update[k] += S*u
1417        ymom_update[k] += S*v
1418
1419
1420
1421##############################
1422#OBSOLETE STUFF
1423
1424def balance_deep_and_shallow_old(domain):
1425    """Compute linear combination between stage as computed by
1426    gradient-limiters and stage computed as constant height above bed.
1427    The former takes precedence when heights are large compared to the
1428    bed slope while the latter takes precedence when heights are
1429    relatively small.  Anything in between is computed as a balanced
1430    linear combination in order to avoid numerical disturbances which
1431    would otherwise appear as a result of hard switching between
1432    modes.
1433    """
1434
1435    #OBSOLETE
1436
1437    #Shortcuts
1438    wc = domain.quantities['stage'].centroid_values
1439    zc = domain.quantities['elevation'].centroid_values
1440    hc = wc - zc
1441
1442    wv = domain.quantities['stage'].vertex_values
1443    zv = domain.quantities['elevation'].vertex_values
1444    hv = wv-zv
1445
1446
1447    #Computed linear combination between constant stages and and
1448    #stages parallel to the bed elevation.
1449    for k in range(domain.number_of_elements):
1450        #Compute maximal variation in bed elevation
1451        #  This quantitiy is
1452        #    dz = max_i abs(z_i - z_c)
1453        #  and it is independent of dimension
1454        #  In the 1d case zc = (z0+z1)/2
1455        #  In the 2d case zc = (z0+z1+z2)/3
1456
1457        dz = max(abs(zv[k,0]-zc[k]),
1458                 abs(zv[k,1]-zc[k]),
1459                 abs(zv[k,2]-zc[k]))
1460
1461
1462        hmin = min( hv[k,:] )
1463
1464        #Create alpha in [0,1], where alpha==0 means using shallow
1465        #first order scheme and alpha==1 means using the stage w as
1466        #computed by the gradient limiter (1st or 2nd order)
1467        #
1468        #If hmin > dz/2 then alpha = 1 and the bed will have no effect
1469        #If hmin < 0 then alpha = 0 reverting to constant height above bed.
1470
1471        if dz > 0.0:
1472            alpha = max( min( 2*hmin/dz, 1.0), 0.0 )
1473        else:
1474            #Flat bed
1475            alpha = 1.0
1476
1477        #Weighted balance between stage parallel to bed elevation
1478        #(wvi = zvi + hc) and stage as computed by 1st or 2nd
1479        #order gradient limiter
1480        #(wvi = zvi + hvi) where i=0,1,2 denotes the vertex ids
1481        #
1482        #It follows that the updated wvi is
1483        #  wvi := (1-alpha)*(zvi+hc) + alpha*(zvi+hvi) =
1484        #  zvi + hc + alpha*(hvi - hc)
1485        #
1486        #Note that hvi = zc+hc-zvi in the first order case (constant).
1487
1488        if alpha < 1:
1489            for i in range(3):
1490                wv[k,i] = zv[k,i] + hc[k] + alpha*(hv[k,i]-hc[k])
1491
1492
1493            #Momentums at centroids
1494            xmomc = domain.quantities['xmomentum'].centroid_values
1495            ymomc = domain.quantities['ymomentum'].centroid_values
1496
1497            #Momentums at vertices
1498            xmomv = domain.quantities['xmomentum'].vertex_values
1499            ymomv = domain.quantities['ymomentum'].vertex_values
1500
1501            # Update momentum as a linear combination of
1502            # xmomc and ymomc (shallow) and momentum
1503            # from extrapolator xmomv and ymomv (deep).
1504            xmomv[k,:] = (1-alpha)*xmomc[k] + alpha*xmomv[k,:]
1505            ymomv[k,:] = (1-alpha)*ymomc[k] + alpha*ymomv[k,:]
1506
1507
1508
1509
1510
1511###########################
1512###########################
1513#Geometries
1514
1515
1516#FIXME: Rethink this way of creating values.
1517
1518
1519class Weir:
1520    """Set a bathymetry for weir with a hole and a downstream gutter
1521    x,y are assumed to be in the unit square
1522    """
1523
1524    def __init__(self, stage):
1525        self.inflow_stage = stage
1526
1527    def __call__(self, x, y):
1528        from Numeric import zeros, Float
1529        from math import sqrt
1530
1531        N = len(x)
1532        assert N == len(y)
1533
1534        z = zeros(N, Float)
1535        for i in range(N):
1536            z[i] = -x[i]/2  #General slope
1537
1538            #Flattish bit to the left
1539            if x[i] < 0.3:
1540                z[i] = -x[i]/10
1541
1542            #Weir
1543            if x[i] >= 0.3 and x[i] < 0.4:
1544                z[i] = -x[i]+0.9
1545
1546            #Dip
1547            x0 = 0.6
1548            #depth = -1.3
1549            depth = -1.0
1550            #plateaux = -0.9
1551            plateaux = -0.6
1552            if y[i] < 0.7:
1553                if x[i] > x0 and x[i] < 0.9:
1554                    z[i] = depth
1555
1556                #RHS plateaux
1557                if x[i] >= 0.9:
1558                    z[i] = plateaux
1559
1560
1561            elif y[i] >= 0.7 and y[i] < 1.5:
1562                #Restrict and deepen
1563                if x[i] >= x0 and x[i] < 0.8:
1564                    z[i] = depth-(y[i]/3-0.3)
1565                    #z[i] = depth-y[i]/5
1566                    #z[i] = depth
1567                elif x[i] >= 0.8:
1568                    #RHS plateaux
1569                    z[i] = plateaux
1570
1571            elif y[i] >= 1.5:
1572                if x[i] >= x0 and x[i] < 0.8 + (y[i]-1.5)/1.2:
1573                    #Widen up and stay at constant depth
1574                    z[i] = depth-1.5/5
1575                elif x[i] >= 0.8 + (y[i]-1.5)/1.2:
1576                    #RHS plateaux
1577                    z[i] = plateaux
1578
1579
1580            #Hole in weir (slightly higher than inflow condition)
1581            if x[i] >= 0.3 and x[i] < 0.4 and y[i] > 0.2 and y[i] < 0.4:
1582                z[i] = -x[i]+self.inflow_stage + 0.02
1583
1584            #Channel behind weir
1585            x0 = 0.5
1586            if x[i] >= 0.4 and x[i] < x0 and y[i] > 0.2 and y[i] < 0.4:
1587                z[i] = -x[i]+self.inflow_stage + 0.02
1588
1589            if x[i] >= x0 and x[i] < 0.6 and y[i] > 0.2 and y[i] < 0.4:
1590                #Flatten it out towards the end
1591                z[i] = -x0+self.inflow_stage + 0.02 + (x0-x[i])/5
1592
1593            #Hole to the east
1594            x0 = 1.1; y0 = 0.35
1595            #if x[i] < -0.2 and y < 0.5:
1596            if sqrt((2*(x[i]-x0))**2 + (2*(y[i]-y0))**2) < 0.2:
1597                z[i] = sqrt(((x[i]-x0))**2 + ((y[i]-y0))**2)-1.0
1598
1599            #Tiny channel draining hole
1600            if x[i] >= 1.14 and x[i] < 1.2 and y[i] >= 0.4 and y[i] < 0.6:
1601                z[i] = -0.9 #North south
1602
1603            if x[i] >= 0.9 and x[i] < 1.18 and y[i] >= 0.58 and y[i] < 0.65:
1604                z[i] = -1.0 + (x[i]-0.9)/3 #East west
1605
1606
1607
1608            #Stuff not in use
1609
1610            #Upward slope at inlet to the north west
1611            #if x[i] < 0.0: # and y[i] > 0.5:
1612            #    #z[i] = -y[i]+0.5  #-x[i]/2
1613            #    z[i] = x[i]/4 - y[i]**2 + 0.5
1614
1615            #Hole to the west
1616            #x0 = -0.4; y0 = 0.35 # center
1617            #if sqrt((2*(x[i]-x0))**2 + (2*(y[i]-y0))**2) < 0.2:
1618            #    z[i] = sqrt(((x[i]-x0))**2 + ((y[i]-y0))**2)-0.2
1619
1620
1621
1622
1623
1624        return z/2
1625
1626class Weir_simple:
1627    """Set a bathymetry for weir with a hole and a downstream gutter
1628    x,y are assumed to be in the unit square
1629    """
1630
1631    def __init__(self, stage):
1632        self.inflow_stage = stage
1633
1634    def __call__(self, x, y):
1635        from Numeric import zeros, Float
1636
1637        N = len(x)
1638        assert N == len(y)
1639
1640        z = zeros(N, Float)
1641        for i in range(N):
1642            z[i] = -x[i]  #General slope
1643
1644            #Flat bit to the left
1645            if x[i] < 0.3:
1646                z[i] = -x[i]/10  #General slope
1647
1648            #Weir
1649            if x[i] > 0.3 and x[i] < 0.4:
1650                z[i] = -x[i]+0.9
1651
1652            #Dip
1653            if x[i] > 0.6 and x[i] < 0.9:
1654                z[i] = -x[i]-0.5  #-y[i]/5
1655
1656            #Hole in weir (slightly higher than inflow condition)
1657            if x[i] > 0.3 and x[i] < 0.4 and y[i] > 0.2 and y[i] < 0.4:
1658                z[i] = -x[i]+self.inflow_stage + 0.05
1659
1660
1661        return z/2
1662
1663
1664
1665class Constant_stage:
1666    """Set an initial condition with constant stage
1667    """
1668    def __init__(self, s):
1669        self.s = s
1670
1671    def __call__(self, x, y):
1672        return self.s
1673
1674class Constant_height:
1675    """Set an initial condition with constant water height, e.g
1676    stage s = z+h
1677    """
1678
1679    def __init__(self, W, h):
1680        self.W = W
1681        self.h = h
1682
1683    def __call__(self, x, y):
1684        if self.W is None:
1685            from Numeric import ones, Float
1686            return self.h*ones(len(x), Float)
1687        else:
1688            return self.W(x,y) + self.h
1689
1690
1691
1692
1693def compute_fluxes_python(domain):
1694    """Compute all fluxes and the timestep suitable for all volumes
1695    in domain.
1696
1697    Compute total flux for each conserved quantity using "flux_function"
1698
1699    Fluxes across each edge are scaled by edgelengths and summed up
1700    Resulting flux is then scaled by area and stored in
1701    explicit_update for each of the three conserved quantities
1702    stage, xmomentum and ymomentum
1703
1704    The maximal allowable speed computed by the flux_function for each volume
1705    is converted to a timestep that must not be exceeded. The minimum of
1706    those is computed as the next overall timestep.
1707
1708    Post conditions:
1709      domain.explicit_update is reset to computed flux values
1710      domain.timestep is set to the largest step satisfying all volumes.
1711    """
1712
1713    import sys
1714    from Numeric import zeros, Float
1715
1716    N = domain.number_of_elements
1717
1718    #Shortcuts
1719    Stage = domain.quantities['stage']
1720    Xmom = domain.quantities['xmomentum']
1721    Ymom = domain.quantities['ymomentum']
1722    Bed = domain.quantities['elevation']
1723
1724    #Arrays
1725    stage = Stage.edge_values
1726    xmom =  Xmom.edge_values
1727    ymom =  Ymom.edge_values
1728    bed =   Bed.edge_values
1729
1730    stage_bdry = Stage.boundary_values
1731    xmom_bdry =  Xmom.boundary_values
1732    ymom_bdry =  Ymom.boundary_values
1733
1734    flux = zeros((N,3), Float) #Work array for summing up fluxes
1735
1736    #Loop
1737    timestep = float(sys.maxint)
1738    for k in range(N):
1739
1740        for i in range(3):
1741            #Quantities inside volume facing neighbour i
1742            ql = [stage[k, i], xmom[k, i], ymom[k, i]]
1743            zl = bed[k, i]
1744
1745            #Quantities at neighbour on nearest face
1746            n = domain.neighbours[k,i]
1747            if n < 0:
1748                m = -n-1 #Convert negative flag to index
1749                qr = [stage_bdry[m], xmom_bdry[m], ymom_bdry[m]]
1750                zr = zl #Extend bed elevation to boundary
1751            else:
1752                m = domain.neighbour_edges[k,i]
1753                qr = [stage[n, m], xmom[n, m], ymom[n, m]]
1754                zr = bed[n, m]
1755
1756
1757            #Outward pointing normal vector
1758            normal = domain.normals[k, 2*i:2*i+2]
1759
1760            #Flux computation using provided function
1761            edgeflux, max_speed = flux_function(normal, ql, qr, zl, zr)
1762
1763            flux[k,:] = edgeflux
1764
1765    return flux
1766
1767
1768
1769
1770
1771
1772
1773##############################################
1774#Initialise module
1775
1776
1777from utilities import compile
1778if compile.can_use_C_extension('shallow_water_ext.c'):
1779    #Replace python version with c implementations
1780
1781    from shallow_water_ext import rotate, assign_windfield_values
1782    compute_fluxes = compute_fluxes_c
1783    extrapolate_second_order_sw=extrapolate_second_order_sw_c
1784    gravity = gravity_c
1785    manning_friction = manning_friction_c
1786    h_limiter = h_limiter_c
1787    balance_deep_and_shallow = balance_deep_and_shallow_c
1788    protect_against_infinitesimal_and_negative_heights = protect_against_infinitesimal_and_negative_heights_c
1789
1790
1791    #distribute_to_vertices_and_edges = distribute_to_vertices_and_edges_c #(like MH's)
1792
1793
1794
1795#Optimisation with psyco
1796from config import use_psyco
1797if use_psyco:
1798    try:
1799        import psyco
1800    except:
1801        import os
1802        if os.name == 'posix' and os.uname()[4] == 'x86_64':
1803            pass
1804            #Psyco isn't supported on 64 bit systems, but it doesn't matter
1805        else:
1806            msg = 'WARNING: psyco (speedup) could not import'+\
1807                  ', you may want to consider installing it'
1808            print msg
1809    else:
1810        psyco.bind(Domain.distribute_to_vertices_and_edges)
1811        psyco.bind(Domain.compute_fluxes)
1812
1813if __name__ == "__main__":
1814    pass
1815
1816# Profiling stuff
1817import profile
1818profiler = profile.Profile()
Note: See TracBrowser for help on using the repository browser.