Source code for firedrake.matrix_free.operators

from collections import OrderedDict
import itertools

from mpi4py import MPI
import numpy

from pyop2.mpi import internal_comm, temp_internal_comm
from firedrake.ufl_expr import adjoint, action
from firedrake.formmanipulation import ExtractSubBlock
from firedrake.bcs import DirichletBC, EquationBCSplit
from firedrake.petsc import PETSc
from firedrake.utils import cached_property


__all__ = ("ImplicitMatrixContext", )


@PETSc.Log.EventDecorator()
def find_sub_block(iset, ises, comm):
    """Determine if iset comes from a concatenation of some subset of
    ises.

    :arg iset: a PETSc IS to find in ``ises``.
    :arg ises: An iterable of PETSc ISes.
    :arg comm: User comm used to construct function space

    :returns: The indices into ``ises`` that when concatenated
        together produces ``iset``.

    :raises LookupError: if ``iset`` could not be found in
        ``ises``.
    """
    found = []
    target_indices = iset.indices
    candidates = OrderedDict(enumerate(ises))
    with temp_internal_comm(comm) as icomm:
        while True:
            match = False
            for i, candidate in list(candidates.items()):
                candidate_indices = candidate.indices
                candidate_size, = candidate_indices.shape
                target_size, = target_indices.shape
                # Does the local part of the candidate IS match a prefix
                # of the target indices?
                lmatch = (candidate_size <= target_size
                          and numpy.array_equal(target_indices[:candidate_size], candidate_indices))
                if icomm.allreduce(lmatch, op=MPI.LAND):
                    # Yes, this candidate matched, so remove it from the
                    # target indices, and list of candidate
                    target_indices = target_indices[candidate_size:]
                    found.append(i)
                    candidates.pop(i)
                    # And keep looking for the remainder in the remaining candidates.
                    match = True
            if not match:
                break
        if icomm.allreduce(len(target_indices), op=MPI.SUM) > 0:
            # We didn't manage to hoover up all the target indices, not a match
            raise LookupError("Unable to find %s in %s" % (iset, ises))
    return found


[docs] class ImplicitMatrixContext(object): # By default, these matrices will represent diagonal blocks (the # (0,0) block of a 1x1 block matrix is on the diagonal). on_diag = True """This class gives the Python context for a PETSc Python matrix. :arg a: The bilinear form defining the matrix :arg row_bcs: An iterable of the :class.`.DirichletBC`s that are imposed on the test space. We distinguish between row and column boundary conditions in the case of submatrices off of the diagonal. :arg col_bcs: An iterable of the :class.`.DirichletBC`s that are imposed on the trial space. :arg fcparams: A dictionary of parameters to pass on to the form compiler. :arg appctx: Any extra user-supplied context, available to preconditioners and the like. """ @PETSc.Log.EventDecorator() def __init__(self, a, row_bcs=[], col_bcs=[], fc_params=None, appctx=None): from firedrake.assemble import get_assembler self.a = a self.aT = adjoint(a) self.comm = a.arguments()[0].function_space().comm self._comm = internal_comm(self.comm, self) self.fc_params = fc_params self.appctx = appctx # Collect all DirichletBC instances including # DirichletBCs applied to an EquationBC. # all bcs (DirichletBC, EquationBCSplit) self.bcs = row_bcs self.bcs_col = col_bcs self.row_bcs = tuple(bc for bc in itertools.chain(*row_bcs) if isinstance(bc, DirichletBC)) self.col_bcs = tuple(bc for bc in itertools.chain(*col_bcs) if isinstance(bc, DirichletBC)) # create functions from test and trial space to help # with 1-form assembly test_space, trial_space = [ a.arguments()[i].function_space() for i in (0, 1) ] from firedrake import function, cofunction # Need a cofunction since y receives the assembled result of Ax self._ystar = cofunction.Cofunction(test_space.dual()) self._y = function.Function(test_space) self._x = function.Function(trial_space) self._xstar = cofunction.Cofunction(trial_space.dual()) # These are temporary storage for holding the BC # values during matvec application. _xbc is for # the action and ._ybc is for transpose. if len(self.bcs) > 0: self._xbc = cofunction.Cofunction(trial_space.dual()) if len(self.col_bcs) > 0: self._ybc = cofunction.Cofunction(test_space.dual()) # Get size information from template vecs on test and trial spaces trial_vec = trial_space.dof_dset.layout_vec test_vec = test_space.dof_dset.layout_vec self.col_sizes = trial_vec.getSizes() self.row_sizes = test_vec.getSizes() self.block_size = (test_vec.getBlockSize(), trial_vec.getBlockSize()) self.action = action(self.a, self._x) self.actionT = action(self.aT, self._y) # For assembling action(f, self._x) self.bcs_action = [] for bc in self.bcs: if isinstance(bc, DirichletBC): self.bcs_action.append(bc) elif isinstance(bc, EquationBCSplit): self.bcs_action.append(bc.reconstruct(action_x=self._x)) self._assemble_action = get_assembler(self.action, bcs=self.bcs_action, form_compiler_parameters=self.fc_params, zero_bc_nodes=True).assemble # For assembling action(adjoint(f), self._y) # Sorted list of equation bcs self.objs_actionT = [] for bc in self.bcs: self.objs_actionT += bc.sorted_equation_bcs() self.objs_actionT.append(self) # Each par_loop is to run with appropriate masks on self._y self._assemble_actionT = [] # Deepest EquationBCs first for bc in self.bcs: for ebc in bc.sorted_equation_bcs(): self._assemble_actionT.append( get_assembler(action(adjoint(ebc.f), self._y), form_compiler_parameters=self.fc_params).assemble) # Domain last self._assemble_actionT.append( get_assembler(self.actionT, form_compiler_parameters=self.fc_params).assemble) @cached_property def _diagonal(self): from firedrake import Cofunction assert self.on_diag return Cofunction(self._x.function_space().dual()) @cached_property def _assemble_diagonal(self): from firedrake.assemble import get_assembler return get_assembler(self.a, form_compiler_parameters=self.fc_params, diagonal=True).assemble
[docs] def getDiagonal(self, mat, vec): self._assemble_diagonal(tensor=self._diagonal) diagonal_func = self._diagonal.riesz_representation(riesz_map="l2") for bc in self.bcs: # Operator is identity on boundary nodes bc.set(diagonal_func, 1) self._diagonal.assign(diagonal_func.riesz_representation(riesz_map="l2")) with self._diagonal.dat.vec_ro as v: v.copy(vec)
[docs] def missingDiagonal(self, mat): return (False, -1)
[docs] @PETSc.Log.EventDecorator() def mult(self, mat, X, Y): with self._x.dat.vec_wo as v: X.copy(v) # if we are a block on the diagonal, then the matrix has an # identity block corresponding to the Dirichlet boundary conditions. # our algorithm in this case is to save the BC values, zero them # out before computing the action so that they don't pollute # anything, and then set the values into the result. # This has the effect of applying # [ A_II 0 ; 0 I ] where A_II is the block corresponding only to # non-fixed dofs and I is the identity block on the fixed dofs. # If we are not, then the matrix just has 0s in the rows and columns. for bc in self.col_bcs: bc.zero(self._x) self._assemble_action(tensor=self._ystar) # This sets the essential boundary condition values on the # result. if self.on_diag: if len(self.row_bcs) > 0: # TODO, can we avoid the copy? with self._xbc.dat.vec_wo as v: X.copy(v) for bc in self.row_bcs: bc.set(self._ystar, self._xbc) else: for bc in self.row_bcs: bc.zero(self._ystar) with self._ystar.dat.vec_ro as v: v.copy(Y)
[docs] @PETSc.Log.EventDecorator() def multTranspose(self, mat, Y, X): """ EquationBC makes multTranspose different from mult. Decompose M^T into bundles of columns associated with the rows of M corresponding to cell, facet, edge, and vertice equations (if exist) and add up their contributions. .. code-block:: text Domain ( a a a a 0 a a ) | ( a a a a 0 a a ) | ( a a a a 0 a a ) | EBC1 M = ( b b b b b b b ) | | EBC2 DBC1 ( 0 0 0 0 1 0 0 ) | | | | ( c c c c 0 c c ) | | ( c c c c 0 c c ) | | Multiplication algorithm: To avoid copys, use same y, and update it from left (deepest ebc) to right (least deep ebc or domain). * below can be any number ( a a a b 0 c c ) ( y0 ) ( a a a b 0 c c ) ( y1 ) ( a a a b 0 c c ) ( y2 ) M^T y = ( a a a b 0 c c ) ( y3 ) ( 0 0 0 0 1 0 0 ) ( y4 ) ( a a a b 0 c c ) ( y5 ) ( a a a b 0 c c ) ( y6 ) ( 0 0 0 0 c c c ) ( * ) Matrix is uniform ( 0 0 0 0 c c c ) ( * ) on facet2 (EBC2) ( 0 0 0 0 c c c ) ( * ) = ( 0 0 0 0 c c c ) ( * ) Initial y ( 0 0 0 0 c c c ) ( 0 ) ( 0 0 0 0 c c c ) ( y5 ) ( 0 0 0 0 c c c ) ( y6 ) ( 0 0 0 b b 0 0 ) ( * ) Matrix is uniform ( 0 0 0 b b 0 0 ) ( * ) on facet1 (EBC1) ( 0 0 0 b b 0 0 ) ( * ) + ( 0 0 0 b b 0 0 ) ( y3 ) Update y ( 0 0 0 b b 0 0 ) ( 0 ) ( 0 0 0 b b 0 0 ) ( * ) ( 0 0 0 b b 0 0 ) ( * ) ( a a a a a a a ) ( y0 ) Matrix is uniform ( a a a a a a a ) ( y1 ) on domain ( a a a a a a a ) ( y2 ) + ( a a a a a a a ) ( 0 ) Update y ( a a a a a a a ) ( 0 ) ( a a a a a a a ) ( 0 ) ( a a a a a a a ) ( 0 ) ( 0 ) ( 0 ) Update y replace at the end (DBC1) ( 0 ) + ( 0 ) ( y4 ) ( 0 ) ( 0 ) """ with self._y.dat.vec_wo as v: Y.copy(v) if len(self.bcs) > 0: # Accumulate values in self._x self._xstar.dat.zero() # Apply actionTs in sorted order for aT, obj in zip(self._assemble_actionT, self.objs_actionT): # zero columns associated with DirichletBCs/EquationBCs for obc in obj.bcs: obc.zero(self._y) aT(tensor=self._xbc) self._xstar += self._xbc else: # No DirichletBC/EquationBC # There is only a single element in the list (for the domain equation). # Save to self._x directly aT, = self._assemble_actionT aT(tensor=self._xstar) if self.on_diag: if len(self.col_bcs) > 0: # TODO, can we avoid the copy? with self._ybc.dat.vec_wo as v: Y.copy(v) for bc in self.col_bcs: bc.set(self._xstar, self._ybc) else: for bc in self.col_bcs: bc.zero(self._xstar) with self._xstar.dat.vec_ro as v: v.copy(X)
[docs] def view(self, mat, viewer=None): if viewer is None: return typ = viewer.getType() if typ != PETSc.Viewer.Type.ASCII: return viewer.printfASCII("Firedrake matrix-free operator %s\n" % type(self).__name__)
[docs] def getInfo(self, mat, info=None): memory = self._x.dat.nbytes + self._y.dat.nbytes if hasattr(self, "_xbc"): memory += self._xbc.dat.nbytes if hasattr(self, "_ybc"): memory += self._ybc.dat.nbytes if info is None: info = PETSc.Mat.InfoType.GLOBAL_SUM if info == PETSc.Mat.InfoType.LOCAL: return {"memory": memory} elif info == PETSc.Mat.InfoType.GLOBAL_SUM: gmem = self._comm.allreduce(memory, op=MPI.SUM) return {"memory": gmem} elif info == PETSc.Mat.InfoType.GLOBAL_MAX: gmem = self._comm.allreduce(memory, op=MPI.MAX) return {"memory": gmem} else: raise ValueError("Unknown info type %s" % info)
# Now, to enable fieldsplit preconditioners, we need to enable submatrix # extraction for our custom matrix type. Note that we are splitting UFL # and index sets rather than an assembled matrix, keeping matrix # assembly deferred as long as possible.
[docs] @PETSc.Log.EventDecorator() def createSubMatrix(self, mat, row_is, col_is, target=None): if target is not None: # Repeat call, just return the matrix, since we don't # actually assemble in here. target.assemble() return target # These are the sets of ISes of which the the row and column # space consist. row_ises = self._y.function_space().dof_dset.field_ises col_ises = self._x.function_space().dof_dset.field_ises row_inds = find_sub_block(row_is, row_ises, comm=self.comm) if row_is == col_is and row_ises == col_ises: col_inds = row_inds else: col_inds = find_sub_block(col_is, col_ises, comm=self.comm) splitter = ExtractSubBlock() asub = splitter.split(self.a, argument_indices=(row_inds, col_inds)) Wrow = asub.arguments()[0].function_space() Wcol = asub.arguments()[1].function_space() row_bcs = [] col_bcs = [] for bc in self.bcs: if isinstance(bc, DirichletBC): bc_temp = bc.reconstruct(field=row_inds, V=Wrow, g=bc.function_arg, sub_domain=bc.sub_domain, use_split=True) elif isinstance(bc, EquationBCSplit): bc_temp = bc.reconstruct(field=row_inds, V=Wrow, row_field=row_inds, col_field=col_inds, use_split=True) if bc_temp is not None: row_bcs.append(bc_temp) if Wrow == Wcol and row_inds == col_inds and self.bcs == self.bcs_col: col_bcs = row_bcs else: for bc in self.bcs_col: if isinstance(bc, DirichletBC): bc_temp = bc.reconstruct(field=col_inds, V=Wcol, g=bc.function_arg, sub_domain=bc.sub_domain, use_split=True) elif isinstance(bc, EquationBCSplit): bc_temp = bc.reconstruct(field=col_inds, V=Wcol, row_field=row_inds, col_field=col_inds, use_split=True) if bc_temp is not None: col_bcs.append(bc_temp) submat_ctx = ImplicitMatrixContext(asub, row_bcs=row_bcs, col_bcs=col_bcs, fc_params=self.fc_params, appctx=self.appctx) submat_ctx.on_diag = self.on_diag and row_inds == col_inds submat = PETSc.Mat().create(comm=self._comm) submat.setType("python") submat.setSizes((submat_ctx.row_sizes, submat_ctx.col_sizes), bsize=submat_ctx.block_size) submat.setPythonContext(submat_ctx) submat.setUp() return submat
[docs] @PETSc.Log.EventDecorator() def duplicate(self, mat, copy): if copy == 0: raise NotImplementedError("We do now know how to duplicate a matrix-free MAT when copy=0") newmat_ctx = ImplicitMatrixContext(self.a, row_bcs=self.bcs, col_bcs=self.bcs_col, fc_params=self.fc_params, appctx=self.appctx) newmat = PETSc.Mat().create(comm=self._comm) newmat.setType("python") newmat.setSizes((newmat_ctx.row_sizes, newmat_ctx.col_sizes), bsize=newmat_ctx.block_size) newmat.setPythonContext(newmat_ctx) newmat.setUp() return newmat