source: inundation/pyvolution/shallow_water.py @ 2300

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

Implemented set_quantities_to_be_stored and used it in bedslope.py

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