source: anuga_work/development/pyvolution-1d/shallow_water_adjust.py @ 4151

"""Class Domain -
1D interval domains for finite-volume computations of
the shallow water wave equation.

This module contains a specialisation of class Domain from module domain.py
consisting of methods specific to the Shallow Water Wave Equation


U_t + E_x = S

where

U = [w, uh]
E = [uh, u^2h + gh^2/2]
S represents source terms forcing the system
(e.g. gravity, friction, wind stress, ...)

and _t, _x, _y denote the derivative with respect to t, x and y respectiely.

The quantities are

symbol    variable name    explanation
x         x                horizontal distance from origin [m]
z         elevation        elevation of bed on which flow is modelled [m]
h         height           water height above z [m]
w         stage            absolute water level, w = z+h [m]
u                          speed in the x direction [m/s]
uh        xmomentum        momentum in the x direction [m^2/s]

eta                        mannings friction coefficient [to appear]
nu                         wind stress coefficient [to appear]

The conserved quantities are w, uh

For details see e.g.
Christopher Zoppou and Stephen Roberts,
Catastrophic Collapse of Water Supply Reservoirs in Urban Areas,
Journal of Hydraulic Engineering, vol. 127, No. 7 July 1999


John Jakeman, Ole Nielsen, Stephen Roberts, Duncan Gray, Christopher Zoppou
Geoscience Australia, 2006
"""

#from domain import *
from domain_adjust import *
Generic_Domain = Domain #Rename

#Shallow water domain
class Domain(Generic_Domain):

    def __init__(self, coordinates, boundary = None, tagged_elements = None,
                 geo_reference = None):

        conserved_quantities = ['stage', 'xmomentum','height']
        other_quantities = ['elevation', 'friction']#, 'height']
        Generic_Domain.__init__(self, coordinates, boundary,
                                conserved_quantities, other_quantities,
                                tagged_elements, geo_reference)
        
        from config import minimum_allowed_height, g
        self.minimum_allowed_height = minimum_allowed_height
        self.g = g

        #forcing terms not included in 1d domain ?WHy?
        self.forcing_terms.append(gravity)
        #self.forcing_terms.append(manning_friction)
        #print "\nI have Removed forcing terms line 64 1dsw"

        #Realtime visualisation
        self.visualiser = None
        self.visualise  = False
        self.visualise_color_stage = False
        self.visualise_stage_range = 1.0
        self.visualise_timer = True
        self.visualise_range_z = None
        
        #Stored output
        self.store = True
        self.format = 'sww'
        self.smooth = True

        #Evolve parametrs
        self.cfl = 1.0
        
        #Reduction operation for get_vertex_values
        from util import mean
        self.reduction = mean
        #self.reduction = min  #Looks better near steep slopes

        self.quantities_to_be_stored = ['stage','xmomentum']


    def set_quantities_to_be_stored(self, q):
        """Specify which quantities will be stored in the sww file.

        q must be either:
          - the name of a quantity
          - a list of quantity names
          - None

        In the two first cases, the named quantities will be stored at each
        yieldstep
        (This is in addition to the quantities elevation and friction)  
        If q is None, storage will be switched off altogether.
        """


        if q is None:
            self.quantities_to_be_stored = []            
            self.store = False
            return

        if isinstance(q, basestring):
            q = [q] # Turn argument into a list

        #Check correcness    
        for quantity_name in q:
            msg = 'Quantity %s is not a valid conserved quantity' %quantity_name
            assert quantity_name in self.conserved_quantities, msg 
        
        self.quantities_to_be_stored = q
        

    def initialise_visualiser(self,scale_z=1.0,rect=None):
        #Realtime visualisation
        if self.visualiser is None:
            from realtime_visualisation_new import Visualiser
            self.visualiser = Visualiser(self,scale_z,rect)
            self.visualiser.setup['elevation']=True
            self.visualiser.updating['stage']=True
        self.visualise = True
        if self.visualise_color_stage == True:
            self.visualiser.coloring['stage'] = True
            self.visualiser.qcolor['stage'] = (0.0, 0.0, 0.8)
        print 'initialise visualiser'
        print self.visualiser.setup
        print self.visualiser.updating

    def check_integrity(self):
        Generic_Domain.check_integrity(self)
        #Check that we are solving the shallow water wave equation

        msg = 'First conserved quantity must be "stage"'
        assert self.conserved_quantities[0] == 'stage', msg
        msg = 'Second conserved quantity must be "xmomentum"'
        assert self.conserved_quantities[1] == 'xmomentum', msg

    def extrapolate_second_order_sw(self):
        #Call correct module function
        #(either from this module or C-extension)
        extrapolate_second_order_sw(self)

    def compute_fluxes(self):
        #Call correct module function
        #(either from this module or C-extension)
        compute_fluxes(self)

    def compute_timestep(self):
        #Call correct module function
        compute_timestep(self)

    def distribute_to_vertices_and_edges(self):
        #Call correct module function
        #(either from this module or C-extension)
        distribute_to_vertices_and_edges(self)


    def distribute_stage_to_height(self):
        #Call correct module function
        #(either from this module or C-extension)
        distribute_stage_to_height(self)
        

    def classify_wet_dry_cells(self):
        #Call correct module function
        classify_wet_dry_cells(self)

    def set_initial_conserved_quanitites(self):
        #Call correct module function
        set_initial_conserved_quanitites(self)
    
        
    def adjust_partially_submerged_cells(self):
        #Call correct module function
        adjust_partially_submerged_cells(self)

    def pre_update(self):
        #Call correct module function
        pre_update(self)

    def post_update(self):
        #Call correct module function
        post_update(self)
        
    def evolve(self, yieldstep = None, finaltime = None,
               skip_initial_step = False):
        """Specialisation of basic evolve method from parent class
        """

        #Call check integrity here rather than from user scripts
        #self.check_integrity()

        #msg = 'Parameter beta_h must be in the interval [0, 1)'
        #assert 0 <= self.beta_h < 1.0, msg
        #msg = 'Parameter beta_w must be in the interval [0, 1)'
        #assert 0 <= self.beta_w < 1.0, msg


        #Initial update of vertex and edge values before any storage
        #and or visualisation

        print "will not work if stage not set at vertices"
        self.set_initial_conserved_quanitites()
        
        self.distribute_to_vertices_and_edges()

        self.distribute_stage_to_height()
        
        #Initialise real time viz if requested
        #if self.visualise is True and self.time == 0.0:
        #    if self.visualiser is None:
        #        self.initialise_visualiser()
        #
        #    self.visualiser.update_timer()
        #    self.visualiser.setup_all()

        #Store model data, e.g. for visualisation
        #if self.store is True and self.time == 0.0:
        #    self.initialise_storage()
        #    #print 'Storing results in ' + self.writer.filename
        #else:
        #    pass
        #    #print 'Results will not be stored.'
        #    #print 'To store results set domain.store = True'
        #    #FIXME: Diagnostic output should be controlled by
        #    # a 'verbose' flag living in domain (or in a parent class)

        #Call basic machinery from parent class
        for t in Generic_Domain.evolve(self, yieldstep, finaltime,
                                       skip_initial_step):
            #Real time viz
        #    if self.visualise is True:
        #        self.visualiser.update_all()
        #        self.visualiser.update_timer()


            #Store model data, e.g. for subsequent visualisation
        #    if self.store is True:
        #        self.store_timestep(self.quantities_to_be_stored)

            #FIXME: Could maybe be taken from specified list
            #of 'store every step' quantities

            #Pass control on to outer loop for more specific actions
            yield(t)

    def initialise_storage(self):
        """Create and initialise self.writer object for storing data.
        Also, save x and bed elevation
        """

        import data_manager

        #Initialise writer
        self.writer = data_manager.get_dataobject(self, mode = 'w')

        #Store vertices and connectivity
        self.writer.store_connectivity()


    def store_timestep(self, name):
        """Store named quantity and time.

        Precondition:
           self.write has been initialised
        """
        self.writer.store_timestep(name)


#=============== End of Shallow Water Domain ===============================

#Rotation of momentum vector
def rotate(q, normal, direction = 1):
    """Rotate the momentum component q (q[1], q[2])
    from x,y coordinates to coordinates based on normal vector.

    If direction is negative the rotation is inverted.

    Input vector is preserved

    This function is specific to the shallow water wave equation
    """

    from Numeric import zeros, Float

    assert len(q) == 3,\
           'Vector of conserved quantities must have length 3'\
           'for 2D shallow water equation'

    try:
        l = len(normal)
    except:
        raise 'Normal vector must be an Numeric array'

    assert l == 2, 'Normal vector must have 2 components'


    n1 = normal[0]
    n2 = normal[1]

    r = zeros(len(q), Float) #Rotated quantities
    r[0] = q[0]              #First quantity, height, is not rotated

    if direction == -1:
        n2 = -n2


    r[1] =  n1*q[1] + n2*q[2]
    r[2] = -n2*q[1] + n1*q[2]

    return r

def flux_function(normal, ql, qr, zl, zr):
    """Compute fluxes between volumes for the shallow water wave equation
    cast in terms of w = h+z using the 'central scheme' as described in

    Kurganov, Noelle, Petrova. 'Semidiscrete Central-Upwind Schemes For
    Hyperbolic Conservation Laws and Hamilton-Jacobi Equations'.
    Siam J. Sci. Comput. Vol. 23, No. 3, pp. 707-740.

    The implemented formula is given in equation (3.15) on page 714

    Conserved quantities w, uh, are stored as elements 0 and 1
    in the numerical vectors ql an qr.

    Bed elevations zl and zr.
    """

    from config import g, epsilon
    from math import sqrt
    from Numeric import array

    #print 'ql',ql
    #print 'qr', qr

    #Align momentums with x-axis
    #q_left  = rotate(ql, normal, direction = 1)
    #q_right = rotate(qr, normal, direction = 1)
    q_left = ql
    q_left[1] = q_left[1]*normal
    q_right = qr
    q_right[1] = q_right[1]*normal

    #z = (zl+zr)/2 #Take average of field values
    z = 0.5*(zl+zr) #Take average of field values

    h_left  = q_left[0]   #w=h+z
    w_left  = h_left+z
    uh_left = q_left[1]

    if h_left < epsilon:
        u_left = 0.0  #Could have been negative
        h_left = 0.0
    else:
        u_left  = uh_left/h_left


    h_right  = q_right[0]  #w=h+z
    w_right  = h_right+z
    uh_right = q_right[1]


    if h_right < epsilon:
        u_right = 0.0  #Could have been negative
        h_right = 0.0
    else:
        u_right  = uh_right/h_right

    #vh_left  = q_left[2]
    #vh_right = q_right[2]

    #print "uright %f ulef %f"%(u_right,u_left)

    soundspeed_left  = sqrt(g*h_left)
    soundspeed_right = sqrt(g*h_right)

    #Maximal wave speed
    s_max = max(u_left+soundspeed_left, u_right+soundspeed_right, 0)
    
    #Minimal wave speed
    s_min = min(u_left-soundspeed_left, u_right-soundspeed_right, 0)
    
    #Flux computation

    #flux_left  = array([u_left*h_left,
    #                    u_left*uh_left + 0.5*g*h_left**2])
    #flux_right = array([u_right*h_right,
    #                    u_right*uh_right + 0.5*g*h_right**2])
    flux_left  = array([u_left*h_left,
                        u_left*uh_left + 0.5*g*h_left*h_left])
    flux_right = array([u_right*h_right,
                        u_right*uh_right + 0.5*g*h_right*h_right])

    denom = s_max-s_min
    if denom == 0.0:
        edgeflux = array([0.0, 0.0])
        max_speed = 0.0
    else:
        edgeflux = (s_max*flux_left - s_min*flux_right)/denom
        edgeflux += s_max*s_min*(q_right-q_left)/denom
        
        edgeflux[1] = edgeflux[1]*normal

        max_speed = max(abs(s_max), abs(s_min))

    return edgeflux, max_speed


def flux_function_split(normal, ql, qr, zl, zr):
    from config import g, epsilon
    from math import sqrt
    from Numeric import array

    #print 'ql',ql

    #Align momentums with x-axis
    #q_left  = rotate(ql, normal, direction = 1)
    #q_right = rotate(qr, normal, direction = 1)
    q_left = ql
    q_left[1] = q_left[1]*normal
    q_right = qr
    q_right[1] = q_right[1]*normal

    #z = (zl+zr)/2 #Take average of field values
    z = 0.5*(zl+zr) #Take average of field values

    w_left  = q_left[0]   #w=h+z
    h_left  = w_left-z
    uh_left = q_left[1]

    if h_left < epsilon:
        u_left = 0.0  #Could have been negative
        h_left = 0.0
    else:
        u_left  = uh_left/h_left


    w_right  = q_right[0]  #w=h+z
    h_right  = w_right-z
    uh_right = q_right[1]


    if h_right < epsilon:
        u_right = 0.0  #Could have been negative
        h_right = 0.0
    else:
        u_right  = uh_right/h_right

    #vh_left  = q_left[2]
    #vh_right = q_right[2]

    #soundspeed_left  = sqrt(g*h_left)
    #soundspeed_right = sqrt(g*h_right)

    #Maximal wave speed
    #s_max = max(u_left+soundspeed_left, u_right+soundspeed_right, 0)
    s_max = max(u_left, u_right, 0)
    
    #Minimal wave speed
    #s_min = min(u_left-soundspeed_left, u_right-soundspeed_right, 0)
    s_min = min(u_left, u_right, 0)
    
    #Flux computation

    #flux_left  = array([u_left*h_left,
    #                    u_left*uh_left + 0.5*g*h_left*h_left])
    #flux_right = array([u_right*h_right,
    #                    u_right*uh_right + 0.5*g*h_right*h_right])
    flux_left  = array([u_left*h_left,
                        u_left*uh_left])# + 0.5*g*h_left*h_left])
    flux_right = array([u_right*h_right,
                        u_right*uh_right])# + 0.5*g*h_right*h_right])
    
    denom = s_max-s_min
    if denom == 0.0:
        edgeflux = array([0.0, 0.0])
        max_speed = 0.0
    else:
        edgeflux = (s_max*flux_left - s_min*flux_right)/denom
        edgeflux += s_max*s_min*(q_right-q_left)/denom
        
        edgeflux[1] = edgeflux[1]*normal

        max_speed = max(abs(s_max), abs(s_min))

    return edgeflux, max_speed    

def compute_timestep(domain):
    import sys
    from Numeric import zeros, Float

    N = domain.number_of_elements

    #Shortcuts
    Stage = domain.quantities['stage']
    Xmom = domain.quantities['xmomentum']
    Bed = domain.quantities['elevation']
    
    stage = Stage.vertex_values
    xmom =  Xmom.vertex_values
    bed =   Bed.vertex_values

    stage_bdry = Stage.boundary_values
    xmom_bdry =  Xmom.boundary_values

    flux = zeros(2, Float) #Work array for summing up fluxes
    ql = zeros(2, Float)
    qr = zeros(2, Float)

    #Loop
    timestep = float(sys.maxint)
    enter = True
    for k in range(N):

        flux[:] = 0.  #Reset work array
        for i in range(2):
            #Quantities inside volume facing neighbour i
            ql = [stage[k, i], xmom[k, i]]
            zl = bed[k, i]

            #Quantities at neighbour on nearest face
            n = domain.neighbours[k,i]
            if n < 0:
                m = -n-1 #Convert negative flag to index
                qr[0] = stage_bdry[m]
                qr[1] = xmom_bdry[m]
                zr = zl #Extend bed elevation to boundary
            else:
                #m = domain.neighbour_edges[k,i]
                m = domain.neighbour_vertices[k,i]
                qr[0] = stage[n, m]
                qr[1] =  xmom[n, m]
                zr = bed[n, m]


            #Outward pointing normal vector
            normal = domain.normals[k, i]

            if domain.split == False:
                edgeflux, max_speed = flux_function(normal, ql, qr, zl, zr)
            elif domain.split == True:
                edgeflux, max_speed = flux_function_split(normal, ql, qr, zl, zr)
            #Update optimal_timestep
            try:
                timestep = min(timestep, domain.cfl*0.5*domain.areas[k]/max_speed)
            except ZeroDivisionError:
                pass

    domain.timestep = timestep

def compute_fluxes(domain):
    """Compute all fluxes and the timestep suitable for all volumes
    in domain.

    Compute total flux for each conserved quantity using "flux_function"

    Fluxes across each edge are scaled by edgelengths and summed up
    Resulting flux is then scaled by area and stored in
    explicit_update for each of the three conserved quantities
    stage, xmomentum and ymomentum

    The maximal allowable speed computed by the flux_function for each volume
    is converted to a timestep that must not be exceeded. The minimum of
    those is computed as the next overall timestep.

    Post conditions:
      domain.explicit_update is reset to computed flux values
      domain.timestep is set to the largest step satisfying all volumes.
    """

    #print "in compute_fluxes"

    import sys
    from Numeric import zeros, Float

    N = domain.number_of_elements

    #Shortcuts
    #Stage = domain.quantities['stage']
    Height = domain.quantities['height']
    Xmom = domain.quantities['xmomentum']
#    Ymom = domain.quantities['ymomentum']
    Bed = domain.quantities['elevation']

    #Arrays
    #stage = Stage.edge_values
    #xmom =  Xmom.edge_values
 #   ymom =  Ymom.edge_values
    #bed =   Bed.edge_values
    
    #stage = Stage.vertex_values
    height = Height.vertex_values
    xmom =  Xmom.vertex_values
    bed =   Bed.vertex_values
    
    #print 'stage edge values', stage
    #print 'xmom edge values', xmom
    #print 'bed values', bed

    #stage_bdry = Stage.boundary_values
    height_bdry = Height.boundary_values
    xmom_bdry =  Xmom.boundary_values
    #print 'stage_bdry',stage_bdry
    #print 'xmom_bdry', xmom_bdry
#    ymom_bdry =  Ymom.boundary_values

#    flux = zeros(3, Float) #Work array for summing up fluxes
    flux = zeros(2, Float) #Work array for summing up fluxes
    ql = zeros(2, Float)
    qr = zeros(2, Float)

    #Loop
    timestep = float(sys.maxint)
    enter = True
    for k in range(N):

        flux[:] = 0.  #Reset work array
        #for i in range(3):
        for i in range(2):
            #Quantities inside volume facing neighbour i
            #ql = [stage[k, i], xmom[k, i]]
            ql = [height[k, i], xmom[k, i]]
            zl = bed[k, i]
            #print "cell",k,"left state",ql[0]

            #Quantities at neighbour on nearest face
            n = domain.neighbours[k,i]
            if n < 0:
                m = -n-1 #Convert negative flag to index
                #qr[0] = stage_bdry[m]
                qr[0] = height_bdry[m]
                qr[1] = xmom_bdry[m]
                zr = zl #Extend bed elevation to boundary
            else:
                #m = domain.neighbour_edges[k,i]
                m = domain.neighbour_vertices[k,i]
                #qr = [stage[n, m], xmom[n, m], ymom[n, m]]
                #qr[0] = stage[n, m]
                qr[0] = height[n, m]
                qr[1] =  xmom[n, m]
                zr = bed[n, m]
                #print "cell",k,"right state",qr[0]


            #Outward pointing normal vector
            normal = domain.normals[k, i]
        
            #Flux computation using provided function
            #edgeflux, max_speed = flux_function(normal, ql, qr, zl, zr)
            #print 'ql',ql
            #print 'qr',qr
            

            if domain.split == False:
                edgeflux, max_speed = flux_function(normal, ql, qr, zl, zr)
                #print edgeflux
            elif domain.split == True:
                edgeflux, max_speed = flux_function_split(normal, ql, qr, zl, zr)
            #print 'edgeflux', edgeflux

            # THIS IS THE LINE TO DEAL WITH LEFT AND RIGHT FLUXES
            # flux = edgefluxleft - edgefluxright
            flux -= edgeflux #* domain.edgelengths[k,i]

            #Update optimal_timestep
            try:
                #timestep = min(timestep, 0.5*domain.radii[k]/max_speed)
                timestep = min(timestep, domain.cfl*0.5*domain.areas[k]/max_speed)
                #print "cell %d wet length %f"%(k,domain.areas[k])
                #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])
            except ZeroDivisionError:
                pass

        #Normalise by area and store for when all conserved
        #quantities get updated
        flux /= domain.areas[k]

        #try:
        #    #if flux[0] > 1.0e-6: unecessary because only min is taken
        #    if Height.centroid_values[k] > 1.0e-3:
        #        timestep = min(timestep, domain.cfl*0.5*Height.centroid_values[k]*domain.areas[k]/abs(flux[0]))
        #except ZeroDivisionError:
        #    pass



        #Stage.explicit_update[k] = flux[0]
        Height.explicit_update[k] = flux[0]
        Xmom.explicit_update[k] = flux[1]
        #Ymom.explicit_update[k] = flux[2]


    domain.timestep = timestep
    #print "flux"
    #print Height.explicit_update

    #print "integral %f"%domain.quantities['height'].get_integral()

####################################

# Module functions for gradient limiting (distribute_to_vertices_and_edges)

def distribute_to_vertices_and_edges(domain):
    """Distribution from centroids to vertices specific to the
    shallow water wave
    equation.

    It will ensure that h (w-z) is always non-negative even in the
    presence of steep bed-slopes by taking a weighted average between shallow
    and deep cases.

    In addition, all conserved quantities get distributed as per either a
    constant (order==1) or a piecewise linear function (order==2).

    FIXME: more explanation about removal of artificial variability etc

    Precondition:
      All quantities defined at centroids and bed elevation defined at
      vertices.

    Postcondition
      Conserved quantities defined at vertices

    """

    #print "in distribute to vertices and edges"
    #print "integral %f"%domain.quantities['height'].get_integral()
    
    #from config import optimised_gradient_limiter

    #Remove very thin layers of water
    #protect_against_infinitesimal_and_negative_heights(domain)
    
    #classify_wet_dry_cells(domain)
    #import copy
    #print domain.wet_nodes
    #StageQ = copy.copy(domain.quantities['stage'].centroid_values)
    #adjust_partially_submerged_cells(domain)
    #StageQ2 = domain.quantities['stage'].centroid_values
    #print StageQ-StageQ2
    
    #Extrapolate all conserved quantities
    #if optimised_gradient_limiter:
    #    #MH090605 if second order,
    #    #perform the extrapolation and limiting on
    #    #all of the conserved quantities

    #    if (domain.order == 1):
    #        for name in domain.conserved_quantities:
    #            Q = domain.quantities[name]
    #            Q.extrapolate_first_order()
    #    elif domain.order == 2:
    #        domain.extrapolate_second_order_sw()
    #    else:
    #        raise 'Unknown order'
    #else:
        #old code:

    for name in domain.conserved_quantities:
        Q = domain.quantities[name]
        if domain.order == 1:
            Q.extrapolate_first_order()
        elif domain.order == 2:
            #print "add extrapolate second order to shallow water"
            if (name != 'height'):#& (name != "xmomentum"):
                Q.extrapolate_second_order()
            #Q.limit()
        else:
            raise 'Unknown order'

    #Take bed elevation into account when water heights are small
    #balance_deep_and_shallow(domain)
    protect_dry_beds(domain)
    
    #Compute edge values by interpolation
    #for name in domain.conserved_quantities:
    #    Q = domain.quantities[name]
    #    Q.interpolate_from_vertices_to_edges()

    #print domain.quantities['stage'].centroid_values
    
def protect_dry_beds(domain):
    #print "in protect dry beds"
    #print "integral %f"%domain.quantities['height'].get_integral()
    
    N = domain.number_of_elements
    wc = domain.quantities['stage'].centroid_values
    zc = domain.quantities['elevation'].centroid_values
    hc = domain.quantities['height'].centroid_values
    wv = domain.quantities['stage'].vertex_values
    zv = domain.quantities['elevation'].vertex_values
    xmomv = domain.quantities['xmomentum'].vertex_values
    xmomc = domain.quantities['xmomentum'].centroid_values
    min_centroid_height = 1e-3

    for k in range(N):
        if hc[k] < domain.minimum_allowed_height:
            if hc[k] < domain.epsilon:
                wc[k] = zc[k]
                wv[k,0] = zv[k,0]
                wv[k,1] = zv[k,1]
                xmomc[k] = 0.0
                xmomv[k,:] = 0.0
        

def set_initial_conserved_quanitites(domain):
    from util import calculate_new_wet_area, calculate_wetted_area
    from Numeric import allclose
    
    N = domain.number_of_elements
    wv = domain.quantities['stage'].vertex_values
    zv = domain.quantities['elevation'].vertex_values
    xv = domain.vertices
    zc = domain.quantities['elevation'].centroid_values
    wc = domain.quantities['stage'].centroid_values
    hc = domain.quantities['height'].centroid_values
    xmomc = domain.quantities['xmomentum'].centroid_values
    
    #A1 = domain.quantities['stage'].get_integral()
    for k in range(N):
        if hc[k] < 0.0:
            print "ERROR!!!!!!!!!!!!!!!!!!"
        #if (domain.wet_nodes[k,0] == 0) | (domain.wet_nodes[k,1] == 0):
        if (wv[k,0] <  zv[k,0]) | (wv[k,1] <  zv[k,1]):
            # cell is dry or wetted
            #print "cell %d is wetted or dry"%k
            w1 = wv[k,0]
            w2 = wv[k,1]
            z1 = zv[k,0]
            z2 = zv[k,1]
            x1 = xv[k,0]
            x2 = xv[k,1]
            #print w1,w2
            A = calculate_wetted_area(x1,x2,z1,z2,w1,w2)
            L = x2-x1
            w_centroid, wet_len = calculate_new_wet_area(x1,x2,z1,z2,A)
            #print "cell",k,"stage",w1,w2
            #print A
            if A > 0.0:
                wc[k] = w_centroid # surface is flat in wetted bed
                #wv[k,0] = w1
                #wv[k,1] = w2
                domain.wet_nodes[k,0] = 2  # stops limiter from limiting these cells
                domain.wet_nodes[k,1] = 2
                hc[k] = A/L
               # print "hc %f" %hc[k]
                #xmomc[k] = 0.0
                #print "cell",k,"stage",wc[k],hc[k],wet_len
                #if wet_len > 0.0:
                    #domain.areas[k] = wet_len
            else:
                #print "cell %d is not wetted but dry"%k
                wc[k] = 0.5*(w1+w2) # bed is completely dry
                wv[k,0] = zv[k,0]
                wv[k,1] = zv[k,1]
                domain.wet_nodes[k,0] = 0
                domain.wet_nodes[k,1] = 0
                hc[k] = 0.0
        else:
            #print "cell",k,"is wet"
            hc[k] = wc[k]-zc[k]
            
    #A2 = domain.quantities['stage'].get_integral()
    #assert allclose(A1,A2)
    print "Must set initial conditions for momentum as well in set_initial_conserved_quanitites and adjust"
    #print hc


def adjust_partially_submerged_cells(domain):
    #normal#
    #print "in adjust wetted_cells"
    #print "integral %f"%domain.quantities['height'].get_integral()

    from util import calculate_new_wet_area, analytic_cannal
    min_centroid_height = 1.0e-3
    N = domain.number_of_elements
    wv = domain.quantities['stage'].vertex_values
    zv = domain.quantities['elevation'].vertex_values
    xv = domain.vertices
    xc = domain.centroids
    zc = domain.quantities['elevation'].centroid_values
    wc = domain.quantities['stage'].centroid_values
    hc = domain.quantities['height'].centroid_values
    xmomc = domain.quantities['xmomentum'].centroid_values
    xmomv = domain.quantities['xmomentum'].vertex_values
    
    domain.wet_nodes[:,:] = 0.0

    #print hc

    for k in range(N):
        domain.areas[k] = xv[k,1]-xv[k,0]
        
        if hc[k] <= domain.epsilon:
            if hc[k] < 0.0:
                if hc[k+1] > domain.epsilon:
                  hc[k+1] += hc[k]
                else: #hc[k-1] > domain.epsilon:
                    hc[k-1] += hc[k]
                    
            hc[k] = 0.0 # if this not here get negative heights means timesteping wrong
            wc[k] = zc[k]
            xmomc[k] = 0.0
            xmomv[k,0] = 0.0
            xmomv[k,1] = 0.0
        else: #hc[k] > domain.epsilon: #min_centroid_height:
            z1 = zv[k,0]
            z2 = zv[k,1]
            x1 = xv[k,0]
            x2 = xv[k,1]
            L = x2-x1
            A = hc[k]*L
            if A > 0.0:
                w_centroid, wet_len = calculate_new_wet_area(x1,x2,z1,z2,A)
                if (w_centroid > max(z1,z2)):
                    wc[k] = w_centroid
                    #xmomc[k] = 0.5*(uh1+uh2)
                    #xmomv[k,0] = uh1
                    #xmomv[k,1] = uh2 # sets xmom to be analytic everwhere
                    
                elif (w_centroid > min(z1,z2)):
                    wc[k] = w_centroid # surface is flat in wetted bed
                    #wv[k,0] = w1
                    #wv[k,1] = w2 # this is done by slope limiter
                    domain.wet_nodes[k,0] = 2  # stops limiter from limiting these cells
                    domain.wet_nodes[k,1] = 2
                    #print "cell",k, hc[k], A,L, w1,w2
                    wcrap,uhc = analytic_cannal(xc[k],domain.time+domain.timestep)
                    w1,uh1 = analytic_cannal(x1,domain.time+domain.timestep)
                    w2,uh2 = analytic_cannal(x2,domain.time+domain.timestep)
            
                    if (w_centroid > z1):
                        #w1,uh1 = analytic_cannal(xc[k-1],domain.time+domain.timestep)
                        #xmomv[k,0] = uh1
                        xmomv[k,0] = 0.0
                        xmomv[k,1] = 0.0
                        #xmomc[k] = uh1
                        xmomc[k] = 0.0
                        
                    else:
                        #w2,uh2 = analytic_cannal(xc[k+1],domain.time+domain.timestep)
                        xmomv[k,0] = 0.0
                        xmomv[k,1] = 0.0
                        #xmomv[k,1] = uh2
                        #xmomc[k] = uh2
                        xmomc[k] = 0.0
                else:
                    wc[k] = zc[k] # bed is completely dry
                    domain.wet_nodes[k,0] = 0
                    domain.wet_nodes[k,1] = 0
                    hc[k] = 0.0
                    xmomc[k] = 0.0
                    xmomv[k,0] = 0.0
                    xmomv[k,1] = 0.0
                #if wet_len > 0.0:
                    #domain.areas[k] = wet_len
                    # if wet_len is 0 then bed is flat and area = x2-x1 (as default)
            else:
                wc[k] = zc[k] # bed is completely dry
                domain.wet_nodes[k,0] = 0
                domain.wet_nodes[k,1] = 0
                hc[k] = 0.0
                xmomc[k] = 0.0
                xmomv[k,0] = 0.0
                xmomv[k,1] = 0.0
                
    #for k in range(N):
    #    
    #    wcrap,uhc = analytic_cannal(xc[int(N/2)],domain.time+domain.timestep)
    #    if (k > 1) & (k<N-2):
    #        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):
    #            #print "TRUE"
    #            xmomc[k] = uhc
    #        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):
    #            #print "FALSE"
    #            xmomc[k] = uhc
                
        
    #print "integral %f"%domain.quantities['height'].get_integral()
    #print domain.quantities['height'].centroid_values
    
def adjust_partially_submerged_cells_analytic(domain):
    #analytic#

    #print "in adjust wetted_cells"
    from util import calculate_new_wet_area_analytic
    min_centroid_height = 1.0e-3
    N = domain.number_of_elements
    wv = domain.quantities['stage'].vertex_values
    zv = domain.quantities['elevation'].vertex_values
    xv = domain.vertices
    zc = domain.quantities['elevation'].centroid_values
    wc = domain.quantities['stage'].centroid_values
    hc = domain.quantities['height'].centroid_values
    xmomc = domain.quantities['xmomentum'].centroid_values
    xmomv = domain.quantities['xmomentum'].vertex_values

    domain.wet_nodes[:,:] = 0.0

    #print hc

    for k in range(N):
        domain.areas[k] = xv[k,1]-xv[k,0]
        if hc[k] > domain.epsilon:#min_centroid_height:
            z1 = zv[k,0]
            z2 = zv[k,1]
            x1 = xv[k,0]
            x2 = xv[k,1]
            L = x2-x1
            A = hc[k]*L
            if A > 0.0:
                #print "cell %d is wetted" %k
                #print "time: %f -- is 1 sec behind real time in adjust. i.e. time has not been updated"%domain.time
                w1,w2, wet_len,uh1,uh2 = calculate_new_wet_area_analytic(x1,x2,z1,z2,A,domain.time+1.0)
                xmomc[k] = 0.5*(uh1+uh2)
                xmomv[k,0] = uh1
                xmomv[k,1] = uh2
                if (w1 > z1) & (w2 > z2):
                    #wc[k] = (w1+w2)*0.5
                    #wv[k,0] = w1
                    #wv[k,1] = w2
                    wc[k] = zc[k]+hc[k]
                    xmomv[k,0] = uh1
                    xmomv[k,1] = uh2
                    xmomc[k] = 0.5*(uh1+uh2)
                    #print w1,w2
                elif (w1 > min(z1,z2)) | (w2 > min(z1,z2)):
                    wc[k] = 0.5*(w1+w2)
                    wv[k,0] = w1
                    wv[k,1] = w2
                    #print "cell %d w1 %f w2 %f" %(k,w1,w2)
                    # stops limiter from limiting these cells
                    domain.wet_nodes[k,0] = 2
                    domain.wet_nodes[k,1] = 2
                    if (w1 > z1):
                        xmomv[k,0] = uh1
                        xmomv[k,1] = 0.0
                    else:
                        xmomv[k,0] = 0.0
                        xmomv[k,1] = uh2
                else:
                    wc[k] = zc[k] # bed is completely dry
                    domain.wet_nodes[k,0] = 0
                    domain.wet_nodes[k,1] = 0
                    hc[k] = 0.0
                    xmomc[k] = 0.0
                if wet_len > 0.0:
                    domain.areas[k] = wet_len
                    #print w1,w2
                    #print "cell %d wet length %f"%(k,wet_len)
                    # if wet_len is 0 then bed is flat and area = x2-x1 (as default)
            else:
                wc[k] = zc[k] # bed is completely dry
                domain.wet_nodes[k,0] = 0
                domain.wet_nodes[k,1] = 0
                hc[k] = 0.0
                #xmomc[k] = 0.0
        else:
            wc[k] = zc[k]
            xmomc[k] = 0.0


def distribute_stage_to_height(domain):

    #print "in distribute stage to height"
    #print "integral %f"%domain.quantities['height'].get_integral()

    #print "hv1",domain.quantities['height'].vertex_values
    #print "hc1",domain.quantities['height'].centroid_values
    
    min_centroid_height = 1.0e-3
    N = domain.number_of_elements
    wv = domain.quantities['stage'].vertex_values
    zv = domain.quantities['elevation'].vertex_values
    hv = domain.quantities['height'].vertex_values
    hc = domain.quantities['height'].centroid_values
    for k in range(N):
        for i in range(2):
            if hc[k] > domain.epsilon:
                hv[k,i] = wv[k,i]-zv[k,i]
            elif hc[k] < 0:
                print "ERROR"
            else:
                hv[k,i] = 0.0
        
    #print "hv2",domain.quantities['height'].vertex_values
    #print "hc2",domain.quantities['height'].centroid_values
    #print "integral %f"%domain.quantities['height'].get_integral()
    
def protect_against_infinitesimal_and_negative_heights(domain):
    """Protect against infinitesimal heights and associated high velocities
    """

    #Shortcuts
    wc = domain.quantities['stage'].centroid_values
    zc = domain.quantities['elevation'].centroid_values
    xmomc = domain.quantities['xmomentum'].centroid_values
#    ymomc = domain.quantities['ymomentum'].centroid_values
    hc = wc - zc  #Water depths at centroids

    zv = domain.quantities['elevation'].vertex_values
    wv = domain.quantities['stage'].vertex_values
    #remove the above two lines and corresponding code below

    #Update
    for k in range(domain.number_of_elements):

        if hc[k] < domain.minimum_allowed_height:
            #Control stage
            if hc[k] < domain.epsilon:
                wc[k] = zc[k] # Contain 'lost mass' error
                wv[k,0] = zv[k,0]
                wv[k,1] = zv[k,1]

            #Control momentum
            #xmomc[k] = ymomc[k] = 0.0
            xmomc[k] = 0.0

def h_limiter(domain):
    """Limit slopes for each volume to eliminate artificial variance
    introduced by e.g. second order extrapolator

    limit on h = w-z

    This limiter depends on two quantities (w,z) so it resides within
    this module rather than within quantity.py
    """

    from Numeric import zeros, Float

    N = domain.number_of_elements
    beta_h = domain.beta_h

    #Shortcuts
    wc = domain.quantities['stage'].centroid_values
    zc = domain.quantities['elevation'].centroid_values
    hc = wc - zc

    wv = domain.quantities['stage'].vertex_values
    zv = domain.quantities['elevation'].vertex_values
    hv = wv-zv

    hvbar = zeros(hv.shape, Float) #h-limited values

    #Find min and max of this and neighbour's centroid values
    hmax = zeros(hc.shape, Float)
    hmin = zeros(hc.shape, Float)

    for k in range(N):
        hmax[k] = hmin[k] = hc[k]
        #for i in range(3):
        for i in range(2):   
            n = domain.neighbours[k,i]
            if n >= 0:
                hn = hc[n] #Neighbour's centroid value

                hmin[k] = min(hmin[k], hn)
                hmax[k] = max(hmax[k], hn)


    #Diffences between centroids and maxima/minima
    dhmax = hmax - hc
    dhmin = hmin - hc

    #Deltas between vertex and centroid values
    dh = zeros(hv.shape, Float)
    #for i in range(3):
    for i in range(2):
        dh[:,i] = hv[:,i] - hc

    #Phi limiter
    for k in range(N):

        #Find the gradient limiter (phi) across vertices
        phi = 1.0
        #for i in range(3):
        for i in range(2):
            r = 1.0
            if (dh[k,i] > 0): r = dhmax[k]/dh[k,i]
            if (dh[k,i] < 0): r = dhmin[k]/dh[k,i]

            phi = min( min(r*beta_h, 1), phi )

        #Then update using phi limiter
        #for i in range(3):
        for i in range(2):
            hvbar[k,i] = hc[k] + phi*dh[k,i]

    return hvbar

def balance_deep_and_shallow(domain):
    """Compute linear combination between stage as computed by
    gradient-limiters limiting using w, and stage computed by
    gradient-limiters limiting using h (h-limiter).
    The former takes precedence when heights are large compared to the
    bed slope while the latter takes precedence when heights are
    relatively small.  Anything in between is computed as a balanced
    linear combination in order to avoid numerical disturbances which
    would otherwise appear as a result of hard switching between
    modes.

    The h-limiter is always applied irrespective of the order.
    """

    #Shortcuts
    wc = domain.quantities['stage'].centroid_values
    zc = domain.quantities['elevation'].centroid_values
    hc = wc - zc

    wv = domain.quantities['stage'].vertex_values
    zv = domain.quantities['elevation'].vertex_values
    hv = wv-zv

    #Limit h
    hvbar = h_limiter(domain)

    for k in range(domain.number_of_elements):
        #Compute maximal variation in bed elevation
        #  This quantitiy is
        #    dz = max_i abs(z_i - z_c)
        #  and it is independent of dimension
        #  In the 1d case zc = (z0+z1)/2
        #  In the 2d case zc = (z0+z1+z2)/3

        dz = max(abs(zv[k,0]-zc[k]),
                 abs(zv[k,1]-zc[k]))#,
#                 abs(zv[k,2]-zc[k]))


        hmin = min( hv[k,:] )

        #Create alpha in [0,1], where alpha==0 means using the h-limited
        #stage and alpha==1 means using the w-limited stage as
        #computed by the gradient limiter (both 1st or 2nd order)

        #If hmin > dz/2 then alpha = 1 and the bed will have no effect
        #If hmin < 0 then alpha = 0 reverting to constant height above bed.

        if dz > 0.0:
            alpha = max( min( 2*hmin/dz, 1.0), 0.0 )
        else:
            #Flat bed
            alpha = 1.0

        alpha  = 0.0
        #Let
        #
        #  wvi be the w-limited stage (wvi = zvi + hvi)
        #  wvi- be the h-limited state (wvi- = zvi + hvi-)
        #
        #
        #where i=0,1,2 denotes the vertex ids
        #
        #Weighted balance between w-limited and h-limited stage is
        #
        #  wvi := (1-alpha)*(zvi+hvi-) + alpha*(zvi+hvi)
        #
        #It follows that the updated wvi is
        #  wvi := zvi + (1-alpha)*hvi- + alpha*hvi
        #
        # Momentum is balanced between constant and limited


        #for i in range(3):
        #    wv[k,i] = zv[k,i] + hvbar[k,i]

        #return

        if alpha < 1:

            #for i in range(3):
            for i in range(2):
                wv[k,i] = zv[k,i] + (1.0-alpha)*hvbar[k,i] + alpha*hv[k,i]

            #Momentums at centroids
            xmomc = domain.quantities['xmomentum'].centroid_values
#            ymomc = domain.quantities['ymomentum'].centroid_values

            #Momentums at vertices
            xmomv = domain.quantities['xmomentum'].vertex_values
#           ymomv = domain.quantities['ymomentum'].vertex_values

            # Update momentum as a linear combination of
            # xmomc and ymomc (shallow) and momentum
            # from extrapolator xmomv and ymomv (deep).
            xmomv[k,:] = (1.0-alpha)*xmomc[k] + alpha*xmomv[k,:]
#            ymomv[k,:] = (1-alpha)*ymomc[k] + alpha*ymomv[k,:]


###############################################
#Boundaries - specific to the shallow water wave equation
class Reflective_boundary(Boundary):
    """Reflective boundary returns same conserved quantities as
    those present in its neighbour volume but reflected.

    This class is specific to the shallow water equation as it
    works with the momentum quantities assumed to be the second
    and third conserved quantities.
    """

    def __init__(self, domain = None):
        Boundary.__init__(self)

        if domain is None:
            msg = 'Domain must be specified for reflective boundary'
            raise msg

        #Handy shorthands
        #self.stage   = domain.quantities['stage'].edge_values
        #self.xmom    = domain.quantities['xmomentum'].edge_values
        #self.ymom    = domain.quantities['ymomentum'].edge_values
        #self.normals = domain.normals
        self.stage   = domain.quantities['stage'].vertex_values
        self.xmom    = domain.quantities['xmomentum'].vertex_values
        self.height  = domain.quantities['height'].vertex_values

        from Numeric import zeros, Float
        self.conserved_quantities = zeros(3, Float)
        #self.conserved_quantities = zeros(2, Float)

    def __repr__(self):
        return 'Reflective_boundary'


    def evaluate(self, vol_id, edge_id):
        """Reflective boundaries reverses the outward momentum
        of the volume they serve.
        """

        q = self.conserved_quantities
        q[0] = self.stage[vol_id, edge_id]
        q[1] = self.xmom[vol_id, edge_id]
        #q[2] = self.ymom[vol_id, edge_id]
        q[2] = self.height[vol_id,edge_id] 
        #normal = self.normals[vol_id, 2*edge_id:2*edge_id+2]
        #normal = self.normals[vol_id, 2*edge_id:2*edge_id+1]


        #r = rotate(q, normal, direction = 1)
        #r[1] = -r[1]
        #q = rotate(r, normal, direction = -1)
        r = q
        r[1] = -q[1]
        q = r
        #For start interval there is no outward momentum so do not need to
        #reverse direction in this case

        return q

class Dirichlet_boundary(Boundary):
    """Dirichlet boundary returns constant values for the
    conserved quantities
    """


    def __init__(self, conserved_quantities=None):
        Boundary.__init__(self)

        if conserved_quantities is None:
            msg = 'Must specify one value for each conserved quantity'
            raise msg

        from Numeric import array, Float
        self.conserved_quantities=array(conserved_quantities).astype(Float)

    def __repr__(self):
        return 'Dirichlet boundary (%s)' %self.conserved_quantities

    def evaluate(self, vol_id=None, edge_id=None):
        return self.conserved_quantities


#########################
#Standard forcing terms:
#
def gravity(domain):
    """Apply gravitational pull in the presence of bed slope
    """
    #print "in gravity"
    #print "integral %f"%domain.quantities['height'].get_integral()
    from util import gradient
    from Numeric import zeros, Float, array, sum

    xmom = domain.quantities['xmomentum'].explicit_update
    stage = domain.quantities['stage'].explicit_update
#    ymom = domain.quantities['ymomentum'].explicit_update

    Stage = domain.quantities['stage']
    Elevation = domain.quantities['elevation']
    Height = domain.quantities['height']
    #h = Stage.edge_values - Elevation.edge_values
    h = Stage.vertex_values - Elevation.vertex_values
    hc = Height.centroid_values
    h = Height.vertex_values
    b = Elevation.vertex_values
    w = Stage.vertex_values

    x = domain.get_vertex_coordinates()
    g = domain.g

    for k in range(domain.number_of_elements):
#        avg_h = sum( h[k,:] )/3
        #avg_h = sum( h[k,:] )/2
        avg_h = hc[k]

        #Compute bed slope
        #x0, y0, x1, y1, x2, y2 = x[k,:]
        x0, x1 = x[k,:]
        #z0, z1, z2 = v[k,:]
        b0, b1 = b[k,:]

        w0, w1 = w[k,:]
        wx = gradient(x0, x1, w0, w1)

        #zx, zy = gradient(x0, y0, x1, y1, x2, y2, z0, z1, z2)
        bx = gradient(x0, x1, b0, b1)
        
        #Update momentum (explicit update is reset to source values)
        if domain.split == False:
            xmom[k] += -g*bx*avg_h
            #xmom[k] = -g*bx*avg_h
            #stage[k] = 0.0
        elif domain.split == True:
            xmom[k] += -g*wx*avg_h
            #xmom[k] = -g*wx*avg_h
        #ymom[k] += -g*zy*avg_h
 
def manning_friction(domain):
    """Apply (Manning) friction to water momentum
    """

    from math import sqrt

    w = domain.quantities['stage'].centroid_values
    z = domain.quantities['elevation'].centroid_values
    h = w-z

    uh = domain.quantities['xmomentum'].centroid_values
    #vh = domain.quantities['ymomentum'].centroid_values
    eta = domain.quantities['friction'].centroid_values

    xmom_update = domain.quantities['xmomentum'].semi_implicit_update
    #ymom_update = domain.quantities['ymomentum'].semi_implicit_update

    N = domain.number_of_elements
    eps = domain.minimum_allowed_height
    g = domain.g

    for k in range(N):
        if eta[k] >= eps:
            if h[k] >= eps:
                #S = -g * eta[k]**2 * sqrt((uh[k]**2 + vh[k]**2))
                S = -g * eta[k]**2 * uh[k]
                S /= h[k]**(7.0/3)

                #Update momentum
                xmom_update[k] += S*uh[k]
                #ymom_update[k] += S*vh[k]

def linear_friction(domain):
    """Apply linear friction to water momentum

    Assumes quantity: 'linear_friction' to be present
    """

    from math import sqrt

    w = domain.quantities['stage'].centroid_values
    z = domain.quantities['elevation'].centroid_values
    h = w-z

    uh = domain.quantities['xmomentum'].centroid_values
#    vh = domain.quantities['ymomentum'].centroid_values
    tau = domain.quantities['linear_friction'].centroid_values

    xmom_update = domain.quantities['xmomentum'].semi_implicit_update
#    ymom_update = domain.quantities['ymomentum'].semi_implicit_update

    N = domain.number_of_elements
    eps = domain.minimum_allowed_height
    g = domain.g #Not necessary? Why was this added?

    for k in range(N):
        if tau[k] >= eps:
            if h[k] >= eps:
                S = -tau[k]/h[k]

                #Update momentum
                xmom_update[k] += S*uh[k]
 #              ymom_update[k] += S*vh[k]



def check_forcefield(f):
    """Check that f is either
    1: a callable object f(t,x,y), where x and y are vectors
       and that it returns an array or a list of same length
       as x and y
    2: a scalar
    """

    from Numeric import ones, Float, array


    if callable(f):
        #N = 3
        N = 2
        #x = ones(3, Float)
        #y = ones(3, Float)
        x = ones(2, Float)
        #y = ones(2, Float)
        
        try:
            #q = f(1.0, x=x, y=y)
            q = f(1.0, x=x)
        except Exception, e:
            msg = 'Function %s could not be executed:\n%s' %(f, e)
            #FIXME: Reconsider this semantics
            raise msg

        try:
            q = array(q).astype(Float)
        except:
            msg = 'Return value from vector function %s could ' %f
            msg += 'not be converted into a Numeric array of floats.\n'
            msg += 'Specified function should return either list or array.'
            raise msg

        #Is this really what we want?
        msg = 'Return vector from function %s ' %f
        msg += 'must have same lenght as input vectors'
        assert len(q) == N, msg

    else:
        try:
            f = float(f)
        except:
            msg = 'Force field %s must be either a scalar' %f
            msg += ' or a vector function'
            raise msg
    return f

class Wind_stress:
    """Apply wind stress to water momentum in terms of
    wind speed [m/s] and wind direction [degrees]
    """

    def __init__(self, *args, **kwargs):
        """Initialise windfield from wind speed s [m/s]
        and wind direction phi [degrees]

        Inputs v and phi can be either scalars or Python functions, e.g.

        W = Wind_stress(10, 178)

        #FIXME - 'normal' degrees are assumed for now, i.e. the
        vector (1,0) has zero degrees.
        We may need to convert from 'compass' degrees later on and also
        map from True north to grid north.

        Arguments can also be Python functions of t,x,y as in

        def speed(t,x,y):
            ...
            return s

        def angle(t,x,y):
            ...
            return phi

        where x and y are vectors.

        and then pass the functions in

        W = Wind_stress(speed, angle)

        The instantiated object W can be appended to the list of
        forcing_terms as in

        Alternatively, one vector valued function for (speed, angle)
        can be applied, providing both quantities simultaneously.
        As in
        W = Wind_stress(F), where returns (speed, angle) for each t.

        domain.forcing_terms.append(W)
        """

        from config import rho_a, rho_w, eta_w
        from Numeric import array, Float

        if len(args) == 2:
            s = args[0]
            phi = args[1]
        elif len(args) == 1:
            #Assume vector function returning (s, phi)(t,x,y)
            vector_function = args[0]
            #s = lambda t,x,y: vector_function(t,x=x,y=y)[0]
            #phi = lambda t,x,y: vector_function(t,x=x,y=y)[1]
            s = lambda t,x: vector_function(t,x=x)[0]
            phi = lambda t,x: vector_function(t,x=x)[1]
        else:
           #Assume info is in 2 keyword arguments

           if len(kwargs) == 2:
               s = kwargs['s']
               phi = kwargs['phi']
           else:
               raise 'Assumes two keyword arguments: s=..., phi=....'

        print 'phi', phi
        self.speed = check_forcefield(s)
        self.phi = check_forcefield(phi)

        self.const = eta_w*rho_a/rho_w


    def __call__(self, domain):
        """Evaluate windfield based on values found in domain
        """

        from math import pi, cos, sin, sqrt
        from Numeric import Float, ones, ArrayType

        xmom_update = domain.quantities['xmomentum'].explicit_update
        #ymom_update = domain.quantities['ymomentum'].explicit_update

        N = domain.number_of_elements
        t = domain.time

        if callable(self.speed):
            xc = domain.get_centroid_coordinates()
            #s_vec = self.speed(t, xc[:,0], xc[:,1])
            s_vec = self.speed(t, xc)
        else:
            #Assume s is a scalar

            try:
                s_vec = self.speed * ones(N, Float)
            except:
                msg = 'Speed must be either callable or a scalar: %s' %self.s
                raise msg


        if callable(self.phi):
            xc = domain.get_centroid_coordinates()
            #phi_vec = self.phi(t, xc[:,0], xc[:,1])
            phi_vec = self.phi(t, xc)
        else:
            #Assume phi is a scalar

            try:
                phi_vec = self.phi * ones(N, Float)
            except:
                msg = 'Angle must be either callable or a scalar: %s' %self.phi
                raise msg

        #assign_windfield_values(xmom_update, ymom_update,
        #                        s_vec, phi_vec, self.const)
        assign_windfield_values(xmom_update, s_vec, phi_vec, self.const)


#def assign_windfield_values(xmom_update, ymom_update,
#                            s_vec, phi_vec, const):
def assign_windfield_values(xmom_update, s_vec, phi_vec, const):
    """Python version of assigning wind field to update vectors.
    A c version also exists (for speed)
    """
    from math import pi, cos, sin, sqrt

    N = len(s_vec)
    for k in range(N):
        s = s_vec[k]
        phi = phi_vec[k]

        #Convert to radians
        phi = phi*pi/180

        #Compute velocity vector (u, v)
        u = s*cos(phi)
        v = s*sin(phi)

        #Compute wind stress
        #S = const * sqrt(u**2 + v**2)
        S = const * u
        xmom_update[k] += S*u
        #ymom_update[k] += S*v