source: trunk/anuga_work/development/2010-projects/anuga_1d/sqpipe/sqpipe_test_domain.py @ 8253

from anuga_1d.sqpipe.sqpipe_domain import Sqpipe_domain
from anuga_1d.sqpipe.sqpipe_forcing_terms import *
from anuga_1d.sqpipe.sqpipe_boundary_conditions import *


class Domain(Sqpipe_domain):

    def __init__(self, coordinates, boundary = None, tagged_elements = None, bulk_modulus = 1.0):
        conserved_quantities = ['area', 'discharge']
        forcing_terms = []
        state = numpy.zeros(len(coordinates)-1, numpy.int)
        Sqpipe_domain.__init__(self,
                               coordinates          = coordinates,
                               conserved_quantities = conserved_quantities,
                               forcing_terms        = forcing_terms,
                               boundary             = boundary,
                               state                = state,
                               bulk_modulus         = bulk_modulus,
        #                       update_state_flag         = False,
                               tagged_elements      = tagged_elements)


        self.__doc__ = 'sqpipe_area_discharge_domain'
            
    def distribute_to_vertices_and_edges(self):
        # Shortcuts
        h0 = self.h0
        epsilon = self.epsilon

        area      = self.quantities['area']
        discharge = self.quantities['discharge']
        bed       = self.quantities['elevation']
        height    = self.quantities['height']
        velocity  = self.quantities['velocity']
        width     = self.quantities['width']
        top       = self.quantities['top']
        stage     = self.quantities['stage']

        #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

        #Calculate height (and fix negatives)better be non-negative!
        h_C[:] = a_C/(b_C + h0/(b_C + h0)) # Do we need to protect against small b? Make a b0 instead of h0 or use epsilon?
        w_C[:] = z_C + h_C
        u_C[:]  = d_C/(h_C + h0/(h_C + h0))

        for name in ['velocity', 'stage' ]:
            Q = self.quantities[name]
            if self.order == 1:
                Q.extrapolate_first_order()
            elif self.order == 2:
                Q.extrapolate_second_order()
            else:
                raise 'Unknown order'

        # These have been extrapolated
        w_V = self.quantities['stage'].vertex_values
        u_V = self.quantities['velocity'].vertex_values

        # These are the given geometry and remain fixed
        z_V  = bed.vertex_values
        b_V  = width.vertex_values
        t_V  = top.vertex_values

        # These need to be updated
        a_V  = area.vertex_values
        d_V  = discharge.vertex_values
        h_V  = height.vertex_values

        h_V[:,:]    = w_V - z_V

        # Fix up small heights
        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

        # It may still be possible that h is small 
        # If we set h to zero, we should also set u to 0
        h_V[:,:] = numpy.where(h_V < epsilon, 0.0, h_V)
        u_V[:,:] = numpy.where(h_V < epsilon, 0.0, u_V)

        # Reconstruct conserved quantities and make everything consistent
        # with new h values
        w_V[:,:] = z_V + h_V
        a_V[:,:] = h_V * b_V
        d_V[:,:] = u_V * a_V


        return