source: trunk/anuga_work/development/2010-projects/anuga_1d/pipe/pipe_domain.py @ 8209

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

This module contains a specialisation of class Generic_domain from module generic_domain.py consisting of methods specific to mixed rectangular Pipe/Channel flow using the Shallow Water Wave Equation.

This particular modification of the Domain class implements the ability to vary the height and width of the 1D pipe that the water flows in. As a result the conserved variables are slightly different than previous implementations and so are the equations. In particular, the conserved quantites generalise the open channel flow conserved quantities.

U_t + E_x = S

where

U = [A, Q] = [b*h, u*b*h]
E = [Q, Q^2/A + A]
S represents source terms forcing the system
(e.g. gravity, boundary_stress, friction, wind stress, ...)
gravity = -g*b*h*z_x
boundary_stress = 1/2*g*b_x*h^2

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

The quantities are

symbol    variable name    explanation
A         area             free surface equivalent wetted area = b*h [m^2]
Q         discharge        flux of water = u*b*h [m^3/s]
x         x                horizontal distance from origin [m]
z         elevation        elevation of bottom of pipe [m]
h         height           free surface equivalent water height above z [m]
w         stage            equivalent absolute water level, w = z+h [m]
u         xvelocity        speed in the x direction [m/s]
uh        xmomentum        momentum in the x direction [m^2/s]
b         width            width of pipe [m]
t         top              height of pipe above z [m]
st        state            state of the cell (pressurised or free surface
eta                        mannings friction coefficient [to appear]

The conserved quantities are A, Q
--------------------------------------------------------------------------
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

Add paper details here

John Jakeman, Ole Nielsen, Stephen Roberts, Duncan Gray, Christopher Zoppou,
Padarn Wilson, Geoscience Australia, 2008
"""


from anuga_1d.base.generic_domain import *
import numpy


#Shallow water pipe domain
class Domain(Generic_domain):

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

        conserved_quantities = ['area', 'discharge']
        evolved_quantities = ['area', 'discharge', 'elevation', 'height', 'velocity','width','top','stage','state']
        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
        self.setstageflag = False
        self.discontinousb = False


        # forcing terms gravity and boundary stress are included in the flux calculation
        #self.forcing_terms.append(gravity)
        #self.forcing_terms.append(boundary_stress)
        #self.forcing_terms.append(manning_friction)
       

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

        
        #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.set_quantities_to_be_stored(['area','discharge'])

        self.__doc__ = 'channel_domain'

        self.check_integrity()


    def check_integrity(self):

        #Check that we are solving the shallow water wave equation

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

        msg = 'First evolved quantity must be "area"'
        assert self.evolved_quantities[0] == 'area', msg
        msg = 'Second evolved quantity must be "discharge"'
        assert self.evolved_quantities[1] == 'discharge', msg
        msg = 'Third evolved quantity must be "elevation"'
        assert self.evolved_quantities[2] == 'elevation', msg
        msg = 'Fourth evolved quantity must be "height"'
        assert self.evolved_quantities[3] == 'height', msg
        msg = 'Fifth evolved quantity must be "velocity"'
        assert self.evolved_quantities[4] == 'velocity', msg
        msg = 'Sixth evolved quantity must be "width"'
        assert self.evolved_quantities[5] == 'width', msg
        msg = 'Seventh evolved quantity must be "top"'
        assert self.evolved_quantities[6] == 'top', msg
        msg = 'Eighth evolved quantity must be "stage"'
        assert self.evolved_quantities[7] == 'stage', msg
        msg = 'Ninth evolved quantity must be "state"'
        assert self.evolved_quantities[8] == 'state', msg

        Generic_domain.check_integrity(self)

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

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

#=============== End of Pipe Domain ===============================

#-----------------------------------
# Compute flux definition with pipe
#-----------------------------------
def compute_fluxes_pipe(domain):
    import sys
    timestep = float(sys.maxint)

    area       = domain.quantities['area']
    discharge  = domain.quantities['discharge']
    bed        = domain.quantities['elevation']
    height     = domain.quantities['height']
    velocity   = domain.quantities['velocity']
    width      = domain.quantities['width']
    top        = domain.quantities['top']
    state      = domain.quantities['state']


    from anuga_1d.pipe.pipe_domain_ext import compute_fluxes_pipe_ext
    domain.flux_timestep = compute_fluxes_pipe_ext(timestep,domain,area,discharge,bed,height,velocity,width,top,state)

#-----------------------------------------------------------------------
# Distribute to verticies with stage, velocity and pipe geometry
# reconstructed and then extrapolated.
#-----------------------------------------------------------------------
def distribute_to_vertices_and_edges_limit_s_v(domain):
    import sys
    from anuga_1d.config import epsilon, h0

    N = domain.number_of_elements

    #Shortcuts
    area      = domain.quantities['area']
    discharge = domain.quantities['discharge']
    bed       = domain.quantities['elevation']
    height    = domain.quantities['height']
    velocity  = domain.quantities['velocity']
    width     = domain.quantities['width']
    top       = domain.quantities['top']
    stage     = domain.quantities['stage']
    state     = domain.quantities['state']

    #Arrays   
    a_C   = area.centroid_values
    d_C   = discharge.centroid_values
    z_C   = bed.centroid_values
    h_C   = height.centroid_values
    u_C   = velocity.centroid_values
    b_C   = width.centroid_values
    t_C   = top.centroid_values
    w_C   = stage.centroid_values
    s_C   = state.centroid_values

    # Calculate height, velocity and stage.
    # Here we assume the conserved quantities and the channel geometry
    # (i.e. bed, top and width) have been accurately computed in the previous
    # timestep. 
    h_C[:] = numpy.where(a_C > 0.0, a_C/(b_C + h0/b_C), 0.0)
    u_C[:] = numpy.where(a_C > 0.0, d_C/(a_C + h0/a_C), 0.0)

    w_C[:] = h_C + z_C    
      
    # Extrapolate velocity and stage as well as channel geometry.
    for name in ['velocity', 'stage', 'elevation', 'width', 'top']:
        Q = domain.quantities[name]
        if domain.order == 1:
            Q.extrapolate_first_order()
        elif domain.order == 2:
            Q.extrapolate_second_order()
        else:
            raise 'Unknown order'

    # Stage, bed, width, top and velocity have been extrapolated.
    # We may need to ensure that top is never 0.
    w_V  = stage.vertex_values
    u_V  = velocity.vertex_values
    z_V  = bed.vertex_values
    b_V  = width.vertex_values
    t_V  = top.vertex_values

    # These quantites need to update vertex_values
    a_V  = area.vertex_values
    h_V  = height.vertex_values
    d_V  = discharge.vertex_values

    # State is true of false so should not be extrapolated
    s_V  = state.vertex_values


    # Calculate height and fix up negatives. The main idea here is
    # fix up the wet/dry interface.
    h_V[:,:] = w_V - z_V

    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


    # Protect against negative and small heights. If we set h to zero
    # we better do the same with velocity (i.e. no water, no velocity).
    h_V[:,:] = numpy.where (h_V <= h0, 0.0, h_V)
    u_V[:,:] = numpy.where (h_V <= h0, 0.0, u_V)


    # Clean up conserved quantities
    w_V[:] = z_V + h_V
    a_V[:] = b_V * h_V
    d_V[:] = u_V * a_V

    
    return

#-----------------------------------------------------------------------
# Distribute to verticies with stage, height and velocity reconstructed 
# and then extrapolated.
# In this method, we extrapolate the stage and height out to the vertices.
# The bed, although given as initial data to the problem, is reconstructed
# from the stage and height. This ensures consistency of the reconstructed
# quantities (i.e. w = z + h) as well as protecting against negative
# heights.
#-----------------------------------------------------------------------
def distribute_to_vertices_and_edges_limit_s_v_h(domain):
    import sys
    from anuga_1d.config import epsilon, h0

    N = domain.number_of_elements

    #Shortcuts
    area      = domain.quantities['area']
    discharge = domain.quantities['discharge']
    bed       = domain.quantities['elevation']
    height    = domain.quantities['height']
    velocity  = domain.quantities['velocity']
    width     = domain.quantities['width']
    top     = domain.quantities['top']
    stage     = domain.quantities['stage']
    state     = domain.quantities['state']

    #Arrays   
    a_C   = area.centroid_values
    d_C   = discharge.centroid_values
    z_C   = bed.centroid_values
    h_C   = height.centroid_values
    u_C   = velocity.centroid_values
    b_C   = width.centroid_values
    t_C   = top.centroid_values
    w_C   = stage.centroid_values
    s_C   = state.centroid_values

    # Construct h,u,w from the conserved quantities after protecting
    # conserved quantities from becoming too small.
    a_C[:] = numpy.where( (a_C>h0), a_C, 0.0 )
    d_C[:] = numpy.where( (a_C>h0), d_C, 0.0 )
    h_C[:] = numpy.where( (b_C>h0), a_C/(b_C + h0/b_C), 0.0 )
    u_C[:] = numpy.where( (a_C>h0), d_C/(a_C + h0/a_C), 0.0 )    

    # Set the stage
    w_C[:] = h_C + z_C          

    # Extrapolate "fundamental" quantities.
    # All other quantities will be reconstructed from these.
    for name in ['velocity', 'stage',  'height', 'width', 'top']:
        Q = domain.quantities[name]
        if domain.order == 1:
            Q.extrapolate_first_order()
        elif domain.order == 2:
            Q.extrapolate_second_order()
        else:
            raise 'Unknown order'

    # These quantities have been extrapolated.
    # We may need to ensure top is never 0.
    u_V  = velocity.vertex_values
    w_V  = stage.vertex_values
    h_V  = height.vertex_values
    b_V  = width.vertex_values
    t_V  = top.vertex_values
        
    # These need to be reconstructed
    a_V  = area.vertex_values
    d_V  = discharge.vertex_values
    z_V  = bed.vertex_values

    # State is true or false so should never be extrapolated.
    s_V = state.vertex_values

    # Reconstruct bed from stage and height.
    z_V[:] = w_V-h_V

    # Now reconstruct our conserved quantities from the above
    # reconstructed quantities.
    a_V[:] = b_V*h_V
    d_V[:] = u_V*a_V

    return


#--------------------------------------------------------
#Boundaries - specific to the pipe_domain
#--------------------------------------------------------
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.normals  = domain.normals
        self.area     = domain.quantities['area'].vertex_values
        self.discharge     = domain.quantities['discharge'].vertex_values
        self.bed      = domain.quantities['elevation'].vertex_values
        self.height   = domain.quantities['height'].vertex_values
        self.velocity = domain.quantities['velocity'].vertex_values
        self.width    = domain.quantities['width'].vertex_values
        self.top    = domain.quantities['top'].vertex_values
        self.stage    = domain.quantities['stage'].vertex_values
        self.state    = domain.quantities['state'].vertex_values

        self.evolved_quantities = numpy.zeros(9, numpy.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.evolved_quantities
        q[0] =  self.area[vol_id, edge_id]
        q[1] = -self.discharge[vol_id, edge_id]
        q[2] =  self.bed[vol_id, edge_id]
        q[3] =  self.height[vol_id, edge_id]
        q[4] = -self.velocity[vol_id, edge_id]
        q[5] =  self.width[vol_id,edge_id]
        q[6] =  self.top[vol_id,edge_id]
        q[7] =  self.stage[vol_id,edge_id]
        q[8] =  self.state[vol_id,edge_id]
        return q

class Dirichlet_boundary(Boundary):
    """Dirichlet boundary returns constant values for the
    conserved quantities
if k>5 and k<15:
            print discharge_ud[k],-g*zx*avg_h*avg_b
        discharge_ud[k] +=-g*zx*avg_h*avg_b    """


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

        if evolved_quantities is None:
            msg = 'Must specify one value for each evolved quantity'
            raise msg

        assert len(evolved_quantities) == 9

        self.evolved_quantities=numpy.array(evolved_quantities,numpy.float)

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

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


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



    Area      = domain.quantities['area']
    Discharge  = domain.quantities['discharge']
    Elevation = domain.quantities['elevation']
    Height    = domain.quantities['height']
    Width     = domain.quantities['width']
    Top       = domain.quantities['top']

    discharge_ud  = Discharge.explicit_update


       
    h = Height.vertex_values
    b = Width.vertex_values
    t = Top.vertex_values
    a = Area.vertex_values
    z = Elevation.vertex_values

    x = domain.get_vertex_coordinates()
    g = domain.g
    for k in range(domain.number_of_elements):
        avg_h = 0.5*(h[k,0] + h[k,1])
        avg_b = 0.5*(b[k,0] + b[k,1])
        avg_t = 0.5*(t[k,0] + t[k,1])
        
        #Compute bed slope
        x0, x1 = x[k,:]
        z0, z1 = z[k,:]
        zx = gradient(x0, x1, z0, z1)
       
        #Update momentum (explicit update is reset to source values)
        discharge_ud[k]+= -g*zx*avg_h*avg_b
      

def boundary_stress(domain):

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

    Area     = domain.quantities['area']
    Discharge  = domain.quantities['discharge']
    Elevation = domain.quantities['elevation']
    Height    = domain.quantities['height']
    Width     = domain.quantities['width']
    Top       = domain.quantities['top']

    discharge_ud  = Discharge.explicit_update
 
    h = Height.vertex_values
    b = Width.vertex_values
    t = Top.vertex_values
    a = Area.vertex_values
    z = Elevation.vertex_values

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

    for k in range(domain.number_of_elements):
        avg_h = 0.5*(h[k,0] + h[k,1])
        

        #Compute bed slope
        x0, x1 = x[k,:]
        b0, b1 = b[k,:]
        bx = gradient(x0, x1, b0, b1)
        
        #Update momentum (explicit update is reset to source values)
        discharge_ud[k] += 0.5*g*bx*avg_h*avg_h
        #stage_ud[k] = 0.0
 
 
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
    eta = domain.quantities['friction'].centroid_values

    xmom_update = domain.quantities['xmomentum'].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
    tau = domain.quantities['linear_friction'].centroid_values

    xmom_update = domain.quantities['xmomentum'].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]




def linearb(domain):

    bC = domain.quantities['width'].vertex_values
    
    for i in range(len(bC)-1):
        temp= 0.5*(bC[i,1]+bC[i+1,0])
        bC[i,1]=temp
        bC[i+1,0]=temp