source: inundation/pyvolution/shallow_water.py @ 2113

Last change on this file since 2113 was 2112, checked in by steve, 19 years ago

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