from typing import Any
from collections.abc import Iterable
import itertools
import ufl
from pyop2.utils import as_tuple
from pyop2 import op2
from pyop2.types.mat import _GlobalMatPayload, _DatMatPayload
from firedrake.petsc import PETSc
from firedrake.bcs import DirichletBC
from firedrake.matrix_free import ImplicitMatrixContext
from firedrake.slate import slate
__all__ = ("MatrixBase", "Matrix", "ImplicitMatrix", "AssembledMatrix")
class DummyOP2Mat:
"""A hashable implementation of M.handle"""
def __init__(self, handle):
self.handle = handle
def _get_mat_type(petscmat: PETSc.Mat) -> str:
"""Maps PETSc matrix types to Firedrake notation"""
from firedrake.interpolation import VomOntoVomMatContext
mat_type = petscmat.getType()
if mat_type == "python":
ctx = petscmat.getPythonContext()
if isinstance(ctx, ImplicitMatrixContext):
return "matfree"
elif isinstance(ctx, _GlobalMatPayload):
return "global"
elif isinstance(ctx, _DatMatPayload):
return "dat"
elif isinstance(ctx, VomOntoVomMatContext):
return "vomtovom"
else:
raise NotImplementedError(
f"Python matrix context type '{type(ctx).__name__}' not recognised"
)
elif mat_type.startswith("seq") or mat_type.startswith("mpi"):
return mat_type[3:]
else:
return mat_type
[docs]
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.
Parameters
----------
a
A UFL BaseForm (with two arguments) that this MatrixBase represents,
or a tuple of the arguments it represents, or a slate TensorBase.
bcs
An optional iterable of boundary conditions to apply to this :class:`MatrixBase`.
Empty tuple by default.
fc_params
A dictionary of form compiler parameters for this matrix.
"""
def __init__(
self,
a: ufl.BaseForm | slate.TensorBase | tuple[ufl.Argument | ufl.Coargument, ufl.Argument | ufl.Coargument],
bcs: Iterable[DirichletBC] = (),
fc_params: dict[str, Any] | None = None,
):
"""Initialise a :class:`MatrixBase`."""
if isinstance(a, tuple):
self.a = None
test, trial = a
arguments = a
elif isinstance(a, ufl.BaseForm | slate.TensorBase):
self.a = a
test, trial = a.arguments()
arguments = None
else:
raise TypeError(f"'a' must be a UFL BaseForm, slate TensorBase, or tuple of arguments, not {type(a).__name__}.")
# Iteration over bcs must be in a parallel consistent order
# (so we can't use a set, since the iteration order may differ
# on different processes)
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
# this property after the fact.
self._analyze_form_arguments()
self._arguments = arguments
self.comm = test.function_space().comm
self.block_shape = (len(test.function_space()),
len(trial.function_space()))
self.bcs = bcs
self.form_compiler_parameters = {} if fc_params is None else fc_params
[docs]
def arguments(self):
if self.a:
return self.a.arguments()
else:
return self._arguments
[docs]
def ufl_domains(self):
return self._domains
@property
def has_bcs(self):
"""Return True if this :class:`MatrixBase` has any boundary
conditions attached to it."""
return self._bcs != ()
@property
def bcs(self):
"""The set of boundary conditions attached to this
:class:`.MatrixBase` (may be empty)."""
return self._bcs
@bcs.setter
def bcs(self, bcs):
"""Attach some boundary conditions to this :class:`MatrixBase`.
:arg bcs: a boundary condition (of type
:class:`.DirichletBC`), or an iterable of boundary
conditions. If bcs is None, erase all boundary conditions
on the :class:`.MatrixBase`.
"""
if bcs is not None:
self._bcs = tuple(itertools.chain(*(as_tuple(bc) for bc in bcs)))
else:
self._bcs = ()
def __repr__(self):
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})"
[docs]
def assign(self, val):
"""Set matrix entries."""
if isinstance(val, MatrixBase):
val.petscmat.copy(self.petscmat)
else:
raise TypeError(f"Cannot assign a {type(val).__name__} to a {type(self).__name__}.")
return self
[docs]
def zero(self):
"""Set all matrix entries to zero."""
self.petscmat.zeroEntries()
return self
[docs]
class Matrix(MatrixBase):
"""A representation of an assembled bilinear form.
Parameters
----------
a
The bilinear form this :class:`Matrix` represents.
mat
The underlying matrix object. Either a PyOP2 Mat or a PETSc Mat.
bcs
An iterable of boundary conditions to apply to this :class:`Matrix`.
Empty tuple by default.
fc_params
A dictionary of form compiler parameters for this matrix, by default None.
options_prefix
PETSc options prefix to apply, by default None.
"""
def __init__(
self,
a: ufl.BaseForm,
mat: op2.Mat | PETSc.Mat,
bcs: Iterable[DirichletBC] = (),
fc_params: dict[str, Any] | None = None,
options_prefix: str | None = None,
):
"""Initialise a :class:`Matrix`."""
super().__init__(a, bcs=bcs, fc_params=fc_params)
if isinstance(mat, op2.Mat):
self.M = mat
elif isinstance(mat, PETSc.Mat):
self.M = DummyOP2Mat(mat)
else:
raise TypeError(f"'mat' must be a PyOP2 Mat or PETSc Mat, not {type(mat).__name__}.")
self.petscmat = self.M.handle
if options_prefix is not None:
self.petscmat.setOptionsPrefix(options_prefix)
self.mat_type = _get_mat_type(self.petscmat)
[docs]
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.
Parameters
----------
a
The bilinear form this :class:`ImplicitMatrix` represents.
ctx
An :class:`firedrake.matrix_free.operators.ImplicitMatrixContext` that
defines the operations of the matrix.
bcs
An iterable of boundary conditions to apply to this :class:`Matrix`.
May be `None` if there are no boundary conditions to apply.
Empty tuple by default.
fc_params
A dictionary of form compiler parameters for this matrix, by default None.
options_prefix
PETSc options prefix to apply, by default None.
"""
def __init__(
self,
a: ufl.BaseForm,
ctx: ImplicitMatrixContext,
bcs: Iterable[DirichletBC] = (),
fc_params: dict[str, Any] | None = None,
options_prefix: str | None = None,
):
"""Initialise a :class:`ImplicitMatrix`."""
super().__init__(a, bcs=bcs, fc_params=fc_params)
self.petscmat = PETSc.Mat().create(comm=self.comm)
self.petscmat.setType("python")
self.petscmat.setSizes((ctx.row_sizes, ctx.col_sizes),
bsize=ctx.block_size)
self.petscmat.setPythonContext(ctx)
if options_prefix is not None:
self.petscmat.setOptionsPrefix(options_prefix)
self.petscmat.setUp()
self.petscmat.assemble()
self.mat_type = "matfree"
[docs]
def assemble(self):
# Bump petsc matrix state by assembling it.
# Ensures that if the matrix changed, the preconditioner is
# updated if necessary.
self.petscmat.assemble()
[docs]
class AssembledMatrix(MatrixBase):
"""A representation of a matrix that doesn't require knowing the underlying form.
Parameters
----------
args
A tuple of the arguments the matrix represents.
petscmat
The PETSc matrix this object wraps.
bcs
an iterable of boundary conditions to apply to this :class:`Matrix`.
May be `None` if there are no boundary conditions to apply.
Empty tuple by default.
options_prefix
PETSc options prefix to apply, by default None.
"""
def __init__(
self,
args: tuple[ufl.Argument | ufl.Coargument, ufl.Argument | ufl.Coargument],
petscmat: PETSc.Mat,
bcs: Iterable[DirichletBC] = (),
options_prefix: str | None = None,
):
"""Initialise an :class:`AssembledMatrix`."""
super().__init__(args, bcs=bcs)
self.petscmat = petscmat
if options_prefix is not None:
self.petscmat.setOptionsPrefix(options_prefix)
self.mat_type = _get_mat_type(self.petscmat)
# this mimics op2.Mat.handle
self.M = DummyOP2Mat(self.petscmat)
