Source code for firedrake.linear_solver

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 pyop2.mpi import internal_comm

__all__ = ["LinearSolver"]


[docs] class LinearSolver(OptionsManager): DEFAULT_KSP_PARAMETERS = solving_utils.DEFAULT_KSP_PARAMETERS @PETSc.Log.EventDecorator() def __init__(self, A, *, P=None, solver_parameters=None, nullspace=None, transpose_nullspace=None, near_nullspace=None, options_prefix=None): """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 nullspace: an optional :class:`~.VectorSpaceBasis` (or :class:`~.MixedVectorSpaceBasis` spanning the null space of the operator. :kwarg transpose_nullspace: as for the nullspace, but used to make the right hand side consistent. :kwarg near_nullspace: as for the nullspace, but used to set the near nullpace. :kwarg options_prefix: an optional prefix used to distinguish PETSc options. If not provided a unique prefix will be 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. .. note:: Any boundary conditions for this solve *must* have been applied when assembling the operator. """ if not isinstance(A, matrix.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): 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) 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)
[docs] @cached_property def test_space(self): return self.A.arguments()[0].function_space()
[docs] @cached_property def trial_space(self): return self.A.arguments()[1].function_space()
@cached_property def _rhs(self): from firedrake.assemble import get_assembler u = function.Function(self.trial_space) b = cofunction.Cofunction(self.test_space.dual()) expr = -action(self.A.a, u) return u, get_assembler(expr).assemble, b def _lifted(self, b): u, update, blift = self._rhs u.dat.zero() for bc in self.A.bcs: bc.apply(u) update(tensor=blift) # blift contains -A u_bc blift += b if isinstance(blift, cofunction.Cofunction): blift_func = blift.riesz_representation(riesz_map="l2") for bc in self.A.bcs: bc.apply(blift_func) blift.assign(blift_func.riesz_representation(riesz_map="l2")) else: for bc in self.A.bcs: bc.apply(blift) # blift is now b - A u_bc, and satisfies the boundary conditions return blift
[docs] @PETSc.Log.EventDecorator() def solve(self, x, b): if not isinstance(x, (function.Function, vector.Vector, cofunction.Cofunction)): raise TypeError("Provided solution is a '%s', not a Function, Vector or Cofunction" % type(x).__name__) if isinstance(b, vector.Vector): b = b.function if not isinstance(b, (function.Function, cofunction.Cofunction)): raise TypeError("Provided RHS is a '%s', not a Function or Cofunction" % type(b).__name__) 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)