source: inundation/ga/storm_surge/pyvolution/quantity.py @ 209

"""Class Quantity - Implements values at each triangular element

To create:

   Quantity(domain, vertex_values)

   domain: Associated domain structure. Required.
   
   vertex_values: N x 3 array of values at each vertex for each element.
                  Default None

   If vertex_values are None Create array of zeros compatible with domain.
   Otherwise check that it is compatible with dimenions of domain.
   Otherwise raise an exception
"""

#FIXME: Make Conserved_quantity a subclass of Quantity 

from mesh import Mesh

class Quantity:

    def __init__(self, domain, vertex_values=None):

        from Numeric import array, zeros, Float

        msg = 'First argument in Quantity.__init__ '
        msg += 'must be of class Mesh (or a subclass thereof)'
        assert isinstance(domain, Mesh), msg

        if vertex_values is None:
            N = domain.number_of_elements            
            self.vertex_values = zeros((N, 3), Float)
        else:    
            self.vertex_values = array(vertex_values)

            N, V = self.vertex_values.shape
            assert V == 3,\
                   'Three vertex values per element must be specified'
            

            msg = 'Number of vertex values (%d) must be consistent with'\
                  %N
            msg += 'number of elements in specified domain (%d).'\
                   %domain.number_of_elements
            
            assert N == domain.number_of_elements, msg

        self.domain = domain    

        #Allocate space for other quantities
        self.centroid_values = zeros(N, Float)
        self.edge_values = zeros((N, 3), Float)        

        #Allocate space for boundary values
        L = len(domain.boundary)
        self.boundary_values = zeros(L, Float)

        #Allocate space for updates of conserved quantities by
        #flux calculations and forcing functions
        self.explicit_update = zeros(N, Float )
        self.semi_implicit_update = zeros(N, Float )
                
        #Intialise centroid and edge_values
        self.interpolate()


    def interpolate(self):
        """Compute interpolated values at edges and centroid
        Pre-condition: vertex_values have been set
        """

        #FIXME: Write using vector operations (or in C)
        N = self.vertex_values.shape[0]        
        for i in range(N):
            v0 = self.vertex_values[i, 0]
            v1 = self.vertex_values[i, 1]
            v2 = self.vertex_values[i, 2]
            
            self.centroid_values[i] = (v0 + v1 + v2)/3
            
            self.edge_values[i, 0] = 0.5*(v1 + v2)
            self.edge_values[i, 1] = 0.5*(v0 + v2)            
            self.edge_values[i, 2] = 0.5*(v0 + v1)


    def update(self, timestep):
        """Update centroid values based on values stored in
        explicit_update and semi_implicit_update as well as given timestep
        """

        from Numeric import sum, equal, ones, Float
        
        N = self.centroid_values.shape[0]
        
        #Explicit updates
        self.centroid_values += timestep*self.explicit_update
            
        #Semi implicit updates
        denominator = ones(N, Float)-timestep*self.semi_implicit_update

        if sum(equal(denominator, 0.0)) > 0.0:
            msg = 'Zero division in semi implicit update. Call Stephen :-)'
            raise msg
        else:
            #Update conserved_quantities from semi implicit updates
            self.centroid_values /= denominator


    def compute_gradients(self):
        """Compute gradients of triangle surfaces defined by centroids of
        neighbouring volumes.
        If one face is on the boundary, use own centroid as neighbour centroid.
        If two or more are on the boundary, fall back to first order scheme.
        
        Also return minimum and maximum of conserved quantities 
        """

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

        N = self.centroid_values.shape[0]

        a = zeros(N, Float)
        b = zeros(N, Float)
        
        for k in range(N):
        
            number_of_boundaries = self.domain.number_of_boundaries[k]

            if number_of_boundaries == 3:                 
                #We have zero neighbouring volumes -
                #Fall back to first order scheme

                pass
            
            elif number_of_boundaries == 2:
                #Special case where we have only one neighbouring volume.

                k0 = k  #Self
                #Find index of the one neighbour
                for k1 in self.domain.neighbours[k,:]:
                    if k1 >= 0:
                        break
                assert k1 != k0
                assert k1 >= 0

                #Get data
                q0 = self.centroid_values[k0]
                q1 = self.centroid_values[k1]

                x0, y0 = self.domain.centroids[k0] #V0 centroid
                x1, y1 = self.domain.centroids[k1] #V1 centroid        

                #Gradient
                det = x0*y1 - x1*y0
                if det != 0.0:
                    a[k] = (y1*q0 - y0*q1)/det
                    b[k] = (x0*q1 - x1*q0)/det


            else:
                #One or zero missing neighbours
                #In case of one boundary - own centroid
                #has been inserted as a surrogate for the one
                #missing neighbour in neighbour_surrogates

                #Get data        
                k0 = self.domain.surrogate_neighbours[k,0]
                k1 = self.domain.surrogate_neighbours[k,1]
                k2 = self.domain.surrogate_neighbours[k,2]

                q0 = self.centroid_values[k0]
                q1 = self.centroid_values[k1]
                q2 = self.centroid_values[k2]                    

                x0, y0 = self.domain.centroids[k0] #V0 centroid
                x1, y1 = self.domain.centroids[k1] #V1 centroid
                x2, y2 = self.domain.centroids[k2] #V2 centroid

                #Gradient
                a[k], b[k] = gradient(x0, y0, x1, y1, x2, y2, q0, q1, q2)
        
        return a, b
        

    def limit(self):
        """Limit slopes for each volume to eliminate artificial variance
        introduced by e.g. second order extrapolator
        
        This is an unsophisticated limiter as it does not take into
        account dependencies among quantities.
        
        precondition:
           vertex values are estimated from gradient
        postcondition:
           vertex values are updated
        """

        from Numeric import zeros, Float

        beta = self.domain.beta
        
        qc = self.centroid_values
        qv = self.vertex_values
        
        #Deltas between vertex and centroid values
        dq = zeros( qv.shape, Float)
        for i in range(3):
            dq[:,i] = qv[:,i] - qc
        
        #Find min and max at this and neighbour's centroids
        qmax = zeros(qc.shape, Float)
        qmin = zeros(qc.shape, Float)        

        for k in range(self.domain.number_of_elements):
            qmax[k] = qmin[k] = qc[k]
            for i in range(3):
                n = self.domain.neighbours[k,i]
                if n >= 0:
                    qn = qc[n] #Neighbour's centroid value
      
                    qmin[k] = min(qmin[k], qn)
                    qmax[k] = max(qmax[k], qn)
      
    
        #Diffences between centroids and maxima/minima
        dqmax = qmax - qc 
        dqmin = qmin - qc
  
        for k in range(self.domain.number_of_elements):

            #Find the gradient limiter (phi) across vertices
            phi = 1.0
            for i in range(3):
                r = 1.0
                if (dq[k,i] > 0): r = dqmax[k]/dq[k,i]
                if (dq[k,i] < 0): r = dqmin[k]/dq[k,i]
                
                phi = min( min(r*beta, 1), phi )    

            #Then update using phi limiter
            for i in range(3):            
                qv[k,i] = qc[k] + phi*dq[k,i]

        
    def extrapolate_first_order(self):
        """Extrapolate conserved quantities from centroid to
        vertices for each volume using
        first order scheme.
        """
        
        qc = self.centroid_values
        qv = self.vertex_values

        for i in range(3):
            qv[:,i] = qc
            
            
    def extrapolate_second_order(self):
        """Extrapolate conserved quantities from centroid to
        vertices for each volume using
        second order scheme.
        """
        
        a, b = self.compute_gradients()

        V = self.domain.get_vertex_coordinates()
        qc = self.centroid_values
        qv = self.vertex_values
        
        #Check each triangle
        for k in range(self.domain.number_of_elements):
            #Centroid coordinates            
            x, y = self.domain.centroids[k]

            #vertex coordinates
            x0, y0, x1, y1, x2, y2 = V[k,:]

            #Extrapolate
            qv[k,0] = qc[k] + a[k]*(x0-x) + b[k]*(y0-y)
            qv[k,1] = qc[k] + a[k]*(x1-x) + b[k]*(y1-y)
            qv[k,2] = qc[k] + a[k]*(x2-x) + b[k]*(y2-y)            

            

    def interpolate_from_vertices_to_edges(self):
        #FIXME: Write using vector operations (or in C)
        for k in range(self.vertex_values.shape[0]):
            q0 = self.vertex_values[k, 0]
            q1 = self.vertex_values[k, 1]
            q2 = self.vertex_values[k, 2]
            
            self.edge_values[k, 0] = 0.5*(q1+q2)
            self.edge_values[k, 1] = 0.5*(q0+q2)  
            self.edge_values[k, 2] = 0.5*(q0+q1)

            
            
    def set_values(self, X, location='vertices'):
        """Set values for quantity

        X: Compatible list, Numeric array (see below), constant or function
        location: Where values are to be stored.
                  Permissible options are: vertices, edges, centroid
                  Default is "vertices"

        In case of location == 'centroid' the dimension values must
        be a list of a Numerical array of length N, N being the number
        of elements in the mesh. Otherwise it must be of dimension Nx3

        The values will be stored in elements following their
        internal ordering.

        If values are described a function, it will be evaluated at specified points

        If selected location is vertices, values for centroid and edges
        will be assigned interpolated values.
        In any other case, only values for the specified locations
        will be assigned and the others will be left undefined.
        """


        #FIXME: Should do location error check here only
        
        if location not in ['vertices', 'centroids', 'edges']:
            msg = 'Invalid location: %s' %location
            raise msg
        
        import types
        
        if callable(X):
            #Use function specific method
            self.set_function_values(X, location)            
        elif type(X) in [types.FloatType, types.IntType, types.LongType]:
            if location == 'centroids':
                self.centroid_values[:] = X
            elif location == 'edges':
                self.edge_values[:] = X
            else:
                self.vertex_values[:] = X                

        else:
            #Use array specific method
            self.set_array_values(X, location)

        if location == 'vertices':
            #Intialise centroid and edge_values
            self.interpolate()
            


    def set_function_values(self, f, location='vertices'):
        """Set values for quantity using specified function

        f: x, y -> z Function where x, y and z are arrays
        location: Where values are to be stored.
                  Permissible options are: vertices, edges, centroid
                  Default is "vertices"
        """

        if location == 'centroids':
            P = self.domain.centroids
            self.set_values(f(P[:,0], P[:,1]), location)
        elif location == 'edges':
            raise 'Not yet implemented: %s' %location
        else:
            #Vertices
            P = self.domain.get_vertex_coordinates()
            for i in range(3):
                self.vertex_values[:,i] = f(P[:,2*i], P[:,2*i+1])

            
    def set_array_values(self, values, location='vertices'):
        """Set values for quantity

        values: Numeric array
        location: Where values are to be stored.
                  Permissible options are: vertices, edges, centroid
                  Default is "vertices"

        In case of location == 'centroid' the dimension values must
        be a list of a Numerical array of length N, N being the number
        of elements in the mesh. Otherwise it must be of dimension Nx3

        The values will be stored in elements following their
        internal ordering.

        If selected location is vertices, values for centroid and edges
        will be assigned interpolated values.
        In any other case, only values for the specified locations
        will be assigned and the others will be left undefined.
        """

        from Numeric import array, Float

        values = array(values).astype(Float)

        N = self.centroid_values.shape[0]
        
        msg = 'Number of values must match number of elements'
        assert values.shape[0] == N, msg
        
        if location == 'centroids':
            assert len(values.shape) == 1, 'Values array must be 1d'
            self.centroid_values = values                        
        elif location == 'edges':
            assert len(values.shape) == 2, 'Values array must be 2d'
            msg = 'Array must be N x 3'            
            assert values.shape[1] == 3, msg
            
            self.edge_values = values
        else:
            assert len(values.shape) == 2, 'Values array must be 2d'
            msg = 'Array must be N x 3'            
            assert values.shape[1] == 3, msg
            
            self.vertex_values = values