Add support for objective function in SNES

docs/source/solving-interface.rst

@@ -725,6 +725,29 @@ solving with.
    # Use the approximate inverse of Jp to precondition solves
    solve(a == L, ..., Jp=Jp)
 
+Providing an objective functional
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+An objective functional, usually the energy,
+can be provided to PETSc's `SNES`_ linesearch and trust-region methods.
+For example, the code block below minimizes the Dirichlet energy functional
+using a trust-region method with the Steihaugh-Toint method for the
+trust-region subproblem and an initial trust region size of ``1e-4``.
+The solver monitor will report residual norm, norm of the update, and
+value of the objective functional.
+
+.. code-block:: python3
+
+   E = 0.5 * inner(grad(u), grad(u))*dx + kappa**2 * cosh(u)*dx - inner(f, u)*dx
+   F = derivative(E, u)
+   # Solve an optimisation problem
+   sp = {'snes_type': 'newtontr',
+         'snes_tr_delta0': 1e-4,
+         "snes_monitor": "::ascii_info_detail",
+         'ksp_type': 'cg',
+         'pc_type': 'hypre'}
+   solve(F == 0, u, ..., objective=E, solver_parameters=sp, pre_apply_bcs=False)
+
 Default solver options
 ~~~~~~~~~~~~~~~~~~~~~~
 

firedrake/bcs.py

@@ -8,7 +8,7 @@
 
 import ufl
 from ufl import as_ufl, as_tensor
-from finat.ufl import VectorElement
+from finat.ufl import MixedElement, TensorElement, VectorElement
 import finat
 
 import pyop2 as op2
@@ -746,3 +746,44 @@ def restricted_function_space(V, ids):
         return spaces[0]
     else:
         return firedrake.MixedFunctionSpace(spaces, name=V.name)
+
+
+def bcdofs(bc, ghost=True):
+    # Return the global dofs fixed by a DirichletBC
+    # in the numbering given by concatenation of all the
+    # subspaces of a mixed function space
+    Z = bc.function_space()
+    while Z.parent is not None:
+        Z = Z.parent
+
+    indices = bc._indices
+    offset = 0
+
+    for (i, idx) in enumerate(indices):
+        if isinstance(Z.ufl_element(), VectorElement):
+            offset += idx
+            assert i == len(indices)-1  # assert we're at the end of the chain
+            assert Z.sub(idx).block_size == 1
+        elif isinstance(Z.ufl_element(), MixedElement):
+            if ghost:
+                offset += sum(Z.sub(j).dof_count for j in range(idx))
+            else:
+                offset += sum(Z.sub(j).dof_dset.size * Z.sub(j).block_size for j in range(idx))
+        else:
+            raise NotImplementedError("How are you taking a .sub?")
+
+        Z = Z.sub(idx)
+
+    if Z.parent is not None and isinstance(Z.parent.ufl_element(), (TensorElement, VectorElement)):
+        bs = Z.parent.block_size
+        start = 0
+        stop = 1
+    else:
+        bs = Z.block_size
+        start = 0
+        stop = bs
+    nodes = bc.nodes
+    if not ghost:
+        nodes = nodes[nodes < Z.dof_dset.size]
+
+    return np.concatenate([nodes*bs + j for j in range(start, stop)]) + offset

firedrake/mg/ufl_utils.py

@@ -221,13 +221,16 @@ def inject_on_restrict(fine, restriction, rscale, injection, coarse):
         coefficient_mapping = {}
 
     bcs = [self(bc, self, coefficient_mapping=coefficient_mapping) for bc in problem.bcs]
+    E = self(problem.E, self, coefficient_mapping=coefficient_mapping)
     F = self(problem.F, self, coefficient_mapping=coefficient_mapping)
     J = self(problem.J, self, coefficient_mapping=coefficient_mapping)
     Jp = self(problem.Jp, self, coefficient_mapping=coefficient_mapping)
     u = coefficient_mapping[problem.u_restrict]
 
     fine = problem
-    problem = firedrake.NonlinearVariationalProblem(F, u, bcs=bcs, J=J, Jp=Jp, is_linear=problem.is_linear,
+    problem = firedrake.NonlinearVariationalProblem(F, u, bcs=bcs, J=J, Jp=Jp,
+                                                    objective=E,
+                                                    is_linear=problem.is_linear,
                                                     form_compiler_parameters=problem.form_compiler_parameters)
     fine._coarse = problem
     return problem

firedrake/preconditioners/bddc.py

@@ -1,7 +1,6 @@
 from itertools import repeat
 
 from firedrake.preconditioners.base import PCBase
-from firedrake.preconditioners.patch import bcdofs
 from firedrake.preconditioners.facet_split import get_restriction_indices
 from firedrake.petsc import PETSc
 from firedrake.dmhooks import get_function_space, get_appctx
@@ -13,7 +12,7 @@
 from functools import cached_property
 
 from firedrake.parloops import par_loop, INC, READ
-from firedrake.bcs import DirichletBC
+from firedrake.bcs import bcdofs, DirichletBC
 from firedrake.mesh import Submesh
 from ufl import Form, H1, H2, JacobianDeterminant, dx, inner, replace
 from finat.ufl import BrokenElement

firedrake/preconditioners/fdm.py

@@ -3,12 +3,12 @@
 from itertools import chain, product
 from firedrake.petsc import PETSc
 from firedrake.preconditioners.base import PCBase
-from firedrake.preconditioners.patch import bcdofs
 from firedrake.preconditioners.pmg import (prolongation_matrix_matfree,
                                            evaluate_dual,
                                            get_permutation_to_nodal_elements,
                                            cache_generate_code)
 from firedrake.preconditioners.facet_split import restricted_dofs, split_dofs
+from firedrake.bcs import bcdofs
 from firedrake.formmanipulation import ExtractSubBlock
 from firedrake.functionspace import FunctionSpace, MixedFunctionSpace
 from firedrake.function import Function

firedrake/preconditioners/patch.py

@@ -7,6 +7,7 @@
 from firedrake.preconditioners.asm import validate_overlap
 from firedrake.petsc import PETSc
 import firedrake.cython.patchimpl
+from firedrake.bcs import bcdofs
 from firedrake.solving_utils import _SNESContext
 from firedrake.utils import complex_mode
 from firedrake.dmhooks import get_appctx, push_appctx, pop_appctx
@@ -22,7 +23,6 @@
 import operator
 from functools import cached_property, partial
 import numpy
-from finat.ufl import VectorElement, MixedElement
 from tsfc.ufl_utils import extract_firedrake_constants
 import weakref
 import petsctools
@@ -587,47 +587,6 @@ def make_patch_callables(form: ufl.Form, state: Function | None) -> tuple[
     return cell_callable, interior_facet_callable, exterior_facet_callable
 
 
-def bcdofs(bc, ghost=True):
-    # Return the global dofs fixed by a DirichletBC
-    # in the numbering given by concatenation of all the
-    # subspaces of a mixed function space
-    Z = bc.function_space()
-    while Z.parent is not None:
-        Z = Z.parent
-
-    indices = bc._indices
-    offset = 0
-
-    for (i, idx) in enumerate(indices):
-        if isinstance(Z.ufl_element(), VectorElement):
-            offset += idx
-            assert i == len(indices)-1  # assert we're at the end of the chain
-            assert Z.sub(idx).block_size == 1
-        elif isinstance(Z.ufl_element(), MixedElement):
-            if ghost:
-                offset += sum(Z.sub(j).dof_count for j in range(idx))
-            else:
-                offset += sum(Z.sub(j).dof_dset.size * Z.sub(j).block_size for j in range(idx))
-        else:
-            raise NotImplementedError("How are you taking a .sub?")
-
-        Z = Z.sub(idx)
-
-    if Z.parent is not None and isinstance(Z.parent.ufl_element(), VectorElement):
-        bs = Z.parent.block_size
-        start = 0
-        stop = 1
-    else:
-        bs = Z.block_size
-        start = 0
-        stop = bs
-    nodes = bc.nodes
-    if not ghost:
-        nodes = nodes[nodes < Z.dof_dset.size]
-
-    return numpy.concatenate([nodes*bs + j for j in range(start, stop)]) + offset
-
-
 def select_entity(p, dm=None, exclude=None):
     """Filter entities based on some label.
 

firedrake/solving.py

@@ -149,7 +149,7 @@ def _solve_varproblem(*args, **kwargs):
     "Solve variational problem a == L or F == 0"
 
     # Extract arguments
-    eq, u, bcs, J, Jp, M, form_compiler_parameters, \
+    eq, u, bcs, J, Jp, objective, form_compiler_parameters, \
         solver_parameters, nullspace, nullspace_T, \
         near_nullspace, \
         options_prefix, restrict, pre_apply_bcs = _extract_args(*args, **kwargs)
@@ -167,6 +167,8 @@ def _solve_varproblem(*args, **kwargs):
         raise TypeError(f"Equation LHS must be a ufl.BaseForm, not a {type(eq.lhs).__name__}")
 
     if len(eq.lhs.arguments()) == 2:
+        if objective is not None:
+            raise ValueError("The objective functional only makes sense for nonlinear problems.")
         # Create linear variational problem
         problem = vs.LinearVariationalProblem(eq.lhs, eq.rhs, u, bcs, Jp,
                                               form_compiler_parameters=form_compiler_parameters,
@@ -177,6 +179,7 @@ def _solve_varproblem(*args, **kwargs):
         if eq.rhs != 0:
             raise ValueError(f"RHS of nonlinear Equation must be `0`, not {eq.rhs}")
         problem = vs.NonlinearVariationalProblem(eq.lhs, u, bcs, J, Jp,
+                                                 objective=objective,
                                                  form_compiler_parameters=form_compiler_parameters,
                                                  restrict=restrict)
         create_solver = vs.NonlinearVariationalSolver
@@ -279,7 +282,7 @@ def _extract_args(*args, **kwargs):
     "Extraction of arguments for _solve_varproblem"
 
     # Check for use of valid kwargs
-    valid_kwargs = ["bcs", "J", "Jp", "M",
+    valid_kwargs = ["bcs", "J", "Jp", "objective",
                     "form_compiler_parameters", "solver_parameters",
                     "nullspace", "transpose_nullspace", "near_nullspace",
                     "options_prefix", "appctx", "restrict", "pre_apply_bcs"]
@@ -314,9 +317,7 @@ def _extract_args(*args, **kwargs):
     Jp = kwargs.get("Jp", None)
 
     # Extract functional
-    M = kwargs.get("M", None)
-    if M is not None and not isinstance(M, ufl.Form):
-        raise RuntimeError("Expecting goal functional M to be a UFL Form")
+    objective = kwargs.get("objective", None)
 
     nullspace = kwargs.get("nullspace", None)
     nullspace_T = kwargs.get("transpose_nullspace", None)
@@ -328,7 +329,7 @@ def _extract_args(*args, **kwargs):
     restrict = kwargs.get("restrict", False)
     pre_apply_bcs = kwargs.get("pre_apply_bcs", True)
 
-    return eq, u, bcs, J, Jp, M, form_compiler_parameters, \
+    return eq, u, bcs, J, Jp, objective, form_compiler_parameters, \
         solver_parameters, nullspace, nullspace_T, near_nullspace, \
         options_prefix, restrict, pre_apply_bcs
 

firedrake/solving_utils.py

@@ -5,6 +5,7 @@
 
 from pyop2 import op2
 from firedrake import dmhooks
+from firedrake.bcs import bcdofs
 from firedrake.function import Function
 from firedrake.cofunction import Cofunction
 from firedrake.matrix import MatrixBase
@@ -231,6 +232,7 @@ def __init__(self, problem,
         self.pmatfree = pmatfree
         self.F = problem.F
         self.J = problem.J
+        self.E = problem.E
 
         # For Jp to equal J, bc.Jp must equal bc.J for all EquationBC objects.
         Jp_eq_J = problem.Jp is None and all(bc.Jp_eq_J for bc in problem.bcs)
@@ -261,6 +263,12 @@ def __init__(self, problem,
 
             self.F -= problem.compute_bc_lifting(self.J, self._bc_residual)
 
+        self._assemble_objective = lambda *args, **kwargs: args
+        if self.E:
+            self._assemble_objective = get_assembler(self.E,
+                                                     form_compiler_parameters=self.fcp,
+                                                     ).assemble
+
         self._assemble_residual = get_assembler(self.F, bcs=self.bcs_F,
                                                 form_compiler_parameters=self.fcp,
                                                 zero_bc_nodes=pre_apply_bcs,
@@ -277,6 +285,25 @@ def __init__(self, problem,
         self._coefficient_mapping = None
         self._transfer_manager = transfer_manager
 
+    @cached_property
+    def bc_iset(self):
+        if self._problem.restrict:
+            return None
+        bcs = tuple(self._problem.dirichlet_bcs())
+        V = self._x.function_space()
+        if len(bcs) > 0:
+            bc_nodes = numpy.unique(
+                numpy.concatenate([bcdofs(bc, ghost=False) for bc in bcs],
+                                  dtype=PETSc.IntType)
+            )
+            bc_nodes = V.dof_dset.lgmap.apply(bc_nodes)
+        else:
+            bc_nodes = numpy.empty(0, dtype=PETSc.IntType)
+
+        bc_is = PETSc.IS().createGeneral(bc_nodes, comm=V.comm)
+        bc_is.sort()
+        return bc_is
+
     def reconstruct(self, problem=None, mat_type=None, pmat_type=None, **kwargs):
         """Reconstruct this _SNESContext instance with new arguments."""
         problem = problem or self._problem
@@ -342,6 +369,12 @@ def transfer_manager(self, manager):
             raise ValueError("Must set transfer manager before first use.")
         self._transfer_manager = manager
 
+    def set_objective(self, snes):
+        if self._problem.E:
+            snes.setObjective(self.form_objective)
+        else:
+            snes.setObjective(None)
+
     def set_function(self, snes):
         r"""Set the residual evaluation function"""
         with self._F.dat.vec_wo as v:
@@ -360,6 +393,9 @@ def set_nullspace(self, nullspace, ises=None, transpose=False, near=False):
         if ises is not None:
             nullspace._apply(ises, transpose=transpose, near=near)
 
+    def set_ksp_postsolve(self, snes):
+        snes.getKSP().setPostSolve(self.ksp_postsolve)
+
     @PETSc.Log.EventDecorator()
     def split(self, fields):
         from firedrake import replace, as_vector, split, zero
@@ -449,6 +485,27 @@ def split(self, fields):
             splits.append(self.reconstruct(new_problem, options_prefix=options_prefix))
         return self._splits.setdefault(tuple(fields), splits)
 
+    @staticmethod
+    def form_objective(snes, X):
+        r"""Form the objective for this problem
+
+        :arg snes: a PETSc SNES object
+        :arg X: the current guess (a Vec)
+        """
+        dm = snes.getDM()
+        ctx = dmhooks.get_appctx(dm)
+        # X may not be the same vector as the vec behind self._x, so
+        # copy guess in from X.
+        with ctx._x.dat.vec_wo as v:
+            X.copy(v)
+
+        # Apply DirichletBC on the solution
+        if not ctx._problem.restrict:
+            for bc in ctx._problem.dirichlet_bcs():
+                bc.apply(ctx._x)
+
+        return ctx._assemble_objective()
+
     @staticmethod
     def form_function(snes, X, F):
         r"""Form the residual for this problem
@@ -547,6 +604,13 @@ def compute_operators(ksp, J, P):
             assert P.handle == ctx._pjac.petscmat.handle
             ctx._assemble_pjac(ctx._pjac)
 
+    @staticmethod
+    def ksp_postsolve(ksp, rhs, sol):
+        dm = ksp.getDM()
+        ctx = dmhooks.get_appctx(dm)
+        if ctx.pre_apply_bcs and not ctx._problem.restrict:
+            sol.isset(ctx.bc_iset, 0)
+
     @cached_property
     def _assembler_jac(self):
         from firedrake.assemble import get_assembler