Source code for firedrake.variational_solver

import ufl
from itertools import chain
from contextlib import ExitStack

from firedrake import dmhooks, slate, solving, solving_utils, ufl_expr, utils
from firedrake import function
from firedrake.petsc import PETSc, OptionsManager, flatten_parameters
from firedrake.bcs import DirichletBC
from firedrake.adjoint_utils import NonlinearVariationalProblemMixin, NonlinearVariationalSolverMixin

__all__ = ["LinearVariationalProblem",
           "LinearVariationalSolver",
           "NonlinearVariationalProblem",
           "NonlinearVariationalSolver"]


def check_pde_args(F, J, Jp):
    if not isinstance(F, (ufl.BaseForm, slate.slate.TensorBase)):
        raise TypeError("Provided residual is a '%s', not a BaseForm or Slate Tensor" % type(F).__name__)
    if len(F.arguments()) != 1:
        raise ValueError("Provided residual is not a linear form")
    if not isinstance(J, (ufl.BaseForm, slate.slate.TensorBase)):
        raise TypeError("Provided Jacobian is a '%s', not a BaseForm or Slate Tensor" % type(J).__name__)
    if len(J.arguments()) != 2:
        raise ValueError("Provided Jacobian is not a bilinear form")
    if Jp is not None and not isinstance(Jp, (ufl.BaseForm, slate.slate.TensorBase)):
        raise TypeError("Provided preconditioner is a '%s', not a BaseForm or Slate Tensor" % type(Jp).__name__)
    if Jp is not None and len(Jp.arguments()) != 2:
        raise ValueError("Provided preconditioner is not a bilinear form")


def is_form_consistent(is_linear, bcs):
    # Check form style consistency
    if not (is_linear == all(bc.is_linear for bc in bcs if not isinstance(bc, DirichletBC))
            or not is_linear == all(not bc.is_linear for bc in bcs if not isinstance(bc, DirichletBC))):
        raise TypeError("Form style mismatch: some forms are given in 'F == 0' style, but others are given in 'A == b' style.")


[docs] class NonlinearVariationalProblem(NonlinearVariationalProblemMixin): r"""Nonlinear variational problem F(u; v) = 0.""" @PETSc.Log.EventDecorator() @NonlinearVariationalProblemMixin._ad_annotate_init def __init__(self, F, u, bcs=None, J=None, Jp=None, form_compiler_parameters=None, is_linear=False): r""" :param F: the nonlinear form :param u: the :class:`.Function` to solve for :param bcs: the boundary conditions (optional) :param J: the Jacobian J = dF/du (optional) :param Jp: a form used for preconditioning the linear system, optional, if not supplied then the Jacobian itself will be used. :param dict form_compiler_parameters: parameters to pass to the form compiler (optional) :is_linear: internally used to check if all domain/bc forms are given either in 'A == b' style or in 'F == 0' style. """ self.bcs = solving._extract_bcs(bcs) # Check form style consistency self.is_linear = is_linear is_form_consistent(self.is_linear, self.bcs) self.Jp_eq_J = Jp is None self.u = u self.F = F self.Jp = Jp if not isinstance(self.u, function.Function): raise TypeError("Provided solution is a '%s', not a Function" % type(self.u).__name__) # Use the user-provided Jacobian. If none is provided, derive # the Jacobian from the residual. self.J = J or ufl_expr.derivative(F, u) # Argument checking check_pde_args(self.F, self.J, self.Jp) # Store form compiler parameters self.form_compiler_parameters = form_compiler_parameters self._constant_jacobian = False
[docs] def dirichlet_bcs(self): for bc in self.bcs: yield from bc.dirichlet_bcs()
[docs] @utils.cached_property def dm(self): return self.u.function_space().dm
[docs] class NonlinearVariationalSolver(OptionsManager, NonlinearVariationalSolverMixin): r"""Solves a :class:`NonlinearVariationalProblem`.""" DEFAULT_SNES_PARAMETERS = {"snes_type": "newtonls", "snes_linesearch_type": "basic"} # Looser default tolerance for KSP inside SNES. DEFAULT_KSP_PARAMETERS = solving_utils.DEFAULT_KSP_PARAMETERS.copy() DEFAULT_KSP_PARAMETERS["ksp_rtol"] = 1e-5 @PETSc.Log.EventDecorator() @NonlinearVariationalSolverMixin._ad_annotate_init def __init__(self, problem, *, solver_parameters=None, options_prefix=None, nullspace=None, transpose_nullspace=None, near_nullspace=None, appctx=None, pre_jacobian_callback=None, post_jacobian_callback=None, pre_function_callback=None, post_function_callback=None): r""" :arg problem: A :class:`NonlinearVariationalProblem` to solve. :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 specify the near nullspace (for multigrid solvers). :kwarg solver_parameters: Solver parameters to pass to PETSc. This should be a dict mapping PETSc options to values. :kwarg appctx: A dictionary containing application context that is passed to the preconditioner if matrix-free. :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. :kwarg pre_jacobian_callback: A user-defined function that will be called immediately before Jacobian assembly. This can be used, for example, to update a coefficient function that has a complicated dependence on the unknown solution. :kwarg post_jacobian_callback: As above, but called after the Jacobian has been assembled. :kwarg pre_function_callback: As above, but called immediately before residual assembly. :kwarg post_function_callback: As above, but called immediately after residual assembly. Example usage of the ``solver_parameters`` option: to set the nonlinear solver type to just use a linear solver, use .. code-block:: python3 {'snes_type': 'ksponly'} PETSc flag options (where the presence of the option means something) should be specified with ``None``. For example: .. code-block:: python3 {'snes_monitor': None} To use the ``pre_jacobian_callback`` or ``pre_function_callback`` functionality, the user-defined function must accept the current solution as a petsc4py Vec. Example usage is given below: .. code-block:: python3 def update_diffusivity(current_solution): with cursol.dat.vec_wo as v: current_solution.copy(v) solve(trial*test*dx == dot(grad(cursol), grad(test))*dx, diffusivity) solver = NonlinearVariationalSolver(problem, pre_jacobian_callback=update_diffusivity) """ assert isinstance(problem, NonlinearVariationalProblem) solver_parameters = flatten_parameters(solver_parameters or {}) solver_parameters = solving_utils.set_defaults(solver_parameters, problem.J.arguments(), ksp_defaults=self.DEFAULT_KSP_PARAMETERS, snes_defaults=self.DEFAULT_SNES_PARAMETERS) super().__init__(solver_parameters, options_prefix) # Now the correct parameters live in self.parameters (via the # OptionsManager mixin) mat_type = self.parameters.get("mat_type") pmat_type = self.parameters.get("pmat_type") ctx = solving_utils._SNESContext(problem, mat_type=mat_type, pmat_type=pmat_type, appctx=appctx, pre_jacobian_callback=pre_jacobian_callback, pre_function_callback=pre_function_callback, post_jacobian_callback=post_jacobian_callback, post_function_callback=post_function_callback, options_prefix=self.options_prefix) self.snes = PETSc.SNES().create(comm=problem.dm.comm) self._problem = problem self._ctx = ctx self._work = problem.u.dof_dset.layout_vec.duplicate() self.snes.setDM(problem.dm) ctx.set_function(self.snes) ctx.set_jacobian(self.snes) ctx.set_nullspace(nullspace, problem.J.arguments()[0].function_space()._ises, transpose=False, near=False) ctx.set_nullspace(transpose_nullspace, problem.J.arguments()[1].function_space()._ises, transpose=True, near=False) ctx.set_nullspace(near_nullspace, problem.J.arguments()[0].function_space()._ises, transpose=False, near=True) ctx._nullspace = nullspace ctx._nullspace_T = transpose_nullspace ctx._near_nullspace = near_nullspace # Set from options now, so that people who want to noodle with # the snes object directly (mostly Patrick), can. We need the # DM with an app context in place so that if the DM is active # on a subKSP the context is available. dm = self.snes.getDM() with dmhooks.add_hooks(dm, self, appctx=self._ctx, save=False): self.set_from_options(self.snes) # Used for custom grid transfer. self._transfer_operators = () self._setup = False
[docs] def set_transfer_manager(self, manager): r"""Set the object that manages transfer between grid levels. Typically a :class:`~.TransferManager` object. :arg manager: Transfer manager, should conform to the TransferManager interface. :raises ValueError: if called after the transfer manager is setup. """ self._ctx.transfer_manager = manager
[docs] @PETSc.Log.EventDecorator() @NonlinearVariationalSolverMixin._ad_annotate_solve def solve(self, bounds=None): r"""Solve the variational problem. :arg bounds: Optional bounds on the solution (lower, upper). ``lower`` and ``upper`` must both be :class:`~.Function`\s. or :class:`~.Vector`\s. .. note:: If bounds are provided the ``snes_type`` must be set to ``vinewtonssls`` or ``vinewtonrsls``. """ # Make sure the DM has this solver's callback functions self._ctx.set_function(self.snes) self._ctx.set_jacobian(self.snes) # Make sure appcontext is attached to every coefficient DM before we solve. problem = self._problem forms = (problem.F, problem.J, problem.Jp) coefficients = utils.unique(chain.from_iterable(form.coefficients() for form in forms if form is not None)) # Make sure the solution dm is visited last solution_dm = self.snes.getDM() problem_dms = [V.dm for V in utils.unique(c.function_space() for c in coefficients) if V.dm != solution_dm] problem_dms.append(solution_dm) for dbc in problem.dirichlet_bcs(): dbc.apply(problem.u) if bounds is not None: lower, upper = bounds with lower.dat.vec_ro as lb, upper.dat.vec_ro as ub: self.snes.setVariableBounds(lb, ub) work = self._work with problem.u.dat.vec as u: u.copy(work) with ExitStack() as stack: # Ensure options database has full set of options (so monitors # work right) for ctx in chain([self.inserted_options()], [dmhooks.add_hooks(dm, self, appctx=self._ctx) for dm in problem_dms], self._transfer_operators): stack.enter_context(ctx) self.snes.solve(None, work) work.copy(u) self._setup = True solving_utils.check_snes_convergence(self.snes) # Grab the comm associated with the `_problem` and call PETSc's garbage cleanup routine comm = self._problem.u.function_space().mesh()._comm PETSc.garbage_cleanup(comm=comm)
[docs] class LinearVariationalProblem(NonlinearVariationalProblem): r"""Linear variational problem a(u, v) = L(v).""" @PETSc.Log.EventDecorator() def __init__(self, a, L, u, bcs=None, aP=None, form_compiler_parameters=None, constant_jacobian=False): r""" :param a: the bilinear form :param L: the linear form :param u: the :class:`.Function` to which the solution will be assigned :param bcs: the boundary conditions (optional) :param aP: an optional operator to assemble to precondition the system (if not provided a preconditioner may be computed from ``a``) :param dict form_compiler_parameters: parameters to pass to the form compiler (optional) :param constant_jacobian: (optional) flag indicating that the Jacobian is constant (i.e. does not depend on varying fields). If your Jacobian does not change, set this flag to ``True``. """ # In the linear case, the Jacobian is the equation LHS. J = a # Jacobian is checked in superclass, but let's check L here. if not isinstance(L, (ufl.BaseForm, slate.slate.TensorBase)) and L == 0: F = ufl_expr.action(J, u) else: if not isinstance(L, (ufl.BaseForm, slate.slate.TensorBase)): raise TypeError("Provided RHS is a '%s', not a Form or Slate Tensor" % type(L).__name__) if len(L.arguments()) != 1: raise ValueError("Provided RHS is not a linear form") F = ufl_expr.action(J, u) - L super(LinearVariationalProblem, self).__init__(F, u, bcs, J, aP, form_compiler_parameters=form_compiler_parameters, is_linear=True) self._constant_jacobian = constant_jacobian
[docs] class LinearVariationalSolver(NonlinearVariationalSolver): r"""Solves a :class:`LinearVariationalProblem`. :arg problem: A :class:`LinearVariationalProblem` to solve. :kwarg solver_parameters: Solver parameters to pass to PETSc. This should be a dict mapping PETSc options to values. :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 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. :kwarg appctx: A dictionary containing application context that is passed to the preconditioner if matrix-free. :kwarg pre_jacobian_callback: A user-defined function that will be called immediately before Jacobian assembly. This can be used, for example, to update a coefficient function that has a complicated dependence on the unknown solution. :kwarg post_jacobian_callback: As above, but called after the Jacobian has been assembled. :kwarg pre_function_callback: As above, but called immediately before residual assembly. :kwarg post_function_callback: As above, but called immediately after residual assembly. See also :class:`NonlinearVariationalSolver` for nonlinear problems. """ DEFAULT_SNES_PARAMETERS = {"snes_type": "ksponly"} DEFAULT_KSP_PARAMETERS = solving_utils.DEFAULT_KSP_PARAMETERS
[docs] def invalidate_jacobian(self): r""" Forces the matrix to be reassembled next time it is required. """ self._ctx._jacobian_assembled = False