source: anuga_work/development/flow_1d/sww_flow/sww_forcing_terms.py @ 7830

Last change on this file since 7830 was 7830, checked in by steve, 14 years ago

Moving channel code to numpy

File size: 4.5 KB
Line 
1#! /usr/bin/python
2
3
4"""
5Forcing terms for the shallow water equations, gravity, friction etc
6"""
7
8__author__="Stephen Roberts"
9__date__ ="$05/06/2010 5:49:35 PM$"
10
11
12def gravity(domain):
13    """Apply gravitational pull in the presence of bed slope
14    """
15
16    from flow_1d.utilities.util import gradient
17
18    xmom  = domain.quantities['xmomentum'].explicit_update
19    stage = domain.quantities['stage'].explicit_update
20#    ymom = domain.quantities['ymomentum'].explicit_update
21
22    Stage = domain.quantities['stage']
23    Elevation = domain.quantities['elevation']
24    #h = Stage.edge_values - Elevation.edge_values
25    h = Stage.vertex_values - Elevation.vertex_values
26    b = Elevation.vertex_values
27    w = Stage.vertex_values
28
29    x = domain.get_vertex_coordinates()
30    g = domain.g
31
32    for k in range(domain.number_of_elements):
33#        avg_h = sum( h[k,:] )/3
34        avg_h = sum( h[k,:] )/2
35
36        #Compute bed slope
37        #x0, y0, x1, y1, x2, y2 = x[k,:]
38        x0, x1 = x[k,:]
39        #z0, z1, z2 = v[k,:]
40        b0, b1 = b[k,:]
41
42        w0, w1 = w[k,:]
43        wx = gradient(x0, x1, w0, w1)
44
45        #zx, zy = gradient(x0, y0, x1, y1, x2, y2, z0, z1, z2)
46        bx = gradient(x0, x1, b0, b1)
47
48        #Update momentum (explicit update is reset to source values)
49        xmom[k] += -g*bx*avg_h
50        #xmom[k] = -g*bx*avg_h
51        #stage[k] = 0.0
52
53
54def manning_friction(domain):
55    """Apply (Manning) friction to water momentum
56    """
57
58    from math import sqrt
59
60    w = domain.quantities['stage'].centroid_values
61    z = domain.quantities['elevation'].centroid_values
62    h = w-z
63
64    uh = domain.quantities['xmomentum'].centroid_values
65    #vh = domain.quantities['ymomentum'].centroid_values
66    eta = domain.quantities['friction'].centroid_values
67
68    xmom_update = domain.quantities['xmomentum'].semi_implicit_update
69    #ymom_update = domain.quantities['ymomentum'].semi_implicit_update
70
71    N = domain.number_of_elements
72    eps = domain.minimum_allowed_height
73    g = domain.g
74
75    for k in range(N):
76        if eta[k] >= eps:
77            if h[k] >= eps:
78                #S = -g * eta[k]**2 * sqrt((uh[k]**2 + vh[k]**2))
79                S = -g * eta[k]**2 * uh[k]
80                S /= h[k]**(7.0/3)
81
82                #Update momentum
83                xmom_update[k] += S*uh[k]
84                #ymom_update[k] += S*vh[k]
85
86def linear_friction(domain):
87    """Apply linear friction to water momentum
88
89    Assumes quantity: 'linear_friction' to be present
90    """
91
92    from math import sqrt
93
94    w = domain.quantities['stage'].centroid_values
95    z = domain.quantities['elevation'].centroid_values
96    h = w-z
97
98    uh = domain.quantities['xmomentum'].centroid_values
99#    vh = domain.quantities['ymomentum'].centroid_values
100    tau = domain.quantities['linear_friction'].centroid_values
101
102    xmom_update = domain.quantities['xmomentum'].semi_implicit_update
103#    ymom_update = domain.quantities['ymomentum'].semi_implicit_update
104
105    N = domain.number_of_elements
106    eps = domain.minimum_allowed_height
107    g = domain.g #Not necessary? Why was this added?
108
109    for k in range(N):
110        if tau[k] >= eps:
111            if h[k] >= eps:
112                S = -tau[k]/h[k]
113
114                #Update momentum
115                xmom_update[k] += S*uh[k]
116 #              ymom_update[k] += S*vh[k]
117
118
119
120def check_forcefield(f):
121    """Check that f is either
122    1: a callable object f(t,x,y), where x and y are vectors
123       and that it returns an array or a list of same length
124       as x and y
125    2: a scalar
126    """
127
128
129    if callable(f):
130        #N = 3
131        N = 2
132        #x = ones(3, numpy.float)
133        #y = ones(3, numpy.float)
134        x = ones(2, numpy.float)
135        #y = ones(2, numpy.float)
136
137        try:
138            #q = f(1.0, x=x, y=y)
139            q = f(1.0, x=x)
140        except Exception, e:
141            msg = 'Function %s could not be executed:\n%s' %(f, e)
142            #FIXME: Reconsider this semantics
143            raise msg
144
145        try:
146            q = numpy.array(q, numpy.float)
147        except:
148            msg = 'Return value from vector function %s could ' %f
149            msg += 'not be converted into a numpy array of numpy.floats.\n'
150            msg += 'Specified function should return either list or array.'
151            raise msg
152
153        #Is this really what we want?
154        msg = 'Return vector from function %s ' %f
155        msg += 'must have same lenght as input vectors'
156        assert len(q) == N, msg
157
158    else:
159        try:
160            f = float(f)
161        except:
162            msg = 'Force field %s must be either a scalar' %f
163            msg += ' or a vector function'
164            raise msg
165    return f
166
167
Note: See TracBrowser for help on using the repository browser.