source: anuga_work/development/anuga_1d/anuga_1d_suggestion1_2_3/shallow_water_domain_nonwellbalanced3.py @ 7576

Last change on this file since 7576 was 7576, checked in by sudi, 14 years ago
File size: 43.8 KB
Line 
1"""Class Domain -
21D interval 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 = S
10
11where
12
13U = [w, uh]
14E = [uh, u^2h + gh^2/2]
15S represents source terms forcing the system
16(e.g. gravity, friction, wind stress, ...)
17
18and _t, _x, _y denote the derivative with respect to t, x and y respectiely.
19
20The quantities are
21
22symbol    variable name    explanation
23x         x                horizontal distance from origin [m]
24z         elevation        elevation of bed on which flow is modelled [m]
25h         height           water height above z [m]
26w         stage            absolute water level, w = z+h [m]
27u                          speed in the x direction [m/s]
28uh        xmomentum        momentum in the x direction [m^2/s]
29
30eta                        mannings friction coefficient [to appear]
31nu                         wind stress coefficient [to appear]
32
33The conserved quantities are w, uh
34
35For details see e.g.
36Christopher Zoppou and Stephen Roberts,
37Catastrophic Collapse of Water Supply Reservoirs in Urban Areas,
38Journal of Hydraulic Engineering, vol. 127, No. 7 July 1999
39
40
41John Jakeman, Ole Nielsen, Stephen Roberts, Duncan Gray, Christopher Zoppou
42Geoscience Australia, 2006
43"""
44
45
46from domain import *
47Generic_Domain = Domain #Rename
48
49#Shallow water domain
50class Domain(Generic_Domain):
51
52    def __init__(self, coordinates, boundary = None, tagged_elements = None):
53
54        conserved_quantities = ['stage', 'xmomentum']
55        evolved_quantities = ['stage', 'xmomentum', 'elevation', 'height', 'velocity']
56        other_quantities = ['friction']
57        Generic_Domain.__init__(self, coordinates, boundary,
58                                conserved_quantities, evolved_quantities, other_quantities,
59                                tagged_elements)
60       
61        from config import minimum_allowed_height, g, h0
62        self.minimum_allowed_height = minimum_allowed_height
63        self.g = g
64        self.h0 = h0
65
66        #forcing terms not included in 1d domain ?WHy?
67        self.forcing_terms.append(gravity)
68        #self.forcing_terms.append(manning_friction)
69        #print "\nI have Removed forcing terms line 64 1dsw"
70
71        #Realtime visualisation
72        self.visualiser = None
73        self.visualise  = False
74        self.visualise_color_stage = False
75        self.visualise_stage_range = 1.0
76        self.visualise_timer = True
77        self.visualise_range_z = None
78       
79        #Stored output
80        self.store = True
81        self.format = 'sww'
82        self.smooth = True
83
84        #Evolve parametrs
85        self.cfl = 1.0
86       
87        #Reduction operation for get_vertex_values
88        from util import mean
89        self.reduction = mean
90        #self.reduction = min  #Looks better near steep slopes
91
92        self.quantities_to_be_stored = ['stage','xmomentum']
93
94        self.__doc__ = 'shallow_water_domain_suggestion2'
95        self.check_integrity()
96
97
98    def set_quantities_to_be_stored(self, q):
99        """Specify which quantities will be stored in the sww file.
100
101        q must be either:
102          - the name of a quantity
103          - a list of quantity names
104          - None
105
106        In the two first cases, the named quantities will be stored at each
107        yieldstep
108        (This is in addition to the quantities elevation and friction) 
109        If q is None, storage will be switched off altogether.
110        """
111
112
113        if q is None:
114            self.quantities_to_be_stored = []           
115            self.store = False
116            return
117
118        if isinstance(q, basestring):
119            q = [q] # Turn argument into a list
120
121        #Check correcness   
122        for quantity_name in q:
123            msg = 'Quantity %s is not a valid conserved quantity' %quantity_name
124            assert quantity_name in self.conserved_quantities, msg
125       
126        self.quantities_to_be_stored = q
127       
128
129    def initialise_visualiser(self,scale_z=1.0,rect=None):
130        #Realtime visualisation
131        if self.visualiser is None:
132            from realtime_visualisation_new import Visualiser
133            self.visualiser = Visualiser(self,scale_z,rect)
134            self.visualiser.setup['elevation']=True
135            self.visualiser.updating['stage']=True
136        self.visualise = True
137        if self.visualise_color_stage == True:
138            self.visualiser.coloring['stage'] = True
139            self.visualiser.qcolor['stage'] = (0.0, 0.0, 0.8)
140        print 'initialise visualiser'
141        print self.visualiser.setup
142        print self.visualiser.updating
143
144    def check_integrity(self):
145        Generic_Domain.check_integrity(self)
146        #Check that we are solving the shallow water wave equation
147
148        msg = 'First conserved quantity must be "stage"'
149        assert self.conserved_quantities[0] == 'stage', msg
150        msg = 'Second conserved quantity must be "xmomentum"'
151        assert self.conserved_quantities[1] == 'xmomentum', msg
152               
153        msg = 'First evolved quantity must be "stage"'
154        assert self.evolved_quantities[0] == 'stage', msg
155        msg = 'Second evolved quantity must be "xmomentum"'
156        assert self.evolved_quantities[1] == 'xmomentum', msg
157        msg = 'Third evolved quantity must be "elevation"'
158        assert self.evolved_quantities[2] == 'elevation', msg
159        msg = 'Fourth evolved quantity must be "height"'
160        assert self.evolved_quantities[3] == 'height', msg
161        msg = 'Fifth evolved quantity must be "velocity"'
162        assert self.evolved_quantities[4] == 'velocity', msg
163
164    def extrapolate_second_order_sw(self):
165        #Call correct module function
166        #(either from this module or C-extension)
167        extrapolate_second_order_sw(self)
168
169    def compute_fluxes(self):
170        #Call correct module function
171        #(either from this module or C-extension)
172        compute_fluxes_C_nonwellbalanced3(self)
173
174    def compute_timestep(self):
175        #Call correct module function
176        compute_timestep(self)
177
178    def distribute_to_vertices_and_edges(self):
179        #Call correct module function
180        #(either from this module or C-extension)
181        distribute_to_vertices_and_edges(self)
182       
183    def evolve(self, yieldstep = None, finaltime = None, duration = None, skip_initial_step = False):
184        #Call basic machinery from parent class
185        for t in Generic_Domain.evolve(self, yieldstep, finaltime, duration, skip_initial_step):
186            yield(t)
187
188    def initialise_storage(self):
189        """Create and initialise self.writer object for storing data.
190        Also, save x and bed elevation
191        """
192
193        import data_manager
194
195        #Initialise writer
196        self.writer = data_manager.get_dataobject(self, mode = 'w')
197
198        #Store vertices and connectivity
199        self.writer.store_connectivity()
200
201
202    def store_timestep(self, name):
203        """Store named quantity and time.
204
205        Precondition:
206           self.write has been initialised
207        """
208        self.writer.store_timestep(name)
209
210
211#=============== End of Shallow Water Domain ===============================
212
213#Rotation of momentum vector
214def rotate(q, normal, direction = 1):
215    """Rotate the momentum component q (q[1], q[2])
216    from x,y coordinates to coordinates based on normal vector.
217
218    If direction is negative the rotation is inverted.
219
220    Input vector is preserved
221
222    This function is specific to the shallow water wave equation
223    """
224
225    from Numeric import zeros, Float
226
227    assert len(q) == 3,\
228           'Vector of conserved quantities must have length 3'\
229           'for 2D shallow water equation'
230
231    try:
232        l = len(normal)
233    except:
234        raise 'Normal vector must be an Numeric array'
235
236    assert l == 2, 'Normal vector must have 2 components'
237
238
239    n1 = normal[0]
240    n2 = normal[1]
241
242    r = zeros(len(q), Float) #Rotated quantities
243    r[0] = q[0]              #First quantity, height, is not rotated
244
245    if direction == -1:
246        n2 = -n2
247
248
249    r[1] =  n1*q[1] + n2*q[2]
250    r[2] = -n2*q[1] + n1*q[2]
251
252    return r
253
254
255def flux_function(normal, ql, qr, zl, zr):
256    """Compute fluxes between volumes for the shallow water wave equation
257    cast in terms of w = h+z using the 'central scheme' as described in
258
259    Kurganov, Noelle, Petrova. 'Semidiscrete Central-Upwind Schemes For
260    Hyperbolic Conservation Laws and Hamilton-Jacobi Equations'.
261    Siam J. Sci. Comput. Vol. 23, No. 3, pp. 707-740.
262
263    The implemented formula is given in equation (3.15) on page 714
264
265    Conserved quantities w, uh, are stored as elements 0 and 1
266    in the numerical vectors ql an qr.
267
268    Bed elevations zl and zr.
269    """
270
271    from config import g, epsilon, h0
272    from math import sqrt
273    from Numeric import array
274
275    #print 'ql',ql
276
277    #Align momentums with x-axis
278    #q_left  = rotate(ql, normal, direction = 1)
279    #q_right = rotate(qr, normal, direction = 1)
280    q_left = ql
281    q_left[1] = q_left[1]*normal
282    q_right = qr
283    q_right[1] = q_right[1]*normal
284       
285    #z = (zl+zr)/2 #Take average of field values
286    z = 0.5*(zl+zr) #Take average of field values
287
288    w_left  = q_left[0]   #w=h+z
289    h_left  = w_left-z
290    uh_left = q_left[1]
291
292    if h_left < epsilon:
293        u_left = 0.0  #Could have been negative
294        h_left = 0.0
295    else:
296        u_left  = uh_left/(h_left +  h0/h_left)
297
298
299    uh_left = u_left*h_left
300
301
302    w_right  = q_right[0]  #w=h+z
303    h_right  = w_right-z
304    uh_right = q_right[1]
305
306    if h_right < epsilon:
307        u_right = 0.0  #Could have been negative
308        h_right = 0.0
309    else:
310        u_right  = uh_right/(h_right + h0/h_right)
311
312    uh_right = u_right*h_right
313       
314
315    #We have got h and u at vertex, then the following is the calculation of fluxes!
316    soundspeed_left  = sqrt(g*h_left)
317    soundspeed_right = sqrt(g*h_right)
318
319    #Maximal wave speed
320    s_max = max(u_left+soundspeed_left, u_right+soundspeed_right, 0)
321   
322    #Minimal wave speed
323    s_min = min(u_left-soundspeed_left, u_right-soundspeed_right, 0)
324   
325    #Flux computation
326    flux_left  = array([u_left*h_left,
327                        u_left*uh_left + 0.5*g*h_left*h_left])
328    flux_right = array([u_right*h_right,
329                        u_right*uh_right + 0.5*g*h_right*h_right])
330
331    denom = s_max-s_min
332    if denom == 0.0:
333        edgeflux = array([0.0, 0.0])
334        max_speed = 0.0
335    else:
336        edgeflux = (s_max*flux_left - s_min*flux_right)/denom
337        edgeflux += s_max*s_min*(q_right-q_left)/denom
338       
339        edgeflux[1] = edgeflux[1]*normal
340
341        max_speed = max(abs(s_max), abs(s_min))
342
343    return edgeflux, max_speed
344#
345def compute_fluxes(domain):
346    """
347        Compute all fluxes and the timestep suitable for all volumes
348    in domain.
349
350    Compute total flux for each conserved quantity using "flux_function"
351
352    Fluxes across each edge are scaled by edgelengths and summed up
353    Resulting flux is then scaled by area and stored in
354    explicit_update for each of the three conserved quantities
355    stage, xmomentum and ymomentum
356
357    The maximal allowable speed computed by the flux_function for each volume
358    is converted to a timestep that must not be exceeded. The minimum of
359    those is computed as the next overall timestep.
360
361    Post conditions:
362      domain.explicit_update is reset to computed flux values
363      domain.timestep is set to the largest step satisfying all volumes.
364   
365    """
366
367    import sys
368    from Numeric import zeros, Float
369    N = domain.number_of_elements
370    Stage = domain.quantities['stage']
371    Xmom  = domain.quantities['xmomentum']
372    Bed   = domain.quantities['elevation']
373
374    stage_bdry = Stage.boundary_values
375    xmom_bdry =  Xmom.boundary_values
376    stage, xmom, bed, h_V, u_V = distribute_to_vertices_and_edges(domain)
377    domain.update_boundary()
378       
379    flux = zeros(2, Float) #Work array for summing up fluxes
380    ql = zeros(2, Float)
381    qr = zeros(2, Float)
382
383    #Loop
384    timestep = float(sys.maxint)
385    enter = True
386    for k in range(N):
387
388        flux[:] = 0.  #Reset work array
389        #for i in range(3):
390        for i in range(2):
391            #Quantities inside volume facing neighbour i
392            #ql[0] = stage[k, i]
393            #ql[1] = xmom[k, i]
394            ql = [stage[k, i], xmom[k, i]]
395            zl = bed[k, i]
396
397            #Quantities at neighbour on nearest face
398            n = domain.neighbours[k,i]
399            if n < 0:
400                m = -n-1 #Convert negative flag to index
401                qr[0] = stage_bdry[m]
402                qr[1] = xmom_bdry[m]
403                zr = zl #Extend bed elevation to boundary
404            else:
405                #m = domain.neighbour_edges[k,i]
406                m = domain.neighbour_vertices[k,i]
407                #qr = [stage[n, m], xmom[n, m], ymom[n, m]]
408                qr[0] = stage[n, m]
409                qr[1] = xmom[n, m]
410                zr = bed[n, m]
411
412
413            #Outward pointing normal vector
414            normal = domain.normals[k, i]
415       
416            #Flux computation using provided function
417     
418            edgeflux, max_speed = flux_function(normal, ql, qr, zl, zr)
419
420            #print 'edgeflux', edgeflux
421
422            # THIS IS THE LINE TO DEAL WITH LEFT AND RIGHT FLUXES
423            # flux = edgefluxleft - edgefluxright
424            flux -= edgeflux #* domain.edgelengths[k,i]
425            #Update optimal_timestep
426            try:
427                #timestep = min(timestep, 0.5*domain.radii[k]/max_speed)
428                timestep = min(timestep, domain.cfl*0.5*domain.areas[k]/max_speed)
429            except ZeroDivisionError:
430                pass
431
432        #Normalise by area and store for when all conserved
433        #quantities get updated
434        flux /= domain.areas[k]
435
436        Stage.explicit_update[k] = flux[0]
437        Xmom.explicit_update[k] =  flux[1]
438        #Ymom.explicit_update[k] = flux[2]
439        #print "flux cell",k,flux[0]
440
441    domain.flux_timestep = timestep
442    #print domain.quantities['stage'].centroid_values
443#
444def compute_timestep(domain):
445    import sys
446    from Numeric import zeros, Float
447
448    N = domain.number_of_elements
449
450    #Shortcuts
451    Stage = domain.quantities['stage']
452    Xmom = domain.quantities['xmomentum']
453    Bed = domain.quantities['elevation']
454   
455    stage = Stage.vertex_values
456    xmom =  Xmom.vertex_values
457    bed =   Bed.vertex_values
458
459    stage_bdry = Stage.boundary_values
460    xmom_bdry =  Xmom.boundary_values
461
462    flux = zeros(2, Float) #Work array for summing up fluxes
463    ql = zeros(2, Float)
464    qr = zeros(2, Float)
465
466    #Loop
467    timestep = float(sys.maxint)
468    enter = True
469    for k in range(N):
470
471        flux[:] = 0.  #Reset work array
472        for i in range(2):
473            #Quantities inside volume facing neighbour i
474            ql = [stage[k, i], xmom[k, i]]
475            zl = bed[k, i]
476
477            #Quantities at neighbour on nearest face
478            n = domain.neighbours[k,i]
479            if n < 0:
480                m = -n-1 #Convert negative flag to index
481                qr[0] = stage_bdry[m]
482                qr[1] = xmom_bdry[m]
483                zr = zl #Extend bed elevation to boundary
484            else:
485                #m = domain.neighbour_edges[k,i]
486                m = domain.neighbour_vertices[k,i]
487                qr[0] = stage[n, m]
488                qr[1] =  xmom[n, m]
489                zr = bed[n, m]
490
491
492            #Outward pointing normal vector
493            normal = domain.normals[k, i]
494
495            edgeflux, max_speed = flux_function(normal, ql, qr, zl, zr)
496
497            #Update optimal_timestep
498            try:
499                timestep = min(timestep, domain.cfl*0.5*domain.areas[k]/max_speed)
500            except ZeroDivisionError:
501                pass
502
503    domain.timestep = timestep
504
505# Compute flux definition
506def compute_fluxes_C_long(domain):
507    from Numeric import zeros, Float
508    import sys
509
510       
511    timestep = float(sys.maxint)
512    #print 'timestep=',timestep
513    #print 'The type of timestep is',type(timestep)
514       
515    epsilon = domain.epsilon
516    #print 'epsilon=',epsilon
517    #print 'The type of epsilon is',type(epsilon)
518       
519    g = domain.g
520    #print 'g=',g
521    #print 'The type of g is',type(g)
522       
523    neighbours = domain.neighbours
524    #print 'neighbours=',neighbours
525    #print 'The type of neighbours is',type(neighbours)
526       
527    neighbour_vertices = domain.neighbour_vertices
528    #print 'neighbour_vertices=',neighbour_vertices
529    #print 'The type of neighbour_vertices is',type(neighbour_vertices)
530       
531    normals = domain.normals
532    #print 'normals=',normals
533    #print 'The type of normals is',type(normals)
534       
535    areas = domain.areas
536    #print 'areas=',areas
537    #print 'The type of areas is',type(areas)
538       
539    stage_edge_values = domain.quantities['stage'].vertex_values
540    #print 'stage_edge_values=',stage_edge_values
541    #print 'The type of stage_edge_values is',type(stage_edge_values)
542       
543    xmom_edge_values = domain.quantities['xmomentum'].vertex_values
544    #print 'xmom_edge_values=',xmom_edge_values
545    #print 'The type of xmom_edge_values is',type(xmom_edge_values)
546       
547    bed_edge_values = domain.quantities['elevation'].vertex_values
548    #print 'bed_edge_values=',bed_edge_values
549    #print 'The type of bed_edge_values is',type(bed_edge_values)
550       
551    stage_boundary_values = domain.quantities['stage'].boundary_values
552    #print 'stage_boundary_values=',stage_boundary_values
553    #print 'The type of stage_boundary_values is',type(stage_boundary_values)
554       
555    xmom_boundary_values = domain.quantities['xmomentum'].boundary_values
556    #print 'xmom_boundary_values=',xmom_boundary_values
557    #print 'The type of xmom_boundary_values is',type(xmom_boundary_values)
558       
559    stage_explicit_update = domain.quantities['stage'].explicit_update
560    #print 'stage_explicit_update=',stage_explicit_update
561    #print 'The type of stage_explicit_update is',type(stage_explicit_update)
562       
563    xmom_explicit_update = domain.quantities['xmomentum'].explicit_update
564    #print 'xmom_explicit_update=',xmom_explicit_update
565    #print 'The type of xmom_explicit_update is',type(xmom_explicit_update)
566       
567    number_of_elements = len(stage_edge_values)
568    #print 'number_of_elements=',number_of_elements
569    #print 'The type of number_of_elements is',type(number_of_elements)
570       
571    max_speed_array = domain.max_speed_array
572    #print 'max_speed_array=',max_speed_array
573    #print 'The type of max_speed_array is',type(max_speed_array)
574       
575       
576    from comp_flux_ext import compute_fluxes_ext
577       
578    domain.flux_timestep = compute_fluxes_ext(timestep,
579                                  epsilon,
580                                  g,
581                                  neighbours,
582                                  neighbour_vertices,
583                                  normals,
584                                  areas, 
585                                  stage_edge_values,
586                                  xmom_edge_values,
587                                  bed_edge_values,
588                                  stage_boundary_values,
589                                  xmom_boundary_values,
590                                  stage_explicit_update,
591                                  xmom_explicit_update,
592                                  number_of_elements,
593                                  max_speed_array)
594
595
596# Compute flux definition
597def compute_fluxes_C_short(domain):
598    from Numeric import zeros, Float
599    import sys
600
601       
602    timestep = float(sys.maxint)
603
604    stage = domain.quantities['stage']
605    xmom = domain.quantities['xmomentum']
606    bed = domain.quantities['elevation']
607
608
609    from comp_flux_ext import compute_fluxes_ext_short
610
611    domain.flux_timestep = compute_fluxes_ext_short(timestep,domain,stage,xmom,bed)
612
613
614
615# ###################################
616def compute_fluxes_C_nonwellbalanced3(domain):
617    import sys
618    from Numeric import zeros, Float
619   
620    N = domain.number_of_elements
621    timestep = float(sys.maxint)
622    epsilon = domain.epsilon
623    g = domain.g
624    neighbours = domain.neighbours
625    neighbour_vertices = domain.neighbour_vertices
626    normals = domain.normals
627    areas = domain.areas
628   
629    Stage    = domain.quantities['stage']
630    Xmom     = domain.quantities['xmomentum']
631    Bed      = domain.quantities['elevation']
632    Height   = domain.quantities['height']
633    Velocity = domain.quantities['velocity']
634
635
636    stage_boundary_values = Stage.boundary_values
637    xmom_boundary_values  = Xmom.boundary_values
638    bed_boundary_values   = Bed.boundary_values
639    height_boundary_values= Height.boundary_values
640    vel_boundary_values   = Velocity.boundary_values
641   
642    stage_explicit_update = Stage.explicit_update
643    xmom_explicit_update  = Xmom.explicit_update
644    bed_explicit_values   = Bed.explicit_update
645    height_explicit_values= Height.explicit_update
646    vel_explicit_values   = Velocity.explicit_update
647   
648    max_speed_array = domain.max_speed_array
649   
650    domain.distribute_to_vertices_and_edges()
651    domain.update_boundary()
652   
653    stage_V = Stage.vertex_values
654    xmom_V  = Xmom.vertex_values
655    bed_V   = Bed.vertex_values
656    height_V= Height.vertex_values
657    vel_V   = Velocity.vertex_values
658
659    number_of_elements = len(stage_V)
660
661
662       
663    #from config import h_min
664    from comp_flux_ext_nonwellbalanced3 import compute_fluxes_ext_nonwellbalanced3
665    domain.flux_timestep = compute_fluxes_ext_nonwellbalanced3(timestep,
666                                  epsilon,
667                                  g,
668                                                           
669                                  neighbours,
670                                  neighbour_vertices,
671                                  normals,
672                                  areas,
673                                                           
674                                  stage_V,
675                                  xmom_V,
676                                  bed_V,
677                                  height_V,
678                                  vel_V,
679                                                           
680                                  stage_boundary_values,
681                                  xmom_boundary_values,
682                                  bed_boundary_values,
683                                  height_boundary_values,
684                                  vel_boundary_values,
685                                                           
686                                  stage_explicit_update,
687                                  xmom_explicit_update,
688                                  bed_explicit_values,
689                                  height_explicit_values,
690                                  vel_explicit_values,
691                                                           
692                                  number_of_elements,
693                                  max_speed_array)
694 
695
696# ###################################
697
698
699
700
701
702
703# Module functions for gradient limiting (distribute_to_vertices_and_edges)
704
705def distribute_to_vertices_and_edges(domain):
706    """Distribution from centroids to vertices specific to the
707    shallow water wave equation.
708
709    It will ensure that h (w-z) is always non-negative even in the
710    presence of steep bed-slopes by taking a weighted average between shallow
711    and deep cases.
712
713    In addition, all conserved quantities get distributed as per either a
714    constant (order==1) or a piecewise linear function (order==2).
715
716    FIXME: more explanation about removal of artificial variability etc
717
718    Precondition:
719      All quantities defined at centroids and bed elevation defined at
720      vertices.
721
722    Postcondition
723      Conserved quantities defined at vertices
724
725    """
726
727    #from config import optimised_gradient_limiter
728
729    #Remove very thin layers of water
730    #protect_against_infinitesimal_and_negative_heights(domain) 
731
732    import sys
733    from Numeric import zeros, Float
734    from config import epsilon, h0, h_min
735
736    N = domain.number_of_elements
737
738    #Shortcuts
739    Stage = domain.quantities['stage']
740    Xmom = domain.quantities['xmomentum']
741    Bed = domain.quantities['elevation']
742    Height = domain.quantities['height']
743    Velocity = domain.quantities['velocity']
744
745    #Arrays   
746    w_C   = Stage.centroid_values   
747    uh_C  = Xmom.centroid_values   
748    z_C   = Bed.centroid_values
749    h_C   = Height.centroid_values
750    u_C   = Velocity.centroid_values
751
752    for i in range(N):
753        h_C[i] = w_C[i] - z_C[i]                                               
754        if h_C[i] < epsilon : #0.01:
755            #h_C[i] = 0.0
756            #w_C[i] = z_C[i]
757            uh_C[i] = 0.0
758            u_C[i] = 0.0
759           
760        if h_C[i] < epsilon:
761            h_C[i] = 0.0
762            w_C[i] = z_C[i]
763        else:
764            u_C[i]  = uh_C[i]/h_C[i] #+  h0/h_C[i])
765
766    for name in ['stage', 'elevation', 'velocity']:
767        Q = domain.quantities[name]
768        if domain.order == 1:
769            Q.extrapolate_first_order()
770        elif domain.order == 2:
771            #print "add extrapolate second order to shallow water"
772            #if name != 'height':
773            Q.extrapolate_second_order()
774            #Q.limit()
775        else:
776            raise 'Unknown order'
777
778    stage_V = domain.quantities['stage'].vertex_values
779    bed_V   = domain.quantities['elevation'].vertex_values
780    h_V     = domain.quantities['height'].vertex_values
781    u_V     = domain.quantities['velocity'].vertex_values               
782    xmom_V  = domain.quantities['xmomentum'].vertex_values
783
784    i = 0
785    while i < len(stage_V.flat)-1:
786        if bed_V.flat[i] > stage_V.flat[i]: # & (bed_V.flat[i+1] < stage_V.flat[i+1]):
787            delta = bed_V.flat[i] - stage_V.flat[i]
788            bed_V.flat[i] = stage_V.flat[i]
789            bed_V.flat[i+1] = bed_V.flat[i+1] + delta
790        if bed_V.flat[i+1] > stage_V.flat[i+1]: #(bed_V.flat[i] < stage_V.flat[i]) & (bed_V.flat[i+1] > stage_V.flat[i+1]):
791            delta = bed_V.flat[i+1] - stage_V.flat[i+1]
792            bed_V.flat[i+1] = stage_V.flat[i+1]
793            bed_V.flat[i] = bed_V.flat[i] + delta
794        i= i+2
795    #print "stage_V.flat = ", stage_V.flat
796    #print "stage_V      = ", stage_V
797    h_V[:] = stage_V - bed_V
798    xmom_V[:] = u_V * h_V
799   
800    return
801
802   
803#
804def protect_against_infinitesimal_and_negative_heights(domain):
805    """Protect against infinitesimal heights and associated high velocities
806    """
807
808    #Shortcuts
809    wc = domain.quantities['stage'].centroid_values
810    zc = domain.quantities['elevation'].centroid_values
811    xmomc = domain.quantities['xmomentum'].centroid_values
812#    ymomc = domain.quantities['ymomentum'].centroid_values
813    hc = wc - zc  #Water depths at centroids
814
815    zv = domain.quantities['elevation'].vertex_values
816    wv = domain.quantities['stage'].vertex_values
817    hv = wv-zv
818    xmomv = domain.quantities['xmomentum'].vertex_values
819    #remove the above two lines and corresponding code below
820
821    #Update
822    for k in range(domain.number_of_elements):
823
824        if hc[k] < domain.minimum_allowed_height:
825            #Control stage
826            if hc[k] < domain.epsilon:
827                wc[k] = zc[k] # Contain 'lost mass' error
828                wv[k,0] = zv[k,0]
829                wv[k,1] = zv[k,1]
830               
831            xmomc[k] = 0.0
832           
833        #N = domain.number_of_elements
834        #if (k == 0) | (k==N-1):
835        #    wc[k] = zc[k] # Contain 'lost mass' error
836        #    wv[k,0] = zv[k,0]
837        #    wv[k,1] = zv[k,1]
838   
839def h_limiter(domain):
840    """Limit slopes for each volume to eliminate artificial variance
841    introduced by e.g. second order extrapolator
842
843    limit on h = w-z
844
845    This limiter depends on two quantities (w,z) so it resides within
846    this module rather than within quantity.py
847    """
848
849    from Numeric import zeros, Float
850
851    N = domain.number_of_elements
852    beta_h = domain.beta_h
853
854    #Shortcuts
855    wc = domain.quantities['stage'].centroid_values
856    zc = domain.quantities['elevation'].centroid_values
857    hc = wc - zc
858
859    wv = domain.quantities['stage'].vertex_values
860    zv = domain.quantities['elevation'].vertex_values
861    hv = wv-zv
862
863    hvbar = zeros(hv.shape, Float) #h-limited values
864
865    #Find min and max of this and neighbour's centroid values
866    hmax = zeros(hc.shape, Float)
867    hmin = zeros(hc.shape, Float)
868
869    for k in range(N):
870        hmax[k] = hmin[k] = hc[k]
871        #for i in range(3):
872        for i in range(2):   
873            n = domain.neighbours[k,i]
874            if n >= 0:
875                hn = hc[n] #Neighbour's centroid value
876
877                hmin[k] = min(hmin[k], hn)
878                hmax[k] = max(hmax[k], hn)
879
880
881    #Diffences between centroids and maxima/minima
882    dhmax = hmax - hc
883    dhmin = hmin - hc
884
885    #Deltas between vertex and centroid values
886    dh = zeros(hv.shape, Float)
887    #for i in range(3):
888    for i in range(2):
889        dh[:,i] = hv[:,i] - hc
890
891    #Phi limiter
892    for k in range(N):
893
894        #Find the gradient limiter (phi) across vertices
895        phi = 1.0
896        #for i in range(3):
897        for i in range(2):
898            r = 1.0
899            if (dh[k,i] > 0): r = dhmax[k]/dh[k,i]
900            if (dh[k,i] < 0): r = dhmin[k]/dh[k,i]
901
902            phi = min( min(r*beta_h, 1), phi )
903
904        #Then update using phi limiter
905        #for i in range(3):
906        for i in range(2):
907            hvbar[k,i] = hc[k] + phi*dh[k,i]
908
909    return hvbar
910
911def balance_deep_and_shallow(domain):
912    """Compute linear combination between stage as computed by
913    gradient-limiters limiting using w, and stage computed by
914    gradient-limiters limiting using h (h-limiter).
915    The former takes precedence when heights are large compared to the
916    bed slope while the latter takes precedence when heights are
917    relatively small.  Anything in between is computed as a balanced
918    linear combination in order to avoid numerical disturbances which
919    would otherwise appear as a result of hard switching between
920    modes.
921
922    The h-limiter is always applied irrespective of the order.
923    """
924
925    #Shortcuts
926    wc = domain.quantities['stage'].centroid_values
927    zc = domain.quantities['elevation'].centroid_values
928    hc = wc - zc
929
930    wv = domain.quantities['stage'].vertex_values
931    zv = domain.quantities['elevation'].vertex_values
932    hv = wv-zv
933
934    #Limit h
935    hvbar = h_limiter(domain)
936
937    for k in range(domain.number_of_elements):
938        #Compute maximal variation in bed elevation
939        #  This quantitiy is
940        #    dz = max_i abs(z_i - z_c)
941        #  and it is independent of dimension
942        #  In the 1d case zc = (z0+z1)/2
943        #  In the 2d case zc = (z0+z1+z2)/3
944
945        dz = max(abs(zv[k,0]-zc[k]),
946                 abs(zv[k,1]-zc[k]))#,
947#                 abs(zv[k,2]-zc[k]))
948
949
950        hmin = min( hv[k,:] )
951
952        #Create alpha in [0,1], where alpha==0 means using the h-limited
953        #stage and alpha==1 means using the w-limited stage as
954        #computed by the gradient limiter (both 1st or 2nd order)
955
956        #If hmin > dz/2 then alpha = 1 and the bed will have no effect
957        #If hmin < 0 then alpha = 0 reverting to constant height above bed.
958
959        if dz > 0.0:
960            alpha = max( min( 2*hmin/dz, 1.0), 0.0 )
961        else:
962            #Flat bed
963            alpha = 1.0
964
965        alpha  = 0.0
966        #Let
967        #
968        #  wvi be the w-limited stage (wvi = zvi + hvi)
969        #  wvi- be the h-limited state (wvi- = zvi + hvi-)
970        #
971        #
972        #where i=0,1,2 denotes the vertex ids
973        #
974        #Weighted balance between w-limited and h-limited stage is
975        #
976        #  wvi := (1-alpha)*(zvi+hvi-) + alpha*(zvi+hvi)
977        #
978        #It follows that the updated wvi is
979        #  wvi := zvi + (1-alpha)*hvi- + alpha*hvi
980        #
981        # Momentum is balanced between constant and limited
982
983
984        #for i in range(3):
985        #    wv[k,i] = zv[k,i] + hvbar[k,i]
986
987        #return
988
989        if alpha < 1:
990
991            #for i in range(3):
992            for i in range(2):
993                wv[k,i] = zv[k,i] + (1.0-alpha)*hvbar[k,i] + alpha*hv[k,i]
994
995            #Momentums at centroids
996            xmomc = domain.quantities['xmomentum'].centroid_values
997#            ymomc = domain.quantities['ymomentum'].centroid_values
998
999            #Momentums at vertices
1000            xmomv = domain.quantities['xmomentum'].vertex_values
1001#           ymomv = domain.quantities['ymomentum'].vertex_values
1002
1003            # Update momentum as a linear combination of
1004            # xmomc and ymomc (shallow) and momentum
1005            # from extrapolator xmomv and ymomv (deep).
1006            xmomv[k,:] = (1.0-alpha)*xmomc[k] + alpha*xmomv[k,:]
1007#            ymomv[k,:] = (1-alpha)*ymomc[k] + alpha*ymomv[k,:]
1008
1009
1010###############################################
1011#Boundaries - specific to the shallow water wave equation
1012class Reflective_boundary(Boundary):
1013    """Reflective boundary returns same conserved quantities as
1014    those present in its neighbour volume but reflected.
1015
1016    This class is specific to the shallow water equation as it
1017    works with the momentum quantities assumed to be the second
1018    and third conserved quantities.
1019    """
1020
1021    def __init__(self, domain = None):
1022        Boundary.__init__(self)
1023
1024        if domain is None:
1025            msg = 'Domain must be specified for reflective boundary'
1026            raise msg
1027
1028        #Handy shorthands
1029        #self.stage   = domain.quantities['stage'].edge_values
1030        #self.xmom    = domain.quantities['xmomentum'].edge_values
1031        #self.ymom    = domain.quantities['ymomentum'].edge_values
1032        self.normals  = domain.normals
1033        self.stage    = domain.quantities['stage'].vertex_values
1034        self.xmom     = domain.quantities['xmomentum'].vertex_values
1035        self.bed      = domain.quantities['elevation'].vertex_values
1036        self.height   = domain.quantities['height'].vertex_values
1037        self.velocity = domain.quantities['velocity'].vertex_values
1038
1039        from Numeric import zeros, Float
1040        #self.conserved_quantities = zeros(3, Float)
1041        #self.conserved_quantities = zeros(2, Float)
1042        self.evolved_quantities = zeros(5, Float)
1043
1044    def __repr__(self):
1045        return 'Reflective_boundary'
1046
1047    """
1048    def evaluate(self, vol_id, edge_id):
1049        #Reflective boundaries reverses the outward momentum
1050        #of the volume they serve.
1051       
1052        #q = self.conserved_quantities
1053        q = self.evolved_quantities
1054        q[0] = self.stage[vol_id, edge_id]
1055        q[1] = self.xmom[vol_id, edge_id]
1056        #q[2] = self.ymom[vol_id, edge_id]
1057        #normal = self.normals[vol_id, 2*edge_id:2*edge_id+2]
1058        #normal = self.normals[vol_id, 2*edge_id:2*edge_id+1]
1059        normal = self.normals[vol_id,edge_id]
1060
1061        #r = rotate(q, normal, direction = 1)
1062        #r[1] = -r[1]
1063        #q = rotate(r, normal, direction = -1)
1064        r = q
1065        r[1] = normal*r[1]
1066        r[1] = -r[1]
1067        r[1] = normal*r[1]
1068        q = r
1069        #For start interval there is no outward momentum so do not need to
1070        #reverse direction in this case
1071
1072        return q
1073    """
1074    def evaluate(self, vol_id, edge_id):
1075        """Reflective boundaries reverses the outward momentum
1076        of the volume they serve.
1077        """
1078        q = self.evolved_quantities
1079        q[0] = self.stage[vol_id, edge_id]
1080        q[1] = -self.xmom[vol_id, edge_id]
1081        q[2] = self.bed[vol_id, edge_id]
1082        q[3] = self.height[vol_id, edge_id]
1083        q[4] = -self.velocity[vol_id, edge_id]
1084        return q
1085
1086
1087   
1088
1089class Dirichlet_boundary(Boundary):
1090    """Dirichlet boundary returns constant values for the
1091    conserved quantities
1092    """
1093
1094
1095    def __init__(self, conserved_quantities=None):
1096        Boundary.__init__(self)
1097
1098        if conserved_quantities is None:
1099            msg = 'Must specify one value for each conserved quantity'
1100            raise msg
1101
1102        from Numeric import array, Float
1103        self.conserved_quantities=array(conserved_quantities).astype(Float)
1104
1105    def __repr__(self):
1106        return 'Dirichlet boundary (%s)' %self.conserved_quantities
1107
1108    def evaluate(self, vol_id=None, edge_id=None):
1109        return self.conserved_quantities
1110
1111
1112#########################
1113#Standard forcing terms:
1114#
1115def gravity(domain):
1116    """Apply gravitational pull in the presence of bed slope
1117    """
1118
1119    from util import gradient
1120    from Numeric import zeros, Float, array, sum
1121
1122    xmom = domain.quantities['xmomentum'].explicit_update
1123    stage = domain.quantities['stage'].explicit_update
1124#    ymom = domain.quantities['ymomentum'].explicit_update
1125
1126    Stage = domain.quantities['stage']
1127    Elevation = domain.quantities['elevation']
1128    #h = Stage.edge_values - Elevation.edge_values
1129    h = Stage.vertex_values - Elevation.vertex_values
1130    b = Elevation.vertex_values
1131    w = Stage.vertex_values
1132
1133    x = domain.get_vertex_coordinates()
1134    g = domain.g
1135
1136    for k in range(domain.number_of_elements):
1137#        avg_h = sum( h[k,:] )/3
1138        avg_h = sum( h[k,:] )/2
1139
1140        #Compute bed slope
1141        #x0, y0, x1, y1, x2, y2 = x[k,:]
1142        x0, x1 = x[k,:]
1143        #z0, z1, z2 = v[k,:]
1144        b0, b1 = b[k,:]
1145
1146        w0, w1 = w[k,:]
1147        wx = gradient(x0, x1, w0, w1)
1148
1149        #zx, zy = gradient(x0, y0, x1, y1, x2, y2, z0, z1, z2)
1150        bx = gradient(x0, x1, b0, b1)
1151       
1152        #Update momentum (explicit update is reset to source values)
1153        xmom[k] += -g*bx*avg_h
1154        #xmom[k] = -g*bx*avg_h
1155        #stage[k] = 0.0
1156 
1157 
1158def manning_friction(domain):
1159    """Apply (Manning) friction to water momentum
1160    """
1161
1162    from math import sqrt
1163
1164    w = domain.quantities['stage'].centroid_values
1165    z = domain.quantities['elevation'].centroid_values
1166    h = w-z
1167
1168    uh = domain.quantities['xmomentum'].centroid_values
1169    #vh = domain.quantities['ymomentum'].centroid_values
1170    eta = domain.quantities['friction'].centroid_values
1171
1172    xmom_update = domain.quantities['xmomentum'].semi_implicit_update
1173    #ymom_update = domain.quantities['ymomentum'].semi_implicit_update
1174
1175    N = domain.number_of_elements
1176    eps = domain.minimum_allowed_height
1177    g = domain.g
1178
1179    for k in range(N):
1180        if eta[k] >= eps:
1181            if h[k] >= eps:
1182                #S = -g * eta[k]**2 * sqrt((uh[k]**2 + vh[k]**2))
1183                S = -g * eta[k]**2 * uh[k]
1184                S /= h[k]**(7.0/3)
1185
1186                #Update momentum
1187                xmom_update[k] += S*uh[k]
1188                #ymom_update[k] += S*vh[k]
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        N = 2
1238        #x = ones(3, Float)
1239        #y = ones(3, Float)
1240        x = ones(2, Float)
1241        #y = ones(2, Float)
1242       
1243        try:
1244            #q = f(1.0, x=x, y=y)
1245            q = f(1.0, x=x)
1246        except Exception, e:
1247            msg = 'Function %s could not be executed:\n%s' %(f, e)
1248            #FIXME: Reconsider this semantics
1249            raise msg
1250
1251        try:
1252            q = array(q).astype(Float)
1253        except:
1254            msg = 'Return value from vector function %s could ' %f
1255            msg += 'not be converted into a Numeric array of floats.\n'
1256            msg += 'Specified function should return either list or array.'
1257            raise msg
1258
1259        #Is this really what we want?
1260        msg = 'Return vector from function %s ' %f
1261        msg += 'must have same lenght as input vectors'
1262        assert len(q) == N, msg
1263
1264    else:
1265        try:
1266            f = float(f)
1267        except:
1268            msg = 'Force field %s must be either a scalar' %f
1269            msg += ' or a vector function'
1270            raise msg
1271    return f
1272
1273class Wind_stress:
1274    """Apply wind stress to water momentum in terms of
1275    wind speed [m/s] and wind direction [degrees]
1276    """
1277
1278    def __init__(self, *args, **kwargs):
1279        """Initialise windfield from wind speed s [m/s]
1280        and wind direction phi [degrees]
1281
1282        Inputs v and phi can be either scalars or Python functions, e.g.
1283
1284        W = Wind_stress(10, 178)
1285
1286        #FIXME - 'normal' degrees are assumed for now, i.e. the
1287        vector (1,0) has zero degrees.
1288        We may need to convert from 'compass' degrees later on and also
1289        map from True north to grid north.
1290
1291        Arguments can also be Python functions of t,x,y as in
1292
1293        def speed(t,x,y):
1294            ...
1295            return s
1296
1297        def angle(t,x,y):
1298            ...
1299            return phi
1300
1301        where x and y are vectors.
1302
1303        and then pass the functions in
1304
1305        W = Wind_stress(speed, angle)
1306
1307        The instantiated object W can be appended to the list of
1308        forcing_terms as in
1309
1310        Alternatively, one vector valued function for (speed, angle)
1311        can be applied, providing both quantities simultaneously.
1312        As in
1313        W = Wind_stress(F), where returns (speed, angle) for each t.
1314
1315        domain.forcing_terms.append(W)
1316        """
1317
1318        from config import rho_a, rho_w, eta_w
1319        from Numeric import array, Float
1320
1321        if len(args) == 2:
1322            s = args[0]
1323            phi = args[1]
1324        elif len(args) == 1:
1325            #Assume vector function returning (s, phi)(t,x,y)
1326            vector_function = args[0]
1327            #s = lambda t,x,y: vector_function(t,x=x,y=y)[0]
1328            #phi = lambda t,x,y: vector_function(t,x=x,y=y)[1]
1329            s = lambda t,x: vector_function(t,x=x)[0]
1330            phi = lambda t,x: vector_function(t,x=x)[1]
1331        else:
1332           #Assume info is in 2 keyword arguments
1333
1334           if len(kwargs) == 2:
1335               s = kwargs['s']
1336               phi = kwargs['phi']
1337           else:
1338               raise 'Assumes two keyword arguments: s=..., phi=....'
1339
1340        print 'phi', phi
1341        self.speed = check_forcefield(s)
1342        self.phi = check_forcefield(phi)
1343
1344        self.const = eta_w*rho_a/rho_w
1345
1346
1347    def __call__(self, domain):
1348        """Evaluate windfield based on values found in domain
1349        """
1350
1351        from math import pi, cos, sin, sqrt
1352        from Numeric import Float, ones, ArrayType
1353
1354        xmom_update = domain.quantities['xmomentum'].explicit_update
1355        #ymom_update = domain.quantities['ymomentum'].explicit_update
1356
1357        N = domain.number_of_elements
1358        t = domain.time
1359
1360        if callable(self.speed):
1361            xc = domain.get_centroid_coordinates()
1362            #s_vec = self.speed(t, xc[:,0], xc[:,1])
1363            s_vec = self.speed(t, xc)
1364        else:
1365            #Assume s is a scalar
1366
1367            try:
1368                s_vec = self.speed * ones(N, Float)
1369            except:
1370                msg = 'Speed must be either callable or a scalar: %s' %self.s
1371                raise msg
1372
1373
1374        if callable(self.phi):
1375            xc = domain.get_centroid_coordinates()
1376            #phi_vec = self.phi(t, xc[:,0], xc[:,1])
1377            phi_vec = self.phi(t, xc)
1378        else:
1379            #Assume phi is a scalar
1380
1381            try:
1382                phi_vec = self.phi * ones(N, Float)
1383            except:
1384                msg = 'Angle must be either callable or a scalar: %s' %self.phi
1385                raise msg
1386
1387        #assign_windfield_values(xmom_update, ymom_update,
1388        #                        s_vec, phi_vec, self.const)
1389        assign_windfield_values(xmom_update, s_vec, phi_vec, self.const)
1390
1391
1392#def assign_windfield_values(xmom_update, ymom_update,
1393#                            s_vec, phi_vec, const):
1394def assign_windfield_values(xmom_update, s_vec, phi_vec, const):
1395    """Python version of assigning wind field to update vectors.
1396    A c version also exists (for speed)
1397    """
1398    from math import pi, cos, sin, sqrt
1399
1400    N = len(s_vec)
1401    for k in range(N):
1402        s = s_vec[k]
1403        phi = phi_vec[k]
1404
1405        #Convert to radians
1406        phi = phi*pi/180
1407
1408        #Compute velocity vector (u, v)
1409        u = s*cos(phi)
1410        v = s*sin(phi)
1411
1412        #Compute wind stress
1413        #S = const * sqrt(u**2 + v**2)
1414        S = const * u
1415        xmom_update[k] += S*u
1416        #ymom_update[k] += S*v
Note: See TracBrowser for help on using the repository browser.