From 34f87edc9dc6d20e296526357547f8055294c11a Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Mon, 15 Dec 2025 15:12:31 +0000 Subject: [PATCH 01/25] add test --- .../firedrake/regression/test_interp_dual.py | 71 ++++++++++--------- 1 file changed, 37 insertions(+), 34 deletions(-) diff --git a/tests/firedrake/regression/test_interp_dual.py b/tests/firedrake/regression/test_interp_dual.py index b0ce971a95..ce91878bae 100644 --- a/tests/firedrake/regression/test_interp_dual.py +++ b/tests/firedrake/regression/test_interp_dual.py @@ -3,6 +3,7 @@ from firedrake import * from firedrake.utils import complex_mode from firedrake.matrix import MatrixBase +from firedrake.matrix import MatrixBase import ufl @@ -361,37 +362,39 @@ def test_assemble_action_adjoint(V1, V2): a = interpolate(u, V2) # V1 x V2^* -> R, equiv. V1 -> V2 assert a.arguments() == (TestFunction(V2.dual()), TrialFunction(V1)) - f_form = inner(1, TestFunction(V2)) * dx - - for f in (f_form, assemble(f_form)): - expr = action(adjoint(assemble(a)), f) - assert isinstance(expr, Action) - res = assemble(expr) - assert isinstance(res, Cofunction) - assert res.function_space() == V1.dual() - - expr2 = action(f, a) # This simplifies into an Interpolate - assert isinstance(expr2, Interpolate) - res2 = assemble(expr2) - assert isinstance(res2, Cofunction) - assert res2.function_space() == V1.dual() - assert np.allclose(res.dat.data, res2.dat.data) - - A = assemble(a) - assert isinstance(A, MatrixBase) - - # This doesn't explicitly assemble the adjoint of A, but uses multHermitian - expr3 = action(f, A) - assert isinstance(expr3, Action) - res3 = assemble(expr3) - assert isinstance(res3, Cofunction) - assert res3.function_space() == V1.dual() - assert np.allclose(res.dat.data, res3.dat.data) - - # This is simplified into action(f, A) to avoid explicit assembly of adjoint(A) - expr4 = action(adjoint(A), f) - assert isinstance(expr4, Action) - res4 = assemble(expr4) - assert isinstance(res4, Cofunction) - assert res4.function_space() == V1.dual() - assert np.allclose(res.dat.data, res4.dat.data) + a_adj = adjoint(a) # V2^* x V1 -> R, equiv. V2^* -> V1^* + assert a_adj.arguments() == (TestFunction(V1), TrialFunction(V2.dual())) + + f = assemble(inner(1, TestFunction(V2)) * dx) + + expr = action(a_adj, f) + assert isinstance(expr, Action) + res = assemble(expr) + assert isinstance(res, Cofunction) + assert res.function_space() == V1.dual() + + expr2 = action(f, a) # This simplifies into an Interpolate + assert isinstance(expr2, Interpolate) + res2 = assemble(expr2) + assert isinstance(res2, Cofunction) + assert res2.function_space() == V1.dual() + assert np.allclose(res.dat.data, res2.dat.data) + + A = assemble(a) + assert isinstance(A, MatrixBase) + + # This explicitly assembles the adjoint of A, then matmults with f + expr3 = action(adjoint(A), f) + assert isinstance(expr3, Action) + res3 = assemble(expr3) + assert isinstance(res3, Cofunction) + assert res3.function_space() == V1.dual() + assert np.allclose(res.dat.data, res3.dat.data) + + # This doesn't explicitly assemble the adjoint of A, but uses multHermitian + expr4 = action(f, A) + assert isinstance(expr4, Action) + res4 = assemble(expr4) + assert isinstance(res4, Cofunction) + assert res4.function_space() == V1.dual() + assert np.allclose(res2.dat.data, res4.dat.data) From 1a4b9305b828b8316cec061653ec8066579735c2 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Mon, 15 Dec 2025 15:12:40 +0000 Subject: [PATCH 02/25] add assembly case update comment --- tests/firedrake/regression/test_interp_dual.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/firedrake/regression/test_interp_dual.py b/tests/firedrake/regression/test_interp_dual.py index ce91878bae..8a4ef67522 100644 --- a/tests/firedrake/regression/test_interp_dual.py +++ b/tests/firedrake/regression/test_interp_dual.py @@ -383,7 +383,7 @@ def test_assemble_action_adjoint(V1, V2): A = assemble(a) assert isinstance(A, MatrixBase) - # This explicitly assembles the adjoint of A, then matmults with f + # This (currently) explicitly assembles the adjoint of A, then matmults with f expr3 = action(adjoint(A), f) assert isinstance(expr3, Action) res3 = assemble(expr3) From 0461be421dbc248125fd9af9345cf79a130a5749 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Mon, 15 Dec 2025 18:08:15 +0000 Subject: [PATCH 03/25] add simplification --- tests/firedrake/regression/test_interp_dual.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/firedrake/regression/test_interp_dual.py b/tests/firedrake/regression/test_interp_dual.py index 8a4ef67522..824f9de03b 100644 --- a/tests/firedrake/regression/test_interp_dual.py +++ b/tests/firedrake/regression/test_interp_dual.py @@ -383,18 +383,18 @@ def test_assemble_action_adjoint(V1, V2): A = assemble(a) assert isinstance(A, MatrixBase) - # This (currently) explicitly assembles the adjoint of A, then matmults with f - expr3 = action(adjoint(A), f) + # This doesn't explicitly assemble the adjoint of A, but uses multHermitian + expr3 = action(f, A) assert isinstance(expr3, Action) res3 = assemble(expr3) assert isinstance(res3, Cofunction) assert res3.function_space() == V1.dual() assert np.allclose(res.dat.data, res3.dat.data) - # This doesn't explicitly assemble the adjoint of A, but uses multHermitian - expr4 = action(f, A) + # This is simplified into action(f, A) to avoid explicit assembly of adjoint(A) + expr4 = action(adjoint(A), f) assert isinstance(expr4, Action) res4 = assemble(expr4) assert isinstance(res4, Cofunction) assert res4.function_space() == V1.dual() - assert np.allclose(res2.dat.data, res4.dat.data) + assert np.allclose(res.dat.data, res4.dat.data) From 6ddd6d7009d7cf19502e28bd295d18ec41a98030 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Tue, 9 Dec 2025 15:03:10 +0000 Subject: [PATCH 04/25] WIP --- firedrake/interpolation.py | 4 +- firedrake/matrix.py | 107 +++++++++++++++++++++---------------- 2 files changed, 63 insertions(+), 48 deletions(-) diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index 6f0954d8ea..36120425ae 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -33,7 +33,7 @@ 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 MatrixBase, Matrix from firedrake.bcs import DirichletBC from firedrake.formmanipulation import split_form from firedrake.functionspace import VectorFunctionSpace, TensorFunctionSpace, FunctionSpace @@ -364,7 +364,7 @@ def assemble( if tensor: result.copy(tensor.petscmat) return tensor - return AssembledMatrix(self.interpolate_args, bcs, result) + return Matrix(self.ufl_interpolate, bcs, mat_type, result) else: assert isinstance(tensor, Function | Cofunction | None) return tensor.assign(result) if tensor else result diff --git a/firedrake/matrix.py b/firedrake/matrix.py index c9f363a581..393caa42ba 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -1,9 +1,13 @@ import itertools -import ufl +from typing import Any, Iterable, Literal +import ufl +from ufl.argument import BaseArgument from pyop2 import op2 from pyop2.utils import as_tuple from firedrake.petsc import PETSc +from firedrake.bcs import DirichletBC +from firedrake.matrix_free.operators import ImplicitMatrixContext class DummyOP2Mat: @@ -13,25 +17,37 @@ def __init__(self, handle): 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 - """ - def __init__(self, a, bcs, mat_type, fc_params=None): + def __init__( + self, + a: ufl.BaseForm | tuple[BaseArgument, BaseArgument], + mat_type: Literal["aij", "baij", "dense", "nest", "matfree"], + bcs: Iterable[DirichletBC] | None = None, + fc_params: dict["str", Any] | None = None, + ): + """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. + + Parameters + ---------- + a + A UFL BaseForm (with two arguments) that this MatrixBase represents, + or a tuple of the arguments it represents. + mat_type + Matrix type used in the assembly of the PETSc matrix: 'aij', 'baij', 'dense' or 'nest', + or 'matfree' for matrix-free. + fc_params + A dictionary of form compiler parameters for this matrix. + bcs + An optional iterable of boundary conditions to apply to this :class:`MatrixBase`. + None by default. + """ if isinstance(a, tuple): self.a = None test, trial = a arguments = a else: + assert isinstance(a, ufl.BaseForm) self.a = a test, trial = a.arguments() arguments = None @@ -39,7 +55,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__(self, 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 @@ -54,10 +70,6 @@ def __init__(self, a, bcs, mat_type, fc_params=None): 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 def arguments(self): @@ -96,13 +108,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}, bcs={self.bcs})" 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): @@ -132,36 +141,43 @@ def zero(self): return self -class Matrix(MatrixBase): - """A representation of an assembled bilinear form. +"""A representation of an assembled bilinear form. - :arg a: the bilinear form this :class:`Matrix` represents. +: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 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. +:arg mat_type: matrix type of assembled matrix. - :kwarg fc_params: a dict of form compiler parameters for this 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. +A ``pyop2.types.mat.Mat`` will be built from the remaining +arguments, for valid values, see ``pyop2.types.mat.Mat`` source code. - .. note:: +.. 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`). + 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) +class Matrix(MatrixBase): + + def __init__( + self, + a: ufl.BaseForm, + mat: op2.Mat | PETSc.Mat, + mat_type: Literal["aij", "baij", "dense", "nest"], + bcs: Iterable[DirichletBC] | None = None, + fc_params: dict["str", Any] | None = None, + options_prefix: str | None = None, + *args, + ): + super().__init__(self, a, bcs, mat_type, fc_params=fc_params) + self.M = mat self.petscmat = self.M.handle if options_prefix is not None: self.petscmat.setOptionsPrefix(options_prefix) @@ -201,7 +217,6 @@ def __init__(self, a, bcs, *args, **kwargs): 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, From 366057a2b48739b2ca9c278031ff75257033934f Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Tue, 9 Dec 2025 16:00:54 +0000 Subject: [PATCH 05/25] WIP: change matrix interface --- firedrake/assemble.py | 27 ++++++++++----- firedrake/interpolation.py | 2 +- firedrake/matrix.py | 67 ++++++++++++++++++++++++-------------- 3 files changed, 63 insertions(+), 33 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index c0590bd9d3..96519ccd1d 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1381,10 +1381,14 @@ 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.Matrix(self._form, op2mat, self._mat_type, + 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 +1587,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 or {} + ) + return matrix.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/interpolation.py b/firedrake/interpolation.py index 36120425ae..dcc7554d99 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -364,7 +364,7 @@ def assemble( if tensor: result.copy(tensor.petscmat) return tensor - return Matrix(self.ufl_interpolate, bcs, mat_type, result) + return Matrix(self.ufl_interpolate, result, mat_type, bcs=bcs) else: assert isinstance(tensor, Function | Cofunction | None) return tensor.assign(result) if tensor else result diff --git a/firedrake/matrix.py b/firedrake/matrix.py index 393caa42ba..72cf4e48e8 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -1,13 +1,15 @@ +from __future__ import annotations +from typing import Any, Iterable, Literal, TYPE_CHECKING import itertools -from typing import Any, Iterable, Literal + +if TYPE_CHECKING: + from firedrake.bcs import DirichletBC import ufl from ufl.argument import BaseArgument from pyop2 import op2 from pyop2.utils import as_tuple from firedrake.petsc import PETSc -from firedrake.bcs import DirichletBC -from firedrake.matrix_free.operators import ImplicitMatrixContext class DummyOP2Mat: @@ -23,7 +25,7 @@ def __init__( a: ufl.BaseForm | tuple[BaseArgument, BaseArgument], mat_type: Literal["aij", "baij", "dense", "nest", "matfree"], bcs: Iterable[DirichletBC] | None = None, - fc_params: dict["str", Any] | None = None, + fc_params: dict[str, Any] | None = None, ): """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. @@ -55,7 +57,7 @@ def __init__( # (so we can't use a set, since the iteration order may differ # on different processes) - super().__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 @@ -141,7 +143,7 @@ def zero(self): return self -"""A representation of an assembled bilinear form. +""" :arg a: the bilinear form this :class:`Matrix` represents. @@ -172,14 +174,34 @@ def __init__( mat: op2.Mat | PETSc.Mat, mat_type: Literal["aij", "baij", "dense", "nest"], bcs: Iterable[DirichletBC] | None = None, - fc_params: dict["str", Any] | None = None, + fc_params: dict[str, Any] | None = None, options_prefix: str | None = None, - *args, ): - super().__init__(self, a, bcs, mat_type, fc_params=fc_params) - self.M = mat + """A representation of an assembled bilinear form. + + Parameters + ---------- + a + The bilinear form this :class:`Matrix` represents. + mat : op2.Mat | PETSc.Mat + _description_ + mat_type + _description_ + bcs : Iterable[DirichletBC] | None, optional + _description_, by default None + fc_params : dict[str, Any] | None, optional + _description_, by default None + options_prefix : str | None, optional + _description_, by default None + """ + super().__init__(a, mat_type, 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: + if options_prefix: self.petscmat.setOptionsPrefix(options_prefix) self.mat_type = mat_type @@ -209,19 +231,16 @@ class ImplicitMatrix(MatrixBase): by a :class:`.Function`). """ - 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", {}) - - ctx = ImplicitMatrixContext(a, - row_bcs=self.bcs, - col_bcs=self.bcs, - fc_params=fc_params, - appctx=appctx) + def __init__( + self, + a, + ctx, + bcs, + fc_params, + options_prefix, + ): + super().__init__(a, bcs, "matfree", fc_params) + self.petscmat = PETSc.Mat().create(comm=self.comm) self.petscmat.setType("python") self.petscmat.setSizes((ctx.row_sizes, ctx.col_sizes), From 6f6e739652b767693dc26ae908f37e83da45eb05 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Wed, 10 Dec 2025 12:36:51 +0000 Subject: [PATCH 06/25] WIP: use implicitmatrix for interpolation --- firedrake/interpolation.py | 17 ++++++++++++----- firedrake/matrix.py | 38 ++++++++++++++++++++------------------ 2 files changed, 32 insertions(+), 23 deletions(-) diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index dcc7554d99..2019797433 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, Matrix +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 @@ -364,7 +365,13 @@ def assemble( if tensor: result.copy(tensor.petscmat) return tensor - return Matrix(self.ufl_interpolate, result, mat_type, bcs=bcs) + if mat_type == "matfree": + ctx = ImplicitMatrixContext( + self.ufl_interpolate, row_bcs=bcs, col_bcs=bcs, + ) + return ImplicitMatrix(self.ufl_interpolate, ctx, bcs=bcs) + else: + return Matrix(self.ufl_interpolate, result, mat_type, bcs=bcs) else: assert isinstance(tensor, Function | Cofunction | None) return tensor.assign(result) if tensor else result @@ -591,7 +598,7 @@ def callable() -> Function | Number: @property def _allowed_mat_types(self): - return {"aij", None} + return {"aij", "matfree", None} class SameMeshInterpolator(Interpolator): @@ -754,7 +761,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 +1649,4 @@ def callable() -> Number: @property def _allowed_mat_types(self): - return {"aij", "nest", None} + return {"aij", "nest", "matfree", None} diff --git a/firedrake/matrix.py b/firedrake/matrix.py index 72cf4e48e8..359b39fe7d 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -210,36 +210,38 @@ def assemble(self): has been removed. Use 'assemble(a, bcs=bcs)', which\ now returns an assembled matrix.") +"""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. -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 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 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. - :kwarg fc_params: a dict of form compiler parameters for this matrix. +.. note:: - .. 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`). - 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`). +""" + + +class ImplicitMatrix(MatrixBase): - """ def __init__( self, a, ctx, - bcs, - fc_params, - options_prefix, + bcs = None, + fc_params = None, + options_prefix = None, ): - super().__init__(a, bcs, "matfree", fc_params) + super().__init__(a, "matfree", bcs=bcs, fc_params=fc_params) self.petscmat = PETSc.Mat().create(comm=self.comm) self.petscmat.setType("python") From 5e2fd5f20314996738547773367c9e1f5ffa7a1d Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Wed, 10 Dec 2025 17:15:36 +0000 Subject: [PATCH 07/25] add test --- firedrake/matrix.py | 127 ++++++++---------- firedrake/matrix_free/operators.py | 25 ++-- .../firedrake/regression/test_interpolate.py | 32 +++++ 3 files changed, 107 insertions(+), 77 deletions(-) diff --git a/firedrake/matrix.py b/firedrake/matrix.py index 359b39fe7d..baa2b3109b 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -4,6 +4,7 @@ if TYPE_CHECKING: from firedrake.bcs import DirichletBC + from firedrake.matrix_free.operators import ImplicitMatrixContext import ufl from ufl.argument import BaseArgument @@ -143,29 +144,6 @@ def zero(self): return self -""" - -: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`). - -""" - class Matrix(MatrixBase): def __init__( @@ -184,15 +162,17 @@ def __init__( a The bilinear form this :class:`Matrix` represents. mat : op2.Mat | PETSc.Mat - _description_ + The underlying matrix object. Either a PyOP2 Mat or a PETSc Mat. mat_type - _description_ + The type of the PETSc matrix. bcs : Iterable[DirichletBC] | None, optional - _description_, by default None + An iterable of boundary conditions to apply to this :class:`Matrix`. + May be `None` if there are no boundary conditions to apply. + By default None. fc_params : dict[str, Any] | None, optional - _description_, by default None + A dictionary of form compiler parameters for this matrix, by default None. options_prefix : str | None, optional - _description_, by default None + PETSc options prefix to apply, by default None. """ super().__init__(a, mat_type, bcs=bcs, fc_params=fc_params) if isinstance(mat, op2.Mat): @@ -210,37 +190,37 @@ def assemble(self): has been removed. Use 'assemble(a, bcs=bcs)', which\ now returns an assembled matrix.") -"""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`). - -""" - class ImplicitMatrix(MatrixBase): def __init__( self, - a, - ctx, - bcs = None, - fc_params = None, - options_prefix = None, + a: ufl.BaseForm, + ctx: ImplicitMatrixContext, + bcs: Iterable[DirichletBC] | None = None, + fc_params: dict[str, Any] | None = None, + options_prefix: str | None = None, ): + """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. + + Parameters + ---------- + a + The bilinear form this :class:`ImplicitMatrix` represents. + ctx + An :class:`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. + By default None. + 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, "matfree", bcs=bcs, fc_params=fc_params) self.petscmat = PETSc.Mat().create(comm=self.comm) @@ -260,27 +240,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. + def __init__( + self, + args: tuple[BaseArgument, BaseArgument], + petscmat: PETSc.Mat, + mat_type: Literal["aij", "baij", "dense", "nest", "matfree"], + bcs: Iterable[DirichletBC] | None = None, + options_prefix: str | None = None, + ): + """A representation of a matrix that doesn't require knowing the underlying form. - :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") + Parameters + ---------- + args + A tuple of the arguments the matrix represents. + petscmat + The PETSc matrix this object wraps. + mat_type + The type of the PETSc matrix. + bcs + an iterable of boundary conditions to apply to this :class:`Matrix`. + May be `None` if there are no boundary conditions to apply. By default None. + options_prefix + PETSc options prefix to apply, by default None. + """ + super().__init__(args, mat_type, bcs=bcs) self.petscmat = petscmat - if options_prefix is not None: + if options_prefix: self.petscmat.setOptionsPrefix(options_prefix) # this mimics op2.Mat.handle self.M = DummyOP2Mat(self.mat()) - - def mat(self): - return self.petscmat diff --git a/firedrake/matrix_free/operators.py b/firedrake/matrix_free/operators.py index 0ecb021e68..fd47c67ac0 100644 --- a/firedrake/matrix_free/operators.py +++ b/firedrake/matrix_free/operators.py @@ -1,8 +1,10 @@ from collections import OrderedDict +from typing import 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 @@ -62,11 +64,6 @@ def find_sub_block(iset, ises, comm): return found -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 @@ -86,9 +83,22 @@ class ImplicitMatrixContext(object): preconditioners and the like. """ + +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 + + @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] | None = None, + col_bcs: Iterable[DirichletBC] | None = None, + fc_params=None, + appctx=None + ): from firedrake.assemble import get_assembler self.a = a @@ -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.py b/tests/firedrake/regression/test_interpolate.py index 66d733c57c..8f25dce018 100644 --- a/tests/firedrake/regression/test_interpolate.py +++ b/tests/firedrake/regression/test_interpolate.py @@ -741,3 +741,35 @@ def test_interpolate_form_mixed(): res3 = assemble(inner(u, q) * dx) # V x W -> R assert mat_equals(res1, res3) + + +@pytest.mark.parallel([1, 3]) +# @pytest.mark.parametrize("mode", ["forward", "adjoint"]) +@pytest.mark.parametrize("shape", ["scalar", "vector", "tensor"]) +def test_interpolate_operator_matfree(shape): + mesh = UnitSquareMesh(4, 4) + x, y = SpatialCoordinate(mesh) + + if shape == "scalar": + fs = FunctionSpace + expr = sin(x + y) + elif shape == "vector": + fs = VectorFunctionSpace + expr = as_vector([sin(x + y), cos(x + y)]) + elif shape == "tensor": + fs = TensorFunctionSpace + expr = as_tensor([[sin(x + y), x * y], [x + y, cos(x + y)]]) + + V1 = fs(mesh, "CG", 1) + V2 = fs(mesh, "CG", 2) + + exact = Function(V1).interpolate(expr) + + u = TrialFunction(V2) + I = assemble(interpolate(u, V1), mat_type="matfree") + + f = Function(V2).interpolate(expr) + res = assemble(I @ f) + + assert np.allclose(res.dat.data, exact.dat.data) + From be9177f9bd3e58e0e97ad49c112fc7a4d7aed0eb Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Thu, 11 Dec 2025 13:45:39 +0000 Subject: [PATCH 08/25] edit tests; fixes --- firedrake/assemble.py | 4 +- firedrake/interpolation.py | 11 +- firedrake/matrix.py | 2 +- .../firedrake/regression/test_interpolate.py | 32 ----- .../regression/test_interpolator_types.py | 127 +++++++++++------- 5 files changed, 91 insertions(+), 85 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 96519ccd1d..a525e08459 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -634,8 +634,8 @@ def base_form_assembly_visitor(self, expr, tensor, bcs, *args): 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) + return matrix.AssembledMatrix(expr.arguments(), petscmat, self._mat_type, + bcs=bcs, options_prefix=self._options_prefix) @staticmethod def base_form_postorder_traversal(expr, visitor, visited={}): diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index 2019797433..38a193d012 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -357,6 +357,12 @@ def assemble( """ self._check_mat_type(mat_type) + if mat_type == "matfree": + 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 @@ -365,11 +371,6 @@ def assemble( if tensor: result.copy(tensor.petscmat) return tensor - if mat_type == "matfree": - ctx = ImplicitMatrixContext( - self.ufl_interpolate, row_bcs=bcs, col_bcs=bcs, - ) - return ImplicitMatrix(self.ufl_interpolate, ctx, bcs=bcs) else: return Matrix(self.ufl_interpolate, result, mat_type, bcs=bcs) else: diff --git a/firedrake/matrix.py b/firedrake/matrix.py index baa2b3109b..9a79101797 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -272,4 +272,4 @@ def __init__( self.petscmat.setOptionsPrefix(options_prefix) # this mimics op2.Mat.handle - self.M = DummyOP2Mat(self.mat()) + self.M = DummyOP2Mat(self.petscmat) diff --git a/tests/firedrake/regression/test_interpolate.py b/tests/firedrake/regression/test_interpolate.py index 8f25dce018..66d733c57c 100644 --- a/tests/firedrake/regression/test_interpolate.py +++ b/tests/firedrake/regression/test_interpolate.py @@ -741,35 +741,3 @@ def test_interpolate_form_mixed(): res3 = assemble(inner(u, q) * dx) # V x W -> R assert mat_equals(res1, res3) - - -@pytest.mark.parallel([1, 3]) -# @pytest.mark.parametrize("mode", ["forward", "adjoint"]) -@pytest.mark.parametrize("shape", ["scalar", "vector", "tensor"]) -def test_interpolate_operator_matfree(shape): - mesh = UnitSquareMesh(4, 4) - x, y = SpatialCoordinate(mesh) - - if shape == "scalar": - fs = FunctionSpace - expr = sin(x + y) - elif shape == "vector": - fs = VectorFunctionSpace - expr = as_vector([sin(x + y), cos(x + y)]) - elif shape == "tensor": - fs = TensorFunctionSpace - expr = as_tensor([[sin(x + y), x * y], [x + y, cos(x + y)]]) - - V1 = fs(mesh, "CG", 1) - V2 = fs(mesh, "CG", 2) - - exact = Function(V1).interpolate(expr) - - u = TrialFunction(V2) - I = assemble(interpolate(u, V1), mat_type="matfree") - - f = Function(V2).interpolate(expr) - res = assemble(I @ f) - - assert np.allclose(res.dat.data, exact.dat.data) - diff --git a/tests/firedrake/regression/test_interpolator_types.py b/tests/firedrake/regression/test_interpolator_types.py index c77900b781..8693b0b983 100644 --- a/tests/firedrake/regression/test_interpolator_types.py +++ b/tests/firedrake/regression/test_interpolator_types.py @@ -3,12 +3,13 @@ MixedInterpolator, SameMeshInterpolator, CrossMeshInterpolator, get_interpolator, VomOntoVomInterpolator, ) +from firedrake.matrix import ImplicitMatrix, Matrix import pytest 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 +18,53 @@ 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) + 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) with pytest.raises(NotImplementedError): # MatNest only implemented for interpolation between MixedFunctionSpaces @@ -50,25 +72,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 +115,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 +132,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 +144,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 +161,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 +175,24 @@ 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" - for (i, j) in [(0, 0), (0, 1), (1, 0), (1, 1)]: - sub_mat = res.petscmat.getNestSubMatrix(i, j) + assert I_mat.petscmat.type == "nest" + for i in range(2): + sub_mat = I_mat.petscmat.getNestSubMatrix(i, i) if value_shape == "scalar": # Always seqaij for scalar assert sub_mat.type == "seqaij" @@ -166,8 +200,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") From b1225edc2374fe1e383953ed3db1f3d8ee034089 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 12 Dec 2025 11:07:18 +0000 Subject: [PATCH 09/25] change type hint --- firedrake/matrix.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/firedrake/matrix.py b/firedrake/matrix.py index 9a79101797..f3f21ccc52 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -3,9 +3,10 @@ import itertools if TYPE_CHECKING: - from firedrake.bcs import DirichletBC + from firedrake.bcs import BCBase from firedrake.matrix_free.operators import ImplicitMatrixContext - + from firedrake.slate.slate import TensorBase + import ufl from ufl.argument import BaseArgument from pyop2 import op2 @@ -13,6 +14,7 @@ from firedrake.petsc import PETSc + class DummyOP2Mat: """A hashable implementation of M.handle""" def __init__(self, handle): @@ -23,9 +25,9 @@ class MatrixBase(ufl.Matrix): def __init__( self, - a: ufl.BaseForm | tuple[BaseArgument, BaseArgument], + a: ufl.BaseForm | TensorBase | tuple[BaseArgument, BaseArgument], mat_type: Literal["aij", "baij", "dense", "nest", "matfree"], - bcs: Iterable[DirichletBC] | None = None, + bcs: Iterable[BCBase] | None = None, fc_params: dict[str, Any] | None = None, ): """A representation of the linear operator associated with a bilinear form and bcs. @@ -44,13 +46,14 @@ def __init__( bcs An optional iterable of boundary conditions to apply to this :class:`MatrixBase`. None by default. - """ + """ + from firedrake.slate.slate import TensorBase if isinstance(a, tuple): self.a = None test, trial = a arguments = a else: - assert isinstance(a, ufl.BaseForm) + assert isinstance(a, ufl.BaseForm | TensorBase) self.a = a test, trial = a.arguments() arguments = None @@ -151,7 +154,7 @@ def __init__( a: ufl.BaseForm, mat: op2.Mat | PETSc.Mat, mat_type: Literal["aij", "baij", "dense", "nest"], - bcs: Iterable[DirichletBC] | None = None, + bcs: Iterable[BCBase] | None = None, fc_params: dict[str, Any] | None = None, options_prefix: str | None = None, ): @@ -197,7 +200,7 @@ def __init__( self, a: ufl.BaseForm, ctx: ImplicitMatrixContext, - bcs: Iterable[DirichletBC] | None = None, + bcs: Iterable[BCBase] | None = None, fc_params: dict[str, Any] | None = None, options_prefix: str | None = None, ): @@ -246,7 +249,7 @@ def __init__( args: tuple[BaseArgument, BaseArgument], petscmat: PETSc.Mat, mat_type: Literal["aij", "baij", "dense", "nest", "matfree"], - bcs: Iterable[DirichletBC] | None = None, + bcs: Iterable[BCBase] | None = None, options_prefix: str | None = None, ): """A representation of a matrix that doesn't require knowing the underlying form. From 46efa2bd2984b6124368c3a68c57e9b9c57579fc Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 12 Dec 2025 13:54:34 +0000 Subject: [PATCH 10/25] changes --- firedrake/assemble.py | 13 ++---- firedrake/external_operators/ml_operator.py | 10 ++-- firedrake/formmanipulation.py | 3 +- firedrake/function.py | 3 ++ firedrake/interpolation.py | 3 +- firedrake/matrix.py | 29 ++++-------- firedrake/matrix_free/operators.py | 51 +++++++++++---------- tests/firedrake/regression/test_matrix.py | 2 +- 8 files changed, 50 insertions(+), 64 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index a525e08459..be1cdc4eeb 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -365,9 +365,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.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") @@ -634,8 +633,7 @@ def base_form_assembly_visitor(self, expr, tensor, bcs, *args): raise TypeError(f"Unrecognised BaseForm instance: {expr}") def assembled_matrix(self, expr, bcs, petscmat): - return matrix.AssembledMatrix(expr.arguments(), petscmat, self._mat_type, - bcs=bcs, options_prefix=self._options_prefix) + return matrix.AssembledMatrix(expr.arguments(), petscmat, bcs=bcs, options_prefix=self._options_prefix) @staticmethod def base_form_postorder_traversal(expr, visitor, visited={}): @@ -1385,8 +1383,7 @@ def allocate(self): sparsity, mat_type=self._mat_type, sub_mat_type=self._sub_mat_type, dtype=ScalarType ) - return matrix.Matrix(self._form, op2mat, self._mat_type, - bcs=self._bcs, + return matrix.Matrix(self._form, op2mat, bcs=self._bcs, fc_params=self._form_compiler_params, options_prefix=self._options_prefix) @@ -1591,7 +1588,7 @@ def allocate(self): ctx = ImplicitMatrixContext( self._form, row_bcs=self._bcs, col_bcs=self._bcs, fc_params=self._form_compiler_params, - appctx=self._appctx or {} + appctx=self._appctx ) return matrix.ImplicitMatrix( self._form, ctx, self._bcs, 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..f24f17519a 100644 --- a/firedrake/formmanipulation.py +++ b/firedrake/formmanipulation.py @@ -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/function.py b/firedrake/function.py index e0b746c663..be70cd694f 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -3,6 +3,7 @@ import sys import ufl import warnings +from ufl.algorithms.analysis import extract_arguments from ufl.duals import is_dual from ufl.formatting.ufl2unicode import ufl2unicode from ufl.domain import extract_unique_domain @@ -381,6 +382,8 @@ def interpolate(self, Returns `self` """ from firedrake import interpolate, assemble + if len(extract_arguments(expression)) > 0: + raise ValueError("Can't interpolate an expression with arguments into a Function.") V = self.function_space() interp = interpolate(expression, V, **kwargs) return assemble(interp, tensor=self, ad_block_tag=ad_block_tag) diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index 38a193d012..bd44a23a59 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -364,6 +364,7 @@ def assemble( 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) @@ -372,7 +373,7 @@ def assemble( result.copy(tensor.petscmat) return tensor else: - return Matrix(self.ufl_interpolate, result, mat_type, bcs=bcs) + return Matrix(self.ufl_interpolate, result, bcs=bcs) else: assert isinstance(tensor, Function | Cofunction | None) return tensor.assign(result) if tensor else result diff --git a/firedrake/matrix.py b/firedrake/matrix.py index f3f21ccc52..ddac0545de 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -26,7 +26,6 @@ class MatrixBase(ufl.Matrix): def __init__( self, a: ufl.BaseForm | TensorBase | tuple[BaseArgument, BaseArgument], - mat_type: Literal["aij", "baij", "dense", "nest", "matfree"], bcs: Iterable[BCBase] | None = None, fc_params: dict[str, Any] | None = None, ): @@ -37,10 +36,7 @@ def __init__( ---------- a A UFL BaseForm (with two arguments) that this MatrixBase represents, - or a tuple of the arguments it represents. - mat_type - Matrix type used in the assembly of the PETSc matrix: 'aij', 'baij', 'dense' or 'nest', - or 'matfree' for matrix-free. + or a tuple of the arguments it represents, or a slate TensorBase. fc_params A dictionary of form compiler parameters for this matrix. bcs @@ -69,14 +65,11 @@ def __init__( self._analyze_form_arguments() self._arguments = arguments - if bcs is None: - bcs = () - self.bcs = bcs + self.bcs = bcs or () self.comm = test.function_space().comm self.block_shape = (len(test.function_space()), len(trial.function_space())) - self.mat_type = mat_type - self.form_compiler_parameters = fc_params + self.form_compiler_parameters = fc_params or {} def arguments(self): if self.a: @@ -153,7 +146,6 @@ def __init__( self, a: ufl.BaseForm, mat: op2.Mat | PETSc.Mat, - mat_type: Literal["aij", "baij", "dense", "nest"], bcs: Iterable[BCBase] | None = None, fc_params: dict[str, Any] | None = None, options_prefix: str | None = None, @@ -166,8 +158,6 @@ def __init__( The bilinear form this :class:`Matrix` represents. mat : op2.Mat | PETSc.Mat The underlying matrix object. Either a PyOP2 Mat or a PETSc Mat. - mat_type - The type of the PETSc matrix. bcs : Iterable[DirichletBC] | None, optional An iterable of boundary conditions to apply to this :class:`Matrix`. May be `None` if there are no boundary conditions to apply. @@ -177,7 +167,7 @@ def __init__( options_prefix : str | None, optional PETSc options prefix to apply, by default None. """ - super().__init__(a, mat_type, bcs=bcs, fc_params=fc_params) + super().__init__(a, bcs=bcs, fc_params=fc_params) if isinstance(mat, op2.Mat): self.M = mat else: @@ -186,7 +176,7 @@ def __init__( self.petscmat = self.M.handle if options_prefix: self.petscmat.setOptionsPrefix(options_prefix) - self.mat_type = mat_type + self.mat_type = self.petscmat.getType() def assemble(self): raise NotImplementedError("API compatibility to apply bcs after 'assemble(a)'\ @@ -224,7 +214,7 @@ def __init__( options_prefix PETSc options prefix to apply, by default None. """ - super().__init__(a, "matfree", bcs=bcs, fc_params=fc_params) + super().__init__(a, bcs=bcs, fc_params=fc_params) self.petscmat = PETSc.Mat().create(comm=self.comm) self.petscmat.setType("python") @@ -234,6 +224,7 @@ def __init__( 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. @@ -248,7 +239,6 @@ def __init__( self, args: tuple[BaseArgument, BaseArgument], petscmat: PETSc.Mat, - mat_type: Literal["aij", "baij", "dense", "nest", "matfree"], bcs: Iterable[BCBase] | None = None, options_prefix: str | None = None, ): @@ -260,19 +250,18 @@ def __init__( A tuple of the arguments the matrix represents. petscmat The PETSc matrix this object wraps. - mat_type - The type of the PETSc matrix. bcs an iterable of boundary conditions to apply to this :class:`Matrix`. May be `None` if there are no boundary conditions to apply. By default None. options_prefix PETSc options prefix to apply, by default None. """ - super().__init__(args, mat_type, bcs=bcs) + super().__init__(args, bcs=bcs) self.petscmat = petscmat if options_prefix: self.petscmat.setOptionsPrefix(options_prefix) + self.mat_type = self.petscmat.getType() # this mimics op2.Mat.handle self.M = DummyOP2Mat(self.petscmat) diff --git a/firedrake/matrix_free/operators.py b/firedrake/matrix_free/operators.py index fd47c67ac0..bfe5b30e5c 100644 --- a/firedrake/matrix_free/operators.py +++ b/firedrake/matrix_free/operators.py @@ -1,5 +1,5 @@ from collections import OrderedDict -from typing import Iterable +from typing import Any, Iterable import itertools from mpi4py import MPI @@ -64,26 +64,6 @@ def find_sub_block(iset, ises, comm): return found - """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. - - """ - class ImplicitMatrixContext(object): # By default, these matrices will represent diagonal blocks (the # (0,0) block of a 1x1 block matrix is on the diagonal). @@ -96,16 +76,37 @@ def __init__( a: ufl.BaseForm, row_bcs: Iterable[DirichletBC] | None = None, col_bcs: Iterable[DirichletBC] | None = None, - fc_params=None, - appctx=None + 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. By default None. + col_bcs + An iterable of the :class.`.DirichletBC`s that are imposed + on the trial space. By default None. + 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 = fc_params or {} + self.appctx = appctx or {} # Collect all DirichletBC instances including # DirichletBCs applied to an EquationBC. 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) From c6d9a72c85ae29a39e2c35fc78af1d2ff193e407 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 12 Dec 2025 13:59:16 +0000 Subject: [PATCH 11/25] lint --- firedrake/assemble.py | 4 +- firedrake/matrix.py | 59 +++++++++++++++--------------- firedrake/matrix_free/operators.py | 11 +++--- 3 files changed, 36 insertions(+), 38 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index be1cdc4eeb..86fdd400f6 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1591,8 +1591,8 @@ def allocate(self): appctx=self._appctx ) return matrix.ImplicitMatrix( - self._form, ctx, self._bcs, - fc_params=self._form_compiler_params, + self._form, ctx, self._bcs, + fc_params=self._form_compiler_params, options_prefix=self._options_prefix ) diff --git a/firedrake/matrix.py b/firedrake/matrix.py index ddac0545de..6759c955e0 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -1,12 +1,12 @@ from __future__ import annotations -from typing import Any, Iterable, Literal, TYPE_CHECKING +from typing import Any, Iterable, TYPE_CHECKING import itertools if TYPE_CHECKING: from firedrake.bcs import BCBase from firedrake.matrix_free.operators import ImplicitMatrixContext from firedrake.slate.slate import TensorBase - + import ufl from ufl.argument import BaseArgument from pyop2 import op2 @@ -14,7 +14,6 @@ from firedrake.petsc import PETSc - class DummyOP2Mat: """A hashable implementation of M.handle""" def __init__(self, handle): @@ -24,11 +23,11 @@ def __init__(self, handle): class MatrixBase(ufl.Matrix): def __init__( - self, - a: ufl.BaseForm | TensorBase | tuple[BaseArgument, BaseArgument], - bcs: Iterable[BCBase] | None = None, - fc_params: dict[str, Any] | None = None, - ): + self, + a: ufl.BaseForm | TensorBase | tuple[BaseArgument, BaseArgument], + bcs: Iterable[BCBase] | None = None, + fc_params: dict[str, Any] | None = None, + ): """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. @@ -143,13 +142,13 @@ def zero(self): class Matrix(MatrixBase): def __init__( - self, - a: ufl.BaseForm, - mat: op2.Mat | PETSc.Mat, - bcs: Iterable[BCBase] | None = None, - fc_params: dict[str, Any] | None = None, - options_prefix: str | None = None, - ): + self, + a: ufl.BaseForm, + mat: op2.Mat | PETSc.Mat, + bcs: Iterable[BCBase] | None = None, + fc_params: dict[str, Any] | None = None, + options_prefix: str | None = None, + ): """A representation of an assembled bilinear form. Parameters @@ -187,15 +186,15 @@ def assemble(self): class ImplicitMatrix(MatrixBase): def __init__( - self, - a: ufl.BaseForm, - ctx: ImplicitMatrixContext, - bcs: Iterable[BCBase] | None = None, - fc_params: dict[str, Any] | None = None, - options_prefix: str | None = None, - ): - """A representation of the action of bilinear form operating without - explicitly assembling the associated matrix. This class wraps the + self, + a: ufl.BaseForm, + ctx: ImplicitMatrixContext, + bcs: Iterable[BCBase] | None = None, + fc_params: dict[str, Any] | None = None, + options_prefix: str | None = None, + ): + """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. Parameters @@ -236,12 +235,12 @@ def assemble(self): class AssembledMatrix(MatrixBase): def __init__( - self, - args: tuple[BaseArgument, BaseArgument], - petscmat: PETSc.Mat, - bcs: Iterable[BCBase] | None = None, - options_prefix: str | None = None, - ): + self, + args: tuple[BaseArgument, BaseArgument], + petscmat: PETSc.Mat, + bcs: Iterable[BCBase] | None = None, + options_prefix: str | None = None, + ): """A representation of a matrix that doesn't require knowing the underlying form. Parameters diff --git a/firedrake/matrix_free/operators.py b/firedrake/matrix_free/operators.py index bfe5b30e5c..27b4b4b4dc 100644 --- a/firedrake/matrix_free/operators.py +++ b/firedrake/matrix_free/operators.py @@ -69,16 +69,15 @@ class ImplicitMatrixContext(object): # (0,0) block of a 1x1 block matrix is on the diagonal). on_diag = True - @PETSc.Log.EventDecorator() def __init__( - self, - a: ufl.BaseForm, + self, + a: ufl.BaseForm, row_bcs: Iterable[DirichletBC] | None = None, col_bcs: Iterable[DirichletBC] | None = None, - fc_params : dict[str, Any] | None = None, + 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 @@ -97,7 +96,7 @@ def __init__( A dictionary of parameters to pass on to the form compiler. By default None. appctx - Any extra user-supplied context, available to preconditioners + Any extra user-supplied context, available to preconditioners and the like. By default None. """ from firedrake.assemble import get_assembler From aabd4d4cf6a40799a92b439842bd68829918b9f3 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 12 Dec 2025 16:23:51 +0000 Subject: [PATCH 12/25] fixes and review suggestions --- firedrake/function.py | 3 +- firedrake/matrix.py | 28 +++++++++++++------ firedrake/matrix_free/operators.py | 4 +-- .../regression/test_interpolator_types.py | 7 +++-- 4 files changed, 28 insertions(+), 14 deletions(-) diff --git a/firedrake/function.py b/firedrake/function.py index be70cd694f..06a0430a73 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -4,6 +4,7 @@ import ufl import warnings from ufl.algorithms.analysis import extract_arguments +from ufl.constantvalue import as_ufl from ufl.duals import is_dual from ufl.formatting.ufl2unicode import ufl2unicode from ufl.domain import extract_unique_domain @@ -382,7 +383,7 @@ def interpolate(self, Returns `self` """ from firedrake import interpolate, assemble - if len(extract_arguments(expression)) > 0: + if len(extract_arguments(as_ufl(expression))) > 0: raise ValueError("Can't interpolate an expression with arguments into a Function.") V = self.function_space() interp = interpolate(expression, V, **kwargs) diff --git a/firedrake/matrix.py b/firedrake/matrix.py index 6759c955e0..995f2d8460 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -20,6 +20,14 @@ 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): def __init__( @@ -64,11 +72,12 @@ def __init__( self._analyze_form_arguments() self._arguments = arguments - self.bcs = bcs or () self.comm = test.function_space().comm self.block_shape = (len(test.function_space()), len(trial.function_space())) - self.form_compiler_parameters = fc_params or {} + + self.bcs = () if bcs is None else bcs + self.form_compiler_parameters = {} if fc_params is None else fc_params def arguments(self): if self.a: @@ -106,10 +115,10 @@ def bcs(self, bcs): self._bcs = () def __repr__(self): - return f"{type(self).__name__}(a={self.a}, bcs={self.bcs})" + return f"{type(self).__name__}(a={repr(self.a)}, bcs={repr(self.bcs)})" def __str__(self): - return f"assembled {type(self).__name__}(a={self.a}, bcs={self.bcs})" + return f"assembled {type(self).__name__}(a={str(self.a)}, bcs={str(self.bcs)})" def __add__(self, other): if isinstance(other, MatrixBase): @@ -173,9 +182,9 @@ def __init__( assert isinstance(mat, PETSc.Mat) self.M = DummyOP2Mat(mat) self.petscmat = self.M.handle - if options_prefix: + if options_prefix is not None: self.petscmat.setOptionsPrefix(options_prefix) - self.mat_type = self.petscmat.getType() + self.mat_type = _get_mat_type(self.petscmat) def assemble(self): raise NotImplementedError("API compatibility to apply bcs after 'assemble(a)'\ @@ -220,7 +229,8 @@ def __init__( 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" @@ -258,9 +268,9 @@ def __init__( super().__init__(args, bcs=bcs) self.petscmat = petscmat - if options_prefix: + if options_prefix is not None: self.petscmat.setOptionsPrefix(options_prefix) - self.mat_type = self.petscmat.getType() + self.mat_type = _get_mat_type(self.petscmat) # this mimics op2.Mat.handle self.M = DummyOP2Mat(self.petscmat) diff --git a/firedrake/matrix_free/operators.py b/firedrake/matrix_free/operators.py index 27b4b4b4dc..49b1271622 100644 --- a/firedrake/matrix_free/operators.py +++ b/firedrake/matrix_free/operators.py @@ -104,8 +104,8 @@ def __init__( self.a = a self.aT = adjoint(a) self.comm = a.arguments()[0].function_space().comm - self.fc_params = fc_params or {} - self.appctx = appctx or {} + 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. diff --git a/tests/firedrake/regression/test_interpolator_types.py b/tests/firedrake/regression/test_interpolator_types.py index 8693b0b983..7c289c4aee 100644 --- a/tests/firedrake/regression/test_interpolator_types.py +++ b/tests/firedrake/regression/test_interpolator_types.py @@ -191,8 +191,11 @@ def test_mixed_same_mesh_mattype(value_shape, mat_type, sub_mat_type): assert I_mat.petscmat.type == "seqaij" else: assert I_mat.petscmat.type == "nest" - for i in range(2): - sub_mat = I_mat.petscmat.getNestSubMatrix(i, i) + for (i, j) in [(0, 0), (0, 1), (1, 0), (1, 1)]: + 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" From 4c9fb3aa0effb83ca1b5bb4485d278b570e1abeb Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 12 Dec 2025 17:12:34 +0000 Subject: [PATCH 13/25] rename test --- ...test_interpolator_types.py => test_interpolation_operators.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/firedrake/regression/{test_interpolator_types.py => test_interpolation_operators.py} (100%) diff --git a/tests/firedrake/regression/test_interpolator_types.py b/tests/firedrake/regression/test_interpolation_operators.py similarity index 100% rename from tests/firedrake/regression/test_interpolator_types.py rename to tests/firedrake/regression/test_interpolation_operators.py From 51d8efc31d89f9dd8dd63bdafe7e65f193ba3ac0 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 12 Dec 2025 17:12:45 +0000 Subject: [PATCH 14/25] fixes --- firedrake/assemble.py | 2 +- firedrake/interpolation.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 86fdd400f6..40433b933f 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -238,7 +238,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, matrix.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): diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index bd44a23a59..c0b89c2d36 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -357,7 +357,7 @@ def assemble( """ self._check_mat_type(mat_type) - if mat_type == "matfree": + if mat_type == "matfree" and self.rank == 2: ctx = ImplicitMatrixContext( self.ufl_interpolate, row_bcs=bcs, col_bcs=bcs, ) From fda5b340c2746efe259e588c24934cc3b633e404 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 19 Dec 2025 15:30:16 +0000 Subject: [PATCH 15/25] extra test --- tests/firedrake/regression/test_interp_dual.py | 1 - tests/firedrake/regression/test_interpolation_operators.py | 7 +++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/firedrake/regression/test_interp_dual.py b/tests/firedrake/regression/test_interp_dual.py index 824f9de03b..ee83d2b9e7 100644 --- a/tests/firedrake/regression/test_interp_dual.py +++ b/tests/firedrake/regression/test_interp_dual.py @@ -3,7 +3,6 @@ from firedrake import * from firedrake.utils import complex_mode from firedrake.matrix import MatrixBase -from firedrake.matrix import MatrixBase import ufl diff --git a/tests/firedrake/regression/test_interpolation_operators.py b/tests/firedrake/regression/test_interpolation_operators.py index 7c289c4aee..6c97264ded 100644 --- a/tests/firedrake/regression/test_interpolation_operators.py +++ b/tests/firedrake/regression/test_interpolation_operators.py @@ -5,6 +5,7 @@ ) from firedrake.matrix import ImplicitMatrix, Matrix import pytest +import numpy as np def params(): @@ -45,6 +46,7 @@ def test_same_mesh_mattype(value_shape, mat_type, mode): 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 @@ -66,6 +68,11 @@ def test_same_mesh_mattype(value_shape, mat_type, mode): 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 assemble(interp, mat_type="nest") From 7456052f3ee9fbbd82daa165907ddaeca5dad458 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Fri, 19 Dec 2025 15:34:43 +0000 Subject: [PATCH 16/25] fix --- .../firedrake/regression/test_interp_dual.py | 70 +++++++++---------- 1 file changed, 34 insertions(+), 36 deletions(-) diff --git a/tests/firedrake/regression/test_interp_dual.py b/tests/firedrake/regression/test_interp_dual.py index ee83d2b9e7..b0ce971a95 100644 --- a/tests/firedrake/regression/test_interp_dual.py +++ b/tests/firedrake/regression/test_interp_dual.py @@ -361,39 +361,37 @@ def test_assemble_action_adjoint(V1, V2): a = interpolate(u, V2) # V1 x V2^* -> R, equiv. V1 -> V2 assert a.arguments() == (TestFunction(V2.dual()), TrialFunction(V1)) - a_adj = adjoint(a) # V2^* x V1 -> R, equiv. V2^* -> V1^* - assert a_adj.arguments() == (TestFunction(V1), TrialFunction(V2.dual())) - - f = assemble(inner(1, TestFunction(V2)) * dx) - - expr = action(a_adj, f) - assert isinstance(expr, Action) - res = assemble(expr) - assert isinstance(res, Cofunction) - assert res.function_space() == V1.dual() - - expr2 = action(f, a) # This simplifies into an Interpolate - assert isinstance(expr2, Interpolate) - res2 = assemble(expr2) - assert isinstance(res2, Cofunction) - assert res2.function_space() == V1.dual() - assert np.allclose(res.dat.data, res2.dat.data) - - A = assemble(a) - assert isinstance(A, MatrixBase) - - # This doesn't explicitly assemble the adjoint of A, but uses multHermitian - expr3 = action(f, A) - assert isinstance(expr3, Action) - res3 = assemble(expr3) - assert isinstance(res3, Cofunction) - assert res3.function_space() == V1.dual() - assert np.allclose(res.dat.data, res3.dat.data) - - # This is simplified into action(f, A) to avoid explicit assembly of adjoint(A) - expr4 = action(adjoint(A), f) - assert isinstance(expr4, Action) - res4 = assemble(expr4) - assert isinstance(res4, Cofunction) - assert res4.function_space() == V1.dual() - assert np.allclose(res.dat.data, res4.dat.data) + f_form = inner(1, TestFunction(V2)) * dx + + for f in (f_form, assemble(f_form)): + expr = action(adjoint(assemble(a)), f) + assert isinstance(expr, Action) + res = assemble(expr) + assert isinstance(res, Cofunction) + assert res.function_space() == V1.dual() + + expr2 = action(f, a) # This simplifies into an Interpolate + assert isinstance(expr2, Interpolate) + res2 = assemble(expr2) + assert isinstance(res2, Cofunction) + assert res2.function_space() == V1.dual() + assert np.allclose(res.dat.data, res2.dat.data) + + A = assemble(a) + assert isinstance(A, MatrixBase) + + # This doesn't explicitly assemble the adjoint of A, but uses multHermitian + expr3 = action(f, A) + assert isinstance(expr3, Action) + res3 = assemble(expr3) + assert isinstance(res3, Cofunction) + assert res3.function_space() == V1.dual() + assert np.allclose(res.dat.data, res3.dat.data) + + # This is simplified into action(f, A) to avoid explicit assembly of adjoint(A) + expr4 = action(adjoint(A), f) + assert isinstance(expr4, Action) + res4 = assemble(expr4) + assert isinstance(res4, Cofunction) + assert res4.function_space() == V1.dual() + assert np.allclose(res.dat.data, res4.dat.data) From 30260c9b0925bf1b463a414823c87cba5e1058b6 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Mon, 19 Jan 2026 17:50:37 +0000 Subject: [PATCH 17/25] code review; pyop2 type hint troubles --- firedrake/assemble.py | 31 +++++++------ firedrake/bcs.py | 7 +-- firedrake/matrix.py | 73 +++++++++++++++++-------------- firedrake/matrix_free/__init__.py | 1 + 4 files changed, 61 insertions(+), 51 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 40433b933f..fc6043a209 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -21,6 +21,7 @@ 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.MatrixBase) 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, @@ -366,7 +367,7 @@ def allocate(self): test, trial = self._form.arguments() sparsity = ExplicitMatrixAssembler._make_sparsity(test, trial, self._mat_type, self._sub_mat_type, self.maps_and_regions) op2mat = op2.Mat(sparsity, mat_type=self._mat_type, sub_mat_type=self._sub_mat_type, dtype=ScalarType) - return matrix.Matrix(self._form, op2mat, bcs=self._bcs, options_prefix=self._options_prefix, fc_params=self._form_compiler_params) + 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") @@ -473,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() @@ -488,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.") @@ -502,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() @@ -583,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") @@ -632,9 +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(), petscmat, bcs=bcs, options_prefix=self._options_prefix) - @staticmethod def base_form_postorder_traversal(expr, visitor, visited={}): if expr in visited: @@ -1383,9 +1383,8 @@ def allocate(self): sparsity, mat_type=self._mat_type, sub_mat_type=self._sub_mat_type, dtype=ScalarType ) - return matrix.Matrix(self._form, op2mat, bcs=self._bcs, - fc_params=self._form_compiler_params, - options_prefix=self._options_prefix) + 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): @@ -1590,7 +1589,7 @@ def allocate(self): fc_params=self._form_compiler_params, appctx=self._appctx ) - return matrix.ImplicitMatrix( + return ImplicitMatrix( self._form, ctx, self._bcs, fc_params=self._form_compiler_params, options_prefix=self._options_prefix diff --git a/firedrake/bcs.py b/firedrake/bcs.py index 0584c5ede1..5a7b5fe394 100644 --- a/firedrake/bcs.py +++ b/firedrake/bcs.py @@ -14,7 +14,6 @@ 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 @@ -186,7 +185,8 @@ def zero(self, r): boundary condition should be applied. """ - if isinstance(r, matrix.MatrixBase): + from firedrake.matrix import MatrixBase + if isinstance(r, MatrixBase): raise NotImplementedError("Zeroing bcs on a Matrix is not supported") for idx in self._indices: @@ -411,7 +411,8 @@ def apply(self, r, u=None): corresponding rows and columns. """ - if isinstance(r, matrix.MatrixBase): + from firedrake.matrix import MatrixBase + if isinstance(r, MatrixBase): 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/matrix.py b/firedrake/matrix.py index 995f2d8460..24b152498a 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -1,18 +1,22 @@ from __future__ import annotations -from typing import Any, Iterable, TYPE_CHECKING -import itertools -if TYPE_CHECKING: - from firedrake.bcs import BCBase - from firedrake.matrix_free.operators import ImplicitMatrixContext - from firedrake.slate.slate import TensorBase +from typing import Any, TYPE_CHECKING +from collections.abc import Iterable +import itertools import ufl -from ufl.argument import BaseArgument -from pyop2 import op2 +import pyop2 +from pyop2.op2 import Mat from pyop2.utils import as_tuple from firedrake.petsc import PETSc +if TYPE_CHECKING: + from firedrake.bcs import DirichletBC + from firedrake.matrix_free import ImplicitMatrixContext + from firedrake.slate import slate + +__all__ = ("MatrixBase", "Matrix", "ImplicitMatrix", "AssembledMatrix") + class DummyOP2Mat: """A hashable implementation of M.handle""" @@ -29,15 +33,17 @@ def _get_mat_type(petscmat: PETSc.Mat) -> str: 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. + """ def __init__( self, - a: ufl.BaseForm | TensorBase | tuple[BaseArgument, BaseArgument], - bcs: Iterable[BCBase] | None = None, + a: ufl.BaseForm | slate.TensorBase | tuple[ufl.argument.BaseArgument, ufl.argument.BaseArgument], + bcs: Iterable[DirichletBC] = (), fc_params: dict[str, Any] | None = None, ): - """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. + """Initialise a :class:`MatrixBase`. Parameters ---------- @@ -50,7 +56,7 @@ def __init__( An optional iterable of boundary conditions to apply to this :class:`MatrixBase`. None by default. """ - from firedrake.slate.slate import TensorBase + from firedrake.slate import TensorBase if isinstance(a, tuple): self.a = None test, trial = a @@ -76,7 +82,7 @@ def __init__( self.block_shape = (len(test.function_space()), len(trial.function_space())) - self.bcs = () if bcs is None else bcs + self.bcs = bcs self.form_compiler_parameters = {} if fc_params is None else fc_params def arguments(self): @@ -118,7 +124,7 @@ def __repr__(self): return f"{type(self).__name__}(a={repr(self.a)}, bcs={repr(self.bcs)})" def __str__(self): - return f"assembled {type(self).__name__}(a={str(self.a)}, bcs={str(self.bcs)})" + return f"assembled {type(self).__name__}(a={self.a}, bcs={self.bcs})" def __add__(self, other): if isinstance(other, MatrixBase): @@ -149,34 +155,35 @@ def zero(self): class Matrix(MatrixBase): + """A representation of an assembled bilinear form.""" def __init__( self, a: ufl.BaseForm, - mat: op2.Mat | PETSc.Mat, - bcs: Iterable[BCBase] | None = None, + mat: pyop2.op2.Mat | PETSc.Mat, + bcs: Iterable[DirichletBC] = (), fc_params: dict[str, Any] | None = None, options_prefix: str | None = None, ): - """A representation of an assembled bilinear form. + """Initialise a :class:`Matrix`. Parameters ---------- a The bilinear form this :class:`Matrix` represents. - mat : op2.Mat | PETSc.Mat + mat The underlying matrix object. Either a PyOP2 Mat or a PETSc Mat. - bcs : Iterable[DirichletBC] | None, optional + bcs An iterable of boundary conditions to apply to this :class:`Matrix`. May be `None` if there are no boundary conditions to apply. By default None. - fc_params : dict[str, Any] | None, optional + fc_params A dictionary of form compiler parameters for this matrix, by default None. - options_prefix : str | None, optional + options_prefix PETSc options prefix to apply, by default None. """ super().__init__(a, bcs=bcs, fc_params=fc_params) - if isinstance(mat, op2.Mat): + if isinstance(mat, Mat): self.M = mat else: assert isinstance(mat, PETSc.Mat) @@ -193,26 +200,27 @@ 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.""" def __init__( self, a: ufl.BaseForm, ctx: ImplicitMatrixContext, - bcs: Iterable[BCBase] | None = None, + bcs: Iterable[DirichletBC] = (), fc_params: dict[str, Any] | None = None, options_prefix: str | None = None, ): - """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. + """Initialise a :class:`ImplicitMatrix`. Parameters ---------- a The bilinear form this :class:`ImplicitMatrix` represents. ctx - An :class:`ImplicitMatrixContext` that defines the operations - of the matrix. + 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. @@ -243,15 +251,16 @@ def assemble(self): class AssembledMatrix(MatrixBase): + """A representation of a matrix that doesn't require knowing the underlying form.""" def __init__( self, - args: tuple[BaseArgument, BaseArgument], + args: tuple[ufl.argument.BaseArgument, ufl.argument.BaseArgument], petscmat: PETSc.Mat, - bcs: Iterable[BCBase] | None = None, + bcs: Iterable[DirichletBC] = (), options_prefix: str | None = None, ): - """A representation of a matrix that doesn't require knowing the underlying form. + """Initialise an :class:`AssembledMatrix`. Parameters ---------- diff --git a/firedrake/matrix_free/__init__.py b/firedrake/matrix_free/__init__.py index e69de29bb2..b0d557a8ec 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 \ No newline at end of file From 2ee118454ce95ab46d1ae9c407d883aef6589e57 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Mon, 19 Jan 2026 18:09:50 +0000 Subject: [PATCH 18/25] lint and tidy --- firedrake/assemble.py | 6 +++--- firedrake/matrix.py | 5 ++--- firedrake/matrix_free/__init__.py | 2 +- 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index fc6043a209..6692cd912a 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -16,7 +16,7 @@ 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 @@ -474,7 +474,7 @@ 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 = Matrix(expr, result, bcs=bcs, + tensor = Matrix(expr, result, bcs=bcs, options_prefix=self._options_prefix, fc_params=self._form_compiler_params) return tensor elif isinstance(expr, ufl.Action): @@ -1383,7 +1383,7 @@ def allocate(self): sparsity, mat_type=self._mat_type, sub_mat_type=self._sub_mat_type, dtype=ScalarType ) - return Matrix(self._form, op2mat, bcs=self._bcs, + return Matrix(self._form, op2mat, bcs=self._bcs, fc_params=self._form_compiler_params, options_prefix=self._options_prefix) @staticmethod diff --git a/firedrake/matrix.py b/firedrake/matrix.py index 24b152498a..b7db405773 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -6,7 +6,6 @@ import ufl import pyop2 -from pyop2.op2 import Mat from pyop2.utils import as_tuple from firedrake.petsc import PETSc @@ -183,7 +182,7 @@ def __init__( PETSc options prefix to apply, by default None. """ super().__init__(a, bcs=bcs, fc_params=fc_params) - if isinstance(mat, Mat): + if isinstance(mat, pyop2.op2.Mat): self.M = mat else: assert isinstance(mat, PETSc.Mat) @@ -219,7 +218,7 @@ def __init__( a The bilinear form this :class:`ImplicitMatrix` represents. ctx - An :class:`firedrake.matrix_free.operators.ImplicitMatrixContext` that + 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`. diff --git a/firedrake/matrix_free/__init__.py b/firedrake/matrix_free/__init__.py index b0d557a8ec..34e1ecb0f2 100644 --- a/firedrake/matrix_free/__init__.py +++ b/firedrake/matrix_free/__init__.py @@ -1 +1 @@ -from firedrake.matrix_free.operators import ImplicitMatrixContext # noqa: F401 \ No newline at end of file +from firedrake.matrix_free.operators import ImplicitMatrixContext # noqa: F401 From ed45fee8718514819c7b6cad7988677cd7c22cfc Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Mon, 19 Jan 2026 18:10:43 +0000 Subject: [PATCH 19/25] suggestion --- firedrake/matrix.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/matrix.py b/firedrake/matrix.py index b7db405773..e58eca3436 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -120,7 +120,7 @@ def bcs(self, bcs): self._bcs = () def __repr__(self): - return f"{type(self).__name__}(a={repr(self.a)}, bcs={repr(self.bcs)})" + return f"{type(self).__name__}(a={self.a!r}, bcs={self.bcs!r})" def __str__(self): return f"assembled {type(self).__name__}(a={self.a}, bcs={self.bcs})" From fa700660847eaac75f97c73c3eb69709901aaa87 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Tue, 20 Jan 2026 12:08:53 +0000 Subject: [PATCH 20/25] fixes --- firedrake/formmanipulation.py | 2 +- firedrake/linear_solver.py | 2 +- firedrake/matrix.py | 26 +++++++++++--------------- 3 files changed, 13 insertions(+), 17 deletions(-) diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py index f24f17519a..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(): 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 e58eca3436..12dd915237 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -1,18 +1,15 @@ -from __future__ import annotations - -from typing import Any, TYPE_CHECKING +from typing import Any from collections.abc import Iterable import itertools import ufl -import pyop2 +from ufl.argument import BaseArgument from pyop2.utils import as_tuple +from pyop2 import op2 from firedrake.petsc import PETSc - -if TYPE_CHECKING: - from firedrake.bcs import DirichletBC - from firedrake.matrix_free import ImplicitMatrixContext - from firedrake.slate import slate +from firedrake.bcs import DirichletBC +from firedrake.matrix_free import ImplicitMatrixContext +from firedrake.slate import slate __all__ = ("MatrixBase", "Matrix", "ImplicitMatrix", "AssembledMatrix") @@ -38,7 +35,7 @@ class MatrixBase(ufl.Matrix): def __init__( self, - a: ufl.BaseForm | slate.TensorBase | tuple[ufl.argument.BaseArgument, ufl.argument.BaseArgument], + a: ufl.BaseForm | slate.TensorBase | tuple[BaseArgument, BaseArgument], bcs: Iterable[DirichletBC] = (), fc_params: dict[str, Any] | None = None, ): @@ -55,13 +52,12 @@ def __init__( An optional iterable of boundary conditions to apply to this :class:`MatrixBase`. None by default. """ - from firedrake.slate import TensorBase if isinstance(a, tuple): self.a = None test, trial = a arguments = a else: - assert isinstance(a, ufl.BaseForm | TensorBase) + assert isinstance(a, ufl.BaseForm | slate.TensorBase) self.a = a test, trial = a.arguments() arguments = None @@ -159,7 +155,7 @@ class Matrix(MatrixBase): def __init__( self, a: ufl.BaseForm, - mat: pyop2.op2.Mat | PETSc.Mat, + mat: op2.Mat | PETSc.Mat, bcs: Iterable[DirichletBC] = (), fc_params: dict[str, Any] | None = None, options_prefix: str | None = None, @@ -182,7 +178,7 @@ def __init__( PETSc options prefix to apply, by default None. """ super().__init__(a, bcs=bcs, fc_params=fc_params) - if isinstance(mat, pyop2.op2.Mat): + if isinstance(mat, op2.Mat): self.M = mat else: assert isinstance(mat, PETSc.Mat) @@ -254,7 +250,7 @@ class AssembledMatrix(MatrixBase): def __init__( self, - args: tuple[ufl.argument.BaseArgument, ufl.argument.BaseArgument], + args: tuple[BaseArgument, BaseArgument], petscmat: PETSc.Mat, bcs: Iterable[DirichletBC] = (), options_prefix: str | None = None, From 66f50b88f2a27430bf1eeb7888c3883ab668299a Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Thu, 22 Jan 2026 14:36:27 +0000 Subject: [PATCH 21/25] docs nearly fixed --- docs/source/conf.py | 1 + firedrake/matrix.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index 5b2cc94c49..ccd3583b24 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -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/matrix.py b/firedrake/matrix.py index 12dd915237..f23f521c74 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -35,7 +35,7 @@ class MatrixBase(ufl.Matrix): def __init__( self, - a: ufl.BaseForm | slate.TensorBase | tuple[BaseArgument, BaseArgument], + a: ufl.BaseForm | slate.TensorBase | tuple[ufl.Argument | ufl.Coargument, ufl.Argument | ufl.Coargument], bcs: Iterable[DirichletBC] = (), fc_params: dict[str, Any] | None = None, ): @@ -250,7 +250,7 @@ class AssembledMatrix(MatrixBase): def __init__( self, - args: tuple[BaseArgument, BaseArgument], + args: tuple[ufl.Argument | ufl.Coargument, ufl.Argument | ufl.Coargument], petscmat: PETSc.Mat, bcs: Iterable[DirichletBC] = (), options_prefix: str | None = None, From b498822df630afdf0316c1f338b1502ddda8165d Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Thu, 22 Jan 2026 14:52:10 +0000 Subject: [PATCH 22/25] review suggestions lint --- firedrake/bcs.py | 14 ++++++-------- firedrake/function.py | 4 ---- firedrake/interpolation.py | 4 ++-- firedrake/matrix.py | 14 +++++++------- firedrake/matrix_free/operators.py | 8 ++++---- 5 files changed, 19 insertions(+), 25 deletions(-) diff --git a/firedrake/bcs.py b/firedrake/bcs.py index 5a7b5fe394..4be7476e76 100644 --- a/firedrake/bcs.py +++ b/firedrake/bcs.py @@ -15,12 +15,12 @@ import firedrake 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'] @@ -185,9 +185,8 @@ def zero(self, r): boundary condition should be applied. """ - from firedrake.matrix import MatrixBase - if isinstance(r, 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,8 +410,7 @@ def apply(self, r, u=None): corresponding rows and columns. """ - from firedrake.matrix import MatrixBase - if isinstance(r, 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/function.py b/firedrake/function.py index 06a0430a73..e0b746c663 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -3,8 +3,6 @@ import sys import ufl import warnings -from ufl.algorithms.analysis import extract_arguments -from ufl.constantvalue import as_ufl from ufl.duals import is_dual from ufl.formatting.ufl2unicode import ufl2unicode from ufl.domain import extract_unique_domain @@ -383,8 +381,6 @@ def interpolate(self, Returns `self` """ from firedrake import interpolate, assemble - if len(extract_arguments(as_ufl(expression))) > 0: - raise ValueError("Can't interpolate an expression with arguments into a Function.") V = self.function_space() interp = interpolate(expression, V, **kwargs) return assemble(interp, tensor=self, ad_block_tag=ad_block_tag) diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index c0b89c2d36..aaa1469e12 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -340,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. diff --git a/firedrake/matrix.py b/firedrake/matrix.py index f23f521c74..4bdfa0d699 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -3,7 +3,6 @@ import itertools import ufl -from ufl.argument import BaseArgument from pyop2.utils import as_tuple from pyop2 import op2 from firedrake.petsc import PETSc @@ -46,11 +45,11 @@ def __init__( a A UFL BaseForm (with two arguments) that this MatrixBase represents, or a tuple of the arguments it represents, or a slate TensorBase. - fc_params - A dictionary of form compiler parameters for this matrix. bcs An optional iterable of boundary conditions to apply to this :class:`MatrixBase`. - None by default. + Empty tuple by default. + fc_params + A dictionary of form compiler parameters for this matrix. """ if isinstance(a, tuple): self.a = None @@ -171,7 +170,7 @@ def __init__( bcs An iterable of boundary conditions to apply to this :class:`Matrix`. May be `None` if there are no boundary conditions to apply. - By default None. + Empty tuple by default. fc_params A dictionary of form compiler parameters for this matrix, by default None. options_prefix @@ -219,7 +218,7 @@ def __init__( bcs An iterable of boundary conditions to apply to this :class:`Matrix`. May be `None` if there are no boundary conditions to apply. - By default None. + Empty tuple by default. fc_params A dictionary of form compiler parameters for this matrix, by default None. options_prefix @@ -265,7 +264,8 @@ def __init__( 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. By default None. + 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. """ diff --git a/firedrake/matrix_free/operators.py b/firedrake/matrix_free/operators.py index 49b1271622..92550485ee 100644 --- a/firedrake/matrix_free/operators.py +++ b/firedrake/matrix_free/operators.py @@ -73,8 +73,8 @@ class ImplicitMatrixContext(object): def __init__( self, a: ufl.BaseForm, - row_bcs: Iterable[DirichletBC] | None = None, - col_bcs: Iterable[DirichletBC] | None = None, + row_bcs: Iterable[DirichletBC] = (), + col_bcs: Iterable[DirichletBC] = (), fc_params: dict[str, Any] | None = None, appctx: dict[str, Any] | None = None ): @@ -88,10 +88,10 @@ def __init__( 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. By default None. + of the diagonal. Empty tuple by default. col_bcs An iterable of the :class.`.DirichletBC`s that are imposed - on the trial space. By default None. + on the trial space. Empty tuple by default. fc_params A dictionary of parameters to pass on to the form compiler. By default None. From 0288460f169cedc014a567651f188eb612de4053 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Tue, 3 Feb 2026 13:57:54 +0000 Subject: [PATCH 23/25] ignore pyop2 mat in docs --- docs/source/conf.py | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/source/conf.py b/docs/source/conf.py index ccd3583b24..a19bdeadeb 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -138,6 +138,7 @@ (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 From 2f0784a647a51173dc8f2dee2fa302a7d220da7f Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Mon, 9 Feb 2026 16:35:17 +0000 Subject: [PATCH 24/25] review suggestions --- docs/source/conf.py | 1 - firedrake/matrix.py | 1 - 2 files changed, 2 deletions(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index a19bdeadeb..da6ae4c19b 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -134,7 +134,6 @@ (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'), diff --git a/firedrake/matrix.py b/firedrake/matrix.py index 4bdfa0d699..2bd96c58b5 100644 --- a/firedrake/matrix.py +++ b/firedrake/matrix.py @@ -169,7 +169,6 @@ def __init__( The underlying matrix object. Either a PyOP2 Mat or a PETSc Mat. 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. From f6cdeb2c95485ecfa57104c58bae99facebd8fe4 Mon Sep 17 00:00:00 2001 From: Leo Collins Date: Tue, 10 Feb 2026 15:32:47 +0000 Subject: [PATCH 25/25] fix --- tests/firedrake/regression/test_interpolate_cross_mesh.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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))