source: inundation/ga/storm_surge/pyvolution/shallow_water.py @ 620

"""Class Domain -
2D triangular 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

FIXME: Write equations here!


Conserved quantities are w (water level or stage), uh (x momentum)
and vh (y momentum).


Ole Nielsen, Stephen Roberts, Duncan Gray, Christopher Zoppou
Geoscience Australia, 2004   
"""

from domain import *
Generic_domain = Domain #Rename

class Domain(Generic_domain):

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

        conserved_quantities = ['level', 'xmomentum', 'ymomentum']
        other_quantities = ['elevation', 'friction'] 
        
        Generic_domain.__init__(self, coordinates, vertices, boundary,
                                conserved_quantities, other_quantities,
                                tagged_elements)

        from config import minimum_allowed_height, g
        self.minimum_allowed_height = minimum_allowed_height
        self.g = g

        self.forcing_terms.append(gravity)
        self.forcing_terms.append(manning_friction)

        #Realtime visualisation
        self.visualise = False

        #Stored output
        self.store = False
        self.format = 'sww'    
        self.smooth = True

        #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 = ['level']

        
        #Establish shortcuts to relevant quantities (for efficiency)
        #self.w = self.quantities['level']
        #self.uh = self.quantities['xmomentum']                 
        #self.vh = self.quantities['ymomentum']                         
        #self.z = self.quantities['elevation']          
        #self.eta = self.quantities['friction']                 

    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 "level"'
        assert self.conserved_quantities[0] == 'level', msg
        msg = 'Second conserved quantity must be "xmomentum"'
        assert self.conserved_quantities[1] == 'xmomentum', msg
        msg = 'Third conserved quantity must be "ymomentum"'
        assert self.conserved_quantities[2] == 'ymomentum', msg
        

        #Check that levels are >= bed elevation
        from Numeric import alltrue, greater_equal
        
        level = self.quantities['level']
        bed = self.quantities['elevation']        

        msg = 'All water levels must be greater than the bed elevation'
        assert alltrue( greater_equal(
            level.vertex_values, bed.vertex_values )), msg

        assert alltrue( greater_equal(
            level.edge_values, bed.edge_values )), msg

        assert alltrue( greater_equal(
            level.centroid_values, bed.centroid_values )), msg        


    def compute_fluxes(self):
        #Call correct module function
        #(either from this module or C-extension)
        compute_fluxes(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)


    #FIXME: Under construction    
#     def set_defaults(self):
#         """Set default values for uninitialised quantities.
#         This is specific to the shallow water wave equation
#         Defaults for 'elevation', 'friction', 'xmomentum' and 'ymomentum'
#         are 0.0. Default for 'level' is whatever the value of 'elevation'.
#         """

#         for name in self.other_quantities + self.conserved_quantities:
#             print name
#             print self.quantities.keys()
#             if not self.quantities.has_key(name):
#                 if name == 'level':
                    
#                     if self.quantities.has_key('elevation'):
#                         z = self.quantities['elevation'].vertex_values
#                         self.set_quantity(name, z) 
#                     else:
#                         self.set_quantity(name, 0.0)                        
#                 else:    
#                     self.set_quantity(name, 0.0)
                    


#         #Lift negative heights up
#         #z = self.quantities['elevation'].vertex_values
#         #w = self.quantities['level'].vertex_values

#         #h = w-z

#         #for k in range(h.shape[0]):
#         #    for i in range(3):
#         #        if h[k, i] < 0.0:
#         #            w[k, i] = z[k, i]

                    
#         #self.quantities['level'].interpolate()
                


    def evolve(self, yieldstep = None, finaltime = None):
        """Specialisation of basic evolve method from parent class
        """

        #Call check integrity here rather than from user scripts 
        #self.check_integrity()
        
        #Initialise real time viz if requested
        if self.visualise is True and self.time == 0.0:
            import realtime_visualisation as visualise
            visualise.create_surface(self)

        #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:
            #print 'Results will not be stored.'
            #print 'To store results set domain.store = True'
            pass
            #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):
            #Real time viz
            if self.visualise is True:
                visualise.update(self)

            #Store model data, e.g. for subsequent visualisation    
            if self.store is True:
                #self.store_timestep(['level', 'xmomentum', 'ymomentum'])
                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,y 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)
        
#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
    """

    #FIXME: Needs to be tested

    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'

    #FIXME: Put this test into C-extension as well
    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



####################################
# Flux computation        
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, vh are stored as elements 0, 1 and 2
    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

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

    z = (zl+zr)/2 #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)
    
    #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,
                        u_left*vh_left])
    flux_right = array([u_right*h_right,
                        u_right*uh_right + 0.5*g*h_right**2,
                        u_right*vh_right])    

    denom = s_max-s_min
    if denom == 0.0:
        edgeflux = array([0.0, 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 = rotate(edgeflux, normal, direction=-1)
        max_speed = max(abs(s_max), abs(s_min))

    return edgeflux, max_speed        


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 
    level, 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. 
    """

    import sys
    from Numeric import zeros, Float

    N = domain.number_of_elements
    
    #Shortcuts
    Level = domain.quantities['level']
    Xmom = domain.quantities['xmomentum']
    Ymom = domain.quantities['ymomentum']
    Bed = domain.quantities['elevation']        

    #Arrays
    level = Level.edge_values
    xmom =  Xmom.edge_values
    ymom =  Ymom.edge_values
    bed =   Bed.edge_values    

    level_bdry = Level.boundary_values
    xmom_bdry =  Xmom.boundary_values
    ymom_bdry =  Ymom.boundary_values
    
    flux = zeros(3, Float) #Work array for summing up fluxes

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

        flux[:] = 0.  #Reset work array
        for i in range(3):
            #Quantities inside volume facing neighbour i
            ql = [level[k, i], xmom[k, i], ymom[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 = [level_bdry[m], xmom_bdry[m], ymom_bdry[m]]
                zr = zl #Extend bed elevation to boundary
            else:    
                m = domain.neighbour_edges[k,i]
                qr = [level[n, m], xmom[n, m], ymom[n, m]]
                zr = bed[n, m]                

                
            #Outward pointing normal vector   
            normal = domain.normals[k, 2*i:2*i+2]

            #Flux computation using provided function
            edgeflux, max_speed = flux_function(normal, ql, qr, zl, zr)
            flux -= edgeflux * domain.edgelengths[k,i]

            #Update optimal_timestep
            try:
                timestep = min(timestep, domain.radii[k]/max_speed)
            except ZeroDivisionError:
                pass

        #Normalise by area and store for when all conserved
        #quantities get updated
        flux /= domain.areas[k]
        Level.explicit_update[k] = flux[0]
        Xmom.explicit_update[k] = flux[1]
        Ymom.explicit_update[k] = flux[2]

    #print 'FLUX l', Level.explicit_update
    #print 'FLUX x', Xmom.explicit_update
    #print 'FLUX y', Ymom.explicit_update    
    
    domain.timestep = timestep    


def compute_fluxes_c(domain):
    """Wrapper calling C version of compute fluxes
    """

    import sys
    from Numeric import zeros, Float

    N = domain.number_of_elements
    
    #Shortcuts
    Level = domain.quantities['level']
    Xmom = domain.quantities['xmomentum']
    Ymom = domain.quantities['ymomentum']
    Bed = domain.quantities['elevation']        

    timestep = float(sys.maxint)    
    from shallow_water_ext import compute_fluxes
    domain.timestep = compute_fluxes(timestep, domain.epsilon, domain.g,
                                     domain.neighbours,
                                     domain.neighbour_edges,
                                     domain.normals,
                                     domain.edgelengths,                       
                                     domain.radii,
                                     domain.areas,
                                     Level.edge_values, 
                                     Xmom.edge_values, 
                                     Ymom.edge_values,  
                                     Bed.edge_values,   
                                     Level.boundary_values,
                                     Xmom.boundary_values,
                                     Ymom.boundary_values,
                                     Level.explicit_update,
                                     Xmom.explicit_update,
                                     Ymom.explicit_update)
                                     

####################################
# 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
    
    """

    #Remove very thin layers of water
    protect_against_infinitesimal_and_negative_heights(domain)

    #Extrapolate all conserved quantities
    for name in domain.conserved_quantities:
        Q = domain.quantities[name]
        if domain.order == 1:
            Q.extrapolate_first_order()
        elif domain.order == 2:
            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)

    #Compute edge values by interpolation
    for name in domain.conserved_quantities:
        Q = domain.quantities[name]
        Q.interpolate_from_vertices_to_edges()
                    


def dry(domain):
    """Protect against infinitesimal heights and associated high velocities
    at vertices
    """

    #FIXME: Experimental
    
    #Shortcuts
    wv = domain.quantities['level'].vertex_values
    zv = domain.quantities['elevation'].vertex_values
    xmomv = domain.quantities['xmomentum'].vertex_values
    ymomv = domain.quantities['ymomentum'].vertex_values            
    hv = wv - zv  #Water depths at vertices

    #Update
    for k in range(domain.number_of_elements):
        hmax = max(hv[k, :])
        
        if hmax < domain.minimum_allowed_height:
            #Control level
            wv[k, :] = zv[k, :]

            #Control momentum
            xmomv[k,:] = ymomv[k,:] = 0.0
    

def protect_against_infinitesimal_and_negative_heights(domain):
    """Protect against infinitesimal heights and associated high velocities
    """
    
    #Shortcuts
    wc = domain.quantities['level'].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    

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

        if hc[k] < domain.minimum_allowed_height: 
            #Control level
            wc[k] = zc[k]

            #Control momentum
            xmomc[k] = ymomc[k] = 0.0
        
        #From 'newstyle
        #if hc[k] < domain.minimum_allowed_height:        
        #     if hc[k] < 0.0:
        #            #Control level and height
        #            wc[k] = zc[k]
        #
        #        #Control momentum
        #        xmomc[k] = ymomc[k] = 0.0
        #else:
            


def protect_against_infinitesimal_and_negative_heights_c(domain):
    """Protect against infinitesimal heights and associated high velocities
    """
    
    #Shortcuts
    wc = domain.quantities['level'].centroid_values
    zc = domain.quantities['elevation'].centroid_values
    xmomc = domain.quantities['xmomentum'].centroid_values
    ymomc = domain.quantities['ymomentum'].centroid_values            

    from shallow_water_ext import protect

    protect(domain.minimum_allowed_height, wc, zc, xmomc, ymomc)
    
    
def balance_deep_and_shallow(domain):
    """Compute linear combination between stage as computed by
    gradient-limiters and stage computed as constant height above bed.
    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.
    """
    
    #Shortcuts
    wc = domain.quantities['level'].centroid_values
    zc = domain.quantities['elevation'].centroid_values
    hc = wc - zc
    
    wv = domain.quantities['level'].vertex_values
    zv = domain.quantities['elevation'].vertex_values
    hv = wv-zv


    #Computed linear combination between constant levels and and
    #levels parallel to the bed elevation.     
    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 shallow 
        #first order scheme and alpha==1 means using the stage w as
        #computed by the gradient limiter (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


        #Weighted balance between stage parallel to bed elevation 
        #(wvi = zvi + hc) and stage as computed by 1st or 2nd 
        #order gradient limiter
        #(wvi = zvi + hvi) where i=0,1,2 denotes the vertex ids
        #
        #It follows that the updated wvi is
        #  wvi := (1-alpha)*(zvi+hc) + alpha*(zvi+hvi) = 
        #  zvi + hc + alpha*(hvi - hc)
        #
        #Note that hvi = zc+hc-zvi in the first order case (constant).

        if alpha < 1:         
            for i in range(3):
                wv[k,i] = zv[k,i] + hc[k] + alpha*(hv[k,i]-hc[k])
            

            #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-alpha)*xmomc[k] + alpha*xmomv[k,:]            
            ymomv[k,:] = (1-alpha)*ymomc[k] + alpha*ymomv[k,:]
                                    
            

def balance_deep_and_shallow_c(domain):
    """Wrapper for C implementation
    """
    
    #Shortcuts
    wc = domain.quantities['level'].centroid_values
    zc = domain.quantities['elevation'].centroid_values
    hc = wc - zc
    
    wv = domain.quantities['level'].vertex_values
    zv = domain.quantities['elevation'].vertex_values
    hv = wv-zv

    #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         

    

    from shallow_water_ext import balance_deep_and_shallow
    balance_deep_and_shallow(wc, zc, hc, wv, zv, hv,
                             xmomc, ymomc, xmomv, ymomv)

    


###############################################
#Boundary - 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.level = domain.quantities['level'].edge_values
        self.xmom = domain.quantities['xmomentum'].edge_values
        self.ymom = domain.quantities['ymomentum'].edge_values          
        self.normals = domain.normals
        
        from Numeric import zeros, Float        
        self.conserved_quantities = zeros(3, 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.level[vol_id, edge_id]
        q[1] = self.xmom[vol_id, edge_id] 
        q[2] = self.ymom[vol_id, edge_id] 
        
        normal = self.normals[vol_id, 2*edge_id:2*edge_id+2]     

        
        r = rotate(q, normal, direction = 1)
        r[1] = -r[1]
        q = rotate(r, normal, direction = -1)

        return q


#########################
#Standard forcing terms:
#
def gravity(domain):
    """Apply gravitational pull in the presence of bed slope
    """

    from util import gradient
    from Numeric import zeros, Float, array, sum

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

    Level = domain.quantities['level']
    Elevation = domain.quantities['elevation']    
    h = Level.edge_values - Elevation.edge_values
    v = Elevation.vertex_values
    
    x = domain.get_vertex_coordinates()
    g = domain.g
    
    for k in range(domain.number_of_elements):    
        avg_h = sum( h[k,:] )/3
        
        #Compute bed slope
        x0, y0, x1, y1, x2, y2 = x[k,:]    
        z0, z1, z2 = v[k,:]

        zx, zy = gradient(x0, y0, x1, y1, x2, y2, z0, z1, z2)

        #Update momentum
        xmom[k] += -g*zx*avg_h
        ymom[k] += -g*zy*avg_h        


def gravity_c(domain):
    """Wrapper calling C version
    """    

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

    Level = domain.quantities['level']    
    Elevation = domain.quantities['elevation']    
    h = Level.edge_values - Elevation.edge_values
    v = Elevation.vertex_values
    
    x = domain.get_vertex_coordinates()
    g = domain.g
    
    
    from shallow_water_ext import gravity
    gravity(g, h, v, x, xmom, ymom)
    
        
def manning_friction(domain):
    """Apply (Manning) friction to water momentum
    """

    from math import sqrt

    w = domain.quantities['level'].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 /= h[k]**(7.0/3)

                #Update momentum
                xmom_update[k] += S*uh[k]
                ymom_update[k] += S*vh[k]
            
        
def manning_friction_c(domain):
    """Wrapper for c version
    """


    xmom = domain.quantities['xmomentum']
    ymom = domain.quantities['ymomentum']    
        
    w = domain.quantities['level'].centroid_values
    z = domain.quantities['elevation'].centroid_values
    
    uh = xmom.centroid_values
    vh = ymom.centroid_values
    eta = domain.quantities['friction'].centroid_values    
    
    xmom_update = xmom.semi_implicit_update
    ymom_update = ymom.semi_implicit_update    
    
    N = domain.number_of_elements
    eps = domain.minimum_allowed_height
    g = domain.g

    from shallow_water_ext import manning_friction
    manning_friction(g, eps, w, z, uh, vh, eta, xmom_update, ymom_update)


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

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

    from math import sqrt

    w = domain.quantities['level'].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
    
    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 of x,y,t, 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
        x = ones(3, Float)
        y = ones(3, Float)            
        try:
            q = f(x, y, 1.0)
        except Exception, e:
            msg = 'Function %s could not be executed:\n%s' %(f, e)
            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

        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
    """

    def __init__(self, s, phi):
        """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 x,y,t as in

        def speed(x,y,t):
            ...
            return s
        
        def angle(x,y,t):
            ...
            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

        domain.forcing_terms.append(W)
        
        """

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

        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(xc[:,0], xc[:,1], t)
        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(xc[:,0], xc[:,1], t)
        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)


def assign_windfield_values(xmom_update, ymom_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)
        xmom_update[k] += S*u
        ymom_update[k] += S*v        
            

class Wind_stress_from_file:
    """Apply wind stress read from a file to water momentum

    At this stage the wind field is assumed to depend on time only, i.e
    no spatial dependency.
    
    FIXME: This class may be incorporated in the generic Wind_stress class
    Superseded
    """

    
    def __init__(self, filename):
        """Initialise windfield from a file with the following format

        time [DD/MM/YY hh:mm:ss], speed [m/s] direction [degrees]
        
        e.g.
        
        03/09/04 19:15:00, 9.53 39
        
        domain.forcing_terms.append(W)
        
        """

        import time
        from Numeric import array
        from config import time_format


        from config import rho_a, rho_w, eta_w
        self.const = eta_w*rho_a/rho_w
        

        try:
            fid = open(filename)
        except Exception, e:
            msg = 'Wind stress file %s could not be opened: %s\n'\
                  %(filename, e)
            raise msg


        line = fid.readline()
        fid.close()
        fields = line.split(',')
        msg = 'File %s must have the format date, values'
        assert len(fields) == 2, msg

        try:
            starttime = time.mktime(time.strptime(fields[0], time_format))
        except ValueError:    
            msg = 'First field in file %s must be' %filename
            msg += ' date-time with format %s.\n' %time_format
            msg += 'I got %s instead.' %fields[0] 
            raise msg

        #Split values
        values = []
        for value in fields[1].split():
            values.append(float(value))

        q = array(values)
        
        msg = 'ERROR: File function must return a 1d list or array '
        assert len(q.shape) == 1, msg
            

        msg = 'Values specified in file must convert to an array of length 2'
        assert len(q) == 2, msg

        self.filename = filename
        self.domain = domain
        self.starttime = starttime


        if domain.starttime is None:
            domain.starttime = starttime
        else:
            msg = 'Start time as specified in domain (%s) is earlier '
            'than the starttime of file %s: %s'\
                  %(domain.starttime, self.filename, self.starttime) 
            if self.starttime > domain.starttime:
                raise msg


        self.read_times() #Now read all times in.

        
    def read_times(self):
        from Numeric import zeros, Float, alltrue
        from config import time_format
        import time
        
        fid = open(self.filename)        
        lines = fid.readlines()
        fid.close()
        
        N = len(lines)

        T = zeros(N, Float)       #Time
        Q = zeros((N, 2), Float)  #Wind field (s, phi)
        
        for i, line in enumerate(lines):
            fields = line.split(',')
            real_time = time.mktime(time.strptime(fields[0], time_format))

            T[i] = real_time - self.start_time

            for j, value in enumerate(fields[1].split()):
                Q[i, j] = float(value)

        msg = 'File %s must list time as a monotonuosly ' %self.filename
        msg += 'increasing sequence'
        assert alltrue( T[1:] - T[:-1] > 0 ), msg

        self.T = T     #Time
        self.Q = Q     #Windfied
        self.index = 0 #Initial index

        
    def __repr__(self):
        return 'Wind field from file'

        

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

        from math import pi, cos, sin, sqrt
        
        xmom_update = domain.quantities['xmomentum'].explicit_update
        ymom_update = domain.quantities['ymomentum'].explicit_update    
        
        N = domain.number_of_elements
        t = self.domain.time
        
        msg = 'Time given in File %s does not match model time'\
              %self.filename
        
        if t < self.T[0]: raise msg
        if t > self.T[-1]: raise msg        

        oldindex = self.index

        #Find slot
        while t > self.T[self.index]: self.index += 1
        while t < self.T[self.index]: self.index -= 1

        #t is now between index and index+1
        ratio = (t - self.T[self.index])/\
                (self.T[self.index+1] - self.T[self.index])

        #Compute interpolated values for s and phi
        q = self.Q[self.index,:] +\
            ratio*(self.Q[self.index+1,:] - self.Q[self.index,:]) 

        s = q[0]
        
        #Convert to radians
        phi = q[1]*pi/180
        
        #Compute velocity vector (u, v)
        u = s*cos(phi)
        v = s*sin(phi)

        #Compute wind stress for this time step
        S = self.const * sqrt(u**2 + v**2)
        xmom_update += S*u
        ymom_update += S*v        
        




###########################
###########################
#Geometries


#FIXME: Rethink this way of creating values.


class Weir:
    """Set a bathymetry for weir with a hole and a downstream gutter
    x,y are assumed to be in the unit square
    """
    
    def __init__(self, stage):
        self.inflow_stage = stage

    def __call__(self, x, y):    
        from Numeric import zeros, Float
        from math import sqrt
    
        N = len(x)
        assert N == len(y)    

        z = zeros(N, Float)
        for i in range(N):
            z[i] = -x[i]/2  #General slope

            #Flattish bit to the left
            if x[i] < 0.3:
                z[i] = -x[i]/10 
            
            #Weir
            if x[i] >= 0.3 and x[i] < 0.4:
                z[i] = -x[i]+0.9

            #Dip
            x0 = 0.6
            #depth = -1.3
            depth = -1.0
            #plateaux = -0.9
            plateaux = -0.6            
            if y[i] < 0.7:
                if x[i] > x0 and x[i] < 0.9:
                    z[i] = depth

                #RHS plateaux
                if x[i] >= 0.9:
                    z[i] = plateaux
                    
                    
            elif y[i] >= 0.7 and y[i] < 1.5:
                #Restrict and deepen
                if x[i] >= x0 and x[i] < 0.8:
                    z[i] = depth-(y[i]/3-0.3)
                    #z[i] = depth-y[i]/5
                    #z[i] = depth
                elif x[i] >= 0.8:    
                    #RHS plateaux
                    z[i] = plateaux
                    
            elif y[i] >= 1.5:
                if x[i] >= x0 and x[i] < 0.8 + (y[i]-1.5)/1.2:
                    #Widen up and stay at constant depth
                    z[i] = depth-1.5/5
                elif x[i] >= 0.8 + (y[i]-1.5)/1.2:                        
                    #RHS plateaux
                    z[i] = plateaux                                    
                    

            #Hole in weir (slightly higher than inflow condition)
            if x[i] >= 0.3 and x[i] < 0.4 and y[i] > 0.2 and y[i] < 0.4:
                z[i] = -x[i]+self.inflow_stage + 0.02

            #Channel behind weir
            x0 = 0.5
            if x[i] >= 0.4 and x[i] < x0 and y[i] > 0.2 and y[i] < 0.4:
                z[i] = -x[i]+self.inflow_stage + 0.02
                
            if x[i] >= x0 and x[i] < 0.6 and y[i] > 0.2 and y[i] < 0.4:
                #Flatten it out towards the end
                z[i] = -x0+self.inflow_stage + 0.02 + (x0-x[i])/5

            #Hole to the east
            x0 = 1.1; y0 = 0.35
            #if x[i] < -0.2 and y < 0.5:
            if sqrt((2*(x[i]-x0))**2 + (2*(y[i]-y0))**2) < 0.2:
                z[i] = sqrt(((x[i]-x0))**2 + ((y[i]-y0))**2)-1.0

            #Tiny channel draining hole
            if x[i] >= 1.14 and x[i] < 1.2 and y[i] >= 0.4 and y[i] < 0.6:
                z[i] = -0.9 #North south
                
            if x[i] >= 0.9 and x[i] < 1.18 and y[i] >= 0.58 and y[i] < 0.65:
                z[i] = -1.0 + (x[i]-0.9)/3 #East west
            
            

            #Stuff not in use
            
            #Upward slope at inlet to the north west
            #if x[i] < 0.0: # and y[i] > 0.5:
            #    #z[i] = -y[i]+0.5  #-x[i]/2
            #    z[i] = x[i]/4 - y[i]**2 + 0.5

            #Hole to the west
            #x0 = -0.4; y0 = 0.35 # center
            #if sqrt((2*(x[i]-x0))**2 + (2*(y[i]-y0))**2) < 0.2:
            #    z[i] = sqrt(((x[i]-x0))**2 + ((y[i]-y0))**2)-0.2





        return z/2

class Weir_simple:
    """Set a bathymetry for weir with a hole and a downstream gutter
    x,y are assumed to be in the unit square
    """
    
    def __init__(self, stage):
        self.inflow_stage = stage

    def __call__(self, x, y):    
        from Numeric import zeros, Float 
    
        N = len(x)
        assert N == len(y)    

        z = zeros(N, Float)
        for i in range(N):
            z[i] = -x[i]  #General slope

            #Flat bit to the left
            if x[i] < 0.3:
                z[i] = -x[i]/10  #General slope
            
            #Weir
            if x[i] > 0.3 and x[i] < 0.4:
                z[i] = -x[i]+0.9

            #Dip
            if x[i] > 0.6 and x[i] < 0.9:
                z[i] = -x[i]-0.5  #-y[i]/5

            #Hole in weir (slightly higher than inflow condition)
            if x[i] > 0.3 and x[i] < 0.4 and y[i] > 0.2 and y[i] < 0.4:
                z[i] = -x[i]+self.inflow_stage + 0.05       
            

        return z/2
        

            
class Constant_height:
    """Set an initial condition with constant water height, e.g
    stage s = z+h
    """
    def __init__(self, W, h):
        self.W = W
        self.h = h

    def __call__(self, x, y):
        if self.W is None:
            from Numeric import ones, Float
            return self.h*ones(len(x), Float)
        else:
            return self.W(x,y) + self.h



##############################################
#Initialise module


import compile
if compile.can_use_C_extension('shallow_water_ext.c'):
    #Replace python version with c implementations
    
    from shallow_water_ext import rotate, assign_windfield_values
    compute_fluxes = compute_fluxes_c
    gravity = gravity_c
    manning_friction = manning_friction_c
    balance_deep_and_shallow = balance_deep_and_shallow_c
    protect_against_infinitesimal_and_negative_heights = protect_against_infinitesimal_and_negative_heights_c

    
    #distribute_to_vertices_and_edges = distribute_to_vertices_and_edges_c


#Optimisation with psyco
from config import use_psyco
if use_psyco:
    try:
        import psyco
    except:
        msg = 'WARNING: psyco (speedup) could not import'+\
              ', you may want to consider installing it'
        print msg
    else:
        psyco.bind(Domain.distribute_to_vertices_and_edges)
        psyco.bind(Domain.compute_fluxes)
        
if __name__ == "__main__":
    pass