diff --git a/docs/source/conf.py b/docs/source/conf.py index 5b2cc94c49..da6ae4c19b 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -134,10 +134,10 @@ (r'py:.*', r'pyadjoint\..*'), (r'py:.*', r'tsfc\..*'), (r'py:.*', r'ufl\..*'), - (r'py:.*', r'PETSc\..*'), (r'py:.*', r'progress\..*'), # Ignore undocumented PyOP2 ('py:class', 'pyop2.caching.Cached'), + ('py:class', 'pyop2.types.mat.Mat'), # Ignore mission docs from Firedrake internal "private" code # Any "Base" class eg: # firedrake.adjoint.checkpointing.CheckpointBase @@ -416,6 +416,7 @@ 'ufl': ('https://docs.fenicsproject.org/ufl/main/', None), 'FIAT': ('https://firedrakeproject.org/fiat', None), 'petsctools': ('https://firedrakeproject.org/petsctools/', None), + 'petsc4py': ('https://petsc.org/release/petsc4py/', None), 'mpi4py': ('https://mpi4py.readthedocs.io/en/stable/', None), 'h5py': ('http://docs.h5py.org/en/latest/', None), 'h5py.h5p': ('https://api.h5py.org/', None), diff --git a/firedrake/assemble.py b/firedrake/assemble.py index c0590bd9d3..6692cd912a 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -16,11 +16,12 @@ from tsfc.ufl_utils import extract_firedrake_constants import ufl import finat.ufl -from firedrake import (extrusion_utils as eutils, matrix, parameters, solving, +from firedrake import (extrusion_utils as eutils, parameters, solving, tsfc_interface, utils) from firedrake.adjoint_utils import annotate_assemble from firedrake.ufl_expr import extract_domains from firedrake.bcs import DirichletBC, EquationBC, EquationBCSplit +from firedrake.matrix import MatrixBase, Matrix, ImplicitMatrix from firedrake.functionspaceimpl import WithGeometry, FunctionSpace, FiredrakeDualSpace from firedrake.functionspacedata import entity_dofs_key, entity_permutations_key from firedrake.interpolation import get_interpolator @@ -238,7 +239,7 @@ def assemble(self, tensor=None, current_state=None): if isinstance(expr, ufl.algebra.Sum): a, b = [assemble(e) for e in expr.ufl_operands] # Only Expr resulting in a Matrix if assembled are BaseFormOperator - if not all(isinstance(op, matrix.AssembledMatrix) for op in (a, b)): + if not all(isinstance(op, MatrixBase) for op in (a, b)): raise TypeError('Mismatching Sum shapes') return assemble(ufl.FormSum((a, 1), (b, 1)), tensor=tensor) elif isinstance(expr, ufl.algebra.Product): @@ -356,7 +357,7 @@ def __init__(self, def allocate(self): rank = len(self._form.arguments()) if rank == 2 and not self._diagonal: - if isinstance(self._form, matrix.MatrixBase): + if isinstance(self._form, MatrixBase): return self._form elif self._mat_type == "matfree": return MatrixFreeAssembler(self._form, bcs=self._bcs, form_compiler_parameters=self._form_compiler_params, @@ -365,9 +366,8 @@ def allocate(self): else: test, trial = self._form.arguments() sparsity = ExplicitMatrixAssembler._make_sparsity(test, trial, self._mat_type, self._sub_mat_type, self.maps_and_regions) - return matrix.Matrix(self._form, self._bcs, self._mat_type, sparsity, ScalarType, - sub_mat_type=self._sub_mat_type, - options_prefix=self._options_prefix) + op2mat = op2.Mat(sparsity, mat_type=self._mat_type, sub_mat_type=self._sub_mat_type, dtype=ScalarType) + return Matrix(self._form, op2mat, bcs=self._bcs, options_prefix=self._options_prefix, fc_params=self._form_compiler_params) else: raise NotImplementedError("Only implemented for rank = 2 and diagonal = False") @@ -474,13 +474,14 @@ def base_form_assembly_visitor(self, expr, tensor, bcs, *args): # Out-of-place Hermitian transpose mat.petscmat.hermitianTranspose(out=result) if tensor is None: - tensor = self.assembled_matrix(expr, bcs, result) + tensor = Matrix(expr, result, bcs=bcs, + options_prefix=self._options_prefix, fc_params=self._form_compiler_params) return tensor elif isinstance(expr, ufl.Action): if len(args) != 2: raise TypeError("Not enough operands for Action") lhs, rhs = args - if isinstance(lhs, matrix.MatrixBase): + if isinstance(lhs, MatrixBase): if isinstance(rhs, (firedrake.Cofunction, firedrake.Function)): petsc_mat = lhs.petscmat (row, col) = lhs.arguments() @@ -489,11 +490,11 @@ def base_form_assembly_visitor(self, expr, tensor, bcs, *args): with rhs.dat.vec_ro as v_vec, res.dat.vec as res_vec: petsc_mat.mult(v_vec, res_vec) return res - elif isinstance(rhs, matrix.MatrixBase): + elif isinstance(rhs, MatrixBase): result = tensor.petscmat if tensor else PETSc.Mat() lhs.petscmat.matMult(rhs.petscmat, result=result) if tensor is None: - tensor = self.assembled_matrix(expr, bcs, result) + tensor = Matrix(expr, result, bcs=bcs, options_prefix=self._options_prefix) return tensor else: raise TypeError("Incompatible RHS for Action.") @@ -503,7 +504,7 @@ def base_form_assembly_visitor(self, expr, tensor, bcs, *args): with lhs.dat.vec_ro as x, rhs.dat.vec_ro as y: res = x.dot(y) return res - elif isinstance(rhs, matrix.MatrixBase): + elif isinstance(rhs, MatrixBase): # Compute action(Cofunc, Mat) => Mat^* @ Cofunc petsc_mat = rhs.petscmat (_, col) = rhs.arguments() @@ -584,7 +585,8 @@ def base_form_assembly_visitor(self, expr, tensor, bcs, *args): op.handle.copy(result=result) result.scale(w) if tensor is None: - tensor = self.assembled_matrix(expr, bcs, result) + tensor = Matrix(expr.arguments(), result, bcs=bcs, + options_prefix=self._options_prefix, fc_params=self._form_compiler_params) return tensor else: raise TypeError("Mismatching FormSum shapes") @@ -633,10 +635,6 @@ def base_form_assembly_visitor(self, expr, tensor, bcs, *args): else: raise TypeError(f"Unrecognised BaseForm instance: {expr}") - def assembled_matrix(self, expr, bcs, petscmat): - return matrix.AssembledMatrix(expr.arguments(), bcs, petscmat, - options_prefix=self._options_prefix) - @staticmethod def base_form_postorder_traversal(expr, visitor, visited={}): if expr in visited: @@ -1381,10 +1379,12 @@ def allocate(self): self._mat_type, self._sub_mat_type, self._make_maps_and_regions()) - return matrix.Matrix(self._form, self._bcs, self._mat_type, sparsity, ScalarType, - sub_mat_type=self._sub_mat_type, - options_prefix=self._options_prefix, - fc_params=self._form_compiler_params) + op2mat = op2.Mat( + sparsity, mat_type=self._mat_type, sub_mat_type=self._sub_mat_type, + dtype=ScalarType + ) + return Matrix(self._form, op2mat, bcs=self._bcs, + fc_params=self._form_compiler_params, options_prefix=self._options_prefix) @staticmethod def _make_sparsity(test, trial, mat_type, sub_mat_type, maps_and_regions): @@ -1583,10 +1583,17 @@ def __init__(self, form, bcs=None, form_compiler_parameters=None, self._appctx = appctx def allocate(self): - return matrix.ImplicitMatrix(self._form, self._bcs, - fc_params=self._form_compiler_params, - options_prefix=self._options_prefix, - appctx=self._appctx or {}) + from firedrake.matrix_free.operators import ImplicitMatrixContext + ctx = ImplicitMatrixContext( + self._form, row_bcs=self._bcs, col_bcs=self._bcs, + fc_params=self._form_compiler_params, + appctx=self._appctx + ) + return ImplicitMatrix( + self._form, ctx, self._bcs, + fc_params=self._form_compiler_params, + options_prefix=self._options_prefix + ) def assemble(self, tensor=None, current_state=None): if tensor is None: diff --git a/firedrake/bcs.py b/firedrake/bcs.py index 0584c5ede1..4be7476e76 100644 --- a/firedrake/bcs.py +++ b/firedrake/bcs.py @@ -14,14 +14,13 @@ from pyop2.utils import as_tuple import firedrake -import firedrake.matrix as matrix import firedrake.utils as utils -from firedrake import ufl_expr -from firedrake import slate -from firedrake import solving +from firedrake import ufl_expr, slate, solving from firedrake.formmanipulation import ExtractSubBlock from firedrake.adjoint_utils.dirichletbc import DirichletBCMixin from firedrake.petsc import PETSc +from firedrake.function import Function +from firedrake.cofunction import Cofunction __all__ = ['DirichletBC', 'homogenize', 'EquationBC'] @@ -186,8 +185,8 @@ def zero(self, r): boundary condition should be applied. """ - if isinstance(r, matrix.MatrixBase): - raise NotImplementedError("Zeroing bcs on a Matrix is not supported") + if not isinstance(r, Function | Cofunction): + raise NotImplementedError(f"Zeroing bcs not supported for {type(r).__name__}") for idx in self._indices: r = r.sub(idx) @@ -411,7 +410,7 @@ def apply(self, r, u=None): corresponding rows and columns. """ - if isinstance(r, matrix.MatrixBase): + if isinstance(r, ufl.Matrix): raise NotImplementedError("Capability to delay bc application has been dropped. Use assemble(a, bcs=bcs, ...) to obtain a fully assembled matrix") fs = self._function_space diff --git a/firedrake/external_operators/ml_operator.py b/firedrake/external_operators/ml_operator.py index ba35366cbc..4cdbd91ddb 100644 --- a/firedrake/external_operators/ml_operator.py +++ b/firedrake/external_operators/ml_operator.py @@ -1,5 +1,5 @@ from firedrake.external_operators import AbstractExternalOperator, assemble_method -from firedrake.matrix import AssembledMatrix +from firedrake.matrix import Matrix class MLOperator(AbstractExternalOperator): @@ -58,20 +58,16 @@ def assemble_jacobian(self, *args, **kwargs): """Assemble the Jacobian using the AD engine of the ML framework.""" # Delegate computation to the ML framework. J = self._jac() - # Set bcs - bcs = () - return AssembledMatrix(self, bcs, J) + return Matrix(self, J) @assemble_method(1, (1, 0)) def assemble_jacobian_adjoint(self, *args, **kwargs): """Assemble the Jacobian Hermitian transpose using the AD engine of the ML framework.""" # Delegate computation to the ML framework. J = self._jac() - # Set bcs - bcs = () # Take the adjoint (Hermitian transpose) J.hermitianTranspose() - return AssembledMatrix(self, bcs, J) + return Matrix(self, J) @assemble_method(1, (0, None)) def assemble_jacobian_action(self, *args, **kwargs): diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py index eb830493e1..458c446be4 100644 --- a/firedrake/formmanipulation.py +++ b/firedrake/formmanipulation.py @@ -15,7 +15,6 @@ from firedrake.functionspace import MixedFunctionSpace from firedrake.cofunction import Cofunction from firedrake.ufl_expr import Coargument -from firedrake.matrix import AssembledMatrix def subspace(V, indices): @@ -161,6 +160,7 @@ def cofunction(self, o): return Cofunction(W, val=MixedDat(o.dat[i] for i in indices)) def matrix(self, o): + from firedrake.matrix import AssembledMatrix ises = [] args = [] for a in o.arguments(): @@ -180,8 +180,7 @@ def matrix(self, o): args.append(asplit) submat = o.petscmat.createSubMatrix(*ises) - bcs = () - return AssembledMatrix(tuple(args), bcs, submat) + return AssembledMatrix(tuple(args), submat) def zero_base_form(self, o): return ZeroBaseForm(tuple(map(self, o.arguments()))) diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index 6f0954d8ea..aaa1469e12 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -33,7 +33,8 @@ from firedrake.petsc import PETSc from firedrake.halo import _get_mtype from firedrake.functionspaceimpl import WithGeometry -from firedrake.matrix import MatrixBase, AssembledMatrix +from firedrake.matrix import ImplicitMatrix, MatrixBase, Matrix +from firedrake.matrix_free.operators import ImplicitMatrixContext from firedrake.bcs import DirichletBC from firedrake.formmanipulation import split_form from firedrake.functionspace import VectorFunctionSpace, TensorFunctionSpace, FunctionSpace @@ -339,8 +340,8 @@ def assemble( specified. By default None. mat_type The PETSc matrix type to use when assembling a rank 2 interpolation. - For cross-mesh interpolation, only ``"aij"`` is supported. For same-mesh - interpolation, ``"aij"`` and ``"baij"`` are supported. For same/cross mesh interpolation + For cross-mesh interpolation, ``"aij"`` and ``"matfree"`` are supported. For same-mesh + interpolation, ``"aij"``, ``"baij"``, and ``"matfree"`` are supported. For same/cross mesh interpolation between :func:`.MixedFunctionSpace`, ``"aij"`` and ``"nest"`` are supported. For interpolation between input-ordering linked :func:`.VertexOnlyMesh`, ``"aij"``, ``"baij"``, and ``"matfree"`` are supported. @@ -356,7 +357,14 @@ def assemble( """ self._check_mat_type(mat_type) + if mat_type == "matfree" and self.rank == 2: + ctx = ImplicitMatrixContext( + self.ufl_interpolate, row_bcs=bcs, col_bcs=bcs, + ) + return ImplicitMatrix(self.ufl_interpolate, ctx, bcs=bcs) + result = self._get_callable(tensor=tensor, bcs=bcs, mat_type=mat_type, sub_mat_type=sub_mat_type)() + if self.rank == 2: # Assembling the operator assert isinstance(tensor, MatrixBase | None) @@ -364,7 +372,8 @@ def assemble( if tensor: result.copy(tensor.petscmat) return tensor - return AssembledMatrix(self.interpolate_args, bcs, result) + else: + return Matrix(self.ufl_interpolate, result, bcs=bcs) else: assert isinstance(tensor, Function | Cofunction | None) return tensor.assign(result) if tensor else result @@ -591,7 +600,7 @@ def callable() -> Function | Number: @property def _allowed_mat_types(self): - return {"aij", None} + return {"aij", "matfree", None} class SameMeshInterpolator(Interpolator): @@ -754,7 +763,7 @@ def callable() -> Function | Cofunction | PETSc.Mat | Number: @property def _allowed_mat_types(self): - return {"aij", "baij", None} + return {"aij", "baij", "matfree", None} class VomOntoVomInterpolator(SameMeshInterpolator): @@ -1642,4 +1651,4 @@ def callable() -> Number: @property def _allowed_mat_types(self): - return {"aij", "nest", None} + return {"aij", "nest", "matfree", None} diff --git a/firedrake/linear_solver.py b/firedrake/linear_solver.py index d4b7e70b88..6fb87d4a37 100644 --- a/firedrake/linear_solver.py +++ b/firedrake/linear_solver.py @@ -1,6 +1,5 @@ from firedrake.function import Function from firedrake.cofunction import Cofunction -from firedrake.matrix import MatrixBase from firedrake.petsc import PETSc from firedrake.variational_solver import LinearVariationalProblem, LinearVariationalSolver @@ -38,6 +37,7 @@ def __init__(self, A, *, P=None, **kwargs): Any boundary conditions for this solve *must* have been applied when assembling the operator. """ + from firedrake.matrix import MatrixBase if not isinstance(A, MatrixBase): raise TypeError("Provided operator is a '%s', not a MatrixBase" % type(A).__name__) if P is not None and not isinstance(P, MatrixBase): diff --git a/firedrake/matrix.py b/firedrake/matrix.py index c9f363a581..2bd96c58b5 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -1,9 +1,16 @@ +from typing import Any +from collections.abc import Iterable import itertools -import ufl -from pyop2 import op2 +import ufl from pyop2.utils import as_tuple +from pyop2 import op2 from firedrake.petsc import PETSc +from firedrake.bcs import DirichletBC +from firedrake.matrix_free import ImplicitMatrixContext +from firedrake.slate import slate + +__all__ = ("MatrixBase", "Matrix", "ImplicitMatrix", "AssembledMatrix") class DummyOP2Mat: @@ -12,26 +19,44 @@ def __init__(self, handle): self.handle = handle +def _get_mat_type(petscmat: PETSc.Mat) -> str: + """Maps PETSc matrix types to Firedrake notation""" + mat_type = petscmat.getType() + if mat_type.startswith("seq") or mat_type.startswith("mpi"): + return mat_type[3:] + return mat_type + + class MatrixBase(ufl.Matrix): - """A representation of the linear operator associated with a - bilinear form and bcs. Explicitly assembled matrices and matrix-free - matrix classes will derive from this - - :arg a: the bilinear form this :class:`MatrixBase` represents - or a tuple of the arguments it represents - - :arg bcs: an iterable of boundary conditions to apply to this - :class:`MatrixBase`. May be `None` if there are no boundary - conditions to apply. - :arg mat_type: matrix type of assembled matrix, or 'matfree' for matrix-free - :kwarg fc_params: a dict of form compiler parameters of this matrix + """A representation of the linear operator associated with a bilinear form and bcs. + Explicitly assembled matrices and matrix-free .matrix classes will derive from this. """ - def __init__(self, a, bcs, mat_type, fc_params=None): + + def __init__( + self, + a: ufl.BaseForm | slate.TensorBase | tuple[ufl.Argument | ufl.Coargument, ufl.Argument | ufl.Coargument], + bcs: Iterable[DirichletBC] = (), + fc_params: dict[str, Any] | None = None, + ): + """Initialise a :class:`MatrixBase`. + + Parameters + ---------- + a + A UFL BaseForm (with two arguments) that this MatrixBase represents, + or a tuple of the arguments it represents, or a slate TensorBase. + bcs + An optional iterable of boundary conditions to apply to this :class:`MatrixBase`. + Empty tuple by default. + fc_params + A dictionary of form compiler parameters for this matrix. + """ if isinstance(a, tuple): self.a = None test, trial = a arguments = a else: + assert isinstance(a, ufl.BaseForm | slate.TensorBase) self.a = a test, trial = a.arguments() arguments = None @@ -39,7 +64,7 @@ def __init__(self, a, bcs, mat_type, fc_params=None): # (so we can't use a set, since the iteration order may differ # on different processes) - ufl.Matrix.__init__(self, test.function_space(), trial.function_space()) + super().__init__(test.function_space(), trial.function_space()) # ufl.Matrix._analyze_form_arguments sets the _arguments attribute to # non-Firedrake objects, which breaks things. To avoid this we overwrite @@ -47,18 +72,12 @@ def __init__(self, a, bcs, mat_type, fc_params=None): self._analyze_form_arguments() self._arguments = arguments - if bcs is None: - bcs = () - self.bcs = bcs self.comm = test.function_space().comm self.block_shape = (len(test.function_space()), len(trial.function_space())) - self.mat_type = mat_type - """Matrix type. - Matrix type used in the assembly of the PETSc matrix: 'aij', 'baij', 'dense' or 'nest', - or 'matfree' for matrix-free.""" - self.form_compiler_parameters = fc_params + self.bcs = bcs + self.form_compiler_parameters = {} if fc_params is None else fc_params def arguments(self): if self.a: @@ -96,13 +115,10 @@ def bcs(self, bcs): self._bcs = () def __repr__(self): - return "%s(a=%r, bcs=%r)" % (type(self).__name__, - self.a, - self.bcs) + return f"{type(self).__name__}(a={self.a!r}, bcs={self.bcs!r})" def __str__(self): - return "assembled %s(a=%s, bcs=%s)" % (type(self).__name__, - self.a, self.bcs) + return f"assembled {type(self).__name__}(a={self.a}, bcs={self.bcs})" def __add__(self, other): if isinstance(other, MatrixBase): @@ -133,39 +149,42 @@ def zero(self): class Matrix(MatrixBase): - """A representation of an assembled bilinear form. - - :arg a: the bilinear form this :class:`Matrix` represents. - - :arg bcs: an iterable of boundary conditions to apply to this - :class:`Matrix`. May be `None` if there are no boundary - conditions to apply. - - :arg mat_type: matrix type of assembled matrix. - - :kwarg fc_params: a dict of form compiler parameters for this matrix. - - A ``pyop2.types.mat.Mat`` will be built from the remaining - arguments, for valid values, see ``pyop2.types.mat.Mat`` source code. - - .. note:: - - This object acts to the right on an assembled :class:`.Function` - and to the left on an assembled cofunction (currently represented - by a :class:`.Function`). - - """ - - def __init__(self, a, bcs, mat_type, *args, **kwargs): - # sets self.a, self.bcs, self.mat_type, and self.fc_params - fc_params = kwargs.pop("fc_params", None) - MatrixBase.__init__(self, a, bcs, mat_type, fc_params=fc_params) - options_prefix = kwargs.pop("options_prefix", None) - self.M = op2.Mat(*args, mat_type=mat_type, **kwargs) + """A representation of an assembled bilinear form.""" + + def __init__( + self, + a: ufl.BaseForm, + mat: op2.Mat | PETSc.Mat, + bcs: Iterable[DirichletBC] = (), + fc_params: dict[str, Any] | None = None, + options_prefix: str | None = None, + ): + """Initialise a :class:`Matrix`. + + Parameters + ---------- + a + The bilinear form this :class:`Matrix` represents. + mat + The underlying matrix object. Either a PyOP2 Mat or a PETSc Mat. + bcs + An iterable of boundary conditions to apply to this :class:`Matrix`. + Empty tuple by default. + fc_params + A dictionary of form compiler parameters for this matrix, by default None. + options_prefix + PETSc options prefix to apply, by default None. + """ + super().__init__(a, bcs=bcs, fc_params=fc_params) + if isinstance(mat, op2.Mat): + self.M = mat + else: + assert isinstance(mat, PETSc.Mat) + self.M = DummyOP2Mat(mat) self.petscmat = self.M.handle if options_prefix is not None: self.petscmat.setOptionsPrefix(options_prefix) - self.mat_type = mat_type + self.mat_type = _get_mat_type(self.petscmat) def assemble(self): raise NotImplementedError("API compatibility to apply bcs after 'assemble(a)'\ @@ -174,47 +193,48 @@ def assemble(self): class ImplicitMatrix(MatrixBase): - """A representation of the action of bilinear form operating - without explicitly assembling the associated matrix. This class - wraps the relevant information for Python PETSc matrix. - - :arg a: the bilinear form this :class:`Matrix` represents. - - :arg bcs: an iterable of boundary conditions to apply to this - :class:`Matrix`. May be `None` if there are no boundary - conditions to apply. - - :kwarg fc_params: a dict of form compiler parameters for this matrix. - - .. note:: - - This object acts to the right on an assembled :class:`.Function` - and to the left on an assembled cofunction (currently represented - by a :class:`.Function`). + """A representation of the action of bilinear form operating without + explicitly assembling the associated matrix. This class wraps the + relevant information for Python PETSc matrix.""" + + def __init__( + self, + a: ufl.BaseForm, + ctx: ImplicitMatrixContext, + bcs: Iterable[DirichletBC] = (), + fc_params: dict[str, Any] | None = None, + options_prefix: str | None = None, + ): + """Initialise a :class:`ImplicitMatrix`. + + Parameters + ---------- + a + The bilinear form this :class:`ImplicitMatrix` represents. + ctx + An :class:`firedrake.matrix_free.operators.ImplicitMatrixContext` that + defines the operations of the matrix. + bcs + An iterable of boundary conditions to apply to this :class:`Matrix`. + May be `None` if there are no boundary conditions to apply. + Empty tuple by default. + fc_params + A dictionary of form compiler parameters for this matrix, by default None. + options_prefix + PETSc options prefix to apply, by default None. + """ + super().__init__(a, bcs=bcs, fc_params=fc_params) - """ - def __init__(self, a, bcs, *args, **kwargs): - # sets self.a, self.bcs, self.mat_type, and self.fc_params - fc_params = kwargs["fc_params"] - super(ImplicitMatrix, self).__init__(a, bcs, "matfree", fc_params) - - options_prefix = kwargs.pop("options_prefix") - appctx = kwargs.get("appctx", {}) - - from firedrake.matrix_free.operators import ImplicitMatrixContext - ctx = ImplicitMatrixContext(a, - row_bcs=self.bcs, - col_bcs=self.bcs, - fc_params=fc_params, - appctx=appctx) self.petscmat = PETSc.Mat().create(comm=self.comm) self.petscmat.setType("python") self.petscmat.setSizes((ctx.row_sizes, ctx.col_sizes), bsize=ctx.block_size) self.petscmat.setPythonContext(ctx) - self.petscmat.setOptionsPrefix(options_prefix) + if options_prefix is not None: + self.petscmat.setOptionsPrefix(options_prefix) self.petscmat.setUp() self.petscmat.assemble() + self.mat_type = "matfree" def assemble(self): # Bump petsc matrix state by assembling it. @@ -224,27 +244,36 @@ def assemble(self): class AssembledMatrix(MatrixBase): - """A representation of a matrix that doesn't require knowing the underlying form. - This class wraps the relevant information for Python PETSc matrix. - - :arg a: A tuple of the arguments the matrix represents - - :arg bcs: an iterable of boundary conditions to apply to this - :class:`Matrix`. May be `None` if there are no boundary - conditions to apply. - - :arg petscmat: the already constructed petsc matrix this object represents. - """ - def __init__(self, a, bcs, petscmat, *args, **kwargs): - options_prefix = kwargs.pop("options_prefix", None) - super(AssembledMatrix, self).__init__(a, bcs, "assembled") + """A representation of a matrix that doesn't require knowing the underlying form.""" + + def __init__( + self, + args: tuple[ufl.Argument | ufl.Coargument, ufl.Argument | ufl.Coargument], + petscmat: PETSc.Mat, + bcs: Iterable[DirichletBC] = (), + options_prefix: str | None = None, + ): + """Initialise an :class:`AssembledMatrix`. + + Parameters + ---------- + args + A tuple of the arguments the matrix represents. + petscmat + The PETSc matrix this object wraps. + bcs + an iterable of boundary conditions to apply to this :class:`Matrix`. + May be `None` if there are no boundary conditions to apply. + Empty tuple by default. + options_prefix + PETSc options prefix to apply, by default None. + """ + super().__init__(args, bcs=bcs) self.petscmat = petscmat if options_prefix is not None: self.petscmat.setOptionsPrefix(options_prefix) + self.mat_type = _get_mat_type(self.petscmat) # this mimics op2.Mat.handle - self.M = DummyOP2Mat(self.mat()) - - def mat(self): - return self.petscmat + self.M = DummyOP2Mat(self.petscmat) diff --git a/firedrake/matrix_free/__init__.py b/firedrake/matrix_free/__init__.py index e69de29bb2..34e1ecb0f2 100644 --- a/firedrake/matrix_free/__init__.py +++ b/firedrake/matrix_free/__init__.py @@ -0,0 +1 @@ +from firedrake.matrix_free.operators import ImplicitMatrixContext # noqa: F401 diff --git a/firedrake/matrix_free/operators.py b/firedrake/matrix_free/operators.py index 0ecb021e68..92550485ee 100644 --- a/firedrake/matrix_free/operators.py +++ b/firedrake/matrix_free/operators.py @@ -1,8 +1,10 @@ from collections import OrderedDict +from typing import Any, Iterable import itertools from mpi4py import MPI import numpy +import ufl from pyop2.mpi import temp_internal_comm from firedrake.ufl_expr import adjoint, action @@ -67,35 +69,43 @@ class ImplicitMatrixContext(object): # (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): + def __init__( + self, + a: ufl.BaseForm, + row_bcs: Iterable[DirichletBC] = (), + col_bcs: Iterable[DirichletBC] = (), + fc_params: dict[str, Any] | None = None, + appctx: dict[str, Any] | None = None + ): + """This class gives the Python context for a PETSc Python matrix. + + Parameters + ---------- + a + The bilinear form defining the matrix. + 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. Empty tuple by default. + col_bcs + An iterable of the :class.`.DirichletBC`s that are imposed + on the trial space. Empty tuple by default. + fc_params + A dictionary of parameters to pass on to the form compiler. + By default None. + appctx + Any extra user-supplied context, available to preconditioners + and the like. By default None. + """ from firedrake.assemble import get_assembler self.a = a self.aT = adjoint(a) self.comm = a.arguments()[0].function_space().comm - self.fc_params = fc_params - self.appctx = appctx + self.fc_params = {} if fc_params is None else fc_params + self.appctx = {} if appctx is None else appctx # Collect all DirichletBC instances including # DirichletBCs applied to an EquationBC. @@ -438,7 +448,6 @@ def createSubMatrix(self, mat, row_is, col_is, target=None): @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, diff --git a/tests/firedrake/regression/test_interpolate_cross_mesh.py b/tests/firedrake/regression/test_interpolate_cross_mesh.py index 2011dc9466..0516a16a7c 100644 --- a/tests/firedrake/regression/test_interpolate_cross_mesh.py +++ b/tests/firedrake/regression/test_interpolate_cross_mesh.py @@ -779,7 +779,7 @@ def test_mixed_interpolator_cross_mesh(): # | V1 -> V4 V2 -> V4 | res = assemble(mixed_interp, mat_type="nest") - assert isinstance(res, AssembledMatrix) + assert isinstance(res, Matrix) assert res.petscmat.type == "nest" split_interp = dict(split_form(mixed_interp)) diff --git a/tests/firedrake/regression/test_interpolator_types.py b/tests/firedrake/regression/test_interpolation_operators.py similarity index 51% rename from tests/firedrake/regression/test_interpolator_types.py rename to tests/firedrake/regression/test_interpolation_operators.py index c77900b781..6c97264ded 100644 --- a/tests/firedrake/regression/test_interpolator_types.py +++ b/tests/firedrake/regression/test_interpolation_operators.py @@ -3,12 +3,14 @@ MixedInterpolator, SameMeshInterpolator, CrossMeshInterpolator, get_interpolator, VomOntoVomInterpolator, ) +from firedrake.matrix import ImplicitMatrix, Matrix import pytest +import numpy as np def params(): params = [] - for mat_type in [None, "aij"]: + for mat_type in [None, "aij", "matfree"]: params.append(pytest.param(mat_type, None, id=f"mat_type={mat_type}")) for sub_mat_type in [None, "aij", "baij"]: params.append(pytest.param("nest", sub_mat_type, id=f"nest_sub_mat_type={sub_mat_type}")) @@ -17,32 +19,59 @@ def params(): @pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("value_shape", ["scalar", "vector"], ids=lambda v: f"fs_type={v}") -@pytest.mark.parametrize("mat_type", [None, "aij", "baij"], ids=lambda v: f"mat_type={v}") -def test_same_mesh_mattype(value_shape, mat_type): +@pytest.mark.parametrize("mat_type", [None, "aij", "baij", "matfree"], ids=lambda v: f"mat_type={v}") +@pytest.mark.parametrize("mode", ["forward", "adjoint"], ids=lambda v: f"mode={v}") +def test_same_mesh_mattype(value_shape, mat_type, mode): if COMM_WORLD.size > 1: prefix = "mpi" else: prefix = "seq" - mesh = UnitSquareMesh(4, 4) + + mesh = UnitSquareMesh(10, 10) + x, y = SpatialCoordinate(mesh) if value_shape == "scalar": fs_type = FunctionSpace + expr = x**2 + y**2 else: fs_type = VectorFunctionSpace + expr = as_vector([x**2, y**2]) + V1 = fs_type(mesh, "CG", 1) V2 = fs_type(mesh, "CG", 2) - u = TrialFunction(V1) + if mode == "forward": + exact = Function(V1).interpolate(expr) + interp = interpolate(TrialFunction(V2), V1) # V1 x V2^* -> R + f = Function(V2).interpolate(expr) + elif mode == "adjoint": + v1 = TestFunction(V1) + exact = assemble(inner(1, sum(v1) if value_shape == "vector" else v1) * dx) + forward_interp = interpolate(TrialFunction(V1), V2) # V1 x V2^* -> R + interp = interpolate(TestFunction(V1), TrialFunction(V2.dual())) # V2^* x V1 -> R + v2 = TestFunction(V2) + f = inner(1, sum(v2) if value_shape == "vector" else v2) * dx - interp = interpolate(u, V2) assert isinstance(get_interpolator(interp), SameMeshInterpolator) - res = assemble(interp, mat_type=mat_type) - if value_shape == "scalar": + I_mat = assemble(interp, mat_type=mat_type) + assert isinstance(I_mat, ImplicitMatrix if mat_type == "matfree" else Matrix) + + if mat_type == "matfree": + assert I_mat.petscmat.type == "python" + elif value_shape == "scalar": # Always seqaij for scalar - assert res.petscmat.type == prefix + "aij" + assert I_mat.petscmat.type == prefix + "aij" else: # Defaults to seqaij - assert res.petscmat.type == prefix + (mat_type if mat_type else "aij") + assert I_mat.petscmat.type == prefix + (mat_type if mat_type else "aij") + + res = assemble(action(I_mat, f)) + assert np.allclose(res.dat.data, exact.dat.data) + + if mode == "adjoint": + forward_I_mat = assemble(forward_interp, mat_type=mat_type) + res2 = assemble(action(adjoint(forward_I_mat), f)) + assert np.allclose(res2.dat.data, exact.dat.data) with pytest.raises(NotImplementedError): # MatNest only implemented for interpolation between MixedFunctionSpaces @@ -50,25 +79,34 @@ def test_same_mesh_mattype(value_shape, mat_type): @pytest.mark.parametrize("value_shape", ["scalar", "vector"], ids=lambda v: f"fs_type={v}") -@pytest.mark.parametrize("mat_type", [None, "aij"], ids=lambda v: f"mat_type={v}") +@pytest.mark.parametrize("mat_type", [None, "aij", "matfree"], ids=lambda v: f"mat_type={v}") def test_cross_mesh_mattype(value_shape, mat_type): - mesh1 = UnitSquareMesh(1, 1) - mesh2 = UnitSquareMesh(1, 1) + mesh1 = UnitSquareMesh(8, 8) + x1, y1 = SpatialCoordinate(mesh1) + mesh2 = UnitSquareMesh(2, 2) + x2, y2 = SpatialCoordinate(mesh2) if value_shape == "scalar": fs_type = FunctionSpace + expr1 = x1**2 + y1**2 + expr2 = x2**2 + y2**2 else: fs_type = VectorFunctionSpace + expr1 = as_vector([x1**2, y1**2]) + expr2 = as_vector([x2**2, y2**2]) V1 = fs_type(mesh1, "CG", 1) V2 = fs_type(mesh2, "CG", 1) - u = TrialFunction(V1) - - interp = interpolate(u, V2) + interp = interpolate(TrialFunction(V1), V2) assert isinstance(get_interpolator(interp), CrossMeshInterpolator) - res = assemble(interp, mat_type=mat_type) - # only aij for cross-mesh - assert res.petscmat.type == "seqaij" + I_mat = assemble(interp, mat_type=mat_type) + assert isinstance(I_mat, ImplicitMatrix if mat_type == "matfree" else Matrix) + assert I_mat.petscmat.type == "python" if mat_type == "matfree" else "seqaij" + + f = Function(V1).interpolate(expr1) + res = assemble(action(I_mat, f)) + exact = Function(V2).interpolate(expr2) + assert np.allclose(res.dat.data, exact.dat.data) @pytest.mark.parametrize("value_shape", ["scalar", "vector"], ids=lambda v: f"fs_type={v}") @@ -84,15 +122,16 @@ def test_vomtovom_mattype(value_shape, mat_type): P0DG = fs_type(vom, "DG", 0) P0DG_io = fs_type(vom.input_ordering, "DG", 0) - u = TrialFunction(P0DG) - interp = interpolate(u, P0DG_io) + interp = interpolate(TrialFunction(P0DG), P0DG_io) assert isinstance(get_interpolator(interp), VomOntoVomInterpolator) + res = assemble(interp, mat_type=mat_type) + assert isinstance(res, ImplicitMatrix if mat_type == "matfree" else Matrix) + if not mat_type or mat_type == "matfree": assert res.petscmat.type == "python" else: if value_shape == "scalar": - # Always seqaij for scalar assert res.petscmat.type == "seqaij" else: # Defaults to seqaij @@ -100,7 +139,7 @@ def test_vomtovom_mattype(value_shape, mat_type): @pytest.mark.parametrize("value_shape", ["scalar", "vector"], ids=lambda v: f"fs_type={v}") -@pytest.mark.parametrize("mat_type", [None, "aij", "baij"], ids=lambda v: f"mat_type={v}") +@pytest.mark.parametrize("mat_type", [None, "aij", "baij", "matfree"], ids=lambda v: f"mat_type={v}") def test_point_eval_mattype(value_shape, mat_type): mesh = UnitSquareMesh(1, 1) points = [[0.1, 0.1], [0.5, 0.5], [0.9, 0.9]] @@ -112,13 +151,14 @@ def test_point_eval_mattype(value_shape, mat_type): P0DG = fs_type(vom, "DG", 0) V = fs_type(mesh, "CG", 1) - u = TrialFunction(V) - interp = interpolate(u, P0DG) + interp = interpolate(TrialFunction(V), P0DG) assert isinstance(get_interpolator(interp), SameMeshInterpolator) res = assemble(interp, mat_type=mat_type) + assert isinstance(res, ImplicitMatrix if mat_type == "matfree" else Matrix) - if value_shape == "scalar": - # Always seqaij for scalar + if mat_type == "matfree": + assert res.petscmat.type == "python" + elif value_shape == "scalar": assert res.petscmat.type == "seqaij" else: # Defaults to seqaij @@ -128,7 +168,8 @@ def test_point_eval_mattype(value_shape, mat_type): @pytest.mark.parametrize("value_shape", ["scalar", "vector"], ids=lambda v: f"fs_type={v}") @pytest.mark.parametrize("mat_type,sub_mat_type", params()) def test_mixed_same_mesh_mattype(value_shape, mat_type, sub_mat_type): - mesh = UnitSquareMesh(1, 1) + mesh = UnitSquareMesh(5, 5) + x, y = SpatialCoordinate(mesh) if value_shape == "scalar": fs_type = FunctionSpace else: @@ -141,24 +182,27 @@ def test_mixed_same_mesh_mattype(value_shape, mat_type, sub_mat_type): W = V1 * V2 U = V3 * V4 - w = TrialFunction(W) - w0, w1 = split(w) if value_shape == "scalar": - expr = as_vector([w0 + w1, w0 + w1]) + expr = as_vector([x**2, y**2]) else: - w00, w01 = split(w0) - w10, w11 = split(w1) - expr = as_vector([w00 + w10, w00 + w10, w01 + w11, w01 + w11]) - interp = interpolate(expr, U) + expr = as_vector([x**2, x**2, y**2, y**2]) + + interp = interpolate(TrialFunction(U), W) assert isinstance(get_interpolator(interp), MixedInterpolator) - res = assemble(interp, mat_type=mat_type, sub_mat_type=sub_mat_type) - if not mat_type or mat_type == "aij": - # Defaults to seqaij - assert res.petscmat.type == "seqaij" + I_mat = assemble(interp, mat_type=mat_type, sub_mat_type=sub_mat_type) + assert isinstance(I_mat, ImplicitMatrix if mat_type == "matfree" else Matrix) + + if mat_type == "matfree": + assert I_mat.petscmat.type == "python" + elif not mat_type or mat_type == "aij": + assert I_mat.petscmat.type == "seqaij" else: - assert res.petscmat.type == "nest" + assert I_mat.petscmat.type == "nest" for (i, j) in [(0, 0), (0, 1), (1, 0), (1, 1)]: - sub_mat = res.petscmat.getNestSubMatrix(i, j) + sub_mat = I_mat.petscmat.getNestSubMatrix(i, j) + if i != j: + assert not sub_mat + continue if value_shape == "scalar": # Always seqaij for scalar assert sub_mat.type == "seqaij" @@ -166,8 +210,11 @@ def test_mixed_same_mesh_mattype(value_shape, mat_type, sub_mat_type): # matnest sub_mat_type defaults to aij assert sub_mat.type == "seq" + (sub_mat_type if sub_mat_type else "aij") - with pytest.raises(NotImplementedError): - assemble(interp, mat_type="baij") + f = Function(U).interpolate(expr) + exact = Function(W).interpolate(expr) + res = assemble(action(I_mat, f)) + for resi, exi in zip(res.subfunctions, exact.subfunctions): + assert np.allclose(resi.dat.data, exi.dat.data) with pytest.raises(NotImplementedError): - assemble(interp, mat_type="matfree") + assemble(interp, mat_type="baij") diff --git a/tests/firedrake/regression/test_matrix.py b/tests/firedrake/regression/test_matrix.py index 32623c418c..872a1783de 100644 --- a/tests/firedrake/regression/test_matrix.py +++ b/tests/firedrake/regression/test_matrix.py @@ -42,7 +42,7 @@ def test_solve_with_assembled_matrix(a): x, = SpatialCoordinate(V.mesh()) f = Function(V).interpolate(x) - A = AssembledMatrix((v, u), bcs=(), petscmat=assemble(a).petscmat) + A = AssembledMatrix((v, u), assemble(a).petscmat) L = inner(f, v) * dx solution = Function(V)