# source:anuga_core/source/anuga/shallow_water/forcing.py@7733

Last change on this file since 7733 was 7733, checked in by hudson, 13 years ago

Fixed unit test failures.

File size: 28.8 KB
Line
1"""
2Environmental forcing functions, such as wind and rainfall.
3
4Constraints: See GPL license in the user guide
5Version: 1.0 (\$Revision: 7731 \$)
6ModifiedBy:
7    \$Author: hudson \$
8    \$Date: 2010-05-18 14:54:05 +1000 (Tue, 18 May 2010) \$
9"""
10
11
12from anuga.abstract_2d_finite_volumes.neighbour_mesh import segment_midpoints
13from anuga.utilities.numerical_tools import ensure_numeric
14from anuga.fit_interpolate.interpolate import Modeltime_too_early, \
15                                              Modeltime_too_late
16from anuga.geometry.polygon import is_inside_polygon, inside_polygon, \
17                                    polygon_area
18from types import IntType, FloatType
19from anuga.geospatial_data.geospatial_data import ensure_geospatial
20
21from warnings import warn
22import numpy as num
23
24
25
26def check_forcefield(f):
27    """Check that force object is as expected.
28
29    Check that f is either:
30    1: a callable object f(t,x,y), where x and y are vectors
31       and that it returns an array or a list of same length
32       as x and y
33    2: a scalar
34    """
35
36    if callable(f):
37        N = 3
38        x = num.ones(3, num.float)
39        y = num.ones(3, num.float)
40        try:
41            q = f(1.0, x=x, y=y)
42        except Exception, e:
43            msg = 'Function %s could not be executed:\n%s' %(f, e)
44            # FIXME: Reconsider this semantics
45            raise Exception, msg
46
47        try:
48            q = num.array(q, num.float)
49        except:
50            msg = ('Return value from vector function %s could not '
51                   'be converted into a numeric array of floats.\nSpecified '
52                   'function should return either list or array.' % f)
53            raise Exception, msg
54
55        # Is this really what we want?
56        # info is "(func name, filename, defining line)"
57        func_info = (f.func_name, f.func_code.co_filename,
58                     f.func_code.co_firstlineno)
59        func_msg = 'Function %s (defined in %s, line %d)' % func_info
60        try:
61            result_len = len(q)
62        except:
63            msg = '%s must return vector' % func_msg
64            self.fail(msg)
65        msg = '%s must return vector of length %d' % (func_msg, N)
66        assert result_len == N, msg
67    else:
68        try:
69            f = float(f)
70        except:
71            msg = ('Force field %s must be a scalar value coercible to float.'
72                   % str(f))
73            raise Exception, msg
74
75    return f
76
77
78
79class Wind_stress:
80    """Apply wind stress to water momentum in terms of
81    wind speed [m/s] and wind direction [degrees]
82    """
83    def __init__(self, *args, **kwargs):
84        """Initialise windfield from wind speed s [m/s]
85        and wind direction phi [degrees]
86
87        Inputs v and phi can be either scalars or Python functions, e.g.
88
89        W = Wind_stress(10, 178)
90
91        #FIXME - 'normal' degrees are assumed for now, i.e. the
92        vector (1,0) has zero degrees.
93        We may need to convert from 'compass' degrees later on and also
94        map from True north to grid north.
95
96        Arguments can also be Python functions of t,x,y as in
97
98        def speed(t,x,y):
99            ...
100            return s
101
102        def angle(t,x,y):
103            ...
104            return phi
105
106        where x and y are vectors.
107
108        and then pass the functions in
109
110        W = Wind_stress(speed, angle)
111
112        The instantiated object W can be appended to the list of
113        forcing_terms as in
114
115        Alternatively, one vector valued function for (speed, angle)
116        can be applied, providing both quantities simultaneously.
117        As in
118        W = Wind_stress(F), where returns (speed, angle) for each t.
119
120        domain.forcing_terms.append(W)
121        """
122
123        from anuga.config import rho_a, rho_w, eta_w
124
125        if len(args) == 2:
126            s = args[0]
127            phi = args[1]
128        elif len(args) == 1:
129            # Assume vector function returning (s, phi)(t,x,y)
130            vector_function = args[0]
131            s = lambda t,x,y: vector_function(t,x=x,y=y)[0]
132            phi = lambda t,x,y: vector_function(t,x=x,y=y)[1]
133        else:
134           # Assume info is in 2 keyword arguments
135           if len(kwargs) == 2:
136               s = kwargs['s']
137               phi = kwargs['phi']
138           else:
139               raise Exception, 'Assumes two keyword arguments: s=..., phi=....'
140
141        self.speed = check_forcefield(s)
142        self.phi = check_forcefield(phi)
143
144        self.const = eta_w*rho_a/rho_w
145
146    ##
147    # @brief 'execute' this class instance.
148    # @param domain
149    def __call__(self, domain):
150        """Evaluate windfield based on values found in domain"""
151
152        from math import pi, cos, sin, sqrt
153
154        xmom_update = domain.quantities['xmomentum'].explicit_update
155        ymom_update = domain.quantities['ymomentum'].explicit_update
156
157        N = len(domain)    # number_of_triangles
158        t = domain.time
159
160        if callable(self.speed):
161            xc = domain.get_centroid_coordinates()
162            s_vec = self.speed(t, xc[:,0], xc[:,1])
163        else:
164            # Assume s is a scalar
165            try:
166                s_vec = self.speed * num.ones(N, num.float)
167            except:
168                msg = 'Speed must be either callable or a scalar: %s' %self.s
169                raise msg
170
171        if callable(self.phi):
172            xc = domain.get_centroid_coordinates()
173            phi_vec = self.phi(t, xc[:,0], xc[:,1])
174        else:
175            # Assume phi is a scalar
176
177            try:
178                phi_vec = self.phi * num.ones(N, num.float)
179            except:
180                msg = 'Angle must be either callable or a scalar: %s' %self.phi
181                raise msg
182
183        assign_windfield_values(xmom_update, ymom_update,
184                                s_vec, phi_vec, self.const)
185
186
187##
188# @brief Assign wind field values
189# @param xmom_update
190# @param ymom_update
191# @param s_vec
192# @param phi_vec
193# @param const
194def assign_windfield_values(xmom_update, ymom_update,
195                            s_vec, phi_vec, const):
196    """Python version of assigning wind field to update vectors.
197    A C version also exists (for speed)
198    """
199
200    from math import pi, cos, sin, sqrt
201
202    N = len(s_vec)
203    for k in range(N):
204        s = s_vec[k]
205        phi = phi_vec[k]
206
208        phi = phi*pi/180
209
210        # Compute velocity vector (u, v)
211        u = s*cos(phi)
212        v = s*sin(phi)
213
214        # Compute wind stress
215        S = const * sqrt(u**2 + v**2)
216        xmom_update[k] += S*u
217        ymom_update[k] += S*v
218
219
220##
221# @brief A class for a general explicit forcing term.
222class General_forcing:
223    """General explicit forcing term for update of quantity
224
225    This is used by Inflow and Rainfall for instance
226
227
228    General_forcing(quantity_name, rate, center, radius, polygon)
229
230    domain:     ANUGA computational domain
231    quantity_name: Name of quantity to update.
232                   It must be a known conserved quantity.
233
234    rate [?/s]: Total rate of change over the specified area.
235                This parameter can be either a constant or a
236                function of time. Positive values indicate increases,
237                negative values indicate decreases.
238                Rate can be None at initialisation but must be specified
239                before forcing term is applied (i.e. simulation has started).
240
241    center [m]: Coordinates at center of flow point
242    radius [m]: Size of circular area
243    polygon:    Arbitrary polygon
244    default_rate: Rate to be used if rate fails (e.g. if model time exceeds its data)
245                  Admissible types: None, constant number or function of t
246
247
248    Either center, radius or polygon can be specified but not both.
249    If neither are specified the entire domain gets updated.
250    All coordinates to be specified in absolute UTM coordinates (x, y) assuming the zone of domain.
251
252    Inflow or Rainfall for examples of use
253    """
254
255
256    # FIXME (AnyOne) : Add various methods to allow spatial variations
257
258    ##
259    # @brief Create an instance of this forcing term.
260    # @param domain
261    # @param quantity_name
262    # @param rate
263    # @param center
265    # @param polygon
266    # @param default_rate
267    # @param verbose
268    def __init__(self,
269                 domain,
270                 quantity_name,
271                 rate=0.0,
272                 center=None,
274                 polygon=None,
275                 default_rate=None,
276                 verbose=False):
277
278        from math import pi, cos, sin
279
280        if center is None:
281            msg = 'I got radius but no center.'
282            assert radius is None, msg
283
285            msg += 'I got center but no radius.'
286            assert center is None, msg
287
288        self.domain = domain
289        self.quantity_name = quantity_name
290        self.rate = rate
291        self.center = ensure_numeric(center)
293        self.polygon = polygon
294        self.verbose = verbose
295        self.value = 0.0    # Can be used to remember value at
296                            # previous timestep in order to obtain rate
297
298        # Get boundary (in absolute coordinates)
299        bounding_polygon = domain.get_boundary_polygon()
300
301        # Update area if applicable
302        if center is not None and radius is not None:
303            assert len(center) == 2
304            msg = 'Polygon cannot be specified when center and radius are'
305            assert polygon is None, msg
306
307            # Check that circle center lies within the mesh.
308            msg = 'Center %s specified for forcing term did not' % str(center)
309            msg += 'fall within the domain boundary.'
310            assert is_inside_polygon(center, bounding_polygon), msg
311
312            # Check that circle periphery lies within the mesh.
313            N = 100
314            periphery_points = []
315            for i in range(N):
316                theta = 2*pi*i/100
317
318                x = center[0] + radius*cos(theta)
319                y = center[1] + radius*sin(theta)
320
321                periphery_points.append([x,y])
322
323            for point in periphery_points:
324                msg = 'Point %s on periphery for forcing term' % str(point)
325                msg += ' did not fall within the domain boundary.'
326                assert is_inside_polygon(point, bounding_polygon), msg
327
328        if polygon is not None:
329            # Check that polygon lies within the mesh.
330            for point in self.polygon:
331                msg = 'Point %s in polygon for forcing term' % str(point)
332                msg += ' did not fall within the domain boundary.'
333                assert is_inside_polygon(point, bounding_polygon), msg
334
335        # Pointer to update vector
336        self.update = domain.quantities[self.quantity_name].explicit_update
337
338        # Determine indices in flow area
339        N = len(domain)
340        points = domain.get_centroid_coordinates(absolute=True)
341
342        # Calculate indices in exchange area for this forcing term
343        self.exchange_indices = None
344        if self.center is not None and self.radius is not None:
345            # Inlet is circular
347
348            self.exchange_indices = []
349            for k in range(N):
350                x, y = points[k,:]    # Centroid
351
352                c = self.center
354                    self.exchange_indices.append(k)
355
356        if self.polygon is not None:
357            # Inlet is polygon
358            inlet_region = 'polygon=%s' % (self.polygon)
359            self.exchange_indices = inside_polygon(points, self.polygon)
360
361        if self.exchange_indices is None:
362            self.exchange_area = polygon_area(bounding_polygon)
363        else:
364            if len(self.exchange_indices) == 0:
365                msg = 'No triangles have been identified in '
366                msg += 'specified region: %s' % inlet_region
367                raise Exception, msg
368
369            # Compute exchange area as the sum of areas of triangles identified
370            # by circle or polygon
371            self.exchange_area = 0.0
372            for i in self.exchange_indices:
373                self.exchange_area += domain.areas[i]
374
375
376        msg = 'Exchange area in forcing term'
377        msg += ' has area = %f' %self.exchange_area
378        assert self.exchange_area > 0.0
379
380
381
382
383        # Check and store default_rate
384        msg = ('Keyword argument default_rate must be either None '
385               'or a function of time.\nI got %s.' % str(default_rate))
386        assert (default_rate is None or
387                type(default_rate) in [IntType, FloatType] or
388                callable(default_rate)), msg
389
390        if default_rate is not None:
391            # If it is a constant, make it a function
392            if not callable(default_rate):
393                tmp = default_rate
394                default_rate = lambda t: tmp
395
396            # Check that default_rate is a function of one argument
397            try:
398                default_rate(0.0)
399            except:
400                raise Exception, msg
401
402        self.default_rate = default_rate
403        self.default_rate_invoked = False    # Flag
404
405    ##
406    # @brief Execute this instance.
407    # @param domain
408    def __call__(self, domain):
409        """Apply inflow function at time specified in domain, update stage"""
410
411        # Call virtual method allowing local modifications
412        t = domain.get_time()
413        try:
414            rate = self.update_rate(t)
415        except Modeltime_too_early, e:
416            raise Modeltime_too_early, e
417        except Modeltime_too_late, e:
418            if self.default_rate is None:
419                msg = '%s: ANUGA is trying to run longer than specified data.\n' %str(e)
420                msg += 'You can specify keyword argument default_rate in the '
421                msg += 'forcing function to tell it what to do in the absence of time data.'
422                raise Modeltime_too_late, msg
423            else:
424                # Pass control to default rate function
425                rate = self.default_rate(t)
426
427                if self.default_rate_invoked is False:
428                    # Issue warning the first time
429                    msg = ('%s\n'
430                           'Instead I will use the default rate: %s\n'
431                           'Note: Further warnings will be supressed'
432                           % (str(e), str(self.default_rate)))
433                    warn(msg)
434
435                    # FIXME (Ole): Replace this crude flag with
436                    # Python's ability to print warnings only once.
437                    # See http://docs.python.org/lib/warning-filter.html
438                    self.default_rate_invoked = True
439
440        if rate is None:
441            msg = ('Attribute rate must be specified in General_forcing '
442                   'or its descendants before attempting to call it')
443            raise Exception, msg
444
445        # Now rate is a number
446        if self.verbose is True:
447            log.critical('Rate of %s at time = %.2f = %f'
448                         % (self.quantity_name, domain.get_time(), rate))
449
450        if self.exchange_indices is None:
451            self.update[:] += rate
452        else:
453            # Brute force assignment of restricted rate
454            for k in self.exchange_indices:
455                self.update[k] += rate
456
457    ##
458    # @brief Update the internal rate.
459    # @param t A callable or scalar used to set the rate.
460    # @return The new rate.
461    def update_rate(self, t):
462        """Virtual method allowing local modifications by writing an
463        overriding version in descendant
464        """
465
466        if callable(self.rate):
467            rate = self.rate(t)
468        else:
469            rate = self.rate
470
471        return rate
472
473    ##
474    # @brief Get values for the specified quantity.
475    # @param quantity_name Name of the quantity of interest.
476    # @return The value(s) of the quantity.
477    # @note If 'quantity_name' is None, use self.quantity_name.
478    def get_quantity_values(self, quantity_name=None):
479        """Return values for specified quantity restricted to opening
480
481        Optionally a quantity name can be specified if values from another
482        quantity is sought
483        """
484
485        if quantity_name is None:
486            quantity_name = self.quantity_name
487
488        q = self.domain.quantities[quantity_name]
489        return q.get_values(location='centroids',
490                            indices=self.exchange_indices)
491
492    ##
493    # @brief Set value for the specified quantity.
494    # @param val The value object used to set value.
495    # @param quantity_name Name of the quantity of interest.
496    # @note If 'quantity_name' is None, use self.quantity_name.
497    def set_quantity_values(self, val, quantity_name=None):
498        """Set values for specified quantity restricted to opening
499
500        Optionally a quantity name can be specified if values from another
501        quantity is sought
502        """
503
504        if quantity_name is None:
505            quantity_name = self.quantity_name
506
507        q = self.domain.quantities[self.quantity_name]
508        q.set_values(val,
509                     location='centroids',
510                     indices=self.exchange_indices)
511
512
513##
514# @brief A class for rainfall forcing function.
515# @note Inherits from General_forcing.
516class Rainfall(General_forcing):
517    """Class Rainfall - general 'rain over entire domain' forcing term.
518
519    Used for implementing Rainfall over the entire domain.
520
521        Current Limited to only One Gauge..
522
523        Need to add Spatial Varying Capability
524        (This module came from copying and amending the Inflow Code)
525
526    Rainfall(rain)
527
528    domain
529    rain [mm/s]:  Total rain rate over the specified domain.
530                  NOTE: Raingauge Data needs to reflect the time step.
531                  IE: if Gauge is mm read at a time step, then the input
532                  here is as mm/(timeStep) so 10mm in 5minutes becomes
533                  10/(5x60) = 0.0333mm/s.
534
535                  This parameter can be either a constant or a
536                  function of time. Positive values indicate inflow,
537                  negative values indicate outflow.
538                  (and be used for Infiltration - Write Seperate Module)
539                  The specified flow will be divided by the area of
540                  the inflow region and then applied to update the
541                  stage quantity.
542
543    polygon: Specifies a polygon to restrict the rainfall.
544
545    Examples
546    How to put them in a run File...
547
548    #------------------------------------------------------------------------
549    # Setup specialised forcing terms
550    #------------------------------------------------------------------------
551    # This is the new element implemented by Ole and Rudy to allow direct
552    # input of Rainfall in mm/s
553
554    catchmentrainfall = Rainfall(rain=file_function('Q100_2hr_Rain.tms'))
555                        # Note need path to File in String.
556                        # Else assumed in same directory
557
558    domain.forcing_terms.append(catchmentrainfall)
559    """
560
561    ##
562    # @brief Create an instance of the class.
563    # @param domain Domain of interest.
564    # @param rate Total rain rate over the specified domain (mm/s).
565    # @param center
567    # @param polygon Polygon  to restrict rainfall.
568    # @param default_rate
569    # @param verbose True if this instance is to be verbose.
570    def __init__(self,
571                 domain,
572                 rate=0.0,
573                 center=None,
575                 polygon=None,
576                 default_rate=None,
577                 verbose=False):
578
579        # Converting mm/s to m/s to apply in ANUGA)
580        if callable(rate):
581            rain = lambda t: rate(t)/1000.0
582        else:
583            rain = rate/1000.0
584
585        if default_rate is not None:
586            if callable(default_rate):
587                default_rain = lambda t: default_rate(t)/1000.0
588            else:
589                default_rain = default_rate/1000.0
590        else:
591            default_rain = None
592
593
594
595        General_forcing.__init__(self,
596                                 domain,
597                                 'stage',
598                                 rate=rain,
599                                 center=center,
601                                 polygon=polygon,
602                                 default_rate=default_rain,
603                                 verbose=verbose)
604
605
606##
607# @brief A class for inflow (rain and drain) forcing function.
608# @note Inherits from General_forcing.
609class Inflow(General_forcing):
610    """Class Inflow - general 'rain and drain' forcing term.
611
612    Useful for implementing flows in and out of the domain.
613
615
616    domain
617    rate [m^3/s]: Total flow rate over the specified area.
618                  This parameter can be either a constant or a
619                  function of time. Positive values indicate inflow,
620                  negative values indicate outflow.
621                  The specified flow will be divided by the area of
622                  the inflow region and then applied to update stage.
623    center [m]: Coordinates at center of flow point
624    radius [m]: Size of circular area
625    polygon:    Arbitrary polygon.
626
627    Either center, radius or polygon must be specified
628
629    Examples
630
631    # Constant drain at 0.003 m^3/s.
632    # The outflow area is 0.07**2*pi=0.0154 m^2
633    # This corresponds to a rate of change of 0.003/0.0154 = 0.2 m/s
634    #
635    Inflow((0.7, 0.4), 0.07, -0.003)
636
637
638    # Tap turning up to a maximum inflow of 0.0142 m^3/s.
639    # The inflow area is 0.03**2*pi = 0.00283 m^2
640    # This corresponds to a rate of change of 0.0142/0.00283 = 5 m/s
641    # over the specified area
642    Inflow((0.5, 0.5), 0.03, lambda t: min(0.01*t, 0.0142))
643
644
645    #------------------------------------------------------------------------
646    # Setup specialised forcing terms
647    #------------------------------------------------------------------------
648    # This is the new element implemented by Ole to allow direct input
649    # of Inflow in m^3/s
650
651    hydrograph = Inflow(center=(320, 300), radius=10,
652                        rate=file_function('Q/QPMF_Rot_Sub13.tms'))
653
654    domain.forcing_terms.append(hydrograph)
655    """
656
657    ##
658    # @brief Create an instance of the class.
659    # @param domain Domain of interest.
660    # @param rate Total rain rate over the specified domain (mm/s).
661    # @param center
663    # @param polygon Polygon  to restrict rainfall.
664    # @param default_rate
665    # @param verbose True if this instance is to be verbose.
666    def __init__(self,
667                 domain,
668                 rate=0.0,
669                 center=None,
671                 polygon=None,
672                 default_rate=None,
673                 verbose=False):
674        # Create object first to make area is available
675        General_forcing.__init__(self,
676                                 domain,
677                                 'stage',
678                                 rate=rate,
679                                 center=center,
681                                 polygon=polygon,
682                                 default_rate=default_rate,
683                                 verbose=verbose)
684
685    ##
686    # @brief Update the instance rate.
687    # @param t New rate object.
688    def update_rate(self, t):
689        """Virtual method allowing local modifications by writing an
690        overriding version in descendant
691
692        This one converts m^3/s to m/s which can be added directly
693        to 'stage' in ANUGA
694        """
695
696        if callable(self.rate):
697            _rate = self.rate(t)/self.exchange_area
698        else:
699            _rate = self.rate/self.exchange_area
700
701        return _rate
702
703
704##
705# @brief A class for creating cross sections.
706# @note Inherits from General_forcing.
707class Cross_section:
708    """Class Cross_section - a class to setup a cross section from
709    which you can then calculate flow and energy through cross section
710
711
712    Cross_section(domain, polyline)
713
714    domain:
715    polyline: Representation of desired cross section - it may contain
716              multiple sections allowing for complex shapes. Assume
717              absolute UTM coordinates.
718              Format [[x0, y0], [x1, y1], ...]
719    verbose:
720    """
721
722    ##
723    # @brief Create an instance of the class.
724    # @param domain Domain of interest.
725    # @param polyline Polyline defining cross section
726    # @param verbose True if this instance is to be verbose.
727    def __init__(self,
728                 domain,
729                 polyline=None,
730                 verbose=False):
731
732        self.domain = domain
733        self.polyline = polyline
734        self.verbose = verbose
735
736        # Find all intersections and associated triangles.
737        self.segments = self.domain.get_intersecting_segments(self.polyline,
738                                                              use_cache=True,
739                                                              verbose=self.verbose)
740
741        # Get midpoints
742        self.midpoints = segment_midpoints(self.segments)
743
744        # Make midpoints Geospatial instances
745        self.midpoints = ensure_geospatial(self.midpoints, self.domain.geo_reference)
746
747    ##
748    # @brief set verbose mode
749    def set_verbose(self,verbose=True):
750        """Set verbose mode true or flase
751        """
752
753        self.verbose=verbose
754
755    ##
756    # @brief calculate current flow through cross section
757    def get_flow_through_cross_section(self):
758        """ Output: Total flow [m^3/s] across cross section.
759        """
760
761        # Get interpolated values
762        xmomentum = self.domain.get_quantity('xmomentum')
763        ymomentum = self.domain.get_quantity('ymomentum')
764
765        uh = xmomentum.get_values(interpolation_points=self.midpoints,
766                                  use_cache=True)
767        vh = ymomentum.get_values(interpolation_points=self.midpoints,
768                                  use_cache=True)
769
770        # Compute and sum flows across each segment
771        total_flow = 0
772        for i in range(len(uh)):
773            # Inner product of momentum vector with segment normal [m^2/s]
774            normal = self.segments[i].normal
775            normal_momentum = uh[i]*normal[0] + vh[i]*normal[1]
776
777            # Flow across this segment [m^3/s]
778            segment_flow = normal_momentum*self.segments[i].length
779
780            # Accumulate
781            total_flow += segment_flow
782
783        return total_flow
784
785
786    ##
787    # @brief calculate current energy flow through cross section
788    def get_energy_through_cross_section(self, kind='total'):
789        """Obtain average energy head [m] across specified cross section.
790
791        Output:
792            E: Average energy [m] across given segments for all stored times.
793
794        The average velocity is computed for each triangle intersected by
795        the polyline and averaged weighted by segment lengths.
796
797        The typical usage of this function would be to get average energy of
798        flow in a channel, and the polyline would then be a cross section
799        perpendicular to the flow.
800
801        #FIXME (Ole) - need name for this energy reflecting that its dimension
802        is [m].
803        """
804
805        from anuga.config import g, epsilon, velocity_protection as h0
806
807        # Get interpolated values
808        stage = self.domain.get_quantity('stage')
809        elevation = self.domain.get_quantity('elevation')
810        xmomentum = self.domain.get_quantity('xmomentum')
811        ymomentum = self.domain.get_quantity('ymomentum')
812
813        w = stage.get_values(interpolation_points=self.midpoints, use_cache=True)
814        z = elevation.get_values(interpolation_points=self.midpoints, use_cache=True)
815        uh = xmomentum.get_values(interpolation_points=self.midpoints,
816                                  use_cache=True)
817        vh = ymomentum.get_values(interpolation_points=self.midpoints,
818                                  use_cache=True)
819        h = w-z                # Depth
820
821        # Compute total length of polyline for use with weighted averages
822        total_line_length = 0.0
823        for segment in self.segments:
824            total_line_length += segment.length
825
826        # Compute and sum flows across each segment
827        average_energy = 0.0
828        for i in range(len(w)):
829            # Average velocity across this segment
830            if h[i] > epsilon:
831                # Use protection against degenerate velocities
832                u = uh[i]/(h[i] + h0/h[i])
833                v = vh[i]/(h[i] + h0/h[i])
834            else:
835                u = v = 0.0
836
837            speed_squared = u*u + v*v
838            kinetic_energy = 0.5*speed_squared/g
839
840            if kind == 'specific':
841                segment_energy = h[i] + kinetic_energy
842            elif kind == 'total':
843                segment_energy = w[i] + kinetic_energy
844            else:
845                msg = 'Energy kind must be either "specific" or "total".'
846                msg += ' I got %s' %kind
847
848            # Add to weighted average
849            weigth = self.segments[i].length/total_line_length
850            average_energy += segment_energy*weigth
851
852        return average_energy
853
Note: See TracBrowser for help on using the repository browser.