source: trunk/anuga_work/development/sudi/sw_1d/shock_detector_shv/sww_domain_shm.py @ 7906

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

Sudi Mungkasi, ANU, 2010
"""


from domain import *
Generic_Domain = Domain #Rename

#Shallow water domain
class Domain(Generic_Domain):

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

        self.forcing_terms.append(gravity_F2)
        #self.forcing_terms.append(manning_friction)

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

        self.__doc__ = 'sww_domain_shm'
        self.check_integrity()


    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
                
        msg = 'First evolved quantity must be "stage"'
        assert self.evolved_quantities[0] == 'stage', msg
        msg = 'Second evolved quantity must be "xmomentum"'
        assert self.evolved_quantities[1] == 'xmomentum', 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

    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_C_wellbalanced(self)
        #compute_fluxes_C_nonwellbalanced2(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_shv(self)
        distribute_to_vertices_and_edges_shm(self)
        
    def evolve(self, yieldstep = None, finaltime = None, duration = None, skip_initial_step = False):
        #Call basic machinery from parent class
        for t in Generic_Domain.evolve(self, yieldstep, finaltime, duration, skip_initial_step):
            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 ===============================

# Compute flux definition
# ###################################
def compute_fluxes_C_wellbalanced(domain):
    import sys
    from Numeric import Float
    from numpy import zeros   
    
    N = domain.number_of_elements
    timestep = float(sys.maxint)
    epsilon = domain.epsilon
    g = domain.g
    neighbours = domain.neighbours
    neighbour_vertices = domain.neighbour_vertices
    normals = domain.normals
    areas = domain.areas
    
    Stage    = domain.quantities['stage']
    Xmom     = domain.quantities['xmomentum']
    Bed      = domain.quantities['elevation']
    Height   = domain.quantities['height']
    Velocity = domain.quantities['velocity']


    stage_boundary_values = Stage.boundary_values
    xmom_boundary_values  = Xmom.boundary_values
    bed_boundary_values   = Bed.boundary_values
    height_boundary_values= Height.boundary_values
    vel_boundary_values   = Velocity.boundary_values
    
    stage_explicit_update = Stage.explicit_update
    xmom_explicit_update  = Xmom.explicit_update
    bed_explicit_values   = Bed.explicit_update
    height_explicit_values= Height.explicit_update
    vel_explicit_values   = Velocity.explicit_update
    
    max_speed_array = domain.max_speed_array
   
    domain.distribute_to_vertices_and_edges()
    domain.update_boundary()
    
    stage_V = Stage.vertex_values
    xmom_V  = Xmom.vertex_values
    bed_V   = Bed.vertex_values
    height_V= Height.vertex_values
    vel_V   = Velocity.vertex_values

    number_of_elements = len(stage_V)

    from comp_flux_ext_wellbalanced import compute_fluxes_ext_wellbalanced
    domain.flux_timestep = compute_fluxes_ext_wellbalanced(timestep,
                                  epsilon,
                                  g,
                                                           
                                  neighbours,
                                  neighbour_vertices,
                                  normals,
                                  areas,
                                                           
                                  stage_V,
                                  xmom_V,
                                  bed_V,
                                  height_V,
                                  vel_V,
                                                           
                                  stage_boundary_values,
                                  xmom_boundary_values,
                                  bed_boundary_values,
                                  height_boundary_values,
                                  vel_boundary_values,
                                                           
                                  stage_explicit_update,
                                  xmom_explicit_update,
                                  bed_explicit_values,
                                  height_explicit_values,
                                  vel_explicit_values,
                                                           
                                  number_of_elements,
                                  max_speed_array)


# ###################################
# Module functions for gradient limiting (distribute_to_vertices_and_edges)
def distribute_to_vertices_and_edges_shv(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 Numeric import Float
    from numpy import zeros    
    from 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

    for i in range(N):
        h_C[i] = w_C[i] - z_C[i]                                                
        if h_C[i] <= epsilon:
            uh_C[i] = 0.0
            u_C[i] = 0.0
            #w_C[i] = z_C[i]
        else:
            u_C[i]  = uh_C[i]/(h_C[i] + h0/h_C[i])

    for name in ['stage', 'height', 'velocity']:
        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_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 

    bed_V[:] = stage_V - h_V
    xmom_V[:] = u_V * h_V
    
    return

def distribute_to_vertices_and_edges_shm(domain):
    # shm stands for STAGE, HEIGHT, MOMENTUM
    """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 Numeric import Float
    from numpy import array, zeros
    from 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

    
    for i in range(N):
        h_C[i] = w_C[i] - z_C[i]                                                
        if h_C[i] <= epsilon:
            uh_C[i] = 0.0
            u_C[i] = 0.0
        else:
            u_C[i] = uh_C[i]/(h_C[i] + h0/h_C[i])

    for name in ['stage', 'height', 'xmomentum']:
        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_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 

    bed_V[:] = stage_V - h_V

    for i in range(N):                                          
        if min(h_V[i]) <= 0.0:
            h_V[i] = array([0.0, 0.0])
            stage_V[i] = bed_V[i]
            xmom_V[i] = array([0.0, 0.0])
    
    u_V[:]  = xmom_V/(h_V + h0/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
    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
            
    



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

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

    xmom = domain.quantities['xmomentum'].explicit_update
    stage = domain.quantities['stage'].explicit_update

    Stage = domain.quantities['stage']
    Elevation = domain.quantities['elevation']
    h = Stage.vertex_values - Elevation.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,:] )/2

        #Compute bed slope
        x0, x1 = x[k,:]
        b0, b1 = b[k,:]

        w0, w1 = w[k,:]
        wx = gradient(x0, x1, w0, w1)
        bx = gradient(x0, x1, b0, b1)
        
        #Update momentum (explicit update is reset to source values)
        xmom[k] += -g*bx*avg_h

def gravity_F2(domain):
    """Apply gravitational pull in the presence of bed slope
    """

    from util import gradient
    from Numeric import Float
    from numpy import zeros, array, sum    
    from parameters import F2#This is an additional friction!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

    xmom = domain.quantities['xmomentum'].explicit_update
    stage = domain.quantities['stage'].explicit_update

    Stage = domain.quantities['stage']
    Elevation = domain.quantities['elevation']
    h = Stage.vertex_values - Elevation.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,:] )/2

        #Compute bed slope
        x0, x1 = x[k,:]
        b0, b1 = b[k,:]

        w0, w1 = w[k,:]
        wx = gradient(x0, x1, w0, w1)
        bx = gradient(x0, x1, b0, b1)
        
        #Update momentum (explicit update is reset to source values)
        xmom[k] += -g*bx*avg_h + avg_h*F2#This is an additional friction!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

 
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 * uh[k]
                S /= h[k]**(7.0/3)

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