Use CachedSolverBlock for project (-> #4638)

firedrake/adjoint_utils/__init__.py

@@ -9,7 +9,7 @@
     FunctionMixin, CofunctionMixin
 )
 from firedrake.adjoint_utils.assembly import annotate_assemble  # noqa F401
-from firedrake.adjoint_utils.projection import annotate_project  # noqa F401
+from firedrake.adjoint_utils.projection import annotate_super_project  # noqa F401
 from firedrake.adjoint_utils.variational_solver import (  # noqa F401
     NonlinearVariationalProblemMixin, NonlinearVariationalSolverMixin
 )

firedrake/adjoint_utils/function.py

@@ -1,10 +1,10 @@
 from functools import wraps
 from pyop2.mpi import temp_internal_comm
 import ufl
 from ufl.domain import extract_unique_domain
 from pyadjoint.overloaded_type import create_overloaded_object, FloatingType
 from pyadjoint.tape import annotate_tape, stop_annotating, get_working_tape
-from firedrake.adjoint_utils.blocks import FunctionAssignBlock, ProjectBlock, SubfunctionBlock, FunctionMergeBlock, SupermeshProjectBlock
+from firedrake.adjoint_utils.blocks import FunctionAssignBlock, SubfunctionBlock, FunctionMergeBlock
 import firedrake
 from .checkpointing import disk_checkpointing, CheckpointFunction, \
     CheckpointBase, checkpoint_init_data, DelegatedFunctionCheckpoint
@@ -27,32 +27,6 @@
             init(self, *args, **kwargs)
         return wrapper
 
-    @staticmethod
-    def _ad_annotate_project(project):
-        @wraps(project)
-        def wrapper(self, b, *args, **kwargs):
-            ad_block_tag = kwargs.pop("ad_block_tag", None)
-            annotate = annotate_tape(kwargs)
-
-            if annotate:
-                bcs = kwargs.get("bcs", [])
-                if isinstance(b, firedrake.Function) and extract_unique_domain(b) != self.function_space().mesh():
-                    block = SupermeshProjectBlock(b, self.function_space(), self, bcs, ad_block_tag=ad_block_tag)
-                else:
-                    block = ProjectBlock(b, self.function_space(), self, bcs, ad_block_tag=ad_block_tag)
-
-                tape = get_working_tape()
-                tape.add_block(block)
-
-            with stop_annotating():
-                output = project(self, b, *args, **kwargs)
-
-            if annotate:
-                block.add_output(output.create_block_variable())
-
-            return output
-        return wrapper
-
     @staticmethod
     def _ad_annotate_subfunctions(subfunctions):
         @wraps(subfunctions)

firedrake/adjoint_utils/projection.py

@@ -1,45 +1,35 @@
 from functools import wraps
 from pyadjoint.tape import annotate_tape, stop_annotating, get_working_tape
 from firedrake.adjoint_utils.blocks import ProjectBlock, SupermeshProjectBlock
 from firedrake import function
 from ufl.domain import extract_unique_domain
 
 
-def annotate_project(project):
-    @wraps(project)
-    def wrapper(*args, **kwargs):
-        """The project call performs an equation solve, and so it too must be annotated so that the
-        adjoint and tangent linear models may be constructed automatically by pyadjoint.
-
-        To disable the annotation of this function, just pass :py:data:`annotate=False`. This is useful in
-        cases where the solve is known to be irrelevant or diagnostic for the purposes of the adjoint
-        computation (such as projecting fields to other function spaces for the purposes of
-        visualisation)."""
+def annotate_super_project(project):
+    """
+    A wrapper for SupermeshProjector.project(), only handles
+    the code path that leads to the creation of a
+    SupermeshProjectBlock.
+    """
 
-        ad_block_tag = kwargs.pop("ad_block_tag", None)
+    @wraps(project)
+    def wrapper(self, **kwargs):
         annotate = annotate_tape(kwargs)
+        V = self.target.function_space()
         if annotate:
             bcs = kwargs.get("bcs", [])
             sb_kwargs = ProjectBlock.pop_kwargs(kwargs)
-            if isinstance(args[1], function.Function):
+            if self._target_is_function:
                 # block should be created before project because output might also be an input that needs checkpointing
-                output = args[1]
-                V = output.function_space()
-                if isinstance(args[0], function.Function) and extract_unique_domain(args[0]) != V.mesh():
-                    block = SupermeshProjectBlock(args[0], V, output, bcs, ad_block_tag=ad_block_tag, **sb_kwargs)
-                else:
-                    block = ProjectBlock(args[0], V, output, bcs, ad_block_tag=ad_block_tag, **sb_kwargs)
+                block = SupermeshProjectBlock(self.source, V, self.target, bcs, ad_block_tag=self.ad_block_tag, **sb_kwargs)
 
         with stop_annotating():
-            output = project(*args, **kwargs)
+            output = project(self, **kwargs)
 
         if annotate:
             tape = get_working_tape()
-            if not isinstance(args[1], function.Function):
-                if isinstance(args[0], function.Function) and extract_unique_domain(args[0]) != args[1].mesh():
-                    block = SupermeshProjectBlock(args[0], args[1], output, bcs, ad_block_tag=ad_block_tag, **sb_kwargs)
-                else:
-                    block = ProjectBlock(args[0], args[1], output, bcs, ad_block_tag=ad_block_tag, **sb_kwargs)
+            if not self._target_is_function:
+                block = SupermeshProjectBlock(self.source, V, output, bcs, ad_block_tag=self.ad_block_tag, **sb_kwargs)
             tape.add_block(block)
             block.add_output(output.create_block_variable())
 

firedrake/function.py

@@ -340,7 +340,6 @@ def sub(self, i):
         return data[i]
 
     @PETSc.Log.EventDecorator()
-    @FunctionMixin._ad_annotate_project
     def project(self, b, *args, **kwargs):
         r"""Project ``b`` onto ``self``. ``b`` must be a :class:`Function` or a
         UFL expression.

firedrake/projection.py

@@ -6,14 +6,14 @@
 from typing import Optional, Union
 
 import firedrake
+from firedrake.adjoint_utils import NonlinearVariationalSolverMixin, annotate_super_project
 from firedrake.bcs import BCBase
 from firedrake.petsc import PETSc
 from functools import cached_property
 
 from firedrake.utils import complex_mode, SLATE_SUPPORTS_COMPLEX
 from firedrake import functionspaceimpl
 from firedrake import function
-from firedrake.adjoint_utils import annotate_project
 
 
 __all__ = ['project', 'Projector']
@@ -51,7 +51,6 @@ def check_meshes(source, target):
 
 
 @PETSc.Log.EventDecorator()
-@annotate_project
 def project(
     v: ufl.core.expr.Expr,
     V: Union[firedrake.functionspaceimpl.WithGeometry, firedrake.Function],
@@ -84,8 +83,6 @@ def project(
         Quadrature degree to use when approximating integrands.
     name
         The name of the resulting :class:`.Function`.
-    ad_block_tag
-        String for tagging the resulting block on the Pyadjoint tape.
 
     Returns
     -------
@@ -111,19 +108,22 @@ def project(
         solver_parameters=solver_parameters,
         form_compiler_parameters=form_compiler_parameters,
         use_slate_for_inverse=use_slate_for_inverse,
-        quadrature_degree=quadrature_degree
+        quadrature_degree=quadrature_degree,
+        ad_block_tag=ad_block_tag,
     ).project()
     val.rename(name)
     return val
 
 
 class Assigner(object):
-    def __init__(self, source, target):
+    def __init__(self, source, target, **kwargs):
         self.source = source
         self.target = target
 
+        self._kwargs = kwargs
+
     def project(self):
-        self.target.assign(self.source)
+        self.target.assign(self.source, **self._kwargs)
         return self.target
 
 
@@ -236,7 +236,7 @@ def project(self):
         return self.target
 
 
-class BasicProjector(ProjectorBase):
+class BasicProjector(ProjectorBase, NonlinearVariationalSolverMixin):
     """
     A basic projector projects a UFL expression into a function space
     and places the result in a function from that function space,
@@ -245,6 +245,32 @@ class BasicProjector(ProjectorBase):
     defined on the same mesh.
     """
 
+    def __init__(self, *args, **kwargs):
+        ad_block_tag = kwargs.pop("ad_block_tag", None)
+        super().__init__(*args, **kwargs)
+
+        V = self.target.function_space()
+        mesh = V.mesh()
+
+        dx = firedrake.dx(mesh)
+        w = firedrake.TestFunction(V)
+        Pv = firedrake.TrialFunction(V)
+        a = firedrake.inner(Pv, w) * dx
+        L = firedrake.inner(self.source, w) * dx
+        p = firedrake.LinearVariationalProblem(a, L, self.target)
+
+        solver_params = firedrake.LinearVariationalSolver.DEFAULT_SNES_PARAMETERS | self.solver_parameters
+
+        self._init_as_solver(p, forward_kwargs={"solver_parameters": solver_params}, ad_block_tag=ad_block_tag)
+
+    @NonlinearVariationalSolverMixin._ad_annotate_init
+    def _init_as_solver(self, problem, **kwargs):
+        return
+
+    @NonlinearVariationalSolverMixin._ad_annotate_solve
+    def project(self):
+        return super().project()
+
     @cached_property
     def rhs_form(self):
         v = firedrake.TestFunction(self.target.function_space())
@@ -285,6 +311,12 @@ def rhs(self):
 
 
 class SupermeshProjector(ProjectorBase):
+    def __init__(self, *args, **kwargs):
+        self._target_is_function = kwargs.pop("target_is_function", True)
+        self.ad_block_tag = kwargs.pop("ad_block_tag", None)
+
+        super().__init__(*args, **kwargs)
+
     @cached_property
     def mixed_mass(self):
         from firedrake.supermeshing import assemble_mixed_mass_matrix
@@ -297,6 +329,10 @@ def rhs(self):
             self.mixed_mass.mult(u, v)
         return self.residual
 
+    @annotate_super_project
+    def project(self):
+        return super().project()
+
 
 @PETSc.Log.EventDecorator()
 def Projector(
@@ -307,7 +343,8 @@ def Projector(
     form_compiler_parameters: Optional[dict] = None,
     constant_jacobian: Optional[bool] = True,
     use_slate_for_inverse: Optional[bool] = False,
-    quadrature_degree: Optional[Union[int, tuple[int]]] = None
+    quadrature_degree: Optional[Union[int, tuple[int]]] = None,
+    ad_block_tag: Optional[str] = None,
 ):
     """ Projection class.
 
@@ -339,6 +376,8 @@ def Projector(
         function spaces)(only valid for DG function spaces).
     quadrature_degree
         Quadrature degree to use when approximating integrands.
+    ad_block_tag
+        String for tagging the resulting block on the Pyadjoint tape.
     """
     target = create_output(v_out)
     source = sanitise_input(v, target.function_space())
@@ -347,14 +386,15 @@ def Projector(
         raise ValueError("Shape mismatch between source %s and target %s in project" %
                          (source.ufl_shape, target.ufl_shape))
     if isinstance(v, function.Function) and not bcs and v.function_space() == target.function_space():
-        return Assigner(source, target)
+        return Assigner(source, target, ad_block_tag=ad_block_tag)
     elif source_mesh == target_mesh:
         return BasicProjector(
             source, target, bcs=bcs, solver_parameters=solver_parameters,
             form_compiler_parameters=form_compiler_parameters,
             constant_jacobian=constant_jacobian,
             use_slate_for_inverse=use_slate_for_inverse,
-            quadrature_degree=quadrature_degree
+            quadrature_degree=quadrature_degree,
+            ad_block_tag=ad_block_tag,
         )
     else:
         if bcs is not None:
@@ -366,5 +406,7 @@ def Projector(
             form_compiler_parameters=form_compiler_parameters,
             constant_jacobian=constant_jacobian,
             use_slate_for_inverse=use_slate_for_inverse,
-            quadrature_degree=quadrature_degree
+            quadrature_degree=quadrature_degree,
+            target_is_function=isinstance(v_out, firedrake.Function),
+            ad_block_tag=ad_block_tag,
         )

tests/firedrake/adjoint/test_tlm.py

@@ -234,7 +234,8 @@ def test_projection_function(rg):
 
     bc = DirichletBC(V, Constant(1.), "on_boundary")
     x, y = SpatialCoordinate(mesh)
-    g = project(sin(x)*sin(y), V, annotate=False)
+    with stop_annotating():
+        g = project(sin(x)*sin(y), V)
     expr = sin(g*x)
     f = project(expr, V)