Automated goal-oriented adaptivity

firedrake/__init__.py

@@ -160,6 +160,7 @@ def init_petsc():
     PointexprOperator, point_expr, MLOperator
 )
 from firedrake.progress_bar import ProgressBar  # noqa: F401
+from firedrake.dwr import dwr_marking_callback  # noqa: F401
 
 from firedrake.logging import (  # noqa: F401
     set_level, set_log_handlers, set_log_level, DEBUG, INFO,

firedrake/dmhooks.py

@@ -492,6 +492,9 @@ def _refine_adaptive(dm):
 
     coefficient_mapping = {}
     refined_ctx = refine(ctx, refine, coefficient_mapping=coefficient_mapping)
+    from firedrake.dwr import DWRMarkingCallback
+    if isinstance(ctx._marking_callback, DWRMarkingCallback):
+        refined_ctx._marking_callback = ctx._marking_callback.reconstruct(coefficient_mapping)
     parent = get_parent(dm)
     coarsener = get_ctx_coarsener(dm)
     # Get all DMs from the refined problem

firedrake/dwr.py

@@ -0,0 +1,236 @@
+import numpy as np
+import ufl
+from finat.ufl import BrokenElement, FiniteElement
+
+from firedrake.assemble import assemble
+from firedrake.function import Function
+from firedrake.functionspace import FunctionSpace, TensorFunctionSpace
+from firedrake.petsc import PETSc
+from firedrake.preconditioners.pmg import PMGPC
+from firedrake.ufl_expr import TestFunction, TrialFunction, adjoint, derivative
+from firedrake.variational_solver import (LinearVariationalProblem,
+                                          LinearVariationalSolver,
+                                          NonlinearVariationalProblem,
+                                          NonlinearVariationalSolver)
+from ufl import avg, dS, ds, dx, inner, replace
+
+
+__all__ = ("dwr_marking_callback",)
+
+
+def _replace_arguments(form, *arguments):
+    return replace(form, dict(zip(form.arguments(), arguments)))
+
+
+def _reconstruct_degree(V, degree):
+    element = PMGPC.reconstruct_degree(V.ufl_element(), degree)
+    return V.reconstruct(element=element)
+
+
+def _homogeneous_bcs(bcs, V):
+    return [bc.reconstruct(V=V, indices=bc._indices, g=0) for bc in bcs]
+
+
+def _both(expr):
+    return expr("+") + expr("-")
+
+
+def _residual_indicators(F, dual_error, residual_degree, options_prefix):
+    """Compute the strong residual representation of Rognes and Logg."""
+    v, = F.arguments()
+    V = v.function_space()
+    mesh = V.mesh().unique()
+    dim = mesh.topological_dimension
+    degree = V.ufl_element().degree() + residual_degree
+    variant = "integral"
+
+    bubble_space = FunctionSpace(mesh, "B", dim + 1, variant=variant)
+    bubble = Function(bubble_space).assign(1)
+    if V.value_shape == ():
+        cell_space = FunctionSpace(mesh, "DG", degree, variant=variant)
+    else:
+        cell_space = TensorFunctionSpace(mesh, "DG", degree,
+                                         shape=V.value_shape, variant=variant)
+    cell_trial = TrialFunction(cell_space)
+    cell_test = TestFunction(cell_space)
+    cell_residual = Function(cell_space)
+    cell_problem = LinearVariationalProblem(
+        inner(cell_trial, bubble * cell_test) * dx,
+        _replace_arguments(F, bubble * cell_test), cell_residual,
+    )
+    cell_solver = LinearVariationalSolver(
+        cell_problem, options_prefix=options_prefix + "dwr_cell_"
+    )
+    cell_solver.solve()
+
+    cone_space = FunctionSpace(mesh, "FB", dim, variant=variant)
+    cone = Function(cone_space).assign(1)
+    element = BrokenElement(FiniteElement("FB", cell=mesh.ufl_cell(),
+                                          degree=degree + dim, variant=variant))
+    if V.value_shape == ():
+        facet_space = FunctionSpace(mesh, element)
+    else:
+        facet_space = TensorFunctionSpace(mesh, element, shape=V.value_shape)
+    facet_trial = TrialFunction(facet_space)
+    facet_test = TestFunction(facet_space)
+    facet_residual_hat = Function(facet_space)
+    facet_rhs = (_replace_arguments(F, facet_test)
+                 - inner(cell_residual, facet_test) * dx)
+    facet_lhs = (_both(inner(facet_trial / cone, facet_test)) * dS
+                 + inner(facet_trial / cone, facet_test) * ds)
+    facet_problem = LinearVariationalProblem(
+        facet_lhs, facet_rhs, facet_residual_hat
+    )
+    facet_solver = LinearVariationalSolver(
+        facet_problem, options_prefix=options_prefix + "dwr_facet_"
+    )
+    facet_solver.solve()
+    facet_residual = facet_residual_hat / cone
+
+    indicator_space = FunctionSpace(mesh, "DG", 0)
+    indicator_test = TestFunction(indicator_space)
+    indicators = assemble(
+        inner(inner(cell_residual, dual_error), indicator_test) * dx
+        + inner(avg(inner(facet_residual, dual_error)),
+                _both(indicator_test)) * dS
+        + inner(inner(facet_residual, dual_error), indicator_test) * ds
+    )
+    with indicators.dat.vec as vec:
+        vec.abs()
+    return indicators
+
+
+def _dorfler_mark(indicators, fraction):
+    if not 0 < fraction <= 1:
+        raise ValueError("marking_fraction must lie in (0, 1]")
+    local = indicators.dat.data_ro.copy()
+    if not np.isfinite(local).all():
+        raise RuntimeError("DWR error indicators contain non-finite values")
+    gathered = indicators.comm.allgather(local)
+    values = np.concatenate(gathered)
+    total = values.sum()
+    markers = Function(indicators.function_space())
+    if total <= 0:
+        return markers.assign(1)
+    ordered = np.sort(values)[::-1]
+    count = np.searchsorted(np.cumsum(ordered), fraction * total) + 1
+    threshold = ordered[min(count - 1, len(ordered) - 1)]
+    markers.dat.data_wo[:] = local >= threshold
+    return markers
+
+
+class DWRMarkingCallback:
+    """Mark cells using an automatically localized dual-weighted residual."""
+
+    def __init__(self, goal_functional, primal=None, enrichment_degree=None):
+        if not isinstance(goal_functional, ufl.BaseForm) or goal_functional.arguments():
+            raise ValueError("goal_functional must be a 0-form")
+        self.goal_functional = goal_functional
+        self._primal = primal
+        self._enrichment_degree = enrichment_degree
+        self._high_space = None
+        if primal is not None:
+            V = primal.function_space()
+            self._high_space = _reconstruct_degree(
+                V, V.ufl_element().degree() + enrichment_degree
+            )
+
+    def setup(self, primal, options_prefix):
+        options = PETSc.Options(options_prefix)
+        self._primal = primal
+        self._enrichment_degree = options.getInt("dwr_enrichment_degree", 1)
+        V = primal.function_space()
+        self._high_space = _reconstruct_degree(
+            V, V.ufl_element().degree() + self._enrichment_degree
+        )
+
+    def reconstruct(self, coefficient_mapping):
+        goal = replace(self.goal_functional, coefficient_mapping)
+        primal = coefficient_mapping[self._primal]
+        return type(self)(goal, primal, self._enrichment_degree)
+
+    def __call__(self, ctx, current_solution):
+        return self._mark(ctx, current_solution)
+
+    def _mark(self, ctx, current_solution):
+        problem = ctx._problem
+        V = current_solution.function_space()
+        prefix = ctx.options_prefix or ""
+        options = PETSc.Options(prefix)
+        residual_degree = options.getInt("dwr_residual_degree", 1)
+        marking_fraction = options.getReal("dwr_marking_fraction", 0.5)
+        high_space = self._high_space
+        if high_space is None:
+            raise RuntimeError("DWR marking callback has not been set up")
+
+        dual_low = Function(V, name="dwr_dual_low")
+        direction = TestFunction(V)
+        goal_derivative = derivative(self.goal_functional, current_solution, direction)
+        rhs = assemble(goal_derivative, bcs=_homogeneous_bcs(problem.bcs, V))
+        ctx.solve_jacobian_transpose(rhs, dual_low)
+        if not np.isfinite(dual_low.dat.data_ro).all():
+            raise RuntimeError("DWR low-order dual solution contains non-finite values")
+
+        primal_high = Function(high_space, name="dwr_primal_high")
+        primal_high.interpolate(current_solution)
+        test, = problem.F.arguments()
+        test_high = test.reconstruct(function_space=high_space)
+        F_high = replace(problem.F, {current_solution: primal_high, test: test_high})
+        bcs_high = [bc.reconstruct(V=high_space, indices=bc._indices)
+                    for bc in problem.bcs]
+        high_problem = NonlinearVariationalProblem(F_high, primal_high, bcs_high)
+        primal_solver = NonlinearVariationalSolver(
+            high_problem,
+            options_prefix=prefix + "dwr_primal_",
+        )
+        primal_solver.solve()
+        if not np.isfinite(primal_high.dat.data_ro).all():
+            raise RuntimeError("DWR enriched primal solution contains non-finite values")
+
+        goal_high = replace(self.goal_functional, {current_solution: primal_high})
+        dual_high = Function(high_space, name="dwr_dual_high")
+        direction_high = TrialFunction(high_space)
+        dual_test_high = TestFunction(high_space)
+        jacobian_high = derivative(F_high, primal_high, direction_high)
+        goal_derivative_high = derivative(goal_high, primal_high, dual_test_high)
+        dual_problem = LinearVariationalProblem(
+            adjoint(jacobian_high), goal_derivative_high, dual_high,
+            bcs=_homogeneous_bcs(problem.bcs, high_space),
+        )
+        dual_solver = LinearVariationalSolver(
+            dual_problem,
+            options_prefix=prefix + "dwr_dual_",
+        )
+        dual_solver.solve()
+        if not np.isfinite(dual_high.dat.data_ro).all():
+            raise RuntimeError("DWR enriched dual solution contains non-finite values")
+
+        dual_error = dual_high - dual_low
+        indicators = _residual_indicators(
+            problem.F, dual_error, residual_degree, prefix
+        )
+        return _dorfler_mark(indicators, marking_fraction)
+
+
+def dwr_marking_callback(goal_functional: ufl.BaseForm):
+    """Construct a dual-weighted residual cell-marking callback.
+
+    Parameters
+    ----------
+    goal_functional
+        A scalar UFL 0-form depending on the primal solution.
+
+    Returns
+    -------
+    DWRMarkingCallback
+        A callback suitable for ``solve(..., marking_callback=...)``.
+
+    Notes
+    -----
+    Options are read from the active solver's PETSc options prefix. The
+    supported options are ``dwr_enrichment_degree`` (default 1),
+    ``dwr_residual_degree`` (default 1), and ``dwr_marking_fraction``
+    (default 0.5). The auxiliary solvers use the ``dwr_primal_``,
+    ``dwr_dual_``, ``dwr_cell_``, and ``dwr_facet_`` sub-prefixes.
+    """
+    return DWRMarkingCallback(goal_functional)

firedrake/solving.py

@@ -161,7 +161,7 @@ def _solve_varproblem(*args, **kwargs):
     eq, u, bcs, J, Jp, M, form_compiler_parameters, \
         solver_parameters, nullspace, nullspace_T, \
         near_nullspace, \
-        options_prefix, restrict, pre_apply_bcs = _extract_args(*args, **kwargs)
+        options_prefix, restrict, pre_apply_bcs, marking_callback = _extract_args(*args, **kwargs)
 
     # Check whether solution is valid
     if not isinstance(u, Function):
@@ -197,7 +197,8 @@ def _solve_varproblem(*args, **kwargs):
                            near_nullspace=near_nullspace,
                            options_prefix=options_prefix,
                            appctx=appctx,
-                           pre_apply_bcs=pre_apply_bcs)
+                           pre_apply_bcs=pre_apply_bcs,
+                           marking_callback=marking_callback)
     return solver.solve()
 
 
@@ -291,7 +292,8 @@ def _extract_args(*args, **kwargs):
     valid_kwargs = ["bcs", "J", "Jp", "M",
                     "form_compiler_parameters", "solver_parameters",
                     "nullspace", "transpose_nullspace", "near_nullspace",
-                    "options_prefix", "appctx", "restrict", "pre_apply_bcs"]
+                    "options_prefix", "appctx", "restrict", "pre_apply_bcs",
+                    "marking_callback"]
     for kwarg in kwargs.keys():
         if kwarg not in valid_kwargs:
             raise RuntimeError("Illegal keyword argument '%s'; valid keywords \
@@ -336,10 +338,11 @@ def _extract_args(*args, **kwargs):
     options_prefix = kwargs.get("options_prefix", None)
     restrict = kwargs.get("restrict", False)
     pre_apply_bcs = kwargs.get("pre_apply_bcs", True)
+    marking_callback = kwargs.get("marking_callback", None)
 
     return eq, u, bcs, J, Jp, M, form_compiler_parameters, \
         solver_parameters, nullspace, nullspace_T, near_nullspace, \
-        options_prefix, restrict, pre_apply_bcs
+        options_prefix, restrict, pre_apply_bcs, marking_callback
 
 
 def _extract_bcs(bcs):

firedrake/solving_utils.py

@@ -218,6 +218,7 @@ def __init__(self, problem,
         self._post_jacobian_callback = post_jacobian_callback
         self._post_function_callback = post_function_callback
         self._marking_callback = marking_callback
+        self._snes_ref = None
 
         self.fcp = problem.form_compiler_parameters
         # Function to hold current guess
@@ -302,7 +303,40 @@ def reconstruct(self, problem=None, mat_type=None, pmat_type=None, **kwargs):
         for k, v in default_options.items():
             if kwargs.get(k) is None:
                 kwargs[k] = v
-        return _SNESContext(problem, mat_type, pmat_type, **kwargs)
+        ctx = _SNESContext(problem, mat_type, pmat_type, **kwargs)
+        ctx._snes_ref = self._snes_ref
+        return ctx
+
+    def solve_jacobian_transpose(self, rhs, solution):
+        """Solve the transpose of the current SNES Jacobian.
+
+        Parameters
+        ----------
+        rhs
+            The dual right-hand side.
+        solution
+            The Function in which to store the solution.
+        """
+        if self._snes_ref is None:
+            raise RuntimeError("The SNES context is not attached to a solver")
+        snes = self._snes_ref()
+        if snes is None:
+            raise RuntimeError("The SNES attached to this context has been destroyed")
+        with rhs.dat.vec_ro as b, solution.dat.vec_wo as x:
+            if b.norm() == 0:
+                x.set(0)
+                return
+            ksp = snes.getKSP()
+            ksp.solveTranspose(b, x)
+            reason = ksp.getConvergedReason()
+            if reason < 0:
+                raise ConvergenceError(
+                    "Transpose Jacobian solve failed to converge after "
+                    f"{ksp.getIterationNumber()} iterations.\nReason:\n   "
+                    f"{KSPReasons.get(reason, reason)} "
+                    f"(PC type {ksp.getPC().getType()}, "
+                    f"failure {ksp.getPC().getFailedReason()})"
+                )
 
     @property
     def transfer_manager(self):

firedrake/variational_solver.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import ufl
+import weakref
 from itertools import chain
 from contextlib import ExitStack
 from types import MappingProxyType
@@ -312,6 +313,7 @@ def update_diffusivity(current_solution):
         self.snes = PETSc.SNES().create(comm=problem.dm.comm)
 
         self._ctx = ctx
+        self._ctx._snes_ref = weakref.ref(self.snes)
         self._work = problem.u_restrict.dof_dset.layout_vec.duplicate()
         self.snes.setDM(problem.dm)
         if marking_callback is not None:
@@ -357,6 +359,10 @@ def set_marking_callback(self, callback):
         """
         if not callable(callback):
             raise TypeError(f"marking callback must be callable, not a {type(callback).__name__}")
+        from firedrake.dwr import DWRMarkingCallback
+        if isinstance(callback, DWRMarkingCallback):
+            with self.inserted_options():
+                callback.setup(self._ctx._problem.u, self.options_prefix)
         self._ctx._marking_callback = callback
 
     def get_solution(self):
@@ -444,6 +450,7 @@ def solve(self, bounds=None):
                 self.snes.solve(None, work)
                 # The appctx might have been refined
                 self._ctx = dmhooks.get_appctx(self.snes.getDM())
+                self._ctx._snes_ref = weakref.ref(self.snes)
         problem = self._ctx._problem
         solution = self.snes.getSolution()
         with problem.u_restrict.dat.vec as u: