source: inundation/ga/storm_surge/pyvolution/shallow_water.py @ 768

Last change on this file since 768 was 768, checked in by ole, 19 years ago

Comment about keyword expansion

  • Property svn:keywords set to LastChangedDate LastChangedRevision LastChangedBy HeadURL Id
File size: 39.1 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
19
20The quantities are
21
22symbol    variable name    explanation
23z         elevation        elevation of bed on which flow is modelled
24h         height           water height above z
25w         stage            water level, w = z+h
26u                          speed in the x direction
27v                          speed in the y direction
28uh        xmomentum        momentum in the x direction
29vh        ymomentum        momentum in the y direction
30
31eta                        mannings friction coefficient
32nu                         wind stress coefficient
33
34The conserved quantities are w, uh, vh
35
36
37For details see e.g.
38Christopher Zoppou and Stephen Roberts,
39Catastrophic Collapse of Water Supply Reservoirs in Urban Areas,
40Journal of Hydraulic Engineering, vol. 127, No. 7 July 1999
41
42
43
44Ole Nielsen, Stephen Roberts, Duncan Gray, Christopher Zoppou
45Geoscience Australia, 2004   
46"""
47
48#To set subversion keyword expansion:
49#
50#svn propset svn:keywords "LastChangedDate LastChangedRevision LastChangedBy HeadURL Id" shallow_water.py
51
52
53from domain import *
54from region import *
55Generic_domain = Domain #Rename
56
57class Domain(Generic_domain):
58
59    def __init__(self, coordinates, vertices, boundary = None,
60                 tagged_elements = None):
61
62        conserved_quantities = ['level', 'xmomentum', 'ymomentum']
63        other_quantities = ['elevation', 'friction'] 
64       
65        Generic_domain.__init__(self, coordinates, vertices, boundary,
66                                conserved_quantities, other_quantities,
67                                tagged_elements)
68
69        from config import minimum_allowed_height, g
70        self.minimum_allowed_height = minimum_allowed_height
71        self.g = g
72
73        self.forcing_terms.append(gravity)
74        self.forcing_terms.append(manning_friction)
75
76        #Realtime visualisation
77        self.visualise = False
78
79        #Stored output
80        self.store = False
81        self.format = 'sww'   
82        self.smooth = True
83
84        #Reduction operation for get_vertex_values             
85        #from util import mean
86        #self.reduction = mean
87        self.reduction = min  #Looks better near steep slopes
88       
89        self.quantities_to_be_stored = ['level']
90
91       
92        #Establish shortcuts to relevant quantities (for efficiency)
93        #self.w = self.quantities['level']
94        #self.uh = self.quantities['xmomentum']                 
95        #self.vh = self.quantities['ymomentum']                         
96        #self.z = self.quantities['elevation']         
97        #self.eta = self.quantities['friction']                 
98
99    def check_integrity(self):
100        Generic_domain.check_integrity(self)
101
102        #Check that we are solving the shallow water wave equation
103
104        msg = 'First conserved quantity must be "level"'
105        assert self.conserved_quantities[0] == 'level', msg
106        msg = 'Second conserved quantity must be "xmomentum"'
107        assert self.conserved_quantities[1] == 'xmomentum', msg
108        msg = 'Third conserved quantity must be "ymomentum"'
109        assert self.conserved_quantities[2] == 'ymomentum', msg
110       
111
112        #Check that levels are >= bed elevation
113        from Numeric import alltrue, greater_equal
114       
115        level = self.quantities['level']
116        bed = self.quantities['elevation']       
117
118
119    def compute_fluxes(self):
120        #Call correct module function
121        #(either from this module or C-extension)
122        compute_fluxes(self)
123
124    def distribute_to_vertices_and_edges(self):
125        #Call correct module function
126        #(either from this module or C-extension)       
127        distribute_to_vertices_and_edges(self)
128
129
130    #FIXME: Under construction   
131#     def set_defaults(self):
132#         """Set default values for uninitialised quantities.
133#         This is specific to the shallow water wave equation
134#         Defaults for 'elevation', 'friction', 'xmomentum' and 'ymomentum'
135#         are 0.0. Default for 'level' is whatever the value of 'elevation'.
136#         """
137
138#         for name in self.other_quantities + self.conserved_quantities:
139#             print name
140#             print self.quantities.keys()
141#             if not self.quantities.has_key(name):
142#                 if name == 'level':
143                   
144#                     if self.quantities.has_key('elevation'):
145#                         z = self.quantities['elevation'].vertex_values
146#                         self.set_quantity(name, z)
147#                     else:
148#                         self.set_quantity(name, 0.0)                       
149#                 else:   
150#                     self.set_quantity(name, 0.0)
151                   
152
153
154#         #Lift negative heights up
155#         #z = self.quantities['elevation'].vertex_values
156#         #w = self.quantities['level'].vertex_values
157
158#         #h = w-z
159
160#         #for k in range(h.shape[0]):
161#         #    for i in range(3):
162#         #        if h[k, i] < 0.0:
163#         #            w[k, i] = z[k, i]
164
165                   
166#         #self.quantities['level'].interpolate()
167               
168
169
170    def evolve(self, yieldstep = None, finaltime = None):
171        """Specialisation of basic evolve method from parent class
172        """
173
174        #Call check integrity here rather than from user scripts
175        #self.check_integrity()
176       
177        #Initialise real time viz if requested
178        if self.visualise is True and self.time == 0.0:
179            import realtime_visualisation as visualise
180            visualise.create_surface(self)
181
182        #Store model data, e.g. for visualisation   
183        if self.store is True and self.time == 0.0:
184            self.initialise_storage()
185            #print 'Storing results in ' + self.writer.filename
186        else:
187            #print 'Results will not be stored.'
188            #print 'To store results set domain.store = True'
189            pass
190            #FIXME: Diagnostic output should be controlled by
191            # a 'verbose' flag living in domain (or in a parent class)
192
193        #Call basic machinery from parent class
194        for t in Generic_domain.evolve(self, yieldstep, finaltime):
195            #Real time viz
196            if self.visualise is True:
197                visualise.update(self)
198
199            #Store model data, e.g. for subsequent visualisation   
200            if self.store is True:
201                #self.store_timestep(['level', 'xmomentum', 'ymomentum'])
202                self.store_timestep(self.quantities_to_be_stored)
203               
204                #FIXME: Could maybe be taken from specified list
205                #of 'store every step' quantities       
206
207            #Pass control on to outer loop for more specific actions   
208            yield(t)
209       
210
211    def initialise_storage(self):       
212        """Create and initialise self.writer object for storing data.
213        Also, save x,y and bed elevation
214        """
215
216        import data_manager
217       
218        #Initialise writer
219        self.writer = data_manager.get_dataobject(self, mode = 'w')
220
221        #Store vertices and connectivity
222        self.writer.store_connectivity()
223
224
225    def store_timestep(self, name):
226        """Store named quantity and time.
227
228        Precondition:
229           self.write has been initialised
230        """
231        self.writer.store_timestep(name)
232       
233
234#Rotation of momentum vector
235def rotate(q, normal, direction = 1):
236    """Rotate the momentum component q (q[1], q[2])
237    from x,y coordinates to coordinates based on normal vector.
238
239    If direction is negative the rotation is inverted.
240   
241    Input vector is preserved
242
243    This function is specific to the shallow water wave equation
244    """
245
246    from Numeric import zeros, Float
247   
248    assert len(q) == 3,\
249           'Vector of conserved quantities must have length 3'\
250           'for 2D shallow water equation'
251
252    try:
253        l = len(normal)
254    except:
255        raise 'Normal vector must be an Numeric array'
256
257    assert l == 2, 'Normal vector must have 2 components'
258   
259 
260    n1 = normal[0]
261    n2 = normal[1]   
262   
263    r = zeros(len(q), Float) #Rotated quantities
264    r[0] = q[0]              #First quantity, height, is not rotated
265
266    if direction == -1:
267        n2 = -n2
268
269
270    r[1] =  n1*q[1] + n2*q[2]
271    r[2] = -n2*q[1] + n1*q[2]
272   
273    return r
274
275
276
277####################################
278# Flux computation       
279def flux_function(normal, ql, qr, zl, zr): 
280    """Compute fluxes between volumes for the shallow water wave equation
281    cast in terms of w = h+z using the 'central scheme' as described in
282
283    Kurganov, Noelle, Petrova. 'Semidiscrete Central-Upwind Schemes For
284    Hyperbolic Conservation Laws and Hamilton-Jacobi Equations'.
285    Siam J. Sci. Comput. Vol. 23, No. 3, pp. 707-740.
286
287    The implemented formula is given in equation (3.15) on page 714
288
289    Conserved quantities w, uh, vh are stored as elements 0, 1 and 2
290    in the numerical vectors ql an qr.
291
292    Bed elevations zl and zr.
293    """
294
295    from config import g, epsilon
296    from math import sqrt
297    from Numeric import array
298
299    #Align momentums with x-axis
300    q_left  = rotate(ql, normal, direction = 1)
301    q_right = rotate(qr, normal, direction = 1)
302
303    z = (zl+zr)/2 #Take average of field values
304
305    w_left  = q_left[0]   #w=h+z
306    h_left  = w_left-z
307    uh_left = q_left[1]
308
309    if h_left < epsilon:
310        u_left = 0.0  #Could have been negative
311        h_left = 0.0
312    else:   
313        u_left  = uh_left/h_left
314
315
316    w_right  = q_right[0]  #w=h+z
317    h_right  = w_right-z
318    uh_right = q_right[1]
319
320
321    if h_right < epsilon:
322        u_right = 0.0  #Could have been negative
323        h_right = 0.0
324    else:   
325        u_right  = uh_right/h_right
326
327    vh_left  = q_left[2]
328    vh_right = q_right[2]       
329
330    soundspeed_left  = sqrt(g*h_left) 
331    soundspeed_right = sqrt(g*h_right)
332
333    #Maximal wave speed
334    s_max = max(u_left+soundspeed_left, u_right+soundspeed_right, 0)
335   
336    #Minimal wave speed       
337    s_min = min(u_left-soundspeed_left, u_right-soundspeed_right, 0)
338
339    #Flux computation
340    flux_left  = array([u_left*h_left,
341                        u_left*uh_left + 0.5*g*h_left**2,
342                        u_left*vh_left])
343    flux_right = array([u_right*h_right,
344                        u_right*uh_right + 0.5*g*h_right**2,
345                        u_right*vh_right])   
346
347    denom = s_max-s_min
348    if denom == 0.0:
349        edgeflux = array([0.0, 0.0, 0.0])
350        max_speed = 0.0
351    else:   
352        edgeflux = (s_max*flux_left - s_min*flux_right)/denom
353        edgeflux += s_max*s_min*(q_right-q_left)/denom
354
355        edgeflux = rotate(edgeflux, normal, direction=-1)
356        max_speed = max(abs(s_max), abs(s_min))
357
358    return edgeflux, max_speed       
359
360
361def compute_fluxes(domain):
362    """Compute all fluxes and the timestep suitable for all volumes
363    in domain.
364
365    Compute total flux for each conserved quantity using "flux_function"
366
367    Fluxes across each edge are scaled by edgelengths and summed up
368    Resulting flux is then scaled by area and stored in
369    explicit_update for each of the three conserved quantities
370    level, xmomentum and ymomentum   
371
372    The maximal allowable speed computed by the flux_function for each volume
373    is converted to a timestep that must not be exceeded. The minimum of
374    those is computed as the next overall timestep.
375
376    Post conditions:
377      domain.explicit_update is reset to computed flux values
378      domain.timestep is set to the largest step satisfying all volumes.
379    """
380
381    import sys
382    from Numeric import zeros, Float
383
384    N = domain.number_of_elements
385   
386    #Shortcuts
387    Level = domain.quantities['level']
388    Xmom = domain.quantities['xmomentum']
389    Ymom = domain.quantities['ymomentum']
390    Bed = domain.quantities['elevation']       
391
392    #Arrays
393    level = Level.edge_values
394    xmom =  Xmom.edge_values
395    ymom =  Ymom.edge_values
396    bed =   Bed.edge_values   
397
398    level_bdry = Level.boundary_values
399    xmom_bdry =  Xmom.boundary_values
400    ymom_bdry =  Ymom.boundary_values
401   
402    flux = zeros(3, Float) #Work array for summing up fluxes
403
404    #Loop
405    timestep = float(sys.maxint)   
406    for k in range(N):
407
408        flux[:] = 0.  #Reset work array
409        for i in range(3):
410            #Quantities inside volume facing neighbour i
411            ql = [level[k, i], xmom[k, i], ymom[k, i]]
412            zl = bed[k, i]
413
414            #Quantities at neighbour on nearest face
415            n = domain.neighbours[k,i] 
416            if n < 0:
417                m = -n-1 #Convert negative flag to index
418                qr = [level_bdry[m], xmom_bdry[m], ymom_bdry[m]]
419                zr = zl #Extend bed elevation to boundary
420            else:   
421                m = domain.neighbour_edges[k,i]
422                qr = [level[n, m], xmom[n, m], ymom[n, m]]
423                zr = bed[n, m]               
424
425               
426            #Outward pointing normal vector   
427            normal = domain.normals[k, 2*i:2*i+2]
428
429            #Flux computation using provided function
430            edgeflux, max_speed = flux_function(normal, ql, qr, zl, zr)
431            flux -= edgeflux * domain.edgelengths[k,i]
432
433            #Update optimal_timestep
434            try:
435                timestep = min(timestep, domain.radii[k]/max_speed)
436            except ZeroDivisionError:
437                pass
438
439        #Normalise by area and store for when all conserved
440        #quantities get updated
441        flux /= domain.areas[k]
442        Level.explicit_update[k] = flux[0]
443        Xmom.explicit_update[k] = flux[1]
444        Ymom.explicit_update[k] = flux[2]
445
446   
447    domain.timestep = timestep   
448
449
450def compute_fluxes_c(domain):
451    """Wrapper calling C version of compute fluxes
452    """
453
454    import sys
455    from Numeric import zeros, Float
456
457    N = domain.number_of_elements
458   
459    #Shortcuts
460    Level = domain.quantities['level']
461    Xmom = domain.quantities['xmomentum']
462    Ymom = domain.quantities['ymomentum']
463    Bed = domain.quantities['elevation']       
464
465    timestep = float(sys.maxint)   
466    from shallow_water_ext import compute_fluxes
467    domain.timestep = compute_fluxes(timestep, domain.epsilon, domain.g,
468                                     domain.neighbours,
469                                     domain.neighbour_edges,
470                                     domain.normals,
471                                     domain.edgelengths,                       
472                                     domain.radii,
473                                     domain.areas,
474                                     Level.edge_values, 
475                                     Xmom.edge_values, 
476                                     Ymom.edge_values, 
477                                     Bed.edge_values,   
478                                     Level.boundary_values,
479                                     Xmom.boundary_values,
480                                     Ymom.boundary_values,
481                                     Level.explicit_update,
482                                     Xmom.explicit_update,
483                                     Ymom.explicit_update)
484                                     
485
486####################################
487# Module functions for gradient limiting (distribute_to_vertices_and_edges)
488
489def distribute_to_vertices_and_edges(domain):
490    """Distribution from centroids to vertices specific to the
491    shallow water wave
492    equation.
493
494    It will ensure that h (w-z) is always non-negative even in the
495    presence of steep bed-slopes by taking a weighted average between shallow
496    and deep cases.
497   
498    In addition, all conserved quantities get distributed as per either a
499    constant (order==1) or a piecewise linear function (order==2).
500   
501    FIXME: more explanation about removal of artificial variability etc
502
503    Precondition:
504      All quantities defined at centroids and bed elevation defined at
505      vertices.
506
507    Postcondition
508      Conserved quantities defined at vertices
509   
510    """
511
512    #Remove very thin layers of water
513    protect_against_infinitesimal_and_negative_heights(domain)
514
515    #Extrapolate all conserved quantities
516    for name in domain.conserved_quantities:
517        Q = domain.quantities[name]
518        if domain.order == 1:
519            Q.extrapolate_first_order()
520        elif domain.order == 2:
521            Q.extrapolate_second_order()
522            Q.limit()                               
523        else:
524            raise 'Unknown order'
525
526    #Take bed elevation into account when water heights are small   
527    balance_deep_and_shallow(domain)
528
529    #Compute edge values by interpolation
530    for name in domain.conserved_quantities:
531        Q = domain.quantities[name]
532        Q.interpolate_from_vertices_to_edges()
533                   
534
535
536def dry(domain):
537    """Protect against infinitesimal heights and associated high velocities
538    at vertices
539    """
540
541    #FIXME: Experimental (from old version). Not in use at the moment
542   
543    #Shortcuts
544    wv = domain.quantities['level'].vertex_values
545    zv = domain.quantities['elevation'].vertex_values
546    xmomv = domain.quantities['xmomentum'].vertex_values
547    ymomv = domain.quantities['ymomentum'].vertex_values           
548    hv = wv - zv  #Water depths at vertices
549
550    #Update
551    for k in range(domain.number_of_elements):
552        hmax = max(hv[k, :])
553       
554        if hmax < domain.minimum_allowed_height:
555            #Control level
556            wv[k, :] = zv[k, :]
557
558            #Control momentum
559            xmomv[k,:] = ymomv[k,:] = 0.0
560   
561
562def protect_against_infinitesimal_and_negative_heights(domain):
563    """Protect against infinitesimal heights and associated high velocities
564    """
565   
566    #Shortcuts
567    wc = domain.quantities['level'].centroid_values
568    zc = domain.quantities['elevation'].centroid_values
569    xmomc = domain.quantities['xmomentum'].centroid_values
570    ymomc = domain.quantities['ymomentum'].centroid_values           
571    hc = wc - zc  #Water depths at centroids   
572
573    #Update
574    for k in range(domain.number_of_elements):
575
576        if hc[k] < domain.minimum_allowed_height: 
577            #Control level
578            wc[k] = zc[k]
579
580            #Control momentum
581            xmomc[k] = ymomc[k] = 0.0
582       
583
584
585def protect_against_infinitesimal_and_negative_heights_c(domain):
586    """Protect against infinitesimal heights and associated high velocities
587    """
588   
589    #Shortcuts
590    wc = domain.quantities['level'].centroid_values
591    zc = domain.quantities['elevation'].centroid_values
592    xmomc = domain.quantities['xmomentum'].centroid_values
593    ymomc = domain.quantities['ymomentum'].centroid_values           
594
595    from shallow_water_ext import protect
596
597    protect(domain.minimum_allowed_height, wc, zc, xmomc, ymomc)
598   
599   
600def balance_deep_and_shallow(domain):
601    """Compute linear combination between stage as computed by
602    gradient-limiters and stage computed as constant height above bed.
603    The former takes precedence when heights are large compared to the
604    bed slope while the latter takes precedence when heights are
605    relatively small.  Anything in between is computed as a balanced
606    linear combination in order to avoid numerical disturbances which
607    would otherwise appear as a result of hard switching between
608    modes.
609    """
610   
611    #Shortcuts
612    wc = domain.quantities['level'].centroid_values
613    zc = domain.quantities['elevation'].centroid_values
614    hc = wc - zc
615   
616    wv = domain.quantities['level'].vertex_values
617    zv = domain.quantities['elevation'].vertex_values
618    hv = wv-zv
619
620
621    #Computed linear combination between constant levels and and
622    #levels parallel to the bed elevation.     
623    for k in range(domain.number_of_elements):
624        #Compute maximal variation in bed elevation
625        #  This quantitiy is
626        #    dz = max_i abs(z_i - z_c)
627        #  and it is independent of dimension
628        #  In the 1d case zc = (z0+z1)/2
629        #  In the 2d case zc = (z0+z1+z2)/3
630
631        dz = max(abs(zv[k,0]-zc[k]),
632                 abs(zv[k,1]-zc[k]),
633                 abs(zv[k,2]-zc[k]))
634
635       
636        hmin = min( hv[k,:] )
637
638       
639        #Create alpha in [0,1], where alpha==0 means using shallow
640        #first order scheme and alpha==1 means using the stage w as
641        #computed by the gradient limiter (1st or 2nd order)
642        #
643        #If hmin > dz/2 then alpha = 1 and the bed will have no effect
644        #If hmin < 0 then alpha = 0 reverting to constant height above bed.
645     
646        if dz > 0.0:
647            alpha = max( min( 2*hmin/dz, 1.0), 0.0 )
648        else:
649            #Flat bed
650            alpha = 1.0
651
652
653        #Weighted balance between stage parallel to bed elevation
654        #(wvi = zvi + hc) and stage as computed by 1st or 2nd
655        #order gradient limiter
656        #(wvi = zvi + hvi) where i=0,1,2 denotes the vertex ids
657        #
658        #It follows that the updated wvi is
659        #  wvi := (1-alpha)*(zvi+hc) + alpha*(zvi+hvi) =
660        #  zvi + hc + alpha*(hvi - hc)
661        #
662        #Note that hvi = zc+hc-zvi in the first order case (constant).
663
664        if alpha < 1:         
665            for i in range(3):
666                wv[k,i] = zv[k,i] + hc[k] + alpha*(hv[k,i]-hc[k])
667           
668
669            #Momentums at centroids
670            xmomc = domain.quantities['xmomentum'].centroid_values
671            ymomc = domain.quantities['ymomentum'].centroid_values       
672
673            #Momentums at vertices
674            xmomv = domain.quantities['xmomentum'].vertex_values
675            ymomv = domain.quantities['ymomentum'].vertex_values         
676
677            # Update momentum as a linear combination of
678            # xmomc and ymomc (shallow) and momentum
679            # from extrapolator xmomv and ymomv (deep).
680            xmomv[k,:] = (1-alpha)*xmomc[k] + alpha*xmomv[k,:]           
681            ymomv[k,:] = (1-alpha)*ymomc[k] + alpha*ymomv[k,:]
682                                   
683           
684
685def balance_deep_and_shallow_c(domain):
686    """Wrapper for C implementation
687    """
688   
689    #Shortcuts
690    wc = domain.quantities['level'].centroid_values
691    zc = domain.quantities['elevation'].centroid_values
692    hc = wc - zc
693   
694    wv = domain.quantities['level'].vertex_values
695    zv = domain.quantities['elevation'].vertex_values
696    hv = wv-zv
697
698    #Momentums at centroids
699    xmomc = domain.quantities['xmomentum'].centroid_values
700    ymomc = domain.quantities['ymomentum'].centroid_values       
701
702    #Momentums at vertices
703    xmomv = domain.quantities['xmomentum'].vertex_values
704    ymomv = domain.quantities['ymomentum'].vertex_values         
705
706   
707
708    from shallow_water_ext import balance_deep_and_shallow
709    balance_deep_and_shallow(wc, zc, hc, wv, zv, hv,
710                             xmomc, ymomc, xmomv, ymomv)
711
712   
713
714
715###############################################
716#Boundary - specific to the shallow water wave equation
717class Reflective_boundary(Boundary):
718    """Reflective boundary returns same conserved quantities as
719    those present in its neighbour volume but reflected.
720
721    This class is specific to the shallow water equation as it
722    works with the momentum quantities assumed to be the second
723    and third conserved quantities.
724    """
725   
726    def __init__(self, domain = None):       
727        Boundary.__init__(self)
728
729        if domain is None:
730            msg = 'Domain must be specified for reflective boundary'
731            raise msg
732       
733        #Handy shorthands
734        self.level = domain.quantities['level'].edge_values
735        self.xmom = domain.quantities['xmomentum'].edge_values
736        self.ymom = domain.quantities['ymomentum'].edge_values         
737        self.normals = domain.normals
738       
739        from Numeric import zeros, Float       
740        self.conserved_quantities = zeros(3, Float) 
741       
742    def __repr__(self):
743        return 'Reflective_boundary'
744
745
746    def evaluate(self, vol_id, edge_id):
747        """Reflective boundaries reverses the outward momentum
748        of the volume they serve.
749        """
750
751        q = self.conserved_quantities
752        q[0] = self.level[vol_id, edge_id]
753        q[1] = self.xmom[vol_id, edge_id] 
754        q[2] = self.ymom[vol_id, edge_id] 
755       
756        normal = self.normals[vol_id, 2*edge_id:2*edge_id+2]     
757
758       
759        r = rotate(q, normal, direction = 1)
760        r[1] = -r[1]
761        q = rotate(r, normal, direction = -1)
762
763        return q
764
765
766#########################
767#Standard forcing terms:
768#
769def gravity(domain):
770    """Apply gravitational pull in the presence of bed slope
771    """
772
773    from util import gradient
774    from Numeric import zeros, Float, array, sum
775
776    xmom = domain.quantities['xmomentum'].explicit_update
777    ymom = domain.quantities['ymomentum'].explicit_update
778
779    Level = domain.quantities['level']
780    Elevation = domain.quantities['elevation']   
781    h = Level.edge_values - Elevation.edge_values
782    v = Elevation.vertex_values
783   
784    x = domain.get_vertex_coordinates()
785    g = domain.g
786   
787    for k in range(domain.number_of_elements):   
788        avg_h = sum( h[k,:] )/3
789       
790        #Compute bed slope
791        x0, y0, x1, y1, x2, y2 = x[k,:]   
792        z0, z1, z2 = v[k,:]
793
794        zx, zy = gradient(x0, y0, x1, y1, x2, y2, z0, z1, z2)
795
796        #Update momentum
797        xmom[k] += -g*zx*avg_h
798        ymom[k] += -g*zy*avg_h       
799
800
801def gravity_c(domain):
802    """Wrapper calling C version
803    """   
804
805    xmom = domain.quantities['xmomentum'].explicit_update
806    ymom = domain.quantities['ymomentum'].explicit_update
807
808    Level = domain.quantities['level']   
809    Elevation = domain.quantities['elevation']   
810    h = Level.edge_values - Elevation.edge_values
811    v = Elevation.vertex_values
812   
813    x = domain.get_vertex_coordinates()
814    g = domain.g
815   
816   
817    from shallow_water_ext import gravity
818    gravity(g, h, v, x, xmom, ymom)
819   
820       
821def manning_friction(domain):
822    """Apply (Manning) friction to water momentum
823    """
824
825    from math import sqrt
826
827    w = domain.quantities['level'].centroid_values
828    z = domain.quantities['elevation'].centroid_values
829    h = w-z
830   
831    uh = domain.quantities['xmomentum'].centroid_values
832    vh = domain.quantities['ymomentum'].centroid_values
833    eta = domain.quantities['friction'].centroid_values   
834   
835    xmom_update = domain.quantities['xmomentum'].semi_implicit_update
836    ymom_update = domain.quantities['ymomentum'].semi_implicit_update   
837   
838    N = domain.number_of_elements
839    eps = domain.minimum_allowed_height
840    g = domain.g
841   
842    for k in range(N):
843        if eta[k] >= eps:
844            if h[k] >= eps:
845                S = -g * eta[k]**2 * sqrt((uh[k]**2 + vh[k]**2))
846                S /= h[k]**(7.0/3)
847
848                #Update momentum
849                xmom_update[k] += S*uh[k]
850                ymom_update[k] += S*vh[k]
851           
852       
853def manning_friction_c(domain):
854    """Wrapper for c version
855    """
856
857
858    xmom = domain.quantities['xmomentum']
859    ymom = domain.quantities['ymomentum']   
860       
861    w = domain.quantities['level'].centroid_values
862    z = domain.quantities['elevation'].centroid_values
863   
864    uh = xmom.centroid_values
865    vh = ymom.centroid_values
866    eta = domain.quantities['friction'].centroid_values   
867   
868    xmom_update = xmom.semi_implicit_update
869    ymom_update = ymom.semi_implicit_update   
870   
871    N = domain.number_of_elements
872    eps = domain.minimum_allowed_height
873    g = domain.g
874
875    from shallow_water_ext import manning_friction
876    manning_friction(g, eps, w, z, uh, vh, eta, xmom_update, ymom_update)
877
878
879def linear_friction(domain):
880    """Apply linear friction to water momentum
881
882    Assumes quantity: 'linear_friction' to be present
883    """
884
885    from math import sqrt
886
887    w = domain.quantities['level'].centroid_values
888    z = domain.quantities['elevation'].centroid_values
889    h = w-z
890
891    uh = domain.quantities['xmomentum'].centroid_values
892    vh = domain.quantities['ymomentum'].centroid_values
893    tau = domain.quantities['linear_friction'].centroid_values   
894   
895    xmom_update = domain.quantities['xmomentum'].semi_implicit_update
896    ymom_update = domain.quantities['ymomentum'].semi_implicit_update   
897   
898    N = domain.number_of_elements
899    eps = domain.minimum_allowed_height
900   
901    for k in range(N):
902        if tau[k] >= eps:
903            if h[k] >= eps:
904                S = -tau[k]/h[k]
905
906                #Update momentum
907                xmom_update[k] += S*uh[k]
908                ymom_update[k] += S*vh[k]
909           
910       
911       
912def check_forcefield(f):
913    """Check that f is either
914    1: a callable object f(t,x,y), where x and y are vectors
915       and that it returns an array or a list of same length
916       as x and y
917    2: a scalar   
918    """
919
920    from Numeric import ones, Float, array
921
922   
923    if callable(f):
924        N = 3
925        x = ones(3, Float)
926        y = ones(3, Float)           
927        try:
928            q = f(1.0, x, y)
929        except Exception, e:
930            msg = 'Function %s could not be executed:\n%s' %(f, e)
931            raise msg
932                   
933        try:
934            q = array(q).astype(Float)
935        except:
936            msg = 'Return value from vector function %s could ' %f
937            msg += 'not be converted into a Numeric array of floats.\n'
938            msg += 'Specified function should return either list or array.' 
939            raise msg
940
941        msg = 'Return vector from function %s' %f
942        msg += 'must have same lenght as input vectors'
943        assert len(q) == N, msg
944       
945    else:       
946        try:
947            f = float(f)
948        except:
949            msg = 'Force field %s must be either a scalar' %f
950            msg += ' or a vector function'
951            raise msg
952    return f   
953
954
955class Wind_stress:
956    """Apply wind stress to water momentum in terms of
957    wind speed [m/s] and wind direction [degrees]   
958    """
959
960    def __init__(self, *args, **kwargs):
961        """Initialise windfield from wind speed s [m/s]
962        and wind direction phi [degrees]
963       
964        Inputs v and phi can be either scalars or Python functions, e.g.
965
966        W = Wind_stress(10, 178)
967
968        #FIXME - 'normal' degrees are assumed for now, i.e. the
969        vector (1,0) has zero degrees.
970        We may need to convert from 'compass' degrees later on and also
971        map from True north to grid north.
972
973        Arguments can also be Python functions of t,x,y as in
974
975        def speed(t,x,y):
976            ...
977            return s
978       
979        def angle(t,x,y):
980            ...
981            return phi
982
983        where x and y are vectors.
984
985        and then pass the functions in
986
987        W = Wind_stress(speed, angle)
988
989        The instantiated object W can be appended to the list of
990        forcing_terms as in
991
992        Alternatively, one vector valued function for (speed, angle)
993        can be applied, providing both quantities simultaneously.
994        As in
995        W = Wind_stress(F), where returns (speed, angle) for each t.         
996
997        domain.forcing_terms.append(W)
998        """
999
1000        from config import rho_a, rho_w, eta_w
1001        from Numeric import array, Float               
1002
1003        if len(args) == 2:
1004            s = args[0]
1005            phi = args[1]
1006        elif len(args) == 1:
1007            #Assume vector function returning (s, phi)(t,x,y)
1008            vector_function = args[0]
1009            s = lambda t,x,y: vector_function(t,x,y)[0]
1010            phi = lambda t,x,y: vector_function(t,x,y)[1]
1011        else:
1012           #Assume info is in 2 keyword arguments
1013       
1014           if len(kwargs) == 2:
1015               s = kwargs['s']
1016               phi = kwargs['phi']
1017           else:
1018               raise 'Assumes two keyword arguments: s=..., phi=....'   
1019           
1020        self.speed = check_forcefield(s)
1021        self.phi = check_forcefield(phi)
1022           
1023        self.const = eta_w*rho_a/rho_w
1024       
1025
1026    def __call__(self, domain):
1027        """Evaluate windfield based on values found in domain
1028        """
1029
1030        from math import pi, cos, sin, sqrt
1031        from Numeric import Float, ones, ArrayType
1032       
1033        xmom_update = domain.quantities['xmomentum'].explicit_update
1034        ymom_update = domain.quantities['ymomentum'].explicit_update   
1035       
1036        N = domain.number_of_elements
1037        t = domain.time       
1038       
1039        if callable(self.speed):
1040            xc = domain.get_centroid_coordinates()           
1041            s_vec = self.speed(t, xc[:,0], xc[:,1])
1042        else:
1043            #Assume s is a scalar
1044
1045            try:
1046                s_vec = self.speed * ones(N, Float)
1047            except:
1048                msg = 'Speed must be either callable or a scalar: %s' %self.s
1049                raise msg
1050
1051
1052        if callable(self.phi):
1053            xc = domain.get_centroid_coordinates()           
1054            phi_vec = self.phi(t, xc[:,0], xc[:,1])
1055        else:
1056            #Assume phi is a scalar
1057
1058            try:
1059                phi_vec = self.phi * ones(N, Float)           
1060            except:
1061                msg = 'Angle must be either callable or a scalar: %s' %self.phi
1062                raise msg
1063
1064        assign_windfield_values(xmom_update, ymom_update,
1065                                s_vec, phi_vec, self.const)
1066
1067
1068def assign_windfield_values(xmom_update, ymom_update,
1069                            s_vec, phi_vec, const):
1070    """Python version of assigning wind field to update vectors.
1071    A c version also exists (for speed)
1072    """
1073    from math import pi, cos, sin, sqrt
1074   
1075    N = len(s_vec)
1076    for k in range(N):
1077        s = s_vec[k]
1078        phi = phi_vec[k]
1079
1080        #Convert to radians
1081        phi = phi*pi/180
1082
1083        #Compute velocity vector (u, v)
1084        u = s*cos(phi)
1085        v = s*sin(phi)
1086       
1087        #Compute wind stress
1088        S = const * sqrt(u**2 + v**2)
1089        xmom_update[k] += S*u
1090        ymom_update[k] += S*v       
1091           
1092
1093
1094###########################
1095###########################
1096#Geometries
1097
1098
1099#FIXME: Rethink this way of creating values.
1100
1101
1102class Weir:
1103    """Set a bathymetry for weir with a hole and a downstream gutter
1104    x,y are assumed to be in the unit square
1105    """
1106   
1107    def __init__(self, stage):
1108        self.inflow_stage = stage
1109
1110    def __call__(self, x, y):   
1111        from Numeric import zeros, Float
1112        from math import sqrt
1113   
1114        N = len(x)
1115        assert N == len(y)   
1116
1117        z = zeros(N, Float)
1118        for i in range(N):
1119            z[i] = -x[i]/2  #General slope
1120
1121            #Flattish bit to the left
1122            if x[i] < 0.3:
1123                z[i] = -x[i]/10 
1124           
1125            #Weir
1126            if x[i] >= 0.3 and x[i] < 0.4:
1127                z[i] = -x[i]+0.9
1128
1129            #Dip
1130            x0 = 0.6
1131            #depth = -1.3
1132            depth = -1.0
1133            #plateaux = -0.9
1134            plateaux = -0.6           
1135            if y[i] < 0.7:
1136                if x[i] > x0 and x[i] < 0.9:
1137                    z[i] = depth
1138
1139                #RHS plateaux
1140                if x[i] >= 0.9:
1141                    z[i] = plateaux
1142                   
1143                   
1144            elif y[i] >= 0.7 and y[i] < 1.5:
1145                #Restrict and deepen
1146                if x[i] >= x0 and x[i] < 0.8:
1147                    z[i] = depth-(y[i]/3-0.3)
1148                    #z[i] = depth-y[i]/5
1149                    #z[i] = depth
1150                elif x[i] >= 0.8:   
1151                    #RHS plateaux
1152                    z[i] = plateaux
1153                   
1154            elif y[i] >= 1.5:
1155                if x[i] >= x0 and x[i] < 0.8 + (y[i]-1.5)/1.2:
1156                    #Widen up and stay at constant depth
1157                    z[i] = depth-1.5/5
1158                elif x[i] >= 0.8 + (y[i]-1.5)/1.2:                       
1159                    #RHS plateaux
1160                    z[i] = plateaux                                   
1161                   
1162
1163            #Hole in weir (slightly higher than inflow condition)
1164            if x[i] >= 0.3 and x[i] < 0.4 and y[i] > 0.2 and y[i] < 0.4:
1165                z[i] = -x[i]+self.inflow_stage + 0.02
1166
1167            #Channel behind weir
1168            x0 = 0.5
1169            if x[i] >= 0.4 and x[i] < x0 and y[i] > 0.2 and y[i] < 0.4:
1170                z[i] = -x[i]+self.inflow_stage + 0.02
1171               
1172            if x[i] >= x0 and x[i] < 0.6 and y[i] > 0.2 and y[i] < 0.4:
1173                #Flatten it out towards the end
1174                z[i] = -x0+self.inflow_stage + 0.02 + (x0-x[i])/5
1175
1176            #Hole to the east
1177            x0 = 1.1; y0 = 0.35
1178            #if x[i] < -0.2 and y < 0.5:
1179            if sqrt((2*(x[i]-x0))**2 + (2*(y[i]-y0))**2) < 0.2:
1180                z[i] = sqrt(((x[i]-x0))**2 + ((y[i]-y0))**2)-1.0
1181
1182            #Tiny channel draining hole
1183            if x[i] >= 1.14 and x[i] < 1.2 and y[i] >= 0.4 and y[i] < 0.6:
1184                z[i] = -0.9 #North south
1185               
1186            if x[i] >= 0.9 and x[i] < 1.18 and y[i] >= 0.58 and y[i] < 0.65:
1187                z[i] = -1.0 + (x[i]-0.9)/3 #East west
1188           
1189           
1190
1191            #Stuff not in use
1192           
1193            #Upward slope at inlet to the north west
1194            #if x[i] < 0.0: # and y[i] > 0.5:
1195            #    #z[i] = -y[i]+0.5  #-x[i]/2
1196            #    z[i] = x[i]/4 - y[i]**2 + 0.5
1197
1198            #Hole to the west
1199            #x0 = -0.4; y0 = 0.35 # center
1200            #if sqrt((2*(x[i]-x0))**2 + (2*(y[i]-y0))**2) < 0.2:
1201            #    z[i] = sqrt(((x[i]-x0))**2 + ((y[i]-y0))**2)-0.2
1202
1203
1204
1205
1206
1207        return z/2
1208
1209class Weir_simple:
1210    """Set a bathymetry for weir with a hole and a downstream gutter
1211    x,y are assumed to be in the unit square
1212    """
1213   
1214    def __init__(self, stage):
1215        self.inflow_stage = stage
1216
1217    def __call__(self, x, y):   
1218        from Numeric import zeros, Float
1219   
1220        N = len(x)
1221        assert N == len(y)   
1222
1223        z = zeros(N, Float)
1224        for i in range(N):
1225            z[i] = -x[i]  #General slope
1226
1227            #Flat bit to the left
1228            if x[i] < 0.3:
1229                z[i] = -x[i]/10  #General slope
1230           
1231            #Weir
1232            if x[i] > 0.3 and x[i] < 0.4:
1233                z[i] = -x[i]+0.9
1234
1235            #Dip
1236            if x[i] > 0.6 and x[i] < 0.9:
1237                z[i] = -x[i]-0.5  #-y[i]/5
1238
1239            #Hole in weir (slightly higher than inflow condition)
1240            if x[i] > 0.3 and x[i] < 0.4 and y[i] > 0.2 and y[i] < 0.4:
1241                z[i] = -x[i]+self.inflow_stage + 0.05       
1242           
1243
1244        return z/2
1245       
1246
1247           
1248class Constant_height:
1249    """Set an initial condition with constant water height, e.g
1250    stage s = z+h
1251    """
1252    def __init__(self, W, h):
1253        self.W = W
1254        self.h = h
1255
1256    def __call__(self, x, y):
1257        if self.W is None:
1258            from Numeric import ones, Float
1259            return self.h*ones(len(x), Float)
1260        else:
1261            return self.W(x,y) + self.h
1262
1263
1264
1265##############################################
1266#Initialise module
1267
1268
1269import compile
1270if compile.can_use_C_extension('shallow_water_ext.c'):
1271    #Replace python version with c implementations
1272   
1273    from shallow_water_ext import rotate, assign_windfield_values
1274    compute_fluxes = compute_fluxes_c
1275    gravity = gravity_c
1276    manning_friction = manning_friction_c
1277    balance_deep_and_shallow = balance_deep_and_shallow_c
1278    protect_against_infinitesimal_and_negative_heights = protect_against_infinitesimal_and_negative_heights_c
1279
1280   
1281    #distribute_to_vertices_and_edges = distribute_to_vertices_and_edges_c
1282
1283
1284#Optimisation with psyco
1285from config import use_psyco
1286if use_psyco:
1287    try:
1288        import psyco
1289    except:
1290        msg = 'WARNING: psyco (speedup) could not import'+\
1291              ', you may want to consider installing it'
1292        print msg
1293    else:
1294        psyco.bind(Domain.distribute_to_vertices_and_edges)
1295        psyco.bind(Domain.compute_fluxes)
1296       
1297if __name__ == "__main__":
1298    pass
Note: See TracBrowser for help on using the repository browser.