source: anuga_work/development/anuga_1d/shallow_water_try.py @ 5562

Last change on this file since 5562 was 5562, checked in by steve, 11 years ago

Adding C extension for flux calculation

File size: 31.6 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 *
46#from domain_order2 import *
47from domain_t2 import *
48from comp_flux_ext import *
49Generic_Domain = Domain #Rename
50
51#Shallow water domain
52class Domain(Generic_Domain):
53
54    def __init__(self, coordinates, boundary = None, tagged_elements = None,
55                 geo_reference = None):
56
57        conserved_quantities = ['stage', 'xmomentum']
58        other_quantities = ['elevation', 'friction']
59        Generic_Domain.__init__(self, coordinates, boundary,
60                                conserved_quantities, other_quantities,
61                                tagged_elements, geo_reference)
62       
63        from config import minimum_allowed_height, g
64        self.minimum_allowed_height = minimum_allowed_height
65        self.g = g
66
67        #forcing terms not included in 1d domain ?WHy?
68        self.forcing_terms.append(gravity)
69        #self.forcing_terms.append(manning_friction)
70        #print "\nI have Removed forcing terms line 64 1dsw"
71
72        #Realtime visualisation
73        self.visualiser = None
74        self.visualise  = False
75        self.visualise_color_stage = False
76        self.visualise_stage_range = 1.0
77        self.visualise_timer = True
78        self.visualise_range_z = None
79       
80        #Stored output
81        self.store = True
82        self.format = 'sww'
83        self.smooth = True
84
85        #Evolve parametrs
86        self.cfl = 1.0
87       
88        #Reduction operation for get_vertex_values
89        from util import mean
90        self.reduction = mean
91        #self.reduction = min  #Looks better near steep slopes
92
93        self.quantities_to_be_stored = ['stage','xmomentum']
94
95
96    def set_quantities_to_be_stored(self, q):
97        """Specify which quantities will be stored in the sww file.
98
99        q must be either:
100          - the name of a quantity
101          - a list of quantity names
102          - None
103
104        In the two first cases, the named quantities will be stored at each
105        yieldstep
106        (This is in addition to the quantities elevation and friction) 
107        If q is None, storage will be switched off altogether.
108        """
109
110
111        if q is None:
112            self.quantities_to_be_stored = []           
113            self.store = False
114            return
115
116        if isinstance(q, basestring):
117            q = [q] # Turn argument into a list
118
119        #Check correcness   
120        for quantity_name in q:
121            msg = 'Quantity %s is not a valid conserved quantity' %quantity_name
122            assert quantity_name in self.conserved_quantities, msg
123       
124        self.quantities_to_be_stored = q
125       
126
127    def initialise_visualiser(self,scale_z=1.0,rect=None):
128        #Realtime visualisation
129        if self.visualiser is None:
130            from realtime_visualisation_new import Visualiser
131            self.visualiser = Visualiser(self,scale_z,rect)
132            self.visualiser.setup['elevation']=True
133            self.visualiser.updating['stage']=True
134        self.visualise = True
135        if self.visualise_color_stage == True:
136            self.visualiser.coloring['stage'] = True
137            self.visualiser.qcolor['stage'] = (0.0, 0.0, 0.8)
138        print 'initialise visualiser'
139        print self.visualiser.setup
140        print self.visualiser.updating
141
142    def check_integrity(self):
143        Generic_Domain.check_integrity(self)
144        #Check that we are solving the shallow water wave equation
145
146        msg = 'First conserved quantity must be "stage"'
147        assert self.conserved_quantities[0] == 'stage', msg
148        msg = 'Second conserved quantity must be "xmomentum"'
149        assert self.conserved_quantities[1] == 'xmomentum', msg
150
151    def extrapolate_second_order_sw(self):
152        #Call correct module function
153        #(either from this module or C-extension)
154        extrapolate_second_order_sw(self)
155
156    def compute_fluxes(self):
157        #Call correct module function
158        #(either from this module or C-extension)
159        compute_fluxes(self)
160
161    def compute_timestep(self):
162        #Call correct module function
163        compute_timestep(self)
164
165    def distribute_to_vertices_and_edges(self):
166        #Call correct module function
167        #(either from this module or C-extension)
168        distribute_to_vertices_and_edges(self)
169       
170    def evolve(self, yieldstep = None, finaltime = None,
171               skip_initial_step = False):
172        """Specialisation of basic evolve method from parent class
173        """
174
175        #Call check integrity here rather than from user scripts
176        #self.check_integrity()
177
178        #msg = 'Parameter beta_h must be in the interval [0, 1)'
179        #assert 0 <= self.beta_h < 1.0, msg
180        #msg = 'Parameter beta_w must be in the interval [0, 1)'
181        #assert 0 <= self.beta_w < 1.0, msg
182
183
184        #Initial update of vertex and edge values before any storage
185        #and or visualisation
186       
187        self.distribute_to_vertices_and_edges()
188       
189        #Call basic machinery from parent class
190        for t in Generic_Domain.evolve(self, yieldstep, finaltime,
191                                       skip_initial_step):
192            yield(t)
193
194    def initialise_storage(self):
195        """Create and initialise self.writer object for storing data.
196        Also, save x and bed elevation
197        """
198
199        import data_manager
200
201        #Initialise writer
202        self.writer = data_manager.get_dataobject(self, mode = 'w')
203
204        #Store vertices and connectivity
205        self.writer.store_connectivity()
206
207
208    def store_timestep(self, name):
209        """Store named quantity and time.
210
211        Precondition:
212           self.write has been initialised
213        """
214        self.writer.store_timestep(name)
215
216
217#=============== End of Shallow Water Domain ===============================
218
219#Rotation of momentum vector
220def rotate(q, normal, direction = 1):
221    """Rotate the momentum component q (q[1], q[2])
222    from x,y coordinates to coordinates based on normal vector.
223
224    If direction is negative the rotation is inverted.
225
226    Input vector is preserved
227
228    This function is specific to the shallow water wave equation
229    """
230
231    from Numeric import zeros, Float
232
233    assert len(q) == 3,\
234           'Vector of conserved quantities must have length 3'\
235           'for 2D shallow water equation'
236
237    try:
238        l = len(normal)
239    except:
240        raise 'Normal vector must be an Numeric array'
241
242    assert l == 2, 'Normal vector must have 2 components'
243
244
245    n1 = normal[0]
246    n2 = normal[1]
247
248    r = zeros(len(q), Float) #Rotated quantities
249    r[0] = q[0]              #First quantity, height, is not rotated
250
251    if direction == -1:
252        n2 = -n2
253
254
255    r[1] =  n1*q[1] + n2*q[2]
256    r[2] = -n2*q[1] + n1*q[2]
257
258    return r
259"""++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++"""
260# Compute flux definition
261def compute_fluxes(domain):
262        from Numeric import zeros, Float
263        import sys
264       
265        timestep = float(sys.maxint)
266        print 'timestep=',timestep
267        print 'The type of timestep is',type(timestep)
268       
269        epsilon = domain.epsilon
270        print 'epsilon=',epsilon
271        print 'The type of epsilon is',type(epsilon)
272       
273        g = domain.g
274        print 'g=',g
275        print 'The type of g is',type(g)
276       
277        neighbours = domain.neighbours
278        print 'neighbours=',neighbours
279        print 'The type of neighbours is',type(neighbours)
280       
281        neighbour_vertices = domain.neighbour_vertices
282        print 'neighbour_vertices=',neighbour_vertices
283        print 'The type of neighbour_vertices is',type(neighbour_vertices)
284       
285        normals = domain.normals
286        print 'normals=',normals
287        print 'The type of normals is',type(normals)
288       
289        areas = domain.areas
290        print 'areas=',areas
291        print 'The type of areas is',type(areas)
292       
293        stage_edge_values = domain.quantities['stage'].vertex_values
294        print 'stage_edge_values=',stage_edge_values
295        print 'The type of stage_edge_values is',type(stage_edge_values)
296       
297        xmom_edge_values = domain.quantities['xmomentum'].vertex_values
298        print 'xmom_edge_values=',xmom_edge_values
299        print 'The type of xmom_edge_values is',type(xmom_edge_values)
300       
301        bed_edge_values = domain.quantities['elevation'].vertex_values
302        print 'bed_edge_values=',bed_edge_values
303        print 'The type of bed_edge_values is',type(bed_edge_values)
304       
305        stage_boundary_values = domain.quantities['stage'].boundary_values
306        print 'stage_boundary_values=',stage_boundary_values
307        print 'The type of stage_boundary_values is',type(stage_boundary_values)
308       
309        xmom_boundary_values = domain.quantities['xmomentum'].boundary_values
310        print 'xmom_boundary_values=',xmom_boundary_values
311        print 'The type of xmom_boundary_values is',type(xmom_boundary_values)
312       
313        stage_explicit_update = domain.quantities['stage'].explicit_update
314        print 'stage_explicit_update=',stage_explicit_update
315        print 'The type of stage_explicit_update is',type(stage_explicit_update)
316       
317        xmom_explicit_update = domain.quantities['xmomentum'].explicit_update
318        print 'xmom_explicit_update=',xmom_explicit_update
319        print 'The type of xmom_explicit_update is',type(xmom_explicit_update)
320       
321        number_of_elements = len(stage_edge_values)
322        print 'number_of_elements=',number_of_elements
323        print 'The type of number_of_elements is',type(number_of_elements)
324       
325        max_speed_array = zeros(number_of_elements, Float)
326        print 'max_speed_array=',max_speed_array
327        print 'The type of max_speed_array is',type(max_speed_array)
328       
329        return compute_fluxes_ext(timestep, epsilon, g, neighbours, neighbour_vertices, normals, areas, 
330        stage_edge_values, xmom_edge_values, bed_edge_values, stage_boundary_values, xmom_boundary_values,
331        number_of_elements, max_speed_array)
332####################################
333
334# Module functions for gradient limiting (distribute_to_vertices_and_edges)
335
336def distribute_to_vertices_and_edges(domain):
337    """Distribution from centroids to vertices specific to the
338    shallow water wave
339    equation.
340
341    It will ensure that h (w-z) is always non-negative even in the
342    presence of steep bed-slopes by taking a weighted average between shallow
343    and deep cases.
344
345    In addition, all conserved quantities get distributed as per either a
346    constant (order==1) or a piecewise linear function (order==2).
347
348    FIXME: more explanation about removal of artificial variability etc
349
350    Precondition:
351      All quantities defined at centroids and bed elevation defined at
352      vertices.
353
354    Postcondition
355      Conserved quantities defined at vertices
356
357    """
358
359    #from config import optimised_gradient_limiter
360
361    #Remove very thin layers of water
362    protect_against_infinitesimal_and_negative_heights(domain)
363   
364
365    #Extrapolate all conserved quantities
366    #if optimised_gradient_limiter:
367    #    #MH090605 if second order,
368    #    #perform the extrapolation and limiting on
369    #    #all of the conserved quantities
370
371    #    if (domain.order == 1):
372    #        for name in domain.conserved_quantities:
373    #            Q = domain.quantities[name]
374    #            Q.extrapolate_first_order()
375    #    elif domain.order == 2:
376    #        domain.extrapolate_second_order_sw()
377    #    else:
378    #        raise 'Unknown order'
379    #else:
380        #old code:
381
382    for name in domain.conserved_quantities:
383        Q = domain.quantities[name]
384        if domain.order == 1:
385            Q.extrapolate_first_order()
386        elif domain.order == 2:
387            #print "add extrapolate second order to shallow water"
388            #if name != 'height':
389            Q.extrapolate_second_order()
390            #Q.limit()
391        else:
392            raise 'Unknown order'
393
394    #Take bed elevation into account when water heights are small
395    #balance_deep_and_shallow(domain)
396    #protect_against_infinitesimal_and_negative_heights(domain)
397
398#Compute edge values by interpolation
399    #for name in domain.conserved_quantities:
400    #    Q = domain.quantities[name]
401    #    Q.interpolate_from_vertices_to_edges()   
402   
403def protect_against_infinitesimal_and_negative_heights(domain):
404    """Protect against infinitesimal heights and associated high velocities
405    """
406
407    #Shortcuts
408    wc = domain.quantities['stage'].centroid_values
409    zc = domain.quantities['elevation'].centroid_values
410    xmomc = domain.quantities['xmomentum'].centroid_values
411#    ymomc = domain.quantities['ymomentum'].centroid_values
412    hc = wc - zc  #Water depths at centroids
413
414    zv = domain.quantities['elevation'].vertex_values
415    wv = domain.quantities['stage'].vertex_values
416    hv = wv-zv
417    xmomv = domain.quantities['xmomentum'].vertex_values
418    #remove the above two lines and corresponding code below
419
420    #Update
421    for k in range(domain.number_of_elements):
422
423        if hc[k] < domain.minimum_allowed_height:
424            #Control stage
425            if hc[k] < domain.epsilon:
426                wc[k] = zc[k] # Contain 'lost mass' error
427                wv[k,0] = zv[k,0]
428                wv[k,1] = zv[k,1]
429               
430            xmomc[k] = 0.0
431           
432        #N = domain.number_of_elements
433        #if (k == 0) | (k==N-1):
434        #    wc[k] = zc[k] # Contain 'lost mass' error
435        #    wv[k,0] = zv[k,0]
436        #    wv[k,1] = zv[k,1]
437   
438def h_limiter(domain):
439    """Limit slopes for each volume to eliminate artificial variance
440    introduced by e.g. second order extrapolator
441
442    limit on h = w-z
443
444    This limiter depends on two quantities (w,z) so it resides within
445    this module rather than within quantity.py
446    """
447
448    from Numeric import zeros, Float
449
450    N = domain.number_of_elements
451    beta_h = domain.beta_h
452
453    #Shortcuts
454    wc = domain.quantities['stage'].centroid_values
455    zc = domain.quantities['elevation'].centroid_values
456    hc = wc - zc
457
458    wv = domain.quantities['stage'].vertex_values
459    zv = domain.quantities['elevation'].vertex_values
460    hv = wv-zv
461
462    hvbar = zeros(hv.shape, Float) #h-limited values
463
464    #Find min and max of this and neighbour's centroid values
465    hmax = zeros(hc.shape, Float)
466    hmin = zeros(hc.shape, Float)
467
468    for k in range(N):
469        hmax[k] = hmin[k] = hc[k]
470        #for i in range(3):
471        for i in range(2):   
472            n = domain.neighbours[k,i]
473            if n >= 0:
474                hn = hc[n] #Neighbour's centroid value
475
476                hmin[k] = min(hmin[k], hn)
477                hmax[k] = max(hmax[k], hn)
478
479
480    #Diffences between centroids and maxima/minima
481    dhmax = hmax - hc
482    dhmin = hmin - hc
483
484    #Deltas between vertex and centroid values
485    dh = zeros(hv.shape, Float)
486    #for i in range(3):
487    for i in range(2):
488        dh[:,i] = hv[:,i] - hc
489
490    #Phi limiter
491    for k in range(N):
492
493        #Find the gradient limiter (phi) across vertices
494        phi = 1.0
495        #for i in range(3):
496        for i in range(2):
497            r = 1.0
498            if (dh[k,i] > 0): r = dhmax[k]/dh[k,i]
499            if (dh[k,i] < 0): r = dhmin[k]/dh[k,i]
500
501            phi = min( min(r*beta_h, 1), phi )
502
503        #Then update using phi limiter
504        #for i in range(3):
505        for i in range(2):
506            hvbar[k,i] = hc[k] + phi*dh[k,i]
507
508    return hvbar
509
510def balance_deep_and_shallow(domain):
511    """Compute linear combination between stage as computed by
512    gradient-limiters limiting using w, and stage computed by
513    gradient-limiters limiting using h (h-limiter).
514    The former takes precedence when heights are large compared to the
515    bed slope while the latter takes precedence when heights are
516    relatively small.  Anything in between is computed as a balanced
517    linear combination in order to avoid numerical disturbances which
518    would otherwise appear as a result of hard switching between
519    modes.
520
521    The h-limiter is always applied irrespective of the order.
522    """
523
524    #Shortcuts
525    wc = domain.quantities['stage'].centroid_values
526    zc = domain.quantities['elevation'].centroid_values
527    hc = wc - zc
528
529    wv = domain.quantities['stage'].vertex_values
530    zv = domain.quantities['elevation'].vertex_values
531    hv = wv-zv
532
533    #Limit h
534    hvbar = h_limiter(domain)
535
536    for k in range(domain.number_of_elements):
537        #Compute maximal variation in bed elevation
538        #  This quantitiy is
539        #    dz = max_i abs(z_i - z_c)
540        #  and it is independent of dimension
541        #  In the 1d case zc = (z0+z1)/2
542        #  In the 2d case zc = (z0+z1+z2)/3
543
544        dz = max(abs(zv[k,0]-zc[k]),
545                 abs(zv[k,1]-zc[k]))#,
546#                 abs(zv[k,2]-zc[k]))
547
548
549        hmin = min( hv[k,:] )
550
551        #Create alpha in [0,1], where alpha==0 means using the h-limited
552        #stage and alpha==1 means using the w-limited stage as
553        #computed by the gradient limiter (both 1st or 2nd order)
554
555        #If hmin > dz/2 then alpha = 1 and the bed will have no effect
556        #If hmin < 0 then alpha = 0 reverting to constant height above bed.
557
558        if dz > 0.0:
559            alpha = max( min( 2*hmin/dz, 1.0), 0.0 )
560        else:
561            #Flat bed
562            alpha = 1.0
563
564        alpha  = 0.0
565        #Let
566        #
567        #  wvi be the w-limited stage (wvi = zvi + hvi)
568        #  wvi- be the h-limited state (wvi- = zvi + hvi-)
569        #
570        #
571        #where i=0,1,2 denotes the vertex ids
572        #
573        #Weighted balance between w-limited and h-limited stage is
574        #
575        #  wvi := (1-alpha)*(zvi+hvi-) + alpha*(zvi+hvi)
576        #
577        #It follows that the updated wvi is
578        #  wvi := zvi + (1-alpha)*hvi- + alpha*hvi
579        #
580        # Momentum is balanced between constant and limited
581
582
583        #for i in range(3):
584        #    wv[k,i] = zv[k,i] + hvbar[k,i]
585
586        #return
587
588        if alpha < 1:
589
590            #for i in range(3):
591            for i in range(2):
592                wv[k,i] = zv[k,i] + (1.0-alpha)*hvbar[k,i] + alpha*hv[k,i]
593
594            #Momentums at centroids
595            xmomc = domain.quantities['xmomentum'].centroid_values
596#            ymomc = domain.quantities['ymomentum'].centroid_values
597
598            #Momentums at vertices
599            xmomv = domain.quantities['xmomentum'].vertex_values
600#           ymomv = domain.quantities['ymomentum'].vertex_values
601
602            # Update momentum as a linear combination of
603            # xmomc and ymomc (shallow) and momentum
604            # from extrapolator xmomv and ymomv (deep).
605            xmomv[k,:] = (1.0-alpha)*xmomc[k] + alpha*xmomv[k,:]
606#            ymomv[k,:] = (1-alpha)*ymomc[k] + alpha*ymomv[k,:]
607
608
609###############################################
610#Boundaries - specific to the shallow water wave equation
611class Reflective_boundary(Boundary):
612    """Reflective boundary returns same conserved quantities as
613    those present in its neighbour volume but reflected.
614
615    This class is specific to the shallow water equation as it
616    works with the momentum quantities assumed to be the second
617    and third conserved quantities.
618    """
619
620    def __init__(self, domain = None):
621        Boundary.__init__(self)
622
623        if domain is None:
624            msg = 'Domain must be specified for reflective boundary'
625            raise msg
626
627        #Handy shorthands
628        #self.stage   = domain.quantities['stage'].edge_values
629        #self.xmom    = domain.quantities['xmomentum'].edge_values
630        #self.ymom    = domain.quantities['ymomentum'].edge_values
631        self.normals = domain.normals
632        self.stage   = domain.quantities['stage'].vertex_values
633        self.xmom    = domain.quantities['xmomentum'].vertex_values
634
635        from Numeric import zeros, Float
636        #self.conserved_quantities = zeros(3, Float)
637        self.conserved_quantities = zeros(2, Float)
638
639    def __repr__(self):
640        return 'Reflective_boundary'
641
642
643    def evaluate(self, vol_id, edge_id):
644        """Reflective boundaries reverses the outward momentum
645        of the volume they serve.
646        """
647
648        q = self.conserved_quantities
649        q[0] = self.stage[vol_id, edge_id]
650        q[1] = self.xmom[vol_id, edge_id]
651        #q[2] = self.ymom[vol_id, edge_id]
652        #normal = self.normals[vol_id, 2*edge_id:2*edge_id+2]
653        #normal = self.normals[vol_id, 2*edge_id:2*edge_id+1]
654        normal = self.normals[vol_id,edge_id]
655
656        #r = rotate(q, normal, direction = 1)
657        #r[1] = -r[1]
658        #q = rotate(r, normal, direction = -1)
659        r = q
660        r[1] = normal*r[1]
661        r[1] = -r[1]
662        r[1] = normal*r[1]
663        q = r
664        #For start interval there is no outward momentum so do not need to
665        #reverse direction in this case
666
667        return q
668
669class Dirichlet_boundary(Boundary):
670    """Dirichlet boundary returns constant values for the
671    conserved quantities
672    """
673
674
675    def __init__(self, conserved_quantities=None):
676        Boundary.__init__(self)
677
678        if conserved_quantities is None:
679            msg = 'Must specify one value for each conserved quantity'
680            raise msg
681
682        from Numeric import array, Float
683        self.conserved_quantities=array(conserved_quantities).astype(Float)
684
685    def __repr__(self):
686        return 'Dirichlet boundary (%s)' %self.conserved_quantities
687
688    def evaluate(self, vol_id=None, edge_id=None):
689        return self.conserved_quantities
690
691
692#########################
693#Standard forcing terms:
694#
695def gravity(domain):
696    """Apply gravitational pull in the presence of bed slope
697    """
698
699    from util import gradient
700    from Numeric import zeros, Float, array, sum
701
702    xmom = domain.quantities['xmomentum'].explicit_update
703    stage = domain.quantities['stage'].explicit_update
704#    ymom = domain.quantities['ymomentum'].explicit_update
705
706    Stage = domain.quantities['stage']
707    Elevation = domain.quantities['elevation']
708    #h = Stage.edge_values - Elevation.edge_values
709    h = Stage.vertex_values - Elevation.vertex_values
710    b = Elevation.vertex_values
711    w = Stage.vertex_values
712
713    x = domain.get_vertex_coordinates()
714    g = domain.g
715
716    for k in range(domain.number_of_elements):
717#        avg_h = sum( h[k,:] )/3
718        avg_h = sum( h[k,:] )/2
719
720        #Compute bed slope
721        #x0, y0, x1, y1, x2, y2 = x[k,:]
722        x0, x1 = x[k,:]
723        #z0, z1, z2 = v[k,:]
724        b0, b1 = b[k,:]
725
726        w0, w1 = w[k,:]
727        wx = gradient(x0, x1, w0, w1)
728
729        #zx, zy = gradient(x0, y0, x1, y1, x2, y2, z0, z1, z2)
730        bx = gradient(x0, x1, b0, b1)
731       
732        #Update momentum (explicit update is reset to source values)
733        if domain.split == False:
734            xmom[k] += -g*bx*avg_h
735            #xmom[k] = -g*bx*avg_h
736            #stage[k] = 0.0
737        elif domain.split == True:
738            xmom[k] += -g*wx*avg_h
739            #xmom[k] = -g*wx*avg_h
740        #ymom[k] += -g*zy*avg_h
741 
742def manning_friction(domain):
743    """Apply (Manning) friction to water momentum
744    """
745
746    from math import sqrt
747
748    w = domain.quantities['stage'].centroid_values
749    z = domain.quantities['elevation'].centroid_values
750    h = w-z
751
752    uh = domain.quantities['xmomentum'].centroid_values
753    #vh = domain.quantities['ymomentum'].centroid_values
754    eta = domain.quantities['friction'].centroid_values
755
756    xmom_update = domain.quantities['xmomentum'].semi_implicit_update
757    #ymom_update = domain.quantities['ymomentum'].semi_implicit_update
758
759    N = domain.number_of_elements
760    eps = domain.minimum_allowed_height
761    g = domain.g
762
763    for k in range(N):
764        if eta[k] >= eps:
765            if h[k] >= eps:
766                #S = -g * eta[k]**2 * sqrt((uh[k]**2 + vh[k]**2))
767                S = -g * eta[k]**2 * uh[k]
768                S /= h[k]**(7.0/3)
769
770                #Update momentum
771                xmom_update[k] += S*uh[k]
772                #ymom_update[k] += S*vh[k]
773
774def linear_friction(domain):
775    """Apply linear friction to water momentum
776
777    Assumes quantity: 'linear_friction' to be present
778    """
779
780    from math import sqrt
781
782    w = domain.quantities['stage'].centroid_values
783    z = domain.quantities['elevation'].centroid_values
784    h = w-z
785
786    uh = domain.quantities['xmomentum'].centroid_values
787#    vh = domain.quantities['ymomentum'].centroid_values
788    tau = domain.quantities['linear_friction'].centroid_values
789
790    xmom_update = domain.quantities['xmomentum'].semi_implicit_update
791#    ymom_update = domain.quantities['ymomentum'].semi_implicit_update
792
793    N = domain.number_of_elements
794    eps = domain.minimum_allowed_height
795    g = domain.g #Not necessary? Why was this added?
796
797    for k in range(N):
798        if tau[k] >= eps:
799            if h[k] >= eps:
800                S = -tau[k]/h[k]
801
802                #Update momentum
803                xmom_update[k] += S*uh[k]
804 #              ymom_update[k] += S*vh[k]
805
806
807
808def check_forcefield(f):
809    """Check that f is either
810    1: a callable object f(t,x,y), where x and y are vectors
811       and that it returns an array or a list of same length
812       as x and y
813    2: a scalar
814    """
815
816    from Numeric import ones, Float, array
817
818
819    if callable(f):
820        #N = 3
821        N = 2
822        #x = ones(3, Float)
823        #y = ones(3, Float)
824        x = ones(2, Float)
825        #y = ones(2, Float)
826       
827        try:
828            #q = f(1.0, x=x, y=y)
829            q = f(1.0, x=x)
830        except Exception, e:
831            msg = 'Function %s could not be executed:\n%s' %(f, e)
832            #FIXME: Reconsider this semantics
833            raise msg
834
835        try:
836            q = array(q).astype(Float)
837        except:
838            msg = 'Return value from vector function %s could ' %f
839            msg += 'not be converted into a Numeric array of floats.\n'
840            msg += 'Specified function should return either list or array.'
841            raise msg
842
843        #Is this really what we want?
844        msg = 'Return vector from function %s ' %f
845        msg += 'must have same lenght as input vectors'
846        assert len(q) == N, msg
847
848    else:
849        try:
850            f = float(f)
851        except:
852            msg = 'Force field %s must be either a scalar' %f
853            msg += ' or a vector function'
854            raise msg
855    return f
856
857class Wind_stress:
858    """Apply wind stress to water momentum in terms of
859    wind speed [m/s] and wind direction [degrees]
860    """
861
862    def __init__(self, *args, **kwargs):
863        """Initialise windfield from wind speed s [m/s]
864        and wind direction phi [degrees]
865
866        Inputs v and phi can be either scalars or Python functions, e.g.
867
868        W = Wind_stress(10, 178)
869
870        #FIXME - 'normal' degrees are assumed for now, i.e. the
871        vector (1,0) has zero degrees.
872        We may need to convert from 'compass' degrees later on and also
873        map from True north to grid north.
874
875        Arguments can also be Python functions of t,x,y as in
876
877        def speed(t,x,y):
878            ...
879            return s
880
881        def angle(t,x,y):
882            ...
883            return phi
884
885        where x and y are vectors.
886
887        and then pass the functions in
888
889        W = Wind_stress(speed, angle)
890
891        The instantiated object W can be appended to the list of
892        forcing_terms as in
893
894        Alternatively, one vector valued function for (speed, angle)
895        can be applied, providing both quantities simultaneously.
896        As in
897        W = Wind_stress(F), where returns (speed, angle) for each t.
898
899        domain.forcing_terms.append(W)
900        """
901
902        from config import rho_a, rho_w, eta_w
903        from Numeric import array, Float
904
905        if len(args) == 2:
906            s = args[0]
907            phi = args[1]
908        elif len(args) == 1:
909            #Assume vector function returning (s, phi)(t,x,y)
910            vector_function = args[0]
911            #s = lambda t,x,y: vector_function(t,x=x,y=y)[0]
912            #phi = lambda t,x,y: vector_function(t,x=x,y=y)[1]
913            s = lambda t,x: vector_function(t,x=x)[0]
914            phi = lambda t,x: vector_function(t,x=x)[1]
915        else:
916           #Assume info is in 2 keyword arguments
917
918           if len(kwargs) == 2:
919               s = kwargs['s']
920               phi = kwargs['phi']
921           else:
922               raise 'Assumes two keyword arguments: s=..., phi=....'
923
924        print 'phi', phi
925        self.speed = check_forcefield(s)
926        self.phi = check_forcefield(phi)
927
928        self.const = eta_w*rho_a/rho_w
929
930
931    def __call__(self, domain):
932        """Evaluate windfield based on values found in domain
933        """
934
935        from math import pi, cos, sin, sqrt
936        from Numeric import Float, ones, ArrayType
937
938        xmom_update = domain.quantities['xmomentum'].explicit_update
939        #ymom_update = domain.quantities['ymomentum'].explicit_update
940
941        N = domain.number_of_elements
942        t = domain.time
943
944        if callable(self.speed):
945            xc = domain.get_centroid_coordinates()
946            #s_vec = self.speed(t, xc[:,0], xc[:,1])
947            s_vec = self.speed(t, xc)
948        else:
949            #Assume s is a scalar
950
951            try:
952                s_vec = self.speed * ones(N, Float)
953            except:
954                msg = 'Speed must be either callable or a scalar: %s' %self.s
955                raise msg
956
957
958        if callable(self.phi):
959            xc = domain.get_centroid_coordinates()
960            #phi_vec = self.phi(t, xc[:,0], xc[:,1])
961            phi_vec = self.phi(t, xc)
962        else:
963            #Assume phi is a scalar
964
965            try:
966                phi_vec = self.phi * ones(N, Float)
967            except:
968                msg = 'Angle must be either callable or a scalar: %s' %self.phi
969                raise msg
970
971        #assign_windfield_values(xmom_update, ymom_update,
972        #                        s_vec, phi_vec, self.const)
973        assign_windfield_values(xmom_update, s_vec, phi_vec, self.const)
974
975
976#def assign_windfield_values(xmom_update, ymom_update,
977#                            s_vec, phi_vec, const):
978def assign_windfield_values(xmom_update, s_vec, phi_vec, const):
979    """Python version of assigning wind field to update vectors.
980    A c version also exists (for speed)
981    """
982    from math import pi, cos, sin, sqrt
983
984    N = len(s_vec)
985    for k in range(N):
986        s = s_vec[k]
987        phi = phi_vec[k]
988
989        #Convert to radians
990        phi = phi*pi/180
991
992        #Compute velocity vector (u, v)
993        u = s*cos(phi)
994        v = s*sin(phi)
995
996        #Compute wind stress
997        #S = const * sqrt(u**2 + v**2)
998        S = const * u
999        xmom_update[k] += S*u
1000        #ymom_update[k] += S*v
Note: See TracBrowser for help on using the repository browser.