source: trunk/anuga_work/development/2010-projects/anuga_1d/sww/sww_domain.py @ 8192

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

import numpy

from anuga_1d.base.generic_domain import Generic_domain
from sww_boundary_conditions import *
from sww_forcing_terms import *

#Shallow water domain
class Domain(Generic_domain):

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

        conserved_quantities = ['stage', 'xmomentum']
        evolved_quantities   = ['stage', 'xmomentum', 'elevation', 'height', 'velocity']
        other_quantities     = ['friction']
        Generic_domain.__init__(self,
                                coordinates          = coordinates,
                                boundary             = boundary,
                                conserved_quantities = conserved_quantities,
                                evolved_quantities   = evolved_quantities,
                                other_quantities     = other_quantities,
                                tagged_elements      = tagged_elements)
        
        from anuga_1d.config import minimum_allowed_height, g, h0
        self.minimum_allowed_height = minimum_allowed_height
        self.g = g
        self.h0 = h0

        #forcing terms
        self.forcing_terms.append(gravity)
        #self.forcing_terms.append(manning_friction)
        
        
        #Stored output
        self.store = True
        self.format = 'sww'
        self.smooth = True

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

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

        self.__doc__ = 'sww_domain'


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

        import sys


        timestep = float(sys.maxint)

        stage    = self.quantities['stage']
        xmom     = self.quantities['xmomentum']
        bed      = self.quantities['elevation']
        height   = self.quantities['height']
        velocity = self.quantities['velocity']


        from anuga_1d.sww.sww_comp_flux_ext import compute_fluxes_ext_short

        #self.flux_timestep = compute_fluxes_ext(timestep,self,stage,xmom,bed,height,velocity)

    
        self.flux_timestep = compute_fluxes_ext_short(timestep,self,stage,xmom,bed)



    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 evolve(self, yieldstep = None, finaltime = None, duration = None,
               skip_initial_step = False):
        """Specialisation of basic evolve method from parent class
        """

        #Call basic machinery from parent class
        for t in Generic_domain.evolve(self, yieldstep, finaltime,duration,
                                       skip_initial_step):

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





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


# 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

    """

    #from config import optimised_gradient_limiter

    #Remove very thin layers of water
    #protect_against_infinitesimal_and_negative_heights(domain)  

    import sys
    from anuga_1d.config import epsilon, h0
        
    N = domain.number_of_elements

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

    #Arrays   
    w_C   = Stage.centroid_values    
    uh_C  = Xmom.centroid_values    
    z_C   = Bed.centroid_values
    h_C   = Height.centroid_values
    u_C   = Velocity.centroid_values
        
    #Calculate height (and fix negatives)better be non-negative!
    h_C[:] = w_C - z_C
    u_C[:]  = uh_C/(h_C + h0/h_C)

    #print 'domain.order', domain.order
    
    for name in [ 'velocity', 'stage' ]:
        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':
            Q.extrapolate_second_order()
            #Q.limit()
        else:
            raise 'Unknown order'


    stage_V = domain.quantities['stage'].vertex_values                  
    bed_V   = domain.quantities['elevation'].vertex_values      
    h_V     = domain.quantities['height'].vertex_values
    u_V     = domain.quantities['velocity'].vertex_values
    xmom_V  = domain.quantities['xmomentum'].vertex_values

    h_V[:,:]    = stage_V - bed_V


    #print 'any' ,numpy.any( h_V[:,0] < 0.0)
    #print 'any' ,numpy.any( h_V[:,1] < 0.0)


    h_0 = numpy.where(h_V[:,0] < 0.0, 0.0, h_V[:,0])
    h_1 = numpy.where(h_V[:,0] < 0.0, h_V[:,1]+h_V[:,0], h_V[:,1])

    h_V[:,0] = h_0
    h_V[:,1] = h_1


    h_0 = numpy.where(h_V[:,1] < 0.0, h_V[:,1]+h_V[:,0], h_V[:,0])
    h_1 = numpy.where(h_V[:,1] < 0.0, 0.0, h_V[:,1])

    h_V[:,0] = h_0
    h_V[:,1] = h_1
    
    #print 'any' ,numpy.any( h_V[:,0] < 0.0)
    #print 'any' ,numpy.any( h_V[:,1] < 0.0)

    #h00 = 1e-12
    #print h00

    h_V[:,:] = numpy.where (h_V <= h0, 0.0, h_V)
    u_V[:,:] = numpy.where (h_V <= h0, 0.0, u_V)

#    # protect from edge values going negative
#    h_V[:,1] = numpy.where(h_V[:,0] < 0.0 , h_V[:,1]-h_V[:,0], h_V[:,1])
#    h_V[:,0] = numpy.where(h_V[:,0] < 0.0 , 0.0, h_V[:,0])
#
#    h_V[:,0] = numpy.where(h_V[:,1] < 0.0 , h_V[:,0]-h_V[:,1], h_V[:,0])
#    h_V[:,1] = numpy.where(h_V[:,1] < 0.0 , 0.0, h_V[:,1])
#
#
    #stage_V[:,:] = bed_V + h_V
    bed_V[:,:] = stage_V - h_V
    xmom_V[:,:] = u_V * h_V
    
    return
#







#
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
    hc = wc - zc  #Water depths at centroids

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

    #Update
    #FIXME replace with numpy.where
    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]
                
            xmomc[k] = 0.0
            
        #N = domain.number_of_elements
        #if (k == 0) | (k==N-1):
        #    wc[k] = zc[k] # Contain 'lost mass' error
        #    wv[k,0] = zv[k,0]
        #    wv[k,1] = zv[k,1]
    
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
    """

    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, numpy.float) #h-limited values

    #Find min and max of this and neighbour's centroid values
    hmax = zeros(hc.shape, numpy.float)
    hmin = zeros(hc.shape, numpy.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, numpy.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 numpyal 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,:]