def parabolic_solve(self, u_in, b, a = None, u_out = None, update_matrix=True, \
                       imax=10000, tol=1.0e-8, atol=1.0e-8,
                       iprint=None, output_stats=False):
        """
        Solve for u in the equation

        ( I + dt div a grad ) u = b

        u | boundary = g


        u_in, u_out, f anf g are Quantity objects

        Dirichlet BC g encoded into u_in boundary_values

        Initial guess for iterative scheme is given by
        centroid values of u_in

        Centroid values of a and b provide diffusivity and rhs

        Solution u is retruned in u_out

        """

        if u_out is None:
            u_out = Quantity(self.domain)

        if update_matrix :
            self.update_elliptic_matrix(a)

        self.update_elliptic_boundary_term(u_in)

        self.set_parabolic_solve(True)


        # Pull out arrays and a matrix operator
        IdtA = self
        rhs = b.centroid_values + (self.dt * self.boundary_term)
        x0 = u_in.centroid_values

        x, stats = conjugate_gradient(IdtA,rhs,x0,imax=imax, tol=tol, atol=atol,
                                      iprint=iprint, output_stats=True)

        self.set_parabolic_solve(False)

        u_out.set_values(x, location='centroids')
        u_out.set_boundary_values(u_in.boundary_values)

        if output_stats:
            return u_out, stats
        else:
            return u_out
Exemple #2
0
    def _parabolic_multiply_quantity(self, quantity_in, quantity_out):
        """
        Assumes that update_elliptic_matrix has been run
        """

        if quantity_out is None:
            quantity_out = Quantity(self.domain)

        array_in = quantity_in.centroid_values
        array_out = quantity_out.centroid_values

        X = self._parabolic_multiply_array(array_in, array_out)

        quantity_out.set_values(X, location='centroids')

        return quantity_out
    def _parabolic_multiply_quantity(self, quantity_in, quantity_out):
        """
        Assumes that update_elliptic_matrix has been run
        """

        if quantity_out is None:
            quantity_out = Quantity(self.domain)

        array_in = quantity_in.centroid_values
        array_out = quantity_out.centroid_values

        X = self._parabolic_multiply_array(array_in, array_out)

        quantity_out.set_values(X, location = 'centroids')

        return quantity_out
    def __init__(self, domain, use_diffusivity=True,
                 indices=None,
                 description = None,
                 label = None,
                 logging = False,
                 verbose = False):
                
        Operator.__init__(self, domain, description, label, logging, verbose)
        
        
        self.use_diffusivity = use_diffusivity
            
        self.xmom = self.domain.quantities['xmomentum'].centroid_values
        self.ymom = self.domain.quantities['ymomentum'].centroid_values
        self.depth = self.domain.quantities['height'].centroid_values
        
        self.num_cells = len(self.depth)
        
        if self.use_diffusivity:
        
            self.diffusivity = num.zeros((self.num_cells,))
            self.mix_length = num.zeros((self.num_cells,))
            
            try:
                diff = self.domain.get_quantity('diffusivity')
            except:
                Quantity(domain, name='diffusivity', register=True)
            
        self.domain.set_use_kinematic_viscosity(self.use_diffusivity)

        self.quantity_flag = False
class Collect_max_stage_operator(Operator):
    """
    Simple operator to collect the max stage during a run

    """
    def __init__(self,
                 domain,
                 description=None,
                 label=None,
                 logging=False,
                 verbose=False):

        Operator.__init__(self, domain, description, label, logging, verbose)

        #------------------------------------------
        # Setup a quantity to store max_stage
        #------------------------------------------
        self.max_stage = Quantity(domain, name='max_stage', register=True)
        self.max_stage.set_values(-1.0e+100)

        #------------------------------------------
        # Aliases for stage quantity
        #------------------------------------------
        self.stage = domain.quantities['stage']

    def __call__(self):
        """
        Calculate max_stage at each timestep
        """

        self.max_stage.maximum(self.stage)

    def parallel_safe(self):
        """Operator is applied independently on each cell and
        so is parallel safe.
        """
        return True

    def statistics(self):

        message = self.label + ': Collect_max_stage operator'
        return message

    def timestepping_statistics(self):
        from anuga import indent

        message = indent + self.label + ': Collecting_max_stage'
        return message

    def save_centroid_data_to_csv(self, filename=None):

        self.max_stage.save_centroid_data_to_csv(filename)

    def plot_quantity(self, filename=None, show=True):

        self.max_stage.plot_quantity(filename=filename, show=show)
    def __init__(self,
                 domain,
                 description=None,
                 label=None,
                 logging=False,
                 verbose=False):

        Operator.__init__(self, domain, description, label, logging, verbose)

        #------------------------------------------
        # Setup a quantity to store max_stage
        #------------------------------------------
        self.max_stage = Quantity(domain, name='max_stage', register=True)
        self.max_stage.set_values(-1.0e+100)

        #------------------------------------------
        # Aliases for stage quantity
        #------------------------------------------
        self.stage = domain.quantities['stage']
    def update_elliptic_matrix(self, a=None):
        """
        Updates the data values of matrix representing

        div ( a grad )

        If a is None then we set a = quantity which is set to 1
        """

        #Array self.operator_data is changed by this call, which should flow
        # through to the Sparse_CSR matrix.

        if a is None:
            a = Quantity(self.domain)
            a.set_values(1.0)
            a.set_boundary_values(1.0)
            
        kinematic_viscosity_operator_ext.update_elliptic_matrix(self, \
                a.centroid_values, \
                a.boundary_values)
Exemple #8
0
    def parabolic_solve(self, u_in, b, a = None, u_out = None, update_matrix=True, \
                       imax=10000, tol=1.0e-8, atol=1.0e-8,
                       iprint=None, output_stats=False):
        """
        Solve for u in the equation

        ( I + dt div a grad ) u = b

        u | boundary = g


        u_in, u_out, f anf g are Quantity objects

        Dirichlet BC g encoded into u_in boundary_values

        Initial guess for iterative scheme is given by
        centroid values of u_in

        Centroid values of a and b provide diffusivity and rhs

        Solution u is retruned in u_out

        """

        if u_out is None:
            u_out = Quantity(self.domain)

        if update_matrix:
            self.update_elliptic_matrix(a)

        self.update_elliptic_boundary_term(u_in)

        self.set_parabolic_solve(True)

        # Pull out arrays and a matrix operator
        IdtA = self
        rhs = b.centroid_values + (self.dt * self.boundary_term)
        x0 = u_in.centroid_values

        x, stats = conjugate_gradient(IdtA,
                                      rhs,
                                      x0,
                                      imax=imax,
                                      tol=tol,
                                      atol=atol,
                                      iprint=iprint,
                                      output_stats=True)

        self.set_parabolic_solve(False)

        u_out.set_values(x, location='centroids')
        u_out.set_boundary_values(u_in.boundary_values)

        if output_stats:
            return u_out, stats
        else:
            return u_out
 def set_veg_quantity(self, name_in, quantity_name=None, convert_file=True, save_file=False, load_interp=False):
     """
     Read raster file and sets a vegetation quantity.
 
     Veg quantities are: veg_diameter
                         veg_spacing
                         
     The values in the rasters should be in meters.                  
     Assumes all NODATA values in raster are zero vegetation
     
     Saves a (Nx3) .npy file of x,y,val of non-zero points on save_file
     """
                             
     Quantity(self.domain, name=quantity_name, register=True)
     points = None
 
     if load_interp:
         print 'Veg Load: loading', name_in + '_interp.npy' 
         z_ = num.load(name_in + '_interp.npy')
 
     else:
     
         if convert_file:
             print 'Veg Load: converting asc file', name_in + '.asc'
             self.generic_asc2dem(name_in + '.asc', quantity_name=quantity_name)
             points = self.generic_dem2npy(name_in + '.dem', quantity_name=quantity_name)
     
         if points is None:
             print 'Veg Load: loading', name_in + '.npy' 
             points = num.load(name_in + '.npy')
         
         
         interp = NearestNDInterpolator(points[:,0:2], points[:,2])
     
         coord = self.domain.get_centroid_coordinates(absolute=True)
         z_ = interp( coord )
         
         if save_file:
             print 'Veg Load: saving interpolated file: ', name_in + '_interp.npy'
             num.save(name_in + '_interp.npy', z_)
     
     print 'Veg Load: setting quantity', quantity_name
     self.domain.quantities[quantity_name].set_values(z_, location = 'centroids') 
    def __init__(self,
                 domain,
                 description = None,
                 label = None,
                 logging = False,
                 verbose = False):


        Operator.__init__(self, domain, description, label, logging, verbose)

        #------------------------------------------
        # Setup a quantity to store max_stage
        #------------------------------------------
        self.max_stage = Quantity(domain, name = 'max_stage', register=True)
        self.max_stage.set_values(-1.0e+100)

        #------------------------------------------
        # Aliases for stage quantity
        #------------------------------------------
        self.stage  = domain.quantities['stage']
Exemple #11
0
    def update_elliptic_matrix(self, a=None):
        """
        Updates the data values of matrix representing

        div ( a grad )

        If a is None then we set a = quantity which is set to 1
        """

        #Array self.operator_data is changed by this call, which should flow
        # through to the Sparse_CSR matrix.

        if a is None:
            a = Quantity(self.domain)
            a.set_values(1.0)
            a.set_boundary_values(1.0)

        kinematic_viscosity_operator_ext.update_elliptic_matrix(self, \
                a.centroid_values, \
                a.boundary_values)
Exemple #12
0
class Elliptic_operator(Operator):
    """
    Class for setting up structures and matrices for elliptic differential
    operator using centroid values.

    div ( diffusivity grad )

    where diffusvity is scalar quantity (defaults to quantity with values = 1)
    boundary values of f are used to setup entries associated with cells with boundaries

    There are procedures to apply this operator, ie

    (1) Calculate div( diffusivity grad u )
    using boundary values stored in u

    (2) Calculate ( u + dt div( diffusivity grad u )
    using boundary values stored in u

    (3) Solve div( diffusivity grad u ) = f
    for quantity f and using boundary values stored in u

    (4) Solve ( u + dt div( diffusivity grad u ) = f
    for quantity f using boundary values stored in u

    As an anuga operator, when the __call__ method is called one step of the parabolic
    step is applied. In particular the x and y velocities are updated using



    """
    def __init__(self, domain, use_triangle_areas=True, verbose=False):
        if verbose:
            log.critical('Kinematic Viscosity: Beginning Initialisation')

        Operator.__init__(self, domain)

        #Expose the domain attributes
        self.mesh = self.domain.mesh
        self.boundary = domain.boundary
        self.boundary_enumeration = domain.boundary_enumeration

        # Setup a quantity as diffusivity
        # FIXME SR: Could/Should pass a quantity which already exists
        self.diffusivity = Quantity(self.domain)
        self.diffusivity.set_values(1.0)
        self.diffusivity.set_boundary_values(1.0)

        self.n = len(self.domain)

        self.dt = 0.0  #Need to set to domain.timestep
        self.dt_apply = 0.0

        self.boundary_len = len(self.domain.boundary)
        self.tot_len = self.n + self.boundary_len

        self.verbose = verbose

        #Geometric Information
        if verbose:
            log.critical('Kinematic Viscosity: Building geometric structure')

        self.geo_structure_indices = num.zeros((self.n, 3), num.int)
        self.geo_structure_values = num.zeros((self.n, 3), num.float)

        # Only needs to built once, doesn't change
        kinematic_viscosity_operator_ext.build_geo_structure(self)

        # Setup type of scaling
        self.set_triangle_areas(use_triangle_areas)

        # FIXME SR: should this really be a matrix?
        temp = Sparse(self.n, self.n)
        for i in range(self.n):
            temp[i, i] = 1.0 / self.mesh.areas[i]

        self.triangle_areas = Sparse_CSR(temp)
        #self.triangle_areas

        # FIXME SR: More to do with solving equation
        self.qty_considered = 1  #1 or 2 (uh or vh respectively)

        #Sparse_CSR.data
        self.operator_data = num.zeros((4 * self.n, ), num.float)
        #Sparse_CSR.colind
        self.operator_colind = num.zeros((4 * self.n, ), num.int)
        #Sparse_CSR.rowptr (4 entries in every row, we know this already) = [0,4,8,...,4*n]
        self.operator_rowptr = 4 * num.arange(self.n + 1)

        # Build matrix self.elliptic_matrix [A B]
        self.build_elliptic_matrix(self.diffusivity)

        self.boundary_term = num.zeros((self.n, ), num.float)

        self.parabolic = False  #Are we doing a parabolic solve at the moment?

        self.u_stats = None
        self.v_stats = None

        if verbose: log.critical('Elliptic Operator: Initialisation Done')

    def statistics(self):

        message = 'Elliptic_operator '
        return message

    def timestepping_statistics(self):

        message = '    Kinematic Viscosity Operator: '
        if self.u_stats is not None:
            message += ' u iterations %.5g, ' % self.u_stats.iter

        if self.v_stats is not None:
            message += ' v iterations %.5g, ' % self.v_stats.iter

        return message

    def set_triangle_areas(self, flag=True):

        self.apply_triangle_areas = flag

    def set_parabolic_solve(self, flag):

        self.parabolic = flag

    def build_elliptic_matrix(self, a):
        """
        Builds matrix representing

        div ( a grad )

        which has the form [ A B ]
        """

        #Arrays self.operator_data, self.operator_colind, self.operator_rowptr
        # are setup via this call
        kinematic_viscosity_operator_ext.build_elliptic_matrix(self, \
                a.centroid_values, \
                a.boundary_values)

        self.elliptic_matrix = Sparse_CSR(None, \
                self.operator_data, self.operator_colind, self.operator_rowptr, \
                self.n, self.tot_len)

    def update_elliptic_matrix(self, a=None):
        """
        Updates the data values of matrix representing

        div ( a grad )

        If a is None then we set a = quantity which is set to 1
        """

        #Array self.operator_data is changed by this call, which should flow
        # through to the Sparse_CSR matrix.

        if a is None:
            a = Quantity(self.domain)
            a.set_values(1.0)
            a.set_boundary_values(1.0)

        kinematic_viscosity_operator_ext.update_elliptic_matrix(self, \
                a.centroid_values, \
                a.boundary_values)

    def update_elliptic_boundary_term(self, boundary):

        if isinstance(boundary, Quantity):

            self._update_elliptic_boundary_term(boundary.boundary_values)

        elif isinstance(boundary, num.ndarray):

            self._update_elliptic_boundary_term(boundary.boundary_values)

        else:

            raise TypeError('expecting quantity or numpy array')

    def _update_elliptic_boundary_term(self, b):
        """
        Operator has form [A B] and u = [ u_1 ; b]

        u_1 associated with centroid values of u
        u_2 associated with boundary_values of u

        This procedure calculates B u_2 which can be calculated as

        [A B] [ 0 ; b]

        Assumes that update_elliptic_matrix has just been run.
        """

        n = self.n
        tot_len = self.tot_len

        X = num.zeros((tot_len, ), num.float)

        X[n:] = b
        self.boundary_term[:] = self.elliptic_matrix * X

        #Tidy up
        if self.apply_triangle_areas:
            self.boundary_term[:] = self.triangle_areas * self.boundary_term

    def elliptic_multiply(self, input, output=None):

        if isinstance(input, Quantity):

            assert isinstance(output, Quantity) or output is None

            output = self._elliptic_multiply_quantity(input, output)

        elif isinstance(input, num.ndarray):

            assert isinstance(output, num.ndarray) or output is None

            output = self._elliptic_multiply_array(input, output)

        else:

            raise TypeError('expecting quantity or numpy array')

        return output

    def _elliptic_multiply_quantity(self, quantity_in, quantity_out):
        """
        Assumes that update_elliptic_matrix has been run
        """

        if quantity_out is None:
            quantity_out = Quantity(self.domain)

        array_in = quantity_in.centroid_values
        array_out = quantity_out.centroid_values

        X = self._elliptic_multiply_array(array_in, array_out)

        quantity_out.set_values(X, location='centroids')

        return quantity_out

    def _elliptic_multiply_array(self, array_in, array_out):
        """
        calculates [A B] [array_in ; 0]
        """

        n = self.n
        tot_len = self.tot_len

        V = num.zeros((tot_len, ), num.float)

        assert len(array_in) == n

        if array_out is None:
            array_out = num.zeros_like(array_in)

        V[0:n] = array_in
        V[n:] = 0.0

        if self.apply_triangle_areas:
            V[0:n] = self.triangle_areas * V[0:n]

        array_out[:] = self.elliptic_matrix * V

        return array_out

    def parabolic_multiply(self, input, output=None):

        if isinstance(input, Quantity):

            assert isinstance(output, Quantity) or output is None

            output = self._parabolic_multiply_quantity(input, output)

        elif isinstance(input, num.ndarray):

            assert isinstance(output, num.ndarray) or output is None

            output = self._parabolic_multiply_array(input, output)

        else:

            raise TypeError('expecting quantity or numpy array')

        return output

    def _parabolic_multiply_quantity(self, quantity_in, quantity_out):
        """
        Assumes that update_elliptic_matrix has been run
        """

        if quantity_out is None:
            quantity_out = Quantity(self.domain)

        array_in = quantity_in.centroid_values
        array_out = quantity_out.centroid_values

        X = self._parabolic_multiply_array(array_in, array_out)

        quantity_out.set_values(X, location='centroids')

        return quantity_out

    def _parabolic_multiply_array(self, array_in, array_out):
        """
        calculates ( [ I 0 ; 0  0] + dt [A B] ) [array_in ; 0]
        """

        n = self.n
        tot_len = self.tot_len

        V = num.zeros((tot_len, ), num.float)

        assert len(array_in) == n

        if array_out is None:
            array_out = num.zeros_like(array_in)

        V[0:n] = array_in
        V[n:] = 0.0

        if self.apply_triangle_areas:
            V[0:n] = self.triangle_areas * V[0:n]

        array_out[:] = array_in - self.dt * (self.elliptic_matrix * V)

        return array_out

    def __mul__(self, vector):

        #Vector
        if self.parabolic:
            R = self.parabolic_multiply(vector)
        else:
            #include_boundary=False is this is *only* used for cg_solve()
            R = self.elliptic_multiply(vector)

        return R

    def __rmul__(self, other):
        #Right multiply with scalar
        try:
            other = float(other)
        except:
            msg = 'Sparse matrix can only "right-multiply" onto a scalar'
            raise TypeError(msg)
        else:
            new = self.elliptic_matrix * new
        return new


    def elliptic_solve(self, u_in, b, a = None, u_out = None, update_matrix=True, \
                       imax=10000, tol=1.0e-8, atol=1.0e-8,
                       iprint=None, output_stats=False):
        """ Solving div ( a grad u ) = b
        u | boundary = g

        u_in, u_out, f anf g are Quantity objects

        Dirichlet BC g encoded into u_in boundary_values

        Initial guess for iterative scheme is given by
        centroid values of u_in

        Centroid values of a and b provide diffusivity and rhs

        Solution u is retruned in u_out
        """

        if u_out is None:
            u_out = Quantity(self.domain)

        if update_matrix:
            self.update_elliptic_matrix(a)

        self.update_elliptic_boundary_term(u_in)

        # Pull out arrays and a matrix operator
        A = self
        rhs = b.centroid_values - self.boundary_term
        x0 = u_in.centroid_values

        x, stats = conjugate_gradient(A,
                                      rhs,
                                      x0,
                                      imax=imax,
                                      tol=tol,
                                      atol=atol,
                                      iprint=iprint,
                                      output_stats=True)

        u_out.set_values(x, location='centroids')
        u_out.set_boundary_values(u_in.boundary_values)

        if output_stats:
            return u_out, stats
        else:
            return u_out



    def parabolic_solve(self, u_in, b, a = None, u_out = None, update_matrix=True, \
                       imax=10000, tol=1.0e-8, atol=1.0e-8,
                       iprint=None, output_stats=False):
        """
        Solve for u in the equation

        ( I + dt div a grad ) u = b

        u | boundary = g


        u_in, u_out, f anf g are Quantity objects

        Dirichlet BC g encoded into u_in boundary_values

        Initial guess for iterative scheme is given by
        centroid values of u_in

        Centroid values of a and b provide diffusivity and rhs

        Solution u is retruned in u_out

        """

        if u_out is None:
            u_out = Quantity(self.domain)

        if update_matrix:
            self.update_elliptic_matrix(a)

        self.update_elliptic_boundary_term(u_in)

        self.set_parabolic_solve(True)

        # Pull out arrays and a matrix operator
        IdtA = self
        rhs = b.centroid_values + (self.dt * self.boundary_term)
        x0 = u_in.centroid_values

        x, stats = conjugate_gradient(IdtA,
                                      rhs,
                                      x0,
                                      imax=imax,
                                      tol=tol,
                                      atol=atol,
                                      iprint=iprint,
                                      output_stats=True)

        self.set_parabolic_solve(False)

        u_out.set_values(x, location='centroids')
        u_out.set_boundary_values(u_in.boundary_values)

        if output_stats:
            return u_out, stats
        else:
            return u_out
def distribute(domain, verbose=False, debug=False, parameters=None):
    """ Distribute the domain to all processes

    parameters allows user to change size of ghost layer
    """

    if not pypar_available or numprocs == 1: return domain  # Bypass

    if myid == 0:
        from sequential_distribute import Sequential_distribute
        partition = Sequential_distribute(domain, verbose, debug, parameters)

        partition.distribute(numprocs)

        kwargs, points, vertices, boundary, quantities, boundary_map, \
                domain_name, domain_dir, domain_store, domain_store_centroids, \
                domain_minimum_storable_height, domain_minimum_allowed_height, \
                domain_flow_algorithm, domain_georef, \
                domain_quantities_to_be_stored, domain_smooth \
                 = partition.extract_submesh(0)

        for p in range(1, numprocs):

            tostore = partition.extract_submesh(p)

            send(tostore, p)

    else:

        kwargs, points, vertices, boundary, quantities, boundary_map, \
            domain_name, domain_dir, domain_store, domain_store_centroids, \
            domain_minimum_storable_height, domain_minimum_allowed_height, \
            domain_flow_algorithm, domain_georef, \
            domain_quantities_to_be_stored, domain_smooth \
             = receive(0)

    #---------------------------------------------------------------------------
    # Now Create parallel domain
    #---------------------------------------------------------------------------
    parallel_domain = Parallel_domain(points, vertices, boundary, **kwargs)

    #------------------------------------------------------------------------
    # Copy in quantity data
    #------------------------------------------------------------------------
    for q in quantities:
        try:
            parallel_domain.set_quantity(q, quantities[q])
        except KeyError:
            #print 'Try to create quantity %s'% q
            from anuga import Quantity
            Q = Quantity(parallel_domain, name=q, register=True)
            parallel_domain.set_quantity(q, quantities[q])

    #------------------------------------------------------------------------
    # Transfer boundary conditions to each subdomain
    #------------------------------------------------------------------------
    boundary_map['ghost'] = None  # Add binding to ghost boundary
    parallel_domain.set_boundary(boundary_map)

    #------------------------------------------------------------------------
    # Transfer other attributes to each subdomain
    #------------------------------------------------------------------------

    parallel_domain.set_flow_algorithm(domain_flow_algorithm)
    parallel_domain.set_name(domain_name)
    parallel_domain.set_datadir(domain_dir)
    parallel_domain.set_store(domain_store)
    parallel_domain.set_store_centroids(domain_store_centroids)
    parallel_domain.set_minimum_storable_height(domain_minimum_storable_height)
    parallel_domain.set_minimum_allowed_height(domain_minimum_allowed_height)
    parallel_domain.geo_reference = domain_georef
    parallel_domain.set_quantities_to_be_stored(domain_quantities_to_be_stored)
    parallel_domain.smooth = domain_smooth

    return parallel_domain
    def test_rate_operator_rate_quantity(self):
        from anuga.config import rho_a, rho_w, eta_w
        from math import pi, cos, sin

        a = [0.0, 0.0]
        b = [0.0, 2.0]
        c = [2.0, 0.0]
        d = [0.0, 4.0]
        e = [2.0, 2.0]
        f = [4.0, 0.0]

        points = [a, b, c, d, e, f]
        #             bac,     bce,     ecf,     dbe
        vertices = [[1, 0, 2], [1, 2, 4], [4, 2, 5], [3, 1, 4]]

        domain = Domain(points, vertices)

        #Flat surface with 1m of water
        domain.set_quantity('elevation', 0.0)
        domain.set_quantity('stage', 1.0)
        domain.set_quantity('friction', 0.0)

        Br = Reflective_boundary(domain)
        domain.set_boundary({'exterior': Br})

        verbose = False

        if verbose:
            print domain.quantities['elevation'].centroid_values
            print domain.quantities['stage'].centroid_values
            print domain.quantities['xmomentum'].centroid_values
            print domain.quantities['ymomentum'].centroid_values

        # Apply operator to these triangles
        indices = [0, 1, 3]
        factor = 10.0

        from anuga import Quantity
        rate_Q = Quantity(domain)
        rate_Q.set_values(1.0)

        operator = Rate_operator(domain, rate=rate_Q, factor=factor, \
                                 indices=indices)

        # Apply Operator
        domain.timestep = 2.0
        operator()
        rate = rate_Q.centroid_values[indices]
        t = operator.get_time()
        Q = operator.get_Q()

        rate = rate * factor
        Q_ex = num.sum(domain.areas[indices] * rate)
        d = operator.get_timestep() * rate + 1

        #print "d"
        #print d
        #print Q_ex
        #print Q
        stage_ex = num.array([1.0, 1.0, 1.0, 1.0])
        stage_ex[indices] = d

        verbose = False

        if verbose:
            print domain.quantities['elevation'].centroid_values
            print domain.quantities['stage'].centroid_values
            print domain.quantities['xmomentum'].centroid_values
            print domain.quantities['ymomentum'].centroid_values

        assert num.allclose(domain.quantities['stage'].centroid_values,
                            stage_ex)
        assert num.allclose(domain.quantities['xmomentum'].centroid_values,
                            0.0)
        assert num.allclose(domain.quantities['ymomentum'].centroid_values,
                            0.0)
        assert num.allclose(Q_ex, Q)
        assert num.allclose(domain.fractional_step_volume_integral,
                            ((d - 1.) * domain.areas[indices]).sum())
class Kinematic_viscosity_operator(Operator):
    """
    Class for setting up structures and matrices for kinematic viscosity differential
    operator using centroid values.


    As an anuga operator, when the __call__ method is called one step of the parabolic
    step is applied. In particular the x and y velocities are updated using

    du/dt = div( h grad u )
    dv/dt = div( h grad v )

    """
    def __init__(self,
                 domain,
                 diffusivity='height',
                 use_triangle_areas=True,
                 add_safety=False,
                 verbose=False):

        if verbose:
            log.critical('Kinematic Viscosity: Beginning Initialisation')

        Operator.__init__(self, domain)

        #Expose the domain attributes
        self.mesh = self.domain.mesh
        self.boundary = domain.boundary
        self.boundary_enumeration = domain.boundary_enumeration

        self.diffusivity = diffusivity
        # Setup a quantity as diffusivity
        if diffusivity is None:
            self.diffusivity = Quantity(self.domain)
            self.diffusivity.set_values(1.0)
            self.diffusivity.set_boundary_values(1.0)

        if isinstance(diffusivity, (int, float)):
            self.diffusivity = Quantity(self.domain)
            self.diffusivity.set_values(diffusivity)
            self.diffusivity.set_boundary_values(diffusivity)

        if isinstance(diffusivity, str):
            self.diffusivity = self.domain.get_quantity(diffusivity)

        self.add_safety = add_safety
        self.smooth = 0.1

        assert isinstance(self.diffusivity, Quantity)

        self.n = len(self.domain)

        self.dt = 0.0  #Need to set to domain.timestep
        self.dt_apply = 0.0

        self.boundary_len = len(self.domain.boundary)
        self.tot_len = self.n + self.boundary_len

        self.verbose = verbose

        #Geometric Information
        if verbose:
            log.critical('Kinematic Viscosity: Building geometric structure')

        self.geo_structure_indices = num.zeros((self.n, 3), num.int)
        self.geo_structure_values = num.zeros((self.n, 3), num.float)

        # Only needs to built once, doesn't change
        kinematic_viscosity_operator_ext.build_geo_structure(self)

        # Setup type of scaling
        self.set_triangle_areas(use_triangle_areas)

        # FIXME SR: should this really be a matrix?
        temp = Sparse(self.n, self.n)
        for i in range(self.n):
            temp[i, i] = 1.0 / self.mesh.areas[i]

        self.triangle_areas = Sparse_CSR(temp)
        #self.triangle_areas

        # FIXME SR: More to do with solving equation
        self.qty_considered = 1  #1 or 2 (uh or vh respectively)

        #Sparse_CSR.data
        self.operator_data = num.zeros((4 * self.n, ), num.float)
        #Sparse_CSR.colind
        self.operator_colind = num.zeros((4 * self.n, ), num.int)
        #Sparse_CSR.rowptr (4 entries in every row, we know this already) = [0,4,8,...,4*n]
        self.operator_rowptr = 4 * num.arange(self.n + 1)

        # Build matrix self.elliptic_matrix [A B]
        self.build_elliptic_matrix(self.diffusivity)

        self.boundary_term = num.zeros((self.n, ), num.float)

        self.parabolic = False  #Are we doing a parabolic solve at the moment?

        self.u_stats = None
        self.v_stats = None

        if verbose: log.critical('Kinematic Viscosity: Initialisation Done')

    def __call__(self):
        """ Parabolic update of x and y velocity

        (I + dt (div d grad) ) U^{n+1} = U^{n}

        """

        domain = self.domain

        self.dt = self.dt + domain.get_timestep()

        if self.dt < self.dt_apply:
            return

        # Setup initial values of velocity
        domain.update_centroids_of_velocities_and_height()

        # diffusivity
        if self.add_safety:
            d = self.diffusivity + 0.1
        else:
            d = self.diffusivity

        assert num.all(d.centroid_values >= 0.0)

        #d.set_values(num.where(h.centroid_values<1.0, 0.0, 1.0), location= 'centroids')
        #d.set_boundary_values(num.where(h.boundary_values<1.0, 0.0, 1.0))

        # Quantities to solve
        # Boundary values are implied from BC set in advection part of code
        u = domain.quantities['xvelocity']
        #u.set_boundary_values(0.0)

        v = domain.quantities['yvelocity']
        #v.set_boundary_values(0.0)

        #Update operator using current height
        self.update_elliptic_matrix(d)

        (u, self.u_stats) = self.parabolic_solve(u,
                                                 u,
                                                 d,
                                                 u_out=u,
                                                 update_matrix=False,
                                                 output_stats=True)

        (v, self.v_stats) = self.parabolic_solve(v,
                                                 v,
                                                 d,
                                                 u_out=v,
                                                 update_matrix=False,
                                                 output_stats=True)

        # Update the conserved quantities
        domain.update_centroids_of_momentum_from_velocity()

        self.dt = 0.0

    def statistics(self):

        message = 'Kinematic_viscosity_operator '
        return message

    def timestepping_statistics(self):

        from anuga import indent

        message = indent + 'Kinematic Viscosity Operator: \n'
        if self.u_stats is not None:
            message += indent + indent + 'u: ' + self.u_stats.__str__() + '\n'

        if self.v_stats is not None:
            message += indent + indent + 'v: ' + self.v_stats.__str__()

        return message

    def set_triangle_areas(self, flag=True):

        self.apply_triangle_areas = flag

    def set_parabolic_solve(self, flag):

        self.parabolic = flag

    def build_elliptic_matrix(self, a):
        """
        Builds matrix representing

        div ( a grad )

        which has the form [ A B ]
        """

        #Arrays self.operator_data, self.operator_colind, self.operator_rowptr
        # are setup via this call
        kinematic_viscosity_operator_ext.build_elliptic_matrix(self, \
                a.centroid_values, \
                a.boundary_values)

        self.elliptic_matrix = Sparse_CSR(None, \
                self.operator_data, self.operator_colind, self.operator_rowptr, \
                self.n, self.tot_len)

    def update_elliptic_matrix(self, a=None):
        """
        Updates the data values of matrix representing

        div ( a grad )

        If a == None then we set a = quantity which is set to 1
        """

        #Array self.operator_data is changed by this call, which should flow
        # through to the Sparse_CSR matrix.

        if a == None:
            a = Quantity(self.domain)
            a.set_values(1.0)
            a.set_boundary_values(1.0)

        kinematic_viscosity_operator_ext.update_elliptic_matrix(self, \
                a.centroid_values, \
                a.boundary_values)

    def update_elliptic_boundary_term(self, boundary):

        if isinstance(boundary, Quantity):

            self._update_elliptic_boundary_term(boundary.boundary_values)

        elif isinstance(boundary, num.ndarray):

            self._update_elliptic_boundary_term(boundary)

        else:

            raise TypeError('expecting quantity or numpy array')

    def _update_elliptic_boundary_term(self, b):
        """
        Operator has form [A B] and u = [ u_1 ; b]

        u_1 associated with centroid values of u
        u_2 associated with boundary_values of u

        This procedure calculates B u_2 which can be calculated as

        [A B] [ 0 ; b]

        Assumes that update_elliptic_matrix has just been run.
        """

        n = self.n
        tot_len = self.tot_len

        X = num.zeros((tot_len, ), num.float)

        X[n:] = b
        self.boundary_term[:] = self.elliptic_matrix * X

        #Tidy up
        if self.apply_triangle_areas:
            self.boundary_term[:] = self.triangle_areas * self.boundary_term

    def elliptic_multiply(self, input, output=None):

        if isinstance(input, Quantity):

            assert isinstance(output, Quantity) or output is None

            output = self._elliptic_multiply_quantity(input, output)

        elif isinstance(input, num.ndarray):

            assert isinstance(output, num.ndarray) or output is None

            output = self._elliptic_multiply_array(input, output)

        else:

            raise TypeError('expecting quantity or numpy array')

        return output

    def _elliptic_multiply_quantity(self, quantity_in, quantity_out):
        """
        Assumes that update_elliptic_matrix has been run
        """

        if quantity_out is None:
            quantity_out = Quantity(self.domain)

        array_in = quantity_in.centroid_values
        array_out = quantity_out.centroid_values

        X = self._elliptic_multiply_array(array_in, array_out)

        quantity_out.set_values(X, location='centroids')

        return quantity_out

    def _elliptic_multiply_array(self, array_in, array_out):
        """
        calculates [A B] [array_in ; 0]
        """

        n = self.n
        tot_len = self.tot_len

        V = num.zeros((tot_len, ), num.float)

        assert len(array_in) == n

        if array_out is None:
            array_out = num.zeros_like(array_in)

        V[0:n] = array_in
        V[n:] = 0.0

        if self.apply_triangle_areas:
            V[0:n] = self.triangle_areas * V[0:n]

        array_out[:] = self.elliptic_matrix * V

        return array_out

    def parabolic_multiply(self, input, output=None):

        if isinstance(input, Quantity):

            assert isinstance(output, Quantity) or output is None

            output = self._parabolic_multiply_quantity(input, output)

        elif isinstance(input, num.ndarray):

            assert isinstance(output, num.ndarray) or output is None

            output = self._parabolic_multiply_array(input, output)

        else:

            raise TypeError('expecting quantity or numpy array')

        return output

    def _parabolic_multiply_quantity(self, quantity_in, quantity_out):
        """
        Assumes that update_elliptic_matrix has been run
        """

        if quantity_out is None:
            quantity_out = Quantity(self.domain)

        array_in = quantity_in.centroid_values
        array_out = quantity_out.centroid_values

        X = self._parabolic_multiply_array(array_in, array_out)

        quantity_out.set_values(X, location='centroids')

        return quantity_out

    def _parabolic_multiply_array(self, array_in, array_out):
        """
        calculates ( [ I 0 ; 0  0] + dt [A B] ) [array_in ; 0]
        """

        n = self.n
        tot_len = self.tot_len

        V = num.zeros((tot_len, ), num.float)

        assert len(array_in) == n

        if array_out is None:
            array_out = num.zeros_like(array_in)

        V[0:n] = array_in
        V[n:] = 0.0

        if self.apply_triangle_areas:
            V[0:n] = self.triangle_areas * V[0:n]

        array_out[:] = array_in - self.dt * (self.elliptic_matrix * V)

        return array_out

    def __mul__(self, vector):

        #Vector
        if self.parabolic:
            R = self.parabolic_multiply(vector)
        else:
            #include_boundary=False is this is *only* used for cg_solve()
            R = self.elliptic_multiply(vector)

        return R

    def __rmul__(self, other):
        #Right multiply with scalar
        try:
            other = float(other)
        except:
            msg = 'Sparse matrix can only "right-multiply" onto a scalar'
            raise TypeError(msg)
        else:
            new = self.elliptic_matrix * new
        return new


    def elliptic_solve(self, u_in, b, a = None, u_out = None, update_matrix=True, \
                       imax=10000, tol=1.0e-8, atol=1.0e-8,
                       iprint=None, output_stats=False):
        """ Solving div ( a grad u ) = b
        u | boundary = g

        u_in, u_out, f anf g are Quantity objects

        Dirichlet BC g encoded into u_in boundary_values

        Initial guess for iterative scheme is given by
        centroid values of u_in

        Centroid values of a and b provide diffusivity and rhs

        Solution u is retruned in u_out
        """

        if u_out == None:
            u_out = Quantity(self.domain)

        if update_matrix:
            self.update_elliptic_matrix(a)

        self.update_elliptic_boundary_term(u_in)

        # Pull out arrays and a matrix operator
        A = self
        rhs = b.centroid_values - self.boundary_term
        x0 = u_in.centroid_values

        x, stats = conjugate_gradient(A,
                                      rhs,
                                      x0,
                                      imax=imax,
                                      tol=tol,
                                      atol=atol,
                                      iprint=iprint,
                                      output_stats=True)

        u_out.set_values(x, location='centroids')
        u_out.set_boundary_values(u_in.boundary_values)

        if output_stats:
            return u_out, stats
        else:
            return u_out



    def parabolic_solve(self, u_in, b, a = None, u_out = None, update_matrix=True, \
                       output_stats=False, use_dt_tol=True, iprint=None, imax=10000):
        """
        Solve for u in the equation

        ( I + dt div a grad ) u = b

        u | boundary = g


        u_in, u_out, f anf g are Quantity objects

        Dirichlet BC g encoded into u_in boundary_values

        Initial guess for iterative scheme is given by
        centroid values of u_in

        Centroid values of a and b provide diffusivity and rhs

        Solution u is retruned in u_out

        """

        if use_dt_tol:
            tol = min(self.dt, 0.001)
            atol = min(self.dt, 0.001)
        else:
            tol = 1.0e-5
            atol = 1.0e-5

        if u_out == None:
            u_out = Quantity(self.domain)

        if update_matrix:
            self.update_elliptic_matrix(a)

        self.update_elliptic_boundary_term(u_in)

        self.set_parabolic_solve(True)

        # Pull out arrays and a matrix operator
        IdtA = self
        rhs = b.centroid_values + (self.dt * self.boundary_term)
        x0 = u_in.centroid_values

        x, stats = conjugate_gradient(IdtA,
                                      rhs,
                                      x0,
                                      imax=imax,
                                      tol=tol,
                                      atol=atol,
                                      iprint=iprint,
                                      output_stats=True)

        self.set_parabolic_solve(False)

        u_out.set_values(x, location='centroids')
        u_out.set_boundary_values(u_in.boundary_values)

        if output_stats:
            return u_out, stats
        else:
            return u_out
    def __init__(self, domain, use_triangle_areas=True, verbose=False):
        if verbose: log.critical('Kinematic Viscosity: Beginning Initialisation')
        

        Operator.__init__(self,domain)

        #Expose the domain attributes
        self.mesh = self.domain.mesh
        self.boundary = domain.boundary
        self.boundary_enumeration = domain.boundary_enumeration
        
        # Setup a quantity as diffusivity
        # FIXME SR: Could/Should pass a quantity which already exists
        self.diffusivity = Quantity(self.domain)
        self.diffusivity.set_values(1.0)
        self.diffusivity.set_boundary_values(1.0)
        

        self.n = len(self.domain)

        self.dt = 0.0 #Need to set to domain.timestep
        self.dt_apply = 0.0

        self.boundary_len = len(self.domain.boundary)
        self.tot_len = self.n + self.boundary_len

        self.verbose = verbose

        #Geometric Information
        if verbose: log.critical('Kinematic Viscosity: Building geometric structure')

        self.geo_structure_indices = num.zeros((self.n, 3), num.int)
        self.geo_structure_values = num.zeros((self.n, 3), num.float)

        # Only needs to built once, doesn't change
        kinematic_viscosity_operator_ext.build_geo_structure(self)

        # Setup type of scaling
        self.set_triangle_areas(use_triangle_areas)        

        # FIXME SR: should this really be a matrix?
        temp  = Sparse(self.n, self.n)
        for i in range(self.n):
            temp[i, i] = 1.0 / self.mesh.areas[i]
            
        self.triangle_areas = Sparse_CSR(temp)
        #self.triangle_areas

        # FIXME SR: More to do with solving equation
        self.qty_considered = 1 #1 or 2 (uh or vh respectively)

        #Sparse_CSR.data
        self.operator_data = num.zeros((4 * self.n, ), num.float)
        #Sparse_CSR.colind
        self.operator_colind = num.zeros((4 * self.n, ), num.int)
        #Sparse_CSR.rowptr (4 entries in every row, we know this already) = [0,4,8,...,4*n]
        self.operator_rowptr = 4 * num.arange(self.n + 1)

        # Build matrix self.elliptic_matrix [A B]
        self.build_elliptic_matrix(self.diffusivity)

        self.boundary_term = num.zeros((self.n, ), num.float)

        self.parabolic = False #Are we doing a parabolic solve at the moment?

        self.u_stats = None
        self.v_stats = None

        if verbose: log.critical('Elliptic Operator: Initialisation Done')
class Elliptic_operator(Operator):
    """
    Class for setting up structures and matrices for elliptic differential
    operator using centroid values.

    div ( diffusivity grad )

    where diffusvity is scalar quantity (defaults to quantity with values = 1)
    boundary values of f are used to setup entries associated with cells with boundaries

    There are procedures to apply this operator, ie

    (1) Calculate div( diffusivity grad u )
    using boundary values stored in u

    (2) Calculate ( u + dt div( diffusivity grad u )
    using boundary values stored in u

    (3) Solve div( diffusivity grad u ) = f
    for quantity f and using boundary values stored in u

    (4) Solve ( u + dt div( diffusivity grad u ) = f
    for quantity f using boundary values stored in u

    As an anuga operator, when the __call__ method is called one step of the parabolic
    step is applied. In particular the x and y velocities are updated using



    """

    def __init__(self, domain, use_triangle_areas=True, verbose=False):
        if verbose: log.critical('Kinematic Viscosity: Beginning Initialisation')
        

        Operator.__init__(self,domain)

        #Expose the domain attributes
        self.mesh = self.domain.mesh
        self.boundary = domain.boundary
        self.boundary_enumeration = domain.boundary_enumeration
        
        # Setup a quantity as diffusivity
        # FIXME SR: Could/Should pass a quantity which already exists
        self.diffusivity = Quantity(self.domain)
        self.diffusivity.set_values(1.0)
        self.diffusivity.set_boundary_values(1.0)
        

        self.n = len(self.domain)

        self.dt = 0.0 #Need to set to domain.timestep
        self.dt_apply = 0.0

        self.boundary_len = len(self.domain.boundary)
        self.tot_len = self.n + self.boundary_len

        self.verbose = verbose

        #Geometric Information
        if verbose: log.critical('Kinematic Viscosity: Building geometric structure')

        self.geo_structure_indices = num.zeros((self.n, 3), num.int)
        self.geo_structure_values = num.zeros((self.n, 3), num.float)

        # Only needs to built once, doesn't change
        kinematic_viscosity_operator_ext.build_geo_structure(self)

        # Setup type of scaling
        self.set_triangle_areas(use_triangle_areas)        

        # FIXME SR: should this really be a matrix?
        temp  = Sparse(self.n, self.n)
        for i in range(self.n):
            temp[i, i] = 1.0 / self.mesh.areas[i]
            
        self.triangle_areas = Sparse_CSR(temp)
        #self.triangle_areas

        # FIXME SR: More to do with solving equation
        self.qty_considered = 1 #1 or 2 (uh or vh respectively)

        #Sparse_CSR.data
        self.operator_data = num.zeros((4 * self.n, ), num.float)
        #Sparse_CSR.colind
        self.operator_colind = num.zeros((4 * self.n, ), num.int)
        #Sparse_CSR.rowptr (4 entries in every row, we know this already) = [0,4,8,...,4*n]
        self.operator_rowptr = 4 * num.arange(self.n + 1)

        # Build matrix self.elliptic_matrix [A B]
        self.build_elliptic_matrix(self.diffusivity)

        self.boundary_term = num.zeros((self.n, ), num.float)

        self.parabolic = False #Are we doing a parabolic solve at the moment?

        self.u_stats = None
        self.v_stats = None

        if verbose: log.critical('Elliptic Operator: Initialisation Done')



    def statistics(self):

        message = 'Elliptic_operator '
        return message


    def timestepping_statistics(self):

        message = '    Kinematic Viscosity Operator: '
        if self.u_stats is not None:
            message  += ' u iterations %.5g, ' % self.u_stats.iter

        if self.v_stats is not None:
            message += ' v iterations %.5g, ' % self.v_stats.iter

        return message


    def set_triangle_areas(self,flag=True):

        self.apply_triangle_areas = flag
        

    def set_parabolic_solve(self,flag):

        self.parabolic = flag


    def build_elliptic_matrix(self, a):
        """
        Builds matrix representing

        div ( a grad )

        which has the form [ A B ]
        """

        #Arrays self.operator_data, self.operator_colind, self.operator_rowptr
        # are setup via this call
        kinematic_viscosity_operator_ext.build_elliptic_matrix(self, \
                a.centroid_values, \
                a.boundary_values)

        self.elliptic_matrix = Sparse_CSR(None, \
                self.operator_data, self.operator_colind, self.operator_rowptr, \
                self.n, self.tot_len)


    def update_elliptic_matrix(self, a=None):
        """
        Updates the data values of matrix representing

        div ( a grad )

        If a is None then we set a = quantity which is set to 1
        """

        #Array self.operator_data is changed by this call, which should flow
        # through to the Sparse_CSR matrix.

        if a is None:
            a = Quantity(self.domain)
            a.set_values(1.0)
            a.set_boundary_values(1.0)
            
        kinematic_viscosity_operator_ext.update_elliptic_matrix(self, \
                a.centroid_values, \
                a.boundary_values)
        




    def update_elliptic_boundary_term(self, boundary):


        if isinstance(boundary, Quantity):

            self._update_elliptic_boundary_term(boundary.boundary_values)

        elif isinstance(boundary, num.ndarray):

            self._update_elliptic_boundary_term(boundary.boundary_values)

        else:

            raise  TypeError('expecting quantity or numpy array')


    def _update_elliptic_boundary_term(self, b):
        """
        Operator has form [A B] and u = [ u_1 ; b]

        u_1 associated with centroid values of u
        u_2 associated with boundary_values of u

        This procedure calculates B u_2 which can be calculated as

        [A B] [ 0 ; b]

        Assumes that update_elliptic_matrix has just been run.
        """

        n = self.n
        tot_len = self.tot_len

        X = num.zeros((tot_len,), num.float)

        X[n:] = b
        self.boundary_term[:] = self.elliptic_matrix * X

        #Tidy up
        if self.apply_triangle_areas:
            self.boundary_term[:] = self.triangle_areas * self.boundary_term



    def elliptic_multiply(self, input, output=None):


        if isinstance(input, Quantity):

            assert isinstance(output, Quantity) or output is None

            output = self._elliptic_multiply_quantity(input, output)

        elif isinstance(input, num.ndarray):

            assert isinstance(output, num.ndarray) or output is None

            output = self._elliptic_multiply_array(input, output)

        else:

            raise TypeError('expecting quantity or numpy array')
        
        return output


    def _elliptic_multiply_quantity(self, quantity_in, quantity_out):
        """
        Assumes that update_elliptic_matrix has been run
        """

        if quantity_out is None:
            quantity_out = Quantity(self.domain)

        array_in = quantity_in.centroid_values
        array_out = quantity_out.centroid_values

        X = self._elliptic_multiply_array(array_in, array_out)

        quantity_out.set_values(X, location = 'centroids')
        
        return quantity_out

    def _elliptic_multiply_array(self, array_in, array_out):
        """
        calculates [A B] [array_in ; 0]
        """

        n = self.n
        tot_len = self.tot_len

        V = num.zeros((tot_len,), num.float)

        assert len(array_in) == n

        if array_out is None:
            array_out = num.zeros_like(array_in)

        V[0:n] = array_in
        V[n:] = 0.0


        if self.apply_triangle_areas:
            V[0:n] = self.triangle_areas * V[0:n]


        array_out[:] = self.elliptic_matrix * V


        return array_out





    def parabolic_multiply(self, input, output=None):


        if isinstance(input, Quantity):

            assert isinstance(output, Quantity) or output is None

            output = self._parabolic_multiply_quantity(input, output)

        elif isinstance(input, num.ndarray):

            assert isinstance(output, num.ndarray) or output is None

            output = self._parabolic_multiply_array(input, output)

        else:

            raise TypeError('expecting quantity or numpy array')

        return output


    def _parabolic_multiply_quantity(self, quantity_in, quantity_out):
        """
        Assumes that update_elliptic_matrix has been run
        """

        if quantity_out is None:
            quantity_out = Quantity(self.domain)

        array_in = quantity_in.centroid_values
        array_out = quantity_out.centroid_values

        X = self._parabolic_multiply_array(array_in, array_out)

        quantity_out.set_values(X, location = 'centroids')

        return quantity_out

    def _parabolic_multiply_array(self, array_in, array_out):
        """
        calculates ( [ I 0 ; 0  0] + dt [A B] ) [array_in ; 0]
        """

        n = self.n
        tot_len = self.tot_len

        V = num.zeros((tot_len,), num.float)

        assert len(array_in) == n

        if array_out is None:
            array_out = num.zeros_like(array_in)

        V[0:n] = array_in
        V[n:] = 0.0


        if self.apply_triangle_areas:
            V[0:n] = self.triangle_areas * V[0:n]


        array_out[:] = array_in - self.dt * (self.elliptic_matrix * V)

        return array_out




    def __mul__(self, vector):
        
        #Vector
        if self.parabolic:
            R = self.parabolic_multiply(vector)
        else:
            #include_boundary=False is this is *only* used for cg_solve()
            R = self.elliptic_multiply(vector)

        return R
    
    def __rmul__(self, other):
        #Right multiply with scalar
        try:
            other = float(other)
        except:
            msg = 'Sparse matrix can only "right-multiply" onto a scalar'
            raise TypeError(msg)
        else:
            new = self.elliptic_matrix * new
        return new

    
    def elliptic_solve(self, u_in, b, a = None, u_out = None, update_matrix=True, \
                       imax=10000, tol=1.0e-8, atol=1.0e-8,
                       iprint=None, output_stats=False):
        """ Solving div ( a grad u ) = b
        u | boundary = g

        u_in, u_out, f anf g are Quantity objects

        Dirichlet BC g encoded into u_in boundary_values

        Initial guess for iterative scheme is given by
        centroid values of u_in

        Centroid values of a and b provide diffusivity and rhs

        Solution u is retruned in u_out
        """

        if u_out is None:
            u_out = Quantity(self.domain)

        if update_matrix :
            self.update_elliptic_matrix(a) 

        self.update_elliptic_boundary_term(u_in)

        # Pull out arrays and a matrix operator
        A = self
        rhs = b.centroid_values - self.boundary_term
        x0 = u_in.centroid_values

        x, stats = conjugate_gradient(A,rhs,x0,imax=imax, tol=tol, atol=atol,
                               iprint=iprint, output_stats=True)

        u_out.set_values(x, location='centroids')
        u_out.set_boundary_values(u_in.boundary_values)

        if output_stats:
            return u_out, stats
        else:
            return u_out

    

    def parabolic_solve(self, u_in, b, a = None, u_out = None, update_matrix=True, \
                       imax=10000, tol=1.0e-8, atol=1.0e-8,
                       iprint=None, output_stats=False):
        """
        Solve for u in the equation

        ( I + dt div a grad ) u = b

        u | boundary = g


        u_in, u_out, f anf g are Quantity objects

        Dirichlet BC g encoded into u_in boundary_values

        Initial guess for iterative scheme is given by
        centroid values of u_in

        Centroid values of a and b provide diffusivity and rhs

        Solution u is retruned in u_out

        """

        if u_out is None:
            u_out = Quantity(self.domain)

        if update_matrix :
            self.update_elliptic_matrix(a)

        self.update_elliptic_boundary_term(u_in)

        self.set_parabolic_solve(True)


        # Pull out arrays and a matrix operator
        IdtA = self
        rhs = b.centroid_values + (self.dt * self.boundary_term)
        x0 = u_in.centroid_values

        x, stats = conjugate_gradient(IdtA,rhs,x0,imax=imax, tol=tol, atol=atol,
                                      iprint=iprint, output_stats=True)

        self.set_parabolic_solve(False)

        u_out.set_values(x, location='centroids')
        u_out.set_boundary_values(u_in.boundary_values)

        if output_stats:
            return u_out, stats
        else:
            return u_out
    def test_rate_operator_rate_quantity(self):
        from anuga.config import rho_a, rho_w, eta_w
        from math import pi, cos, sin

        a = [0.0, 0.0]
        b = [0.0, 2.0]
        c = [2.0, 0.0]
        d = [0.0, 4.0]
        e = [2.0, 2.0]
        f = [4.0, 0.0]

        points = [a, b, c, d, e, f]
        #             bac,     bce,     ecf,     dbe
        vertices = [[1,0,2], [1,2,4], [4,2,5], [3,1,4]]

        domain = Domain(points, vertices)

        #Flat surface with 1m of water
        domain.set_quantity('elevation', 0.0)
        domain.set_quantity('stage', 1.0)
        domain.set_quantity('friction', 0.0)

        Br = Reflective_boundary(domain)
        domain.set_boundary({'exterior': Br})

        verbose = False

        if verbose:
            print domain.quantities['elevation'].centroid_values
            print domain.quantities['stage'].centroid_values
            print domain.quantities['xmomentum'].centroid_values
            print domain.quantities['ymomentum'].centroid_values

        # Apply operator to these triangles
        indices = [0,1,3]
        factor = 10.0


        from anuga import Quantity
        rate_Q = Quantity(domain)
        rate_Q.set_values(1.0)

        operator = Rate_operator(domain, rate=rate_Q, factor=factor, \
                                 indices=indices)


        # Apply Operator
        domain.timestep = 2.0
        operator()
        rate = rate_Q.centroid_values[indices]
        t = operator.get_time()
        Q = operator.get_Q()

        rate = rate*factor
        Q_ex = num.sum(domain.areas[indices]*rate)
        d = operator.get_timestep()*rate + 1


        #print "d"
        #print d
        #print Q_ex
        #print Q
        stage_ex = num.array([ 1.0,  1.0,   1.0,  1.0])
        stage_ex[indices] = d
        
        verbose = False
        
        if verbose:
            print domain.quantities['elevation'].centroid_values
            print domain.quantities['stage'].centroid_values
            print domain.quantities['xmomentum'].centroid_values
            print domain.quantities['ymomentum'].centroid_values

        assert num.allclose(domain.quantities['stage'].centroid_values, stage_ex)
        assert num.allclose(domain.quantities['xmomentum'].centroid_values, 0.0)
        assert num.allclose(domain.quantities['ymomentum'].centroid_values, 0.0)
        assert num.allclose(Q_ex, Q)
        assert num.allclose(domain.fractional_step_volume_integral, ((d-1.)*domain.areas[indices]).sum())
class Kinematic_viscosity_operator(Operator):
    """
    Class for setting up structures and matrices for kinematic viscosity differential
    operator using centroid values.


    As an anuga operator, when the __call__ method is called one step of the parabolic
    step is applied. In particular the x and y velocities are updated using

    du/dt = div( h grad u )
    dv/dt = div( h grad v )

    """

    def __init__(self,
                 domain, diffusivity='height',
                 use_triangle_areas=True,
                 add_safety = False,
                 verbose=False):

        if verbose: log.critical('Kinematic Viscosity: Beginning Initialisation')
        

        Operator.__init__(self,domain)

        #Expose the domain attributes
        self.mesh = self.domain.mesh
        self.boundary = domain.boundary
        self.boundary_enumeration = domain.boundary_enumeration

        self.diffusivity = diffusivity
        # Setup a quantity as diffusivity
        if diffusivity is None:
            self.diffusivity = Quantity(self.domain)
            self.diffusivity.set_values(1.0)
            self.diffusivity.set_boundary_values(1.0)

        if isinstance(diffusivity, (int, float)):
            self.diffusivity = Quantity(self.domain)
            self.diffusivity.set_values(diffusivity)
            self.diffusivity.set_boundary_values(diffusivity)

        if isinstance(diffusivity, str):
            self.diffusivity = self.domain.get_quantity(diffusivity)
            

        self.add_safety = add_safety
        self.smooth = 0.1
        
        assert isinstance(self.diffusivity, Quantity)
        

        self.n = len(self.domain)

        self.dt = 0.0 #Need to set to domain.timestep
        self.dt_apply = 0.0

        self.boundary_len = len(self.domain.boundary)
        self.tot_len = self.n + self.boundary_len

        self.verbose = verbose

        #Geometric Information
        if verbose: log.critical('Kinematic Viscosity: Building geometric structure')

        self.geo_structure_indices = num.zeros((self.n, 3), num.int)
        self.geo_structure_values = num.zeros((self.n, 3), num.float)

        # Only needs to built once, doesn't change
        kinematic_viscosity_operator_ext.build_geo_structure(self)

        # Setup type of scaling
        self.set_triangle_areas(use_triangle_areas)        

        # FIXME SR: should this really be a matrix?
        temp  = Sparse(self.n, self.n)
        for i in range(self.n):
            temp[i, i] = 1.0 / self.mesh.areas[i]
            
        self.triangle_areas = Sparse_CSR(temp)
        #self.triangle_areas

        # FIXME SR: More to do with solving equation
        self.qty_considered = 1 #1 or 2 (uh or vh respectively)

        #Sparse_CSR.data
        self.operator_data = num.zeros((4 * self.n, ), num.float)
        #Sparse_CSR.colind
        self.operator_colind = num.zeros((4 * self.n, ), num.int)
        #Sparse_CSR.rowptr (4 entries in every row, we know this already) = [0,4,8,...,4*n]
        self.operator_rowptr = 4 * num.arange(self.n + 1)

        # Build matrix self.elliptic_matrix [A B]
        self.build_elliptic_matrix(self.diffusivity)

        self.boundary_term = num.zeros((self.n, ), num.float)

        self.parabolic = False #Are we doing a parabolic solve at the moment?

        self.u_stats = None
        self.v_stats = None

        if verbose: log.critical('Kinematic Viscosity: Initialisation Done')


    def __call__(self):
        """ Parabolic update of x and y velocity

        (I + dt (div d grad) ) U^{n+1} = U^{n}

        """

        domain = self.domain

        self.dt = self.dt + domain.get_timestep()


        if self.dt < self.dt_apply:
            return


        # Setup initial values of velocity
        domain.update_centroids_of_velocities_and_height()

        # diffusivity
        if self.add_safety:
            d = self.diffusivity + 0.1
        else:
            d = self.diffusivity


        assert num.all(d.centroid_values >= 0.0)

        #d.set_values(num.where(h.centroid_values<1.0, 0.0, 1.0), location= 'centroids')
        #d.set_boundary_values(num.where(h.boundary_values<1.0, 0.0, 1.0))

        # Quantities to solve
        # Boundary values are implied from BC set in advection part of code
        u = domain.quantities['xvelocity']
        #u.set_boundary_values(0.0)

        v = domain.quantities['yvelocity']
        #v.set_boundary_values(0.0)

        #Update operator using current height
        self.update_elliptic_matrix(d)

        (u, self.u_stats) = self.parabolic_solve(u, u, d, u_out=u, update_matrix=False, output_stats=True)

        (v, self.v_stats) = self.parabolic_solve(v, v, d, u_out=v, update_matrix=False, output_stats=True)

        # Update the conserved quantities
        domain.update_centroids_of_momentum_from_velocity()


        self.dt = 0.0



    def statistics(self):

        message = 'Kinematic_viscosity_operator '
        return message


    def timestepping_statistics(self):

        from anuga import indent

        message = indent+'Kinematic Viscosity Operator: \n'
        if self.u_stats is not None:
            message  += indent + indent + 'u: ' + self.u_stats.__str__() +'\n'

        if self.v_stats is not None:
            message  += indent + indent + 'v: ' + self.v_stats.__str__()

        return message


    def set_triangle_areas(self,flag=True):

        self.apply_triangle_areas = flag
        

    def set_parabolic_solve(self,flag):

        self.parabolic = flag


    def build_elliptic_matrix(self, a):
        """
        Builds matrix representing

        div ( a grad )

        which has the form [ A B ]
        """

        #Arrays self.operator_data, self.operator_colind, self.operator_rowptr
        # are setup via this call
        kinematic_viscosity_operator_ext.build_elliptic_matrix(self, \
                a.centroid_values, \
                a.boundary_values)

        self.elliptic_matrix = Sparse_CSR(None, \
                self.operator_data, self.operator_colind, self.operator_rowptr, \
                self.n, self.tot_len)


    def update_elliptic_matrix(self, a=None):
        """
        Updates the data values of matrix representing

        div ( a grad )

        If a == None then we set a = quantity which is set to 1
        """

        #Array self.operator_data is changed by this call, which should flow
        # through to the Sparse_CSR matrix.

        if a == None:
            a = Quantity(self.domain)
            a.set_values(1.0)
            a.set_boundary_values(1.0)
            
        kinematic_viscosity_operator_ext.update_elliptic_matrix(self, \
                a.centroid_values, \
                a.boundary_values)
        




    def update_elliptic_boundary_term(self, boundary):


        if isinstance(boundary, Quantity):

            self._update_elliptic_boundary_term(boundary.boundary_values)

        elif isinstance(boundary, num.ndarray):

            self._update_elliptic_boundary_term(boundary)

        else:

            raise  TypeError('expecting quantity or numpy array')


    def _update_elliptic_boundary_term(self, b):
        """
        Operator has form [A B] and u = [ u_1 ; b]

        u_1 associated with centroid values of u
        u_2 associated with boundary_values of u

        This procedure calculates B u_2 which can be calculated as

        [A B] [ 0 ; b]

        Assumes that update_elliptic_matrix has just been run.
        """

        n = self.n
        tot_len = self.tot_len

        X = num.zeros((tot_len,), num.float)

        X[n:] = b
        self.boundary_term[:] = self.elliptic_matrix * X

        #Tidy up
        if self.apply_triangle_areas:
            self.boundary_term[:] = self.triangle_areas * self.boundary_term



    def elliptic_multiply(self, input, output=None):


        if isinstance(input, Quantity):

            assert isinstance(output, Quantity) or output is None

            output = self._elliptic_multiply_quantity(input, output)

        elif isinstance(input, num.ndarray):

            assert isinstance(output, num.ndarray) or output is None

            output = self._elliptic_multiply_array(input, output)

        else:

            raise TypeError('expecting quantity or numpy array')
        
        return output


    def _elliptic_multiply_quantity(self, quantity_in, quantity_out):
        """
        Assumes that update_elliptic_matrix has been run
        """

        if quantity_out is None:
            quantity_out = Quantity(self.domain)

        array_in = quantity_in.centroid_values
        array_out = quantity_out.centroid_values

        X = self._elliptic_multiply_array(array_in, array_out)

        quantity_out.set_values(X, location = 'centroids')
        
        return quantity_out

    def _elliptic_multiply_array(self, array_in, array_out):
        """
        calculates [A B] [array_in ; 0]
        """

        n = self.n
        tot_len = self.tot_len

        V = num.zeros((tot_len,), num.float)

        assert len(array_in) == n

        if array_out is None:
            array_out = num.zeros_like(array_in)

        V[0:n] = array_in
        V[n:] = 0.0


        if self.apply_triangle_areas:
            V[0:n] = self.triangle_areas * V[0:n]


        array_out[:] = self.elliptic_matrix * V


        return array_out





    def parabolic_multiply(self, input, output=None):


        if isinstance(input, Quantity):

            assert isinstance(output, Quantity) or output is None

            output = self._parabolic_multiply_quantity(input, output)

        elif isinstance(input, num.ndarray):

            assert isinstance(output, num.ndarray) or output is None

            output = self._parabolic_multiply_array(input, output)

        else:

            raise TypeError('expecting quantity or numpy array')

        return output


    def _parabolic_multiply_quantity(self, quantity_in, quantity_out):
        """
        Assumes that update_elliptic_matrix has been run
        """

        if quantity_out is None:
            quantity_out = Quantity(self.domain)

        array_in = quantity_in.centroid_values
        array_out = quantity_out.centroid_values

        X = self._parabolic_multiply_array(array_in, array_out)

        quantity_out.set_values(X, location = 'centroids')

        return quantity_out

    def _parabolic_multiply_array(self, array_in, array_out):
        """
        calculates ( [ I 0 ; 0  0] + dt [A B] ) [array_in ; 0]
        """

        n = self.n
        tot_len = self.tot_len

        V = num.zeros((tot_len,), num.float)

        assert len(array_in) == n

        if array_out is None:
            array_out = num.zeros_like(array_in)

        V[0:n] = array_in
        V[n:] = 0.0


        if self.apply_triangle_areas:
            V[0:n] = self.triangle_areas * V[0:n]


        array_out[:] = array_in - self.dt * (self.elliptic_matrix * V)

        return array_out




    def __mul__(self, vector):
        
        #Vector
        if self.parabolic:
            R = self.parabolic_multiply(vector)
        else:
            #include_boundary=False is this is *only* used for cg_solve()
            R = self.elliptic_multiply(vector)

        return R
    
    def __rmul__(self, other):
        #Right multiply with scalar
        try:
            other = float(other)
        except:
            msg = 'Sparse matrix can only "right-multiply" onto a scalar'
            raise TypeError(msg)
        else:
            new = self.elliptic_matrix * new
        return new

    
    def elliptic_solve(self, u_in, b, a = None, u_out = None, update_matrix=True, \
                       imax=10000, tol=1.0e-8, atol=1.0e-8,
                       iprint=None, output_stats=False):
        """ Solving div ( a grad u ) = b
        u | boundary = g

        u_in, u_out, f anf g are Quantity objects

        Dirichlet BC g encoded into u_in boundary_values

        Initial guess for iterative scheme is given by
        centroid values of u_in

        Centroid values of a and b provide diffusivity and rhs

        Solution u is retruned in u_out
        """

        if u_out == None:
            u_out = Quantity(self.domain)

        if update_matrix :
            self.update_elliptic_matrix(a) 

        self.update_elliptic_boundary_term(u_in)

        # Pull out arrays and a matrix operator
        A = self
        rhs = b.centroid_values - self.boundary_term
        x0 = u_in.centroid_values

        x, stats = conjugate_gradient(A,rhs,x0,imax=imax, tol=tol, atol=atol,
                               iprint=iprint, output_stats=True)

        u_out.set_values(x, location='centroids')
        u_out.set_boundary_values(u_in.boundary_values)

        if output_stats:
            return u_out, stats
        else:
            return u_out

    

    def parabolic_solve(self, u_in, b, a = None, u_out = None, update_matrix=True, \
                       output_stats=False, use_dt_tol=True, iprint=None, imax=10000):
        """
        Solve for u in the equation

        ( I + dt div a grad ) u = b

        u | boundary = g


        u_in, u_out, f anf g are Quantity objects

        Dirichlet BC g encoded into u_in boundary_values

        Initial guess for iterative scheme is given by
        centroid values of u_in

        Centroid values of a and b provide diffusivity and rhs

        Solution u is retruned in u_out

        """


        if use_dt_tol:
            tol  = min(self.dt,0.001)
            atol = min(self.dt,0.001)
        else:
            tol  =  1.0e-5
            atol = 1.0e-5



        
        if u_out == None:
            u_out = Quantity(self.domain)

        if update_matrix :
            self.update_elliptic_matrix(a)

        self.update_elliptic_boundary_term(u_in)

        self.set_parabolic_solve(True)


        # Pull out arrays and a matrix operator
        IdtA = self
        rhs = b.centroid_values + (self.dt * self.boundary_term)
        x0 = u_in.centroid_values

        x, stats = conjugate_gradient(IdtA,rhs,x0,imax=imax, tol=tol, atol=atol,
                                      iprint=iprint, output_stats=True)

        self.set_parabolic_solve(False)

        u_out.set_values(x, location='centroids')
        u_out.set_boundary_values(u_in.boundary_values)

        if output_stats:
            return u_out, stats
        else:
            return u_out
Exemple #20
0
    def __init__(self, domain, use_triangle_areas=True, verbose=False):
        if verbose:
            log.critical('Kinematic Viscosity: Beginning Initialisation')

        Operator.__init__(self, domain)

        #Expose the domain attributes
        self.mesh = self.domain.mesh
        self.boundary = domain.boundary
        self.boundary_enumeration = domain.boundary_enumeration

        # Setup a quantity as diffusivity
        # FIXME SR: Could/Should pass a quantity which already exists
        self.diffusivity = Quantity(self.domain)
        self.diffusivity.set_values(1.0)
        self.diffusivity.set_boundary_values(1.0)

        self.n = len(self.domain)

        self.dt = 0.0  #Need to set to domain.timestep
        self.dt_apply = 0.0

        self.boundary_len = len(self.domain.boundary)
        self.tot_len = self.n + self.boundary_len

        self.verbose = verbose

        #Geometric Information
        if verbose:
            log.critical('Kinematic Viscosity: Building geometric structure')

        self.geo_structure_indices = num.zeros((self.n, 3), num.int)
        self.geo_structure_values = num.zeros((self.n, 3), num.float)

        # Only needs to built once, doesn't change
        kinematic_viscosity_operator_ext.build_geo_structure(self)

        # Setup type of scaling
        self.set_triangle_areas(use_triangle_areas)

        # FIXME SR: should this really be a matrix?
        temp = Sparse(self.n, self.n)
        for i in range(self.n):
            temp[i, i] = 1.0 / self.mesh.areas[i]

        self.triangle_areas = Sparse_CSR(temp)
        #self.triangle_areas

        # FIXME SR: More to do with solving equation
        self.qty_considered = 1  #1 or 2 (uh or vh respectively)

        #Sparse_CSR.data
        self.operator_data = num.zeros((4 * self.n, ), num.float)
        #Sparse_CSR.colind
        self.operator_colind = num.zeros((4 * self.n, ), num.int)
        #Sparse_CSR.rowptr (4 entries in every row, we know this already) = [0,4,8,...,4*n]
        self.operator_rowptr = 4 * num.arange(self.n + 1)

        # Build matrix self.elliptic_matrix [A B]
        self.build_elliptic_matrix(self.diffusivity)

        self.boundary_term = num.zeros((self.n, ), num.float)

        self.parabolic = False  #Are we doing a parabolic solve at the moment?

        self.u_stats = None
        self.v_stats = None

        if verbose: log.critical('Elliptic Operator: Initialisation Done')
class Collect_max_stage_operator(Operator):
    """
    Simple operator to collect the max stage during a run

    """

    def __init__(self,
                 domain,
                 description = None,
                 label = None,
                 logging = False,
                 verbose = False):


        Operator.__init__(self, domain, description, label, logging, verbose)

        #------------------------------------------
        # Setup a quantity to store max_stage
        #------------------------------------------
        self.max_stage = Quantity(domain, name = 'max_stage', register=True)
        self.max_stage.set_values(-1.0e+100)

        #------------------------------------------
        # Aliases for stage quantity
        #------------------------------------------
        self.stage  = domain.quantities['stage']

        

    def __call__(self):
        """
        Calculate max_stage at each timestep
        """

        self.max_stage.maximum(self.stage)


    def parallel_safe(self):
        """Operator is applied independently on each cell and
        so is parallel safe.
        """
        return True

    def statistics(self):

        message = self.label + ': Collect_max_stage operator'
        return message


    def timestepping_statistics(self):
        from anuga import indent

        message  = indent + self.label + ': Collecting_max_stage'
        return message


    def save_centroid_data_to_csv(self, filename=None):

        self.max_stage.save_centroid_data_to_csv(filename)


    def plot_quantity(self, filename=None, show=True):

        self.max_stage.plot_quantity(filename=filename, show=show)