Make LinearSolver a subclass of LinearVariationalSolver

firedrake/assemble.py

@@ -344,7 +344,9 @@ def __init__(self,
     def allocate(self):
         rank = len(self._form.arguments())
         if rank == 2 and not self._diagonal:
-            if self._mat_type == "matfree":
+            if isinstance(self._form, matrix.MatrixBase):
+                return self._form
+            elif self._mat_type == "matfree":
                 return MatrixFreeAssembler(self._form, bcs=self._bcs, form_compiler_parameters=self._form_compiler_params,
                                            options_prefix=self._options_prefix,
                                            appctx=self._appctx).allocate()
@@ -1358,7 +1360,8 @@ def allocate(self):
                                                           self._sub_mat_type,
                                                           self._make_maps_and_regions())
         return matrix.Matrix(self._form, self._bcs, self._mat_type, sparsity, ScalarType,
-                             options_prefix=self._options_prefix)
+                             options_prefix=self._options_prefix,
+                             fc_params=self._form_compiler_params)
 
     @staticmethod
     def _make_sparsity(test, trial, mat_type, sub_mat_type, maps_and_regions):

firedrake/formmanipulation.py

@@ -14,6 +14,7 @@
 from firedrake.petsc import PETSc
 from firedrake.functionspace import MixedFunctionSpace
 from firedrake.cofunction import Cofunction
+from firedrake.matrix import AssembledMatrix
 
 
 def subspace(V, indices):
@@ -147,6 +148,29 @@ def cofunction(self, o):
         else:
             return Cofunction(W, val=MixedDat(o.dat[i] for i in indices))
 
+    def matrix(self, o):
+        ises = []
+        args = []
+        for a in o.arguments():
+            V = a.function_space()
+            iset = PETSc.IS()
+            if a.number() in self.blocks:
+                asplit = self._subspace_argument(a)
+                for f in self.blocks[a.number()]:
+                    fset = V.dof_dset.field_ises[f]
+                    iset = iset.expand(fset)
+            else:
+                asplit = a
+                for fset in V.dof_dset.field_ises:
+                    iset = iset.expand(fset)
+
+            ises.append(iset)
+            args.append(asplit)
+
+        submat = o.petscmat.createSubMatrix(*ises)
+        bcs = ()
+        return AssembledMatrix(tuple(args), bcs, submat)
+
 
 SplitForm = collections.namedtuple("SplitForm", ["indices", "form"])
 

firedrake/linear_solver.py

@@ -1,34 +1,24 @@
-from firedrake.exceptions import ConvergenceError
-import firedrake.function as function
-import firedrake.cofunction as cofunction
-import firedrake.vector as vector
-import firedrake.matrix as matrix
-import firedrake.solving_utils as solving_utils
-from firedrake import dmhooks
-from firedrake.petsc import PETSc, OptionsManager, flatten_parameters
-from firedrake.utils import cached_property
-from firedrake.ufl_expr import action
+from firedrake.function import Function
+from firedrake.cofunction import Cofunction
+from firedrake.matrix import MatrixBase
+from firedrake.petsc import PETSc
 from pyop2.mpi import internal_comm
+from firedrake.variational_solver import LinearVariationalProblem, LinearVariationalSolver
 
 __all__ = ["LinearSolver"]
 
 
-class LinearSolver(OptionsManager):
-
-    DEFAULT_KSP_PARAMETERS = solving_utils.DEFAULT_KSP_PARAMETERS
+class LinearSolver(LinearVariationalSolver):
 
     @PETSc.Log.EventDecorator()
-    def __init__(self, A, *, P=None, solver_parameters=None,
-                 nullspace=None, transpose_nullspace=None,
-                 near_nullspace=None, options_prefix=None,
-                 pre_apply_bcs=True):
+    def __init__(self, A, *, P=None, **kwargs):
         """A linear solver for assembled systems (Ax = b).
 
         :arg A: a :class:`~.MatrixBase` (the operator).
         :arg P: an optional :class:`~.MatrixBase` to construct any
              preconditioner from; if none is supplied ``A`` is
              used to construct the preconditioner.
-        :kwarg parameters: (optional) dict of solver parameters.
+        :kwarg solver_parameters: (optional) dict of solver parameters.
         :kwarg nullspace: an optional :class:`~.VectorSpaceBasis` (or
             :class:`~.MixedVectorSpaceBasis` spanning the null space
             of the operator.
@@ -41,153 +31,59 @@ def __init__(self, A, *, P=None, solver_parameters=None,
                created.  Use this option if you want to pass options
                to the solver from the command line in addition to
                through the ``solver_parameters`` dict.
-        :kwarg pre_apply_bcs: Ignored by this class.
+        :kwarg pre_apply_bcs: If `True`, the bcs are applied before the solve.
+               Otherwise, the bcs are included as part of the linear system.
 
         .. note::
 
           Any boundary conditions for this solve *must* have been
           applied when assembling the operator.
         """
-        if not isinstance(A, matrix.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, matrix.MatrixBase):
+        if P is not None and not isinstance(P, MatrixBase):
             raise TypeError("Provided preconditioner is a '%s', not a MatrixBase" % type(P).__name__)
 
-        solver_parameters = flatten_parameters(solver_parameters or {})
-        solver_parameters = solving_utils.set_defaults(solver_parameters,
-                                                       A.arguments(),
-                                                       ksp_defaults=self.DEFAULT_KSP_PARAMETERS)
+        test, trial = A.arguments()
+        self.x = Function(trial.function_space())
+        self.b = Cofunction(test.function_space().dual())
+
+        problem = LinearVariationalProblem(A, self.b, self.x, bcs=A.bcs, aP=P,
+                                           form_compiler_parameters=A.fc_params,
+                                           constant_jacobian=True)
+        super().__init__(problem, **kwargs)
+
         self.A = A
         self.comm = A.comm
         self._comm = internal_comm(self.comm, self)
         self.P = P if P is not None else A
 
-        # Set up parameters mixin
-        super().__init__(solver_parameters, options_prefix)
-
-        self.A.petscmat.setOptionsPrefix(self.options_prefix)
-        self.P.petscmat.setOptionsPrefix(self.options_prefix)
-
-        # If preconditioning matrix is matrix-free, then default to jacobi
-        if isinstance(self.P, matrix.ImplicitMatrix):
-            self.set_default_parameter("pc_type", "jacobi")
-
-        self.ksp = PETSc.KSP().create(comm=self._comm)
-
-        W = self.test_space
-        # DM provides fieldsplits (but not operators)
-        self.ksp.setDM(W.dm)
-        self.ksp.setDMActive(False)
-
-        if nullspace is not None:
-            nullspace._apply(self.A)
-            if P is not None:
-                nullspace._apply(self.P)
-
-        if transpose_nullspace is not None:
-            transpose_nullspace._apply(self.A, transpose=True)
-            if P is not None:
-                transpose_nullspace._apply(self.P, transpose=True)
-
-        if near_nullspace is not None:
-            near_nullspace._apply(self.A, near=True)
-            if P is not None:
-                near_nullspace._apply(self.P, near=True)
-
-        self.nullspace = nullspace
-        self.transpose_nullspace = transpose_nullspace
-        self.near_nullspace = near_nullspace
-        # Operator setting must come after null space has been
-        # applied
-        self.ksp.setOperators(A=self.A.petscmat, P=self.P.petscmat)
-        # Set from options now (we're not allowed to change parameters
-        # anyway).
-        self.set_from_options(self.ksp)
-
-    @cached_property
-    def test_space(self):
-        return self.A.arguments()[0].function_space()
-
-    @cached_property
-    def trial_space(self):
-        return self.A.arguments()[1].function_space()
-
-    @cached_property
-    def _ulift(self):
-        return function.Function(self.trial_space)
-
-    @cached_property
-    def _blift(self):
-        return cofunction.Cofunction(self.test_space.dual())
-
-    @cached_property
-    def _update_rhs(self):
-        from firedrake.assemble import get_assembler
-        expr = -action(self.A.a, self._ulift)
-        # needs_zeroing = False adds -A * ulift to the exisiting value of the output tensor
-        # zero_bc_nodes = True writes the BC data on the boundary nodes of the output tensor
-        return get_assembler(expr, bcs=self.A.bcs, zero_bc_nodes=False, needs_zeroing=False).assemble
-
-    def _lifted(self, b):
-        self._ulift.dat.zero()
-        for bc in self.A.bcs:
-            bc.apply(self._ulift)
-        # Copy the input to the internal Cofunction so we do not modify the input
-        self._blift.assign(b)
-        self._update_rhs(tensor=self._blift)
-        # self._blift is now b - A u_bc, and satisfies the boundary conditions
-        return self._blift
+        self.ksp = self.snes.ksp
 
     @PETSc.Log.EventDecorator()
     def solve(self, x, b):
         """Solve the linear system with RHS ``b`` and store the solution in ``x``.
 
         Parameters
         ----------
-        x : firedrake.function.Function or firedrake.vector.Vector
-            Existing Function or Vector to place the solution to the linear system in.
-        b : firedrake.cofunction.Cofunction or firedrake.vector.Vector
-            A Cofunction or Vector with the right-hand side of the linear system.
+        x : firedrake.function.Function
+            Existing Function to place the solution to the linear system in.
+        b : firedrake.cofunction.Cofunction
+            A Cofunction with the right-hand side of the linear system.
         """
-        if not isinstance(x, (function.Function, vector.Vector)):
-            raise TypeError(f"Provided solution is a '{type(x).__name__}', not a Function or Vector")
-        if isinstance(x, vector.Vector):
-            x = x.function
-        if not isinstance(b, (cofunction.Cofunction, vector.Vector)):
-            raise TypeError(f"Provided RHS is a '{type(b).__name__}', not a Cofunction or Vector")
-        if isinstance(b, vector.Vector):
-            b = b.function
+        if not isinstance(x, Function):
+            raise TypeError(f"Provided solution is a '{type(x).__name__}', not a Function")
+        if not isinstance(b, Cofunction):
+            raise TypeError(f"Provided RHS is a '{type(b).__name__}', not a Cofunction")
 
         # When solving `Ax = b`, with A: V x U -> R, or equivalently A: V -> U*,
         # we need to make sure that x and b belong to V and U*, respectively.
-        if x.function_space() != self.trial_space:
-            raise ValueError(f"x must be a Function in {self.trial_space}.")
-        if b.function_space() != self.test_space.dual():
-            raise ValueError(f"b must be a Cofunction in {self.test_space.dual()}.")
-
-        if len(self.trial_space) > 1 and self.nullspace is not None:
-            self.nullspace._apply(self.trial_space.dof_dset.field_ises)
-        if len(self.test_space) > 1 and self.transpose_nullspace is not None:
-            self.transpose_nullspace._apply(self.test_space.dof_dset.field_ises,
-                                            transpose=True)
-        if len(self.trial_space) > 1 and self.near_nullspace is not None:
-            self.near_nullspace._apply(self.trial_space.dof_dset.field_ises, near=True)
-
-        if self.A.has_bcs:
-            b = self._lifted(b)
-
-        if self.ksp.getInitialGuessNonzero():
-            acc = x.dat.vec
-        else:
-            acc = x.dat.vec_wo
-
-        with self.inserted_options(), b.dat.vec_ro as rhs, acc as solution, dmhooks.add_hooks(self.ksp.dm, self):
-            self.ksp.solve(rhs, solution)
-
-        r = self.ksp.getConvergedReason()
-        if r < 0:
-            raise ConvergenceError("LinearSolver failed to converge after %d iterations with reason: %s", self.ksp.getIterationNumber(), solving_utils.KSPReasons[r])
-
-        # Grab the comm associated with `x` and call PETSc's garbage cleanup routine
-        comm = x.function_space().mesh()._comm
-        PETSc.garbage_cleanup(comm=comm)
+        if x.function_space() != self.x.function_space():
+            raise ValueError(f"x must be a Function in {self.x.function_space()}.")
+        if b.function_space() != self.b.function_space():
+            raise ValueError(f"b must be a Cofunction in {self.b.function_space()}.")
+
+        self.x.assign(x)
+        self.b.assign(b)
+        super().solve()
+        x.assign(self.x)

firedrake/matrix.py

@@ -20,8 +20,9 @@ class MatrixBase(ufl.Matrix):
         :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):
+    def __init__(self, a, bcs, mat_type, fc_params=None):
         if isinstance(a, tuple):
             self.a = None
             test, trial = a
@@ -52,6 +53,7 @@ def __init__(self, a, bcs, mat_type):
 
         Matrix type used in the assembly of the PETSc matrix: 'aij', 'baij', 'dense' or 'nest',
         or 'matfree' for matrix-free."""
+        self.fc_params = fc_params
 
     def arguments(self):
         if self.a:
@@ -109,6 +111,8 @@ class Matrix(MatrixBase):
 
     :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.
 
@@ -121,8 +125,9 @@ class Matrix(MatrixBase):
     """
 
     def __init__(self, a, bcs, mat_type, *args, **kwargs):
-        # sets self._a, self._bcs, and self._mat_type
-        MatrixBase.__init__(self, a, bcs, mat_type)
+        # 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")
         self.M = op2.Mat(*args, mat_type=mat_type, **kwargs)
         self.petscmat = self.M.handle
@@ -146,6 +151,7 @@ class ImplicitMatrix(MatrixBase):
         :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::
 
@@ -155,8 +161,9 @@ class ImplicitMatrix(MatrixBase):
 
     """
     def __init__(self, a, bcs, *args, **kwargs):
-        # sets self._a, self._bcs, and self._mat_type
-        super(ImplicitMatrix, self).__init__(a, bcs, "matfree")
+        # 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", {})
@@ -165,7 +172,7 @@ def __init__(self, a, bcs, *args, **kwargs):
         ctx = ImplicitMatrixContext(a,
                                     row_bcs=self.bcs,
                                     col_bcs=self.bcs,
-                                    fc_params=kwargs["fc_params"],
+                                    fc_params=fc_params,
                                     appctx=appctx)
         self.petscmat = PETSc.Mat().create(comm=self._comm)
         self.petscmat.setType("python")