Index: /trunk/anuga_core/source/anuga/operators/set_w_uh_vh_operators.py
===================================================================
--- /trunk/anuga_core/source/anuga/operators/set_w_uh_vh_operators.py	(revision 8619)
+++ /trunk/anuga_core/source/anuga/operators/set_w_uh_vh_operators.py	(revision 8619)
@@ -0,0 +1,242 @@
+"""
+Set value operators
+
+Constraints: See GPL license in the user guide
+Version: 1.0 ($Revision: 7731 $)
+"""
+
+__author__="steve"
+__date__ ="$09/03/2012 4:46:39 PM$"
+
+
+from anuga import Domain
+from anuga import Quantity
+import numpy as num
+import anuga.utilities.log as log
+
+from anuga.geometry.polygon import inside_polygon
+
+from anuga.operators.base_operator import Operator
+from anuga.fit_interpolate.interpolate import Modeltime_too_early, \
+                                              Modeltime_too_late
+from anuga import indent
+
+
+
+class Set_w_uh_vh_operator(Operator):
+    """
+    Set the w, uh and vh in a region
+
+    indices: None == all triangles, Empty list [] no triangles
+
+    rate can be a function of time.
+
+    """
+
+    def __init__(self,
+                 domain,
+                 w_uh_vh=None,
+                 indices=None,
+                 description = None,
+                 label = None,
+                 logging = False,
+                 verbose = False):
+
+
+        Operator.__init__(self, domain, description, label, logging, verbose)
+
+        #------------------------------------------
+        # Local variables
+        #------------------------------------------
+        self.w_uh_vh = w_uh_vh
+        self.indices = indices
+
+
+    def __call__(self):
+        """
+        Apply w_uh_vh to those triangles defined in indices
+
+        indices == [], then don't apply anywhere
+        indices == None, then apply everywhere
+        otherwise apply for the specific indices
+        """
+
+
+        if self.indices is []:
+            return
+
+        w_uh_vh = self.get_w_uh_vh()
+
+        if w_uh_vh is None:
+            return
+
+        if self.verbose is True:
+            log.critical('w_uh_vh of %s at time = %.2f = %f'
+                         % (self.quantity_name, domain.get_time(), stage))
+
+        if self.indices is None:
+            self.stage_c[:] = w_uh_vh[0]
+            self.xmom_c[:]  = w_uh_vh[1]
+            self.ymom_c[:]  = w_uh_vh[2]
+        else:
+            self.stage_c[self.indices] = w_uh_vh[0]
+            self.xmom_c[self.indices]  = w_uh_vh[1]
+            self.ymom_c[self.indices]  = w_uh_vh[2]
+
+
+    def get_w_uh_vh(self, t=None):
+        """Get value of w_uh_vh at time t.
+        If t not specified, return stage at current domain time
+        """
+
+        if t is None:
+            t = self.domain.get_time()
+
+        if callable(self.w_uh_vh):
+            try:
+                w_uh_vh = self.w_uh_vh(t)
+            except Modeltime_too_early, e:
+                raise Modeltime_too_early(e)
+            except Modeltime_too_late, e:
+                msg = '%s: ANUGA is trying to run longer than specified data.\n' %str(e)
+                msg += 'You can specify keyword argument default_rate in the '
+                msg += 'function to tell it what to do in the absence of time data.'
+                raise Modeltime_too_late(msg)
+        else:
+            w_uh_vh = self.w_uh_vh
+
+
+        if w_uh_vh  is None:
+            msg = ('Attribute w_uh_vh must be specified in '+self.__name__+
+                   ' before attempting to call it')
+            raise Exception(msg)
+
+        return w_uh_vh
+
+
+
+    def parallel_safe(self):
+        """Operator is applied independently on each cell and
+        so is parallel safe.
+        """
+        return True
+
+    def statistics(self):
+
+        message = self.label + ': Set_w_uh_vh_operator'
+        message = message + ' on triangles '+ str(self.indices)
+        return message
+
+
+    def timestepping_statistics(self):
+
+        message  = indent + self.label + ': Set_w_uh_vh = ' + str(self.get_w_uh_vh())
+        return message
+
+
+
+
+#===============================================================================
+# Specific Stage Operators for circular region.
+#===============================================================================
+class Circular_set_w_uh_vh_operator(Set_w_uh_vh_operator):
+    """
+    Set w_uh_vh over a circular region
+
+    """
+
+    def __init__(self, domain,
+                 w_uh_vh=[0.0, 0.0, 0.0],
+                 center=None,
+                 radius=None,
+                 verbose=False):
+
+        assert center is not None
+        assert radius is not None
+
+
+        # Determine indices in update region
+        N = domain.get_number_of_triangles()
+        points = domain.get_centroid_coordinates(absolute=True)
+
+
+        indices = []
+
+        c = center
+        r = radius
+
+        self.center = center
+        self.radius = radius
+
+        intersect = False
+        for k in range(N):
+            x, y = points[k,:]    # Centroid
+
+            if ((x-c[0])**2+(y-c[1])**2) < r**2:
+                intersect = True
+                indices.append(k)
+
+
+        msg = 'No centroids intersect circle center'+str(center)+' radius '+str(radius)
+        assert intersect, msg
+
+
+
+
+        # It is possible that circle doesn't intersect with mesh (as can happen
+        # for parallel runs
+
+
+        Set_w_uh_vh_operator.__init__(self,
+                                    domain,
+                                    w_uh_vh=w_uh_vh,
+                                    indices=indices,
+                                    verbose=verbose)
+
+
+
+
+
+#===============================================================================
+# Specific Stage Operators for polygonal region.
+#===============================================================================
+class Polygonal_set_w_uh_vh_operator(Set_w_uh_vh_operator):
+    """
+    Add water at certain rate (ms^{-1} = vol/Area/sec) over a
+    polygonal region
+
+    rate can be a function of time.
+
+    """
+
+    def __init__(self, domain,
+                 w_uh_vh=[0.0, 0,0, 0.0],
+                 polygon=None,
+                 verbose=False):
+
+
+        # Determine indices in update region
+        N = domain.get_number_of_triangles()
+        points = domain.get_centroid_coordinates(absolute=True)
+
+
+        indices = inside_polygon(points, polygon)
+        self.polygon = polygon
+
+        # It is possible that circle doesn't intersect with mesh (as can happen
+        # for parallel runs
+
+
+        Set_w_uh_vh_operator.__init__(self,
+                               domain,
+                               w_uh_vh=w_uh_vh,
+                               indices=indices,
+                               verbose=verbose)
+
+
+
+        
+
+
+
+
Index: /trunk/anuga_core/source/anuga/operators/test_set_w_uh_vh_operators.py
===================================================================
--- /trunk/anuga_core/source/anuga/operators/test_set_w_uh_vh_operators.py	(revision 8619)
+++ /trunk/anuga_core/source/anuga/operators/test_set_w_uh_vh_operators.py	(revision 8619)
@@ -0,0 +1,310 @@
+"""  Test set operators - w_uh_vh elevation erosion.
+"""
+
+import unittest, os
+import anuga
+from anuga import Domain
+from anuga import Reflective_boundary
+from anuga import rectangular_cross_domain
+from anuga import file_function
+
+from anuga.config import netcdf_mode_r, netcdf_mode_w, netcdf_mode_a
+from anuga.file_conversion.file_conversion import timefile2netcdf
+from anuga.config import time_format
+
+from set_w_uh_vh_operators import *
+
+import numpy as num
+import warnings
+import time
+
+
+
+class Test_set_w_uh_vh_operators(unittest.TestCase):
+    def setUp(self):
+        pass
+
+    def tearDown(self):
+        pass
+
+
+    def test_set_w_uh_vh_operator_simple(self):
+        from anuga.config import rho_a, rho_w, eta_w
+        from math import pi, cos, sin
+
+        a = [0.0, 0.0]
+        b = [0.0, 2.0]
+        c = [2.0, 0.0]
+        d = [0.0, 4.0]
+        e = [2.0, 2.0]
+        f = [4.0, 0.0]
+
+        points = [a, b, c, d, e, f]
+        #             bac,     bce,     ecf,     dbe
+        vertices = [[1,0,2], [1,2,4], [4,2,5], [3,1,4]]
+
+        domain = Domain(points, vertices)
+
+        #Flat surface with 1m of water
+        domain.set_quantity('elevation', 0)
+        domain.set_quantity('stage', 1.0)
+        domain.set_quantity('xmomentum', 7.0)
+        domain.set_quantity('ymomentum', 8.0)
+        domain.set_quantity('friction', 0)
+
+        Br = Reflective_boundary(domain)
+        domain.set_boundary({'exterior': Br})
+
+
+#        print domain.quantities['w_uh_vh'].centroid_values
+#        print domain.quantities['xmomentum'].centroid_values
+#        print domain.quantities['ymomentum'].centroid_values
+
+        # Apply operator to these triangles
+        indices = [0,1,3]
+
+        w_uh_vh = [3.0, 4.0, 5.0]
+
+
+        operator = Set_w_uh_vh_operator(domain, w_uh_vh=w_uh_vh, indices=indices)
+        
+        # Apply Operator
+        domain.timestep = 2.0
+        operator()
+
+        stage_ex = [ 3.,  3.,   1.,  3.]
+        xmom_ex = [ 4.,  4.,   7.,  4.]
+        ymom_ex = [ 5.,  5.,   8.,  5.]
+
+
+        #print domain.quantities['stage'].centroid_values
+        #print domain.quantities['xmomentum'].centroid_values
+        #print domain.quantities['ymomentum'].centroid_values
+
+        assert num.allclose(domain.quantities['stage'].centroid_values, stage_ex)
+        assert num.allclose(domain.quantities['xmomentum'].centroid_values, xmom_ex)
+        assert num.allclose(domain.quantities['ymomentum'].centroid_values, ymom_ex)
+
+ 
+    def xtest_set_w_uh_vh_operator_negative(self):
+        from anuga.config import rho_a, rho_w, eta_w
+        from math import pi, cos, sin
+
+        a = [0.0, 0.0]
+        b = [0.0, 2.0]
+        c = [2.0, 0.0]
+        d = [0.0, 4.0]
+        e = [2.0, 2.0]
+        f = [4.0, 0.0]
+
+        points = [a, b, c, d, e, f]
+        #             bac,     bce,     ecf,     dbe
+        vertices = [[1,0,2], [1,2,4], [4,2,5], [3,1,4]]
+
+        domain = Domain(points, vertices)
+
+        #Flat surface with 1m of water
+        domain.set_quantity('elevation', lambda x,y : -2*x)
+        domain.set_quantity('w_uh_vh', 1.0)
+        domain.set_quantity('friction', 0)
+
+        Br = Reflective_boundary(domain)
+        domain.set_boundary({'exterior': Br})
+
+#        print domain.quantities['elevation'].centroid_values
+#        print domain.quantities['w_uh_vh'].centroid_values
+#        print domain.quantities['xmomentum'].centroid_values
+#        print domain.quantities['ymomentum'].centroid_values
+
+        # Apply operator to these triangles
+        indices = [0,1,3]
+
+
+
+        #Catchment_Rain_Polygon = read_polygon(join('CatchmentBdy.csv'))
+        #rainfall = file_function(join('1y120m.tms'), quantities=['rainfall'])
+        w_uh_vh = -5.0
+
+
+        operator = Set_w_uh_vh_operator(domain, w_uh_vh=w_uh_vh, indices=indices)
+
+
+        # Apply Operator
+        domain.timestep = 2.0
+        operator()
+
+        w_uh_vh_ex = [ -5.,  -5.,   1.,  -5.]
+
+        #print domain.quantities['elevation'].centroid_values
+        #print domain.quantities['w_uh_vh'].centroid_values
+        #print domain.quantities['xmomentum'].centroid_values
+        #print domain.quantities['ymomentum'].centroid_values
+
+        assert num.allclose(domain.quantities['w_uh_vh'].centroid_values, w_uh_vh_ex)
+        assert num.allclose(domain.quantities['xmomentum'].centroid_values, 0.0)
+        assert num.allclose(domain.quantities['ymomentum'].centroid_values, 0.0)
+
+
+    def xtest_set_w_uh_vh_operator_function(self):
+        from anuga.config import rho_a, rho_w, eta_w
+        from math import pi, cos, sin
+
+        a = [0.0, 0.0]
+        b = [0.0, 2.0]
+        c = [2.0, 0.0]
+        d = [0.0, 4.0]
+        e = [2.0, 2.0]
+        f = [4.0, 0.0]
+
+
+
+
+
+        points = [a, b, c, d, e, f]
+        #             bac,     bce,     ecf,     dbe
+        vertices = [[1,0,2], [1,2,4], [4,2,5], [3,1,4]]
+
+        domain = Domain(points, vertices)
+
+        #Flat surface with 1m of water
+        domain.set_quantity('elevation', 0)
+        domain.set_quantity('w_uh_vh', 1.0)
+        domain.set_quantity('friction', 0)
+
+        Br = Reflective_boundary(domain)
+        domain.set_boundary({'exterior': Br})
+
+
+#        print domain.quantities['w_uh_vh'].centroid_values
+#        print domain.quantities['xmomentum'].centroid_values
+#        print domain.quantities['ymomentum'].centroid_values
+
+        # Apply operator to these triangles
+        indices = [0,1,3]
+
+
+        def w_uh_vh(t):
+            if t < 10.0:
+                return 5.0
+            else:
+                return 10.0
+
+        operator = Set_w_uh_vh_operator(domain, w_uh_vh=w_uh_vh, indices=indices)
+
+        # Apply Operator at time t=1.0
+        domain.set_time(1.0)
+        operator()
+
+        w_uh_vh_ex = [ 5.,  5.,   1.,  5.]
+
+
+#        print domain.quantities['w_uh_vh'].centroid_values
+#        print domain.quantities['xmomentum'].centroid_values
+#        print domain.quantities['ymomentum'].centroid_values
+
+        assert num.allclose(domain.quantities['w_uh_vh'].centroid_values, w_uh_vh_ex)
+        assert num.allclose(domain.quantities['xmomentum'].centroid_values, 0.0)
+        assert num.allclose(domain.quantities['ymomentum'].centroid_values, 0.0)
+
+        # Apply Operator at time t=15.0
+        domain.set_time(15.0)
+        operator()
+
+        w_uh_vh_ex = [ 10.,  10.,   1.,  10.]
+
+
+#        print domain.quantities['w_uh_vh'].centroid_values
+#        print domain.quantities['xmomentum'].centroid_values
+#        print domain.quantities['ymomentum'].centroid_values
+
+        assert num.allclose(domain.quantities['w_uh_vh'].centroid_values, w_uh_vh_ex)
+        assert num.allclose(domain.quantities['xmomentum'].centroid_values, 0.0)
+        assert num.allclose(domain.quantities['ymomentum'].centroid_values, 0.0)
+
+
+    def xtest_set_w_uh_vh_operator_large_function(self):
+        from anuga.config import rho_a, rho_w, eta_w
+        from math import pi, cos, sin
+
+
+        length = 2.0
+        width = 2.0
+        dx = dy = 0.5
+        domain = rectangular_cross_domain(int(length/dx), int(width/dy),
+                                              len1=length, len2=width)
+
+
+        #Flat surface with 1m of water
+        domain.set_quantity('elevation', 0.0)
+        domain.set_quantity('w_uh_vh', 1.0)
+        domain.set_quantity('friction', 0)
+
+        R = Reflective_boundary(domain)
+        domain.set_boundary( {'left': R, 'right': R, 'bottom': R, 'top': R} )
+
+
+#        print domain.quantities['w_uh_vh'].centroid_values
+#        print domain.quantities['xmomentum'].centroid_values
+#        print domain.quantities['ymomentum'].centroid_values
+
+
+
+        def w_uh_vh(t):
+            if t < 10.0:
+                return 5.0
+            else:
+                return 7.0
+
+        polygon = [(0.5,0.5), (1.5,0.5), (1.5,1.5), (0.5,1.5)]
+        operator = Polygonal_set_w_uh_vh_operator(domain, w_uh_vh=w_uh_vh, polygon=polygon)
+
+        # Apply Operator at time t=1.0
+        domain.set_time(1.0)
+        operator()
+
+        w_uh_vh_ex = [ 1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,
+                     1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,
+                     5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  5.0,  1.0,  1.0,
+                     1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  5.0,  5.0,  5.0,  5.0,
+                     5.0,  5.0,  5.0,  5.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,
+                     1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,
+                     1.0,  1.0,  1.0,  1.0]
+
+
+#        print domain.quantities['elevation'].centroid_values
+#        print domain.quantities['w_uh_vh'].centroid_values
+#        print domain.quantities['xmomentum'].centroid_values
+#        print domain.quantities['ymomentum'].centroid_values
+
+        assert num.allclose(domain.quantities['w_uh_vh'].centroid_values, w_uh_vh_ex)
+        assert num.allclose(domain.quantities['xmomentum'].centroid_values, 0.0)
+        assert num.allclose(domain.quantities['ymomentum'].centroid_values, 0.0)
+
+        # Apply Operator at time t=15.0
+        domain.set_time(15.0)
+        operator()
+
+        w_uh_vh_ex = [ 1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,
+                     1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,
+                     7.0,  7.0,  7.0,  7.0,  7.0,  7.0,  7.0,  7.0,  1.0,  1.0,
+                     1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  7.0,  7.0,  7.0,  7.0,
+                     7.0,  7.0,  7.0,  7.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,
+                     1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,  1.0,
+                     1.0,  1.0,  1.0,  1.0]
+
+
+#        print domain.quantities['w_uh_vh'].centroid_values
+#        print domain.quantities['xmomentum'].centroid_values
+#        print domain.quantities['ymomentum'].centroid_values
+
+        assert num.allclose(domain.quantities['w_uh_vh'].centroid_values, w_uh_vh_ex)
+        assert num.allclose(domain.quantities['xmomentum'].centroid_values, 0.0)
+        assert num.allclose(domain.quantities['ymomentum'].centroid_values, 0.0)
+
+
+
+
+if __name__ == "__main__":
+    suite = unittest.makeSuite(Test_set_w_uh_vh_operators, 'test')
+    runner = unittest.TextTestRunner(verbosity=1)
+    runner.run(suite)
