From 31376963894b47eadb71024aa1329392235bbb31 Mon Sep 17 00:00:00 2001 From: Angus Gibson Date: Tue, 14 Jul 2026 13:32:43 +1000 Subject: [PATCH 1/5] Use CachedSolverBlock for project() The ProjectBlock was essentially an extension of SolveVarFormBlock that sets up the mass matrix system for the adjoint machinery. For the original (forward) project() call, the "proper" form is used, but the forward solver on the tape uses a simplified mass matrix. The approach here replicates that behaviour, but it's routed through the new CachedSolverBlock. As a consequence however, we have annotation on the Projector factory, which means that adjoint-friendly projection can occur through either the bare project() call, or by creating a Projector object. The implementation is a bit ugly: because a lot of the caching machinery belongs to the NonlinearVariationalSolverMixin, we essentially dispatch to a dummy initialiser when creating a BasicProjector to ensure the solver caches are created correctly. A call to project() annotates as though a solve occurred, but uses the underlying projection implementation. The SupermeshSolveBlock is actually a standalone object, so we just create that when SupermeshProjector.project() is called. Again, this allows for both bare and object-based projection to be taped. The final case is the "null" project from one function to another on the same space. This creates an Assigner, and I assume it will "just work" through the annotation on assign(). --- firedrake/adjoint_utils/__init__.py | 2 +- firedrake/adjoint_utils/projection.py | 37 ++++++++++--------------- firedrake/projection.py | 39 ++++++++++++++++++++++++--- 3 files changed, 50 insertions(+), 28 deletions(-) diff --git a/firedrake/adjoint_utils/__init__.py b/firedrake/adjoint_utils/__init__.py index e9096735a4..f9cec5e329 100644 --- a/firedrake/adjoint_utils/__init__.py +++ b/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 ) diff --git a/firedrake/adjoint_utils/projection.py b/firedrake/adjoint_utils/projection.py index e8114b4145..8415f1b6f9 100644 --- a/firedrake/adjoint_utils/projection.py +++ b/firedrake/adjoint_utils/projection.py @@ -5,41 +5,32 @@ 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. + """ + @wraps(project) + def wrapper(self, **kwargs): ad_block_tag = kwargs.pop("ad_block_tag", None) 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=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=ad_block_tag, **sb_kwargs) tape.add_block(block) block.add_output(output.create_block_variable()) diff --git a/firedrake/projection.py b/firedrake/projection.py index 5a2d83f371..5e63f07613 100644 --- a/firedrake/projection.py +++ b/firedrake/projection.py @@ -6,6 +6,7 @@ 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 @@ -13,7 +14,6 @@ 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], @@ -236,7 +235,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 +244,29 @@ class BasicProjector(ProjectorBase): defined on the same mesh. """ + def __init__(self, *args, **kwargs): + 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) + + self._init_as_solver(p) + + @NonlinearVariationalSolverMixin._ad_annotate_init + def _init_as_solver(self, problem): + 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 +307,10 @@ def rhs(self): class SupermeshProjector(ProjectorBase): + def __init__(self, *args, **kwargs): + self._target_is_function = kwargs.pop("target_is_function", True) + super().__init__(*args, **kwargs) + @cached_property def mixed_mass(self): from firedrake.supermeshing import assemble_mixed_mass_matrix @@ -297,6 +323,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( @@ -366,5 +396,6 @@ 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), ) From f05824b9390a3ebc77ef34d4d4a72972a56a5a28 Mon Sep 17 00:00:00 2001 From: Angus Gibson Date: Tue, 14 Jul 2026 14:37:20 +1000 Subject: [PATCH 2/5] Pass ad_block_tag through to Projector objects --- firedrake/adjoint_utils/projection.py | 5 ++--- firedrake/projection.py | 27 ++++++++++++++++++--------- 2 files changed, 20 insertions(+), 12 deletions(-) diff --git a/firedrake/adjoint_utils/projection.py b/firedrake/adjoint_utils/projection.py index 8415f1b6f9..c4f3a05b90 100644 --- a/firedrake/adjoint_utils/projection.py +++ b/firedrake/adjoint_utils/projection.py @@ -14,7 +14,6 @@ def annotate_super_project(project): @wraps(project) def wrapper(self, **kwargs): - ad_block_tag = kwargs.pop("ad_block_tag", None) annotate = annotate_tape(kwargs) V = self.target.function_space() if annotate: @@ -22,7 +21,7 @@ def wrapper(self, **kwargs): sb_kwargs = ProjectBlock.pop_kwargs(kwargs) if self._target_is_function: # block should be created before project because output might also be an input that needs checkpointing - block = SupermeshProjectBlock(self.source, V, self.target, 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(self, **kwargs) @@ -30,7 +29,7 @@ def wrapper(self, **kwargs): if annotate: tape = get_working_tape() if not self._target_is_function: - block = SupermeshProjectBlock(self.source, V, output, bcs, ad_block_tag=ad_block_tag, **sb_kwargs) + 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()) diff --git a/firedrake/projection.py b/firedrake/projection.py index 5e63f07613..8f5ea93ef6 100644 --- a/firedrake/projection.py +++ b/firedrake/projection.py @@ -83,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 ------- @@ -110,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 @@ -245,6 +246,7 @@ class BasicProjector(ProjectorBase, NonlinearVariationalSolverMixin): """ def __init__(self, *args, **kwargs): + ad_block_tag = kwargs.pop("ad_block_tag", None) super().__init__(*args, **kwargs) V = self.target.function_space() @@ -257,7 +259,7 @@ def __init__(self, *args, **kwargs): L = firedrake.inner(self.source, w) * dx p = firedrake.LinearVariationalProblem(a, L, self.target) - self._init_as_solver(p) + self._init_as_solver(p, ad_block_tag=ad_block_tag) @NonlinearVariationalSolverMixin._ad_annotate_init def _init_as_solver(self, problem): @@ -309,6 +311,8 @@ 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 @@ -337,7 +341,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. @@ -369,6 +374,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()) @@ -377,14 +384,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: @@ -398,4 +406,5 @@ def Projector( use_slate_for_inverse=use_slate_for_inverse, quadrature_degree=quadrature_degree, target_is_function=isinstance(v_out, firedrake.Function), + ad_block_tag=ad_block_tag, ) From eede792ebae03fe77623352980bb116bd8fe2720 Mon Sep 17 00:00:00 2001 From: Angus Gibson Date: Tue, 14 Jul 2026 14:43:40 +1000 Subject: [PATCH 3/5] Remove annotate=False from project() in test_tlm I guess this makes the changes API-breaking, so they should be smoothed over a bit. But for the moment, this makes the tests pass. --- tests/firedrake/adjoint/test_tlm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/firedrake/adjoint/test_tlm.py b/tests/firedrake/adjoint/test_tlm.py index 5eea04f253..625d8afc8a 100644 --- a/tests/firedrake/adjoint/test_tlm.py +++ b/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) From bcd7e10d0ae41ea05cd8a7ee7659296c279a5ca7 Mon Sep 17 00:00:00 2001 From: Angus Gibson Date: Tue, 14 Jul 2026 15:00:46 +1000 Subject: [PATCH 4/5] Remove Function.project() annotation --- firedrake/adjoint_utils/function.py | 28 +--------------------------- firedrake/function.py | 1 - 2 files changed, 1 insertion(+), 28 deletions(-) diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py index a4b83b9828..e588ce38f7 100644 --- a/firedrake/adjoint_utils/function.py +++ b/firedrake/adjoint_utils/function.py @@ -4,7 +4,7 @@ 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 @@ def wrapper(self, *args, **kwargs): 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) diff --git a/firedrake/function.py b/firedrake/function.py index 9d8a219fb7..08b39c816b 100644 --- a/firedrake/function.py +++ b/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. From c47e9b94f1f2b1960bea083031ce29f813ddc8fb Mon Sep 17 00:00:00 2001 From: Angus Gibson Date: Tue, 14 Jul 2026 14:44:37 +1000 Subject: [PATCH 5/5] Use solver_parameters from forward solver on cache It stands to reason (maybe), that the cached forward solver should just duplicate the solver parameters from the original forward solver. The TLM/adjoint solvers can probably use linear parameters with a mechanism for overriding. This way we can at least get the recomputation solve for project() to converge. --- firedrake/projection.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/firedrake/projection.py b/firedrake/projection.py index 8f5ea93ef6..10cc32eb95 100644 --- a/firedrake/projection.py +++ b/firedrake/projection.py @@ -259,10 +259,12 @@ def __init__(self, *args, **kwargs): L = firedrake.inner(self.source, w) * dx p = firedrake.LinearVariationalProblem(a, L, self.target) - self._init_as_solver(p, ad_block_tag=ad_block_tag) + 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): + def _init_as_solver(self, problem, **kwargs): return @NonlinearVariationalSolverMixin._ad_annotate_solve