source: anuga_work/development/shallow_water_1d/shallow_water_adjust.py @ 4960

Last change on this file since 4960 was 4959, checked in by steve, 16 years ago

Copied John Jakeman's and Sudi Mungkasi's 1d shallow water code

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