From 4dfc06408923ebe26f4389a79e0412979ab62eb7 Mon Sep 17 00:00:00 2001 From: Angus Gibson Date: Thu, 12 Nov 2026 19:19:20 +1100 Subject: [PATCH 1/6] 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 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 | 4 +- firedrake/adjoint_utils/function.py | 29 +----- firedrake/adjoint_utils/projection.py | 139 ++++++++++++++++---------- firedrake/function.py | 1 - firedrake/projection.py | 44 ++++++-- 5 files changed, 122 insertions(+), 95 deletions(-) diff --git a/firedrake/adjoint_utils/__init__.py b/firedrake/adjoint_utils/__init__.py index 7d78643036..80fe9386fa 100644 --- a/firedrake/adjoint_utils/__init__.py +++ b/firedrake/adjoint_utils/__init__.py @@ -9,7 +9,9 @@ FunctionMixin, CofunctionMixin ) from firedrake.adjoint_utils.assembly import annotate_assemble # noqa F401 -from firedrake.adjoint_utils.projection import annotate_project, ProjectorMixin # noqa F401 +from firedrake.adjoint_utils.projection import ( # noqa F401 + annotate_project, annotate_super_project, ProjectorMixin +) from firedrake.adjoint_utils.variational_solver import ( # noqa F401 NonlinearVariationalProblemMixin, NonlinearVariationalSolverMixin ) diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py index 032a289b3d..cdd20cc101 100644 --- a/firedrake/adjoint_utils/function.py +++ b/firedrake/adjoint_utils/function.py @@ -1,9 +1,8 @@ from functools import wraps 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 @@ -26,32 +25,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/adjoint_utils/projection.py b/firedrake/adjoint_utils/projection.py index e847c74258..ab9e323946 100644 --- a/firedrake/adjoint_utils/projection.py +++ b/firedrake/adjoint_utils/projection.py @@ -1,45 +1,66 @@ +import firedrake 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 +from firedrake.adjoint_utils.variational_solver import NonlinearVariationalSolverMixin def annotate_project(project): + """A wrapper for the bare project() call to handle switching annotation + on/off without dispatching into all the types of projectors. + """ + @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. + annotate = annotate_tape(kwargs) + if annotate: + return project(*args, **kwargs) - 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).""" + with stop_annotating(): + return project(*args, **kwargs) - ad_block_tag = kwargs.pop("ad_block_tag", None) + return wrapper + + +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): 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): - # 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) + 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=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()) @@ -48,37 +69,45 @@ def wrapper(*args, **kwargs): return wrapper -class ProjectorMixin: +class ProjectorMixin(NonlinearVariationalSolverMixin): + """A wrapper over the nonlinear variational solver mixin for projectors. + + This establishes the form used for adjoint evaluation of the projector.""" + + def __init__(self, ad_block_tag=None): + V = self.target.function_space() + mesh = V.mesh() + ad_target = firedrake.Function(V) + + 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, ad_target) + + # If we project from an expression containing u (e.g. u + c) onto u, the + # resulting form looks nonlinear. However, the projection is in effect a + # solve onto a temporary function (here ad_target), followed by an implicit + # assignment. Instead of representing this as an explicit solve and assign + # block, we solve for a temporary function and modify _ad_u on the problem. + # This is used for creating the output variable (and thus the dependency + # resolution), in effect merging the implicit assignment into the block + # output phase. + p._ad_u = 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 + @staticmethod def _ad_annotate_project(project): - @wraps(project) - def wrapper(self, *args, **kwargs): - """Annotate the projection performed by a (possibly reused) - Projector, so that repeatedly calling ``project`` tapes one - ProjectBlock per call against the current value of the source. - - To disable the annotation, just pass :py:data:`annotate=False`.""" - ad_block_tag = kwargs.pop("ad_block_tag", None) - annotate = annotate_tape(kwargs) - if annotate: - bcs = self.bcs or [] - output = self.target - V = output.function_space() - # The block is created before projecting because the output - # might also be an input that needs checkpointing. - if isinstance(self.source, function.Function) and extract_unique_domain(self.source) != V.mesh(): - block = SupermeshProjectBlock(self.source, V, output, bcs, ad_block_tag=ad_block_tag) - else: - block = ProjectBlock(self.source, V, output, bcs, ad_block_tag=ad_block_tag) - tape = get_working_tape() - tape.add_block(block) - - with stop_annotating(): - output = project(self, *args, **kwargs) - - if annotate: - block.add_output(output.create_block_variable()) - - return output - - return wrapper + return NonlinearVariationalSolverMixin._ad_annotate_solve(project) diff --git a/firedrake/function.py b/firedrake/function.py index eef7ccce97..9783185b1a 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -356,7 +356,6 @@ def sub(self, indices: tuple[int] | int) -> "Function": return self._components[indices] @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. diff --git a/firedrake/projection.py b/firedrake/projection.py index 57c7fa4d8a..a4cb87570b 100644 --- a/firedrake/projection.py +++ b/firedrake/projection.py @@ -12,7 +12,7 @@ from firedrake import functionspaceimpl from firedrake import function -from firedrake.adjoint_utils import annotate_project, ProjectorMixin +from firedrake.adjoint_utils import annotate_project, annotate_super_project, ProjectorMixin __all__ = ['project', 'Projector'] @@ -110,19 +110,21 @@ 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: - 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 @@ -228,13 +230,12 @@ def residual(self): def rhs(self): pass - @ProjectorMixin._ad_annotate_project def project(self): self.apply_massinv(self.target, self.rhs) return self.target -class BasicProjector(ProjectorBase): +class BasicProjector(ProjectorBase, ProjectorMixin): """ A basic projector projects a UFL expression into a function space and places the result in a function from that function space, @@ -243,6 +244,10 @@ class BasicProjector(ProjectorBase): defined on the same mesh. """ + def __init__(self, *args, ad_block_tag=None, **kwargs): + ProjectorBase.__init__(self, *args, **kwargs) + ProjectorMixin.__init__(self, ad_block_tag=ad_block_tag) + @cached_property def rhs_form(self): v = firedrake.TestFunction(self.target.function_space()) @@ -270,6 +275,10 @@ def rhs_form(self): form = firedrake.inner(self.source, v)*firedrake.dx return form + @ProjectorMixin._ad_annotate_project + def project(self): + return super().project() + @cached_property def assembler(self): from firedrake.assemble import get_assembler @@ -283,6 +292,11 @@ 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 @@ -295,6 +309,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( @@ -305,7 +323,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. @@ -337,6 +356,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()) @@ -345,14 +366,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: @@ -364,5 +386,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, ) From b820eb2f199044ec9e8a1b356277c775964a8cde Mon Sep 17 00:00:00 2001 From: Angus Gibson Date: Fri, 13 Nov 2026 14:56:05 +1100 Subject: [PATCH 2/6] Create ProjectBlock through projection annotation The NLVS mixin has a static method _ad_block_type() that determines what type of block to create on the tape. This is overridden to ProjectBlock for projection operations. An accompanying test for the get_solve_blocks() method has been added. --- firedrake/adjoint_utils/blocks/solving.py | 18 ++---------------- firedrake/adjoint_utils/projection.py | 4 ++++ firedrake/adjoint_utils/variational_solver.py | 6 +++++- tests/firedrake/adjoint/test_projection.py | 11 +++++++++++ 4 files changed, 22 insertions(+), 17 deletions(-) diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py index b2f22f5d9c..a5e39b574c 100644 --- a/firedrake/adjoint_utils/blocks/solving.py +++ b/firedrake/adjoint_utils/blocks/solving.py @@ -1073,22 +1073,8 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, return dFdm -class ProjectBlock(SolveVarFormBlock): - def __init__(self, v, V, output, bcs=[], *args, **kwargs): - mesh = kwargs.pop("mesh", None) - if mesh is None: - mesh = V.mesh() - dx = firedrake.dx(mesh) - w = firedrake.TestFunction(V) - Pv = firedrake.TrialFunction(V) - a = firedrake.inner(Pv, w) * dx - L = firedrake.inner(v, w) * dx - - super().__init__(a == L, output, bcs, *args, **kwargs) - - def _init_solver_parameters(self, args, kwargs): - super()._init_solver_parameters(args, kwargs) - solve_init_params(self, args, kwargs, varform=True) +class ProjectBlock(CachedSolverBlock): + pass class SupermeshProjectBlock(Block): diff --git a/firedrake/adjoint_utils/projection.py b/firedrake/adjoint_utils/projection.py index ab9e323946..98cf717fdf 100644 --- a/firedrake/adjoint_utils/projection.py +++ b/firedrake/adjoint_utils/projection.py @@ -108,6 +108,10 @@ def __init__(self, ad_block_tag=None): def _init_as_solver(self, problem, **kwargs): return + @staticmethod + def _ad_block_type(): + return ProjectBlock + @staticmethod def _ad_annotate_project(project): return NonlinearVariationalSolverMixin._ad_annotate_solve(project) diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py index 4de5d04777..9e3aedbcb7 100644 --- a/firedrake/adjoint_utils/variational_solver.py +++ b/firedrake/adjoint_utils/variational_solver.py @@ -483,6 +483,10 @@ def _ad_hessian_cache(self): d2Fdudm_forms=d2Fdudm_forms, ) + @staticmethod + def _ad_block_type(): + return CachedSolverBlock + @staticmethod def _ad_annotate_solve(solve): @wraps(solve) @@ -497,7 +501,7 @@ def wrapper(self, **kwargs): raise ValueError( "MissingMathsError: we do not know how to differentiate through a variational inequality") - block = CachedSolverBlock( + block = self._ad_block_type()( self._ad_forward_cache, self._ad_tangent_cache, self._ad_adjoint_cache, diff --git a/tests/firedrake/adjoint/test_projection.py b/tests/firedrake/adjoint/test_projection.py index c059dd3dac..079b03a9da 100644 --- a/tests/firedrake/adjoint/test_projection.py +++ b/tests/firedrake/adjoint/test_projection.py @@ -179,3 +179,14 @@ def test_project_to_function_space(): J = assemble(w**2*dx) rf = ReducedFunctional(J, Control(c)) assert taylor_test(rf, Function(R, val=1.), Constant(0.1)) > 1.9 + + +@pytest.mark.skipcomplex +def test_project_blocks(): + mesh = UnitSquareMesh(12, 1) + V = FunctionSpace(mesh, "CG", 1) + R = FunctionSpace(mesh, "R", 0) + u = Function(V) + c = Function(R, val=1.) + project(u+c, u) + assert len(get_solve_blocks()) == 0 From 27fa63fcc598dd7e3d7505c9798f23953f6e64f4 Mon Sep 17 00:00:00 2001 From: Angus Gibson Date: Fri, 13 Nov 2026 16:50:08 +1100 Subject: [PATCH 3/6] Pass BCs into ProjectorMixin Also adds a test that verifies this actually works. --- firedrake/adjoint_utils/projection.py | 4 +++- tests/firedrake/adjoint/test_projection.py | 26 ++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/firedrake/adjoint_utils/projection.py b/firedrake/adjoint_utils/projection.py index 98cf717fdf..9732b6da75 100644 --- a/firedrake/adjoint_utils/projection.py +++ b/firedrake/adjoint_utils/projection.py @@ -84,7 +84,9 @@ def __init__(self, ad_block_tag=None): Pv = firedrake.TrialFunction(V) a = firedrake.inner(Pv, w) * dx L = firedrake.inner(self.source, w) * dx - p = firedrake.LinearVariationalProblem(a, L, ad_target) + p = firedrake.LinearVariationalProblem( + a, L, ad_target, bcs=self.bcs, + ) # If we project from an expression containing u (e.g. u + c) onto u, the # resulting form looks nonlinear. However, the projection is in effect a diff --git a/tests/firedrake/adjoint/test_projection.py b/tests/firedrake/adjoint/test_projection.py index 079b03a9da..6951184b3a 100644 --- a/tests/firedrake/adjoint/test_projection.py +++ b/tests/firedrake/adjoint/test_projection.py @@ -190,3 +190,29 @@ def test_project_blocks(): c = Function(R, val=1.) project(u+c, u) assert len(get_solve_blocks()) == 0 + + +@pytest.mark.skipcomplex +def test_project_bcs(): + mesh = UnitIntervalMesh(8) + V = FunctionSpace(mesh, "CG", 1) + m = Function(V).assign(1.0) + bc = DirichletBC(V, 3.0, "on_boundary") + u = Function(V) + + project(m, u, bcs=bc) + J = assemble(u**2 * dx) + rf = ReducedFunctional(J, Control(m)) + + # create new reference value independent of tape evaluation + with stop_annotating(): + m_new = Function(V).assign(2.0) + u_exp = Function(V) + project(m_new, u_exp, bcs=bc) + J_exp = assemble(u_exp**2 * dx) + + # ensure BCs are passed through the tape + assert rf(m_new) == pytest.approx( + J_exp, rel=1e-12, abs=1e-12, + ) + assert taylor_test(rf, m_new, Constant(0.1)) > 1.9 From a0eb14e4cf93670402a27a838da4a6634e8210b9 Mon Sep 17 00:00:00 2001 From: Angus Gibson Date: Fri, 13 Nov 2026 17:00:22 +1100 Subject: [PATCH 4/6] Clarify seemingly-redundant calls in annotate_project Here, the top-level annotate_project is acting more like a boundary on the annotate=False argument. --- firedrake/adjoint_utils/projection.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/firedrake/adjoint_utils/projection.py b/firedrake/adjoint_utils/projection.py index 9732b6da75..c739083b22 100644 --- a/firedrake/adjoint_utils/projection.py +++ b/firedrake/adjoint_utils/projection.py @@ -14,9 +14,14 @@ def annotate_project(project): def wrapper(*args, **kwargs): annotate = annotate_tape(kwargs) if annotate: + # annotation enabled: let the Projector or Assigner record + # its own block through the ProjectorMixin path return project(*args, **kwargs) with stop_annotating(): + # if global annotation was enabled, but annotation=False + # was passed to this call: take the same path, but ensuring + # all nested annotation is disabled return project(*args, **kwargs) return wrapper From a226b4dc5297f0906da248c272406f514d683a82 Mon Sep 17 00:00:00 2001 From: Angus Gibson Date: Mon, 3 Aug 2026 18:35:56 +1000 Subject: [PATCH 5/6] Remove non-used blocks --- firedrake/adjoint_utils/blocks/__init__.py | 4 +- firedrake/adjoint_utils/blocks/solving.py | 738 +----------------- firedrake/adjoint_utils/projection.py | 2 +- firedrake/adjoint_utils/solving.py | 4 +- firedrake/adjoint_utils/variational_solver.py | 63 +- 5 files changed, 11 insertions(+), 800 deletions(-) diff --git a/firedrake/adjoint_utils/blocks/__init__.py b/firedrake/adjoint_utils/blocks/__init__.py index fa10ff29fa..1e27c51a80 100644 --- a/firedrake/adjoint_utils/blocks/__init__.py +++ b/firedrake/adjoint_utils/blocks/__init__.py @@ -1,8 +1,6 @@ from firedrake.adjoint_utils.blocks.assembly import AssembleBlock # noqa F401 from firedrake.adjoint_utils.blocks.solving import ( # noqa F401 - CachedSolverBlock, GenericSolveBlock, - ProjectBlock, SupermeshProjectBlock, SolveVarFormBlock, - NonlinearVariationalSolveBlock + CachedSolverBlock, ProjectBlock, SupermeshProjectBlock ) from firedrake.adjoint_utils.blocks.function import ( # noqa F401 FunctionAssignBlock, FunctionMergeBlock, SubfunctionBlock diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py index a5e39b574c..1867ff34d5 100644 --- a/firedrake/adjoint_utils/blocks/solving.py +++ b/firedrake/adjoint_utils/blocks/solving.py @@ -1,12 +1,6 @@ -import numpy -import ufl -from ufl.domain import extract_domains, extract_unique_domain -from ufl import replace -from ufl.formatting.ufl2unicode import ufl2unicode from enum import Enum -from pyadjoint import Block, stop_annotating, get_working_tape -from pyadjoint.enlisting import Enlist +from pyadjoint import Block, stop_annotating import firedrake from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint @@ -305,520 +299,6 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_v return hessian_output -class GenericSolveBlock(Block): - pop_kwargs_keys = ["adj_cb", "adj_bdy_cb", "adj2_cb", "adj2_bdy_cb", - "forward_args", "forward_kwargs", "adj_args", - "adj_kwargs"] - - def __init__(self, lhs, rhs, func, bcs, *args, **kwargs): - super().__init__(ad_block_tag=kwargs.pop('ad_block_tag', None)) - self.adj_cb = kwargs.pop("adj_cb", None) - self.adj_bdy_cb = kwargs.pop("adj_bdy_cb", None) - self.adj2_cb = kwargs.pop("adj2_cb", None) - self.adj2_bdy_cb = kwargs.pop("adj2_bdy_cb", None) - - self.forward_args = [] - self.forward_kwargs = {} - self.adj_args = [] - self.adj_kwargs = {} - self.assemble_kwargs = {} - - # Equation LHS - self.lhs = lhs - # Equation RHS - self.rhs = rhs - # Solution function - self.func = func - self.function_space = self.func.function_space() - # Boundary conditions - self.bcs = [] - if bcs is not None: - self.bcs = Enlist(bcs) - - if isinstance(self.lhs, ufl.Form) and isinstance(self.rhs, (ufl.Form, ufl.Cofunction)): - self.linear = True - for c in self.rhs.coefficients(): - self.add_dependency(c, no_duplicates=True) - else: - self.linear = False - - for c in self.lhs.coefficients(): - self.add_dependency(c, no_duplicates=True) - - for bc in self.bcs: - self.add_dependency(bc, no_duplicates=True) - - try: # add all meshes as dependency - for mesh in extract_domains(self.lhs): - self.add_dependency(mesh, no_duplicates=True) - except AttributeError: - pass - - if isinstance(self.rhs, ufl.BaseForm): - # add all meshes as dependency - for mesh in extract_domains(self.rhs): - self.add_dependency(mesh, no_duplicates=True) - - self._init_solver_parameters(args, kwargs) - - def _init_solver_parameters(self, args, kwargs): - self.forward_args = kwargs.pop("forward_args", []) - self.forward_kwargs = kwargs.pop("forward_kwargs", {}) - self.adj_args = kwargs.pop("adj_args", []) - self.adj_kwargs = kwargs.pop("adj_kwargs", {}) - self.assemble_kwargs = {} - - def __str__(self): - try: - lhs_string = ufl2unicode(self.lhs) - except AttributeError: - lhs_string = str(self.lhs) - try: - rhs_string = ufl2unicode(self.rhs) - except AttributeError: - rhs_string = str(self.rhs) - return "solve({} = {})".format(lhs_string, rhs_string) - - def _create_F_form(self): - # Process the equation forms, replacing values with checkpoints, - # and gathering lhs and rhs in one single form. - if self.linear: - tmp_u = firedrake.Function(self.function_space) - F_form = firedrake.action(self.lhs, tmp_u) - self.rhs - else: - tmp_u = self.func - F_form = self.lhs - - replace_map = self._replace_map(F_form) - replace_map[tmp_u] = self.get_outputs()[0].saved_output - return ufl.replace(F_form, replace_map) - - def _homogenize_bcs(self): - bcs = [] - for bc in self.bcs: - if isinstance(bc, firedrake.DirichletBC): - bc = bc.reconstruct(g=0) - bcs.append(bc) - return bcs - - def _create_initial_guess(self): - return firedrake.Function(self.function_space) - - def _recover_bcs(self): - bcs = [] - for block_variable in self.get_dependencies(): - c = block_variable.output - c_rep = block_variable.saved_output - - if isinstance(c, firedrake.DirichletBC): - bcs.append(c_rep) - return bcs - - def _replace_map(self, form): - replace_coeffs = {} - for block_variable in self.get_dependencies(): - coeff = block_variable.output - if coeff in form.coefficients(): - replace_coeffs[coeff] = block_variable.saved_output - return replace_coeffs - - def _replace_form(self, form, func=None): - """Replace the form coefficients with checkpointed values - - func represents the initial guess if relevant. - """ - replace_map = self._replace_map(form) - if func is not None and self.func in replace_map: - firedrake.Function.assign(func, replace_map[self.func]) - replace_map[self.func] = func - return ufl.replace(form, replace_map) - - def _should_compute_boundary_adjoint(self, relevant_dependencies): - # Check if DirichletBC derivative is relevant - bdy = False - for _, dep in relevant_dependencies: - if isinstance(dep.output, firedrake.DirichletBC): - bdy = True - break - return bdy - - @property - def adj_sol(self): - return self.adj_state - - @adj_sol.setter - def adj_sol(self, value): - if self.adj_state is None: - self.adj_state = value.copy(deepcopy=True) - else: - self.adj_state.assign(value) - - def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies): - fwd_block_variable = self.get_outputs()[0] - u = fwd_block_variable.output - - dJdu = adj_inputs[0] - - F_form = self._create_F_form() - - dFdu = firedrake.derivative( - F_form, - fwd_block_variable.saved_output, - firedrake.TrialFunction( - u.function_space() - ) - ) - dFdu_form = firedrake.adjoint(dFdu) - dJdu = dJdu.copy() - - compute_bdy = self._should_compute_boundary_adjoint( - relevant_dependencies - ) - adj_sol, adj_sol_bdy = self._assemble_and_solve_adj_eq( - dFdu_form, dJdu, compute_bdy - ) - self.adj_sol = adj_sol - if self.adj_cb is not None: - self.adj_cb(adj_sol) - if self.adj_bdy_cb is not None and compute_bdy: - self.adj_bdy_cb(adj_sol_bdy) - - r = {} - r["form"] = F_form - r["adj_sol"] = adj_sol - r["adj_sol_bdy"] = adj_sol_bdy - return r - - def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, - prepared=None): - if not self.linear and self.func == block_variable.output: - # We are not able to calculate derivatives wrt initial guess. - return None - F_form = prepared["form"] - adj_sol = prepared["adj_sol"] - adj_sol_bdy = prepared["adj_sol_bdy"] - c = block_variable.output - c_rep = block_variable.saved_output - - if isinstance(c, (firedrake.Function, firedrake.Cofunction)): - trial_function = firedrake.TrialFunction(c.function_space()) - elif isinstance(c, firedrake.DirichletBC): - tmp_bc = c.reconstruct( - g=extract_subfunction(adj_sol_bdy, c.function_space()) - ) - return [tmp_bc] - elif isinstance(c, firedrake.MeshGeometry): - # Using CoordinateDerivative requires us to do action before - # differentiating, might change in the future. - F_form_tmp = firedrake.action(F_form, adj_sol) - X = firedrake.SpatialCoordinate(c_rep) - dFdm = firedrake.derivative( - -F_form_tmp, X, - firedrake.TestFunction(c._ad_function_space()) - ) - - if dFdm == 0: - return firedrake.Function(c._ad_function_space().dual()) - - dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs) - return dFdm - - dFdm = -firedrake.derivative(F_form, c_rep, trial_function) - dFdm = firedrake.adjoint(dFdm) - dFdm = firedrake.action(dFdm, adj_sol) - dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs) - return dFdm - - def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs): - return firedrake.assemble(dFdu_adj_form, **kwargs) - - def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy): - dJdu_copy = dJdu.copy() - # Homogenize and apply boundary conditions on adj_dFdu. - bcs = self._homogenize_bcs() - dFdu = firedrake.assemble(dFdu_adj_form, bcs=bcs, **self.assemble_kwargs) - - adj_sol = firedrake.Function(self.function_space) - firedrake.solve( - dFdu, adj_sol, dJdu, *self.adj_args, **self.adj_kwargs - ) - - adj_sol_bdy = None - if compute_bdy: - adj_sol_bdy = self._compute_adj_bdy( - adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu_copy) - - return adj_sol, adj_sol_bdy - - def _compute_adj_bdy(self, adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu): - adj_sol_bdy = firedrake.assemble(dJdu - firedrake.action(dFdu_adj_form, adj_sol)) - return adj_sol_bdy.riesz_representation("l2") - - def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs): - pass - - def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, - prepared=None): - fwd_block_variable = self.get_outputs()[0] - u = fwd_block_variable.output - - F_form = self._create_F_form() - - # Obtain dFdu. - dFdu = firedrake.derivative( - F_form, - fwd_block_variable.saved_output, - firedrake.TrialFunction(u.function_space()) - ) - V = self.get_outputs()[idx].output.function_space() - - bcs = [] - dFdm = 0. - for block_variable in self.get_dependencies(): - tlm_value = block_variable.tlm_value - c = block_variable.output - c_rep = block_variable.saved_output - - if isinstance(c, firedrake.DirichletBC): - if tlm_value is None: - bcs.append(c.reconstruct(g=0)) - else: - bcs.append(tlm_value) - continue - elif isinstance(c, firedrake.MeshGeometry): - X = firedrake.SpatialCoordinate(c) - c_rep = X - - if tlm_value is None: - continue - - if c == self.func and not self.linear: - continue - - dFdm += firedrake.derivative(-F_form, c_rep, tlm_value) - - if isinstance(dFdm, float): - v = dFdu.arguments()[0] - dFdm = firedrake.inner( - firedrake.Constant(numpy.zeros(v.ufl_shape)), v - ) * firedrake.dx - - dFdm = ufl.algorithms.expand_derivatives(dFdm) - dFdm = firedrake.assemble(dFdm) - dudm = firedrake.Function(V) - result = self._assemble_and_solve_tlm_eq( - firedrake.assemble(dFdu, bcs=bcs, **self.assemble_kwargs), - dFdm, dudm, bcs - ) - return result - - def _assemble_and_solve_tlm_eq(self, dFdu, dFdm, dudm, bcs): - return self._assembled_solve(dFdu, dFdm, dudm, bcs) - - def _assemble_soa_eq_rhs(self, dFdu_form, adj_sol, hessian_input, d2Fdu2): - # Start piecing together the rhs of the soa equation - b = hessian_input.copy() - if not d2Fdu2.empty(): - b_form = firedrake.action(firedrake.adjoint(d2Fdu2), adj_sol) - else: - b_form = d2Fdu2 - - for bo in self.get_dependencies(): - c = bo.output - c_rep = bo.saved_output - tlm_input = bo.tlm_value - - if (c == self.func and not self.linear) or tlm_input is None: - continue - - if isinstance(c, firedrake.MeshGeometry): - X = firedrake.SpatialCoordinate(c) - dFdu_adj = firedrake.action(firedrake.adjoint(dFdu_form), - adj_sol) - d2Fdudm = ufl.algorithms.expand_derivatives( - firedrake.derivative(dFdu_adj, X, tlm_input)) - if not d2Fdudm.empty(): - b_form += d2Fdudm - elif not isinstance(c, firedrake.DirichletBC): - dFdu_adj = firedrake.action(firedrake.adjoint(dFdu_form), - adj_sol) - # b_form += firedrake.derivative(dFdu_adj, c_rep, tlm_input) - bo_form = ufl.algorithms.expand_derivatives( - firedrake.derivative(dFdu_adj, c_rep, tlm_input)) - b_form += bo_form - - b_form = ufl.algorithms.expand_derivatives(b_form) - if not b_form.empty(): - b -= firedrake.assemble(b_form) - - return b - - def _assemble_and_solve_soa_eq(self, dFdu_form, adj_sol, hessian_input, - d2Fdu2, compute_bdy): - b = self._assemble_soa_eq_rhs(dFdu_form, adj_sol, hessian_input, - d2Fdu2) - dFdu_form = firedrake.adjoint(dFdu_form) - adj_sol2, adj_sol2_bdy = self._assemble_and_solve_adj_eq(dFdu_form, b, - compute_bdy) - if self.adj2_cb is not None: - self.adj2_cb(adj_sol2) - if self.adj2_bdy_cb is not None and compute_bdy: - self.adj2_bdy_cb(adj_sol2_bdy) - return adj_sol2, adj_sol2_bdy - - def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, - relevant_dependencies): - # First fetch all relevant values - fwd_block_variable = self.get_outputs()[0] - hessian_input = hessian_inputs[0] - tlm_output = fwd_block_variable.tlm_value - - if hessian_input is None: - return - - if tlm_output is None: - return - - F_form = self._create_F_form() - - # Using the equation Form derive dF/du, d^2F/du^2 * du/dm * direction. - dFdu_form = firedrake.derivative(F_form, - fwd_block_variable.saved_output) - d2Fdu2 = ufl.algorithms.expand_derivatives( - firedrake.derivative(dFdu_form, fwd_block_variable.saved_output, - tlm_output)) - - adj_sol = self.adj_sol - if adj_sol is None: - raise RuntimeError("Hessian computation was run before adjoint.") - bdy = self._should_compute_boundary_adjoint(relevant_dependencies) - adj_sol2, adj_sol2_bdy = self._assemble_and_solve_soa_eq( - dFdu_form, adj_sol, hessian_input, d2Fdu2, bdy - ) - - r = {} - r["adj_sol2"] = adj_sol2 - r["adj_sol2_bdy"] = adj_sol2_bdy - r["form"] = F_form - r["adj_sol"] = adj_sol - - return r - - def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, - block_variable, idx, relevant_dependencies, - prepared=None): - c = block_variable.output - if c == self.func and not self.linear: - return None - - adj_sol2 = prepared["adj_sol2"] - adj_sol2_bdy = prepared["adj_sol2_bdy"] - F_form = prepared["form"] - adj_sol = prepared["adj_sol"] - fwd_block_variable = self.get_outputs()[0] - tlm_output = fwd_block_variable.tlm_value - - c_rep = block_variable.saved_output - - # If m = DirichletBC then d^2F(u,m)/dm^2 = 0 and d^2F(u,m)/dudm = 0, - # so we only have the term dF(u,m)/dm * adj_sol2 - if isinstance(c, firedrake.DirichletBC): - tmp_bc = c.reconstruct( - g=extract_subfunction(adj_sol2_bdy, c.function_space()) - ) - return [tmp_bc] - - if isinstance(c, firedrake.MeshGeometry): - X = firedrake.SpatialCoordinate(c) - W = c._ad_function_space() - else: - W = c.function_space() - - dc = firedrake.TestFunction(W) - form_adj = firedrake.action(F_form, adj_sol) - form_adj2 = firedrake.action(F_form, adj_sol2) - if isinstance(c, firedrake.MeshGeometry): - dFdm_adj = firedrake.derivative(form_adj, X, dc) - dFdm_adj2 = firedrake.derivative(form_adj2, X, dc) - else: - dFdm_adj = firedrake.derivative(form_adj, c_rep, dc) - dFdm_adj2 = firedrake.derivative(form_adj2, c_rep, dc) - - # TODO: Old comment claims this might break on split. Confirm if true - # or not. - d2Fdudm = ufl.algorithms.expand_derivatives( - firedrake.derivative(dFdm_adj, fwd_block_variable.saved_output, - tlm_output)) - - d2Fdm2 = 0 - # We need to add terms from every other dependency - # i.e. the terms d^2F/dm_1dm_2 - for _, bv in relevant_dependencies: - c2 = bv.output - c2_rep = bv.saved_output - - if isinstance(c2, firedrake.DirichletBC): - continue - - tlm_input = bv.tlm_value - if tlm_input is None: - continue - - if c2 == self.func and not self.linear: - continue - - # TODO: If tlm_input is a Sum, this crashes in some instances? - if isinstance(c2_rep, firedrake.MeshGeometry): - X = firedrake.SpatialCoordinate(c2_rep) - d2Fdm2 += ufl.algorithms.expand_derivatives( - firedrake.derivative(dFdm_adj, X, tlm_input) - ) - else: - d2Fdm2 += ufl.algorithms.expand_derivatives( - firedrake.derivative(dFdm_adj, c2_rep, tlm_input) - ) - - hessian_form = ufl.algorithms.expand_derivatives( - d2Fdm2 + dFdm_adj2 + d2Fdudm - ) - hessian_output = 0 - if not hessian_form.empty(): - hessian_output = firedrake.assemble(hessian_form) - hessian_output *= -1. - - return hessian_output - - def prepare_recompute_component(self, inputs, relevant_outputs): - return self._replace_recompute_form() - - def _replace_recompute_form(self): - func = self._create_initial_guess() - - bcs = self._recover_bcs() - lhs = self._replace_form(self.lhs, func=func) - rhs = 0 - if self.linear: - rhs = self._replace_form(self.rhs) - - return lhs, rhs, func, bcs - - def _forward_solve(self, lhs, rhs, func, bcs): - firedrake.solve(lhs == rhs, func, bcs, *self.forward_args, - **self.forward_kwargs) - return func - - def _assembled_solve(self, lhs, rhs, func, bcs, **kwargs): - firedrake.solve(lhs, func, rhs, **kwargs) - return func - - def recompute_component(self, inputs, block_variable, idx, prepared): - lhs, rhs, func, bcs = prepared - result = self._forward_solve(lhs, rhs, func, bcs) - if isinstance(block_variable.checkpoint, firedrake.Function): - result = block_variable.checkpoint.assign(result) - return maybe_disk_checkpoint(result) - - def solve_init_params(self, args, kwargs, varform): if len(self.forward_args) <= 0: self.forward_args = args @@ -862,217 +342,6 @@ def solve_init_params(self, args, kwargs, varform): self.assemble_kwargs["appctx"] = kwargs["appctx"] -class SolveVarFormBlock(GenericSolveBlock): - def __init__(self, equation, func, bcs=[], *args, **kwargs): - lhs = equation.lhs - rhs = equation.rhs - super().__init__(lhs, rhs, func, bcs, *args, **kwargs) - - def _init_solver_parameters(self, args, kwargs): - super()._init_solver_parameters(args, kwargs) - solve_init_params(self, args, kwargs, varform=True) - - -class NonlinearVariationalSolveBlock(GenericSolveBlock): - def __init__(self, equation, func, bcs, adj_cache, problem_J, - solver_kwargs, **kwargs): - lhs = equation.lhs - rhs = equation.rhs - - self._adj_cache = adj_cache - self._dFdm_cache = adj_cache.setdefault("dFdm_cache", {}) - self.problem_J = problem_J - self.solver_kwargs = solver_kwargs - - super().__init__(lhs, rhs, func, bcs, **{**solver_kwargs, **kwargs}) - - if self.problem_J is not None: - for coeff in self.problem_J.coefficients(): - self.add_dependency(coeff, no_duplicates=True) - - def _init_solver_parameters(self, args, kwargs): - super()._init_solver_parameters(args, kwargs) - solve_init_params(self, args, kwargs, varform=True) - - def recompute_component(self, inputs, block_variable, idx, prepared): - tape = get_working_tape() - if self._ad_solvers["recompute_count"] == tape.recompute_count - 1: - # Update how many times the block has been recomputed. - self._ad_solvers["recompute_count"] = tape.recompute_count - if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian: - self._ad_solvers["forward_nlvs"].invalidate_jacobian() - self._ad_solvers["update_adjoint"] = True - return super().recompute_component(inputs, block_variable, idx, prepared) - - def _forward_solve(self, lhs, rhs, func, bcs, **kwargs): - self._ad_solver_replace_forms() - self._ad_solvers["forward_nlvs"].solve() - func.assign(self._ad_solvers["forward_nlvs"]._problem.u) - return func - - def _adjoint_solve(self, dJdu, compute_bdy): - dJdu_copy = dJdu.copy() - # Homogenize and apply boundary conditions on adj_dFdu and dJdu. - for bc in self.bcs: - bc.zero(dJdu) - - if ( - self._ad_solvers["forward_nlvs"]._problem._constant_jacobian - and self._ad_solvers["update_adjoint"] - ): - # Update left hand side of the adjoint equation. - self._ad_solver_replace_forms(SolverType.ADJOINT) - self._ad_solvers["adjoint_lvs"].invalidate_jacobian() - self._ad_solvers["update_adjoint"] = False - elif not self._ad_solvers["forward_nlvs"]._problem._constant_jacobian: - # Update left hand side of the adjoint equation. - self._ad_solver_replace_forms(SolverType.ADJOINT) - - # Update the right hand side of the adjoint equation. - # problem.F._component[1] is the right hand side of the adjoint. - self._ad_solvers["adjoint_lvs"]._problem.F._components[1].assign(dJdu) - - # Solve the adjoint linear variational solver. - self._ad_solvers["adjoint_lvs"].solve() - u_sol = self._ad_solvers["adjoint_lvs"]._problem.u - - adj_sol_bdy = None - if compute_bdy: - jac_adj = self._ad_solvers["adjoint_lvs"]._problem.J - adj_sol_bdy = self._compute_adj_bdy( - u_sol, adj_sol_bdy, jac_adj, dJdu_copy) - return u_sol, adj_sol_bdy - - def _ad_assign_map(self, form, solver): - if solver == SolverType.FORWARD: - count_map = self._ad_solvers["forward_nlvs"]._problem._ad_count_map - else: - count_map = self._ad_solvers["adjoint_lvs"]._problem._ad_count_map - assign_map = {} - form_ad_count_map = dict((count_map[coeff], coeff) - for coeff in form.coefficients()) - for block_variable in self.get_dependencies(): - coeff = block_variable.output - if isinstance(coeff, - (firedrake.Coefficient, firedrake.Constant, - firedrake.Cofunction)): - coeff_count = coeff.count() - if coeff_count in form_ad_count_map: - assign_map[form_ad_count_map[coeff_count]] = \ - block_variable.saved_output - - if ( - solver == SolverType.ADJOINT - and not self._ad_solvers["forward_nlvs"]._problem._constant_jacobian - ): - block_variable = self.get_outputs()[0] - coeff_count = block_variable.output.count() - if coeff_count in form_ad_count_map: - assign_map[form_ad_count_map[coeff_count]] = \ - block_variable.saved_output - return assign_map - - def _ad_assign_coefficients(self, form, solver): - assign_map = self._ad_assign_map(form, solver) - for coeff, value in assign_map.items(): - coeff.assign(value) - - def _ad_assign_bcs(self, problem): - for block_variable in self.get_dependencies(): - bc = block_variable.output - if bc in problem._ad_bc_map: - problem._ad_bc_map[bc].function_arg.assign( - block_variable.saved_output.function_arg) - - def _ad_solver_replace_forms(self, solver=SolverType.FORWARD): - if solver == SolverType.FORWARD: - problem = self._ad_solvers["forward_nlvs"]._problem - self._ad_assign_coefficients(problem.F, solver) - self._ad_assign_coefficients(problem.J, solver) - self._ad_assign_bcs(problem) - else: - self._ad_assign_coefficients( - self._ad_solvers["adjoint_lvs"]._problem.J, solver) - - def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies): - compute_bdy = self._should_compute_boundary_adjoint( - relevant_dependencies - ) - adj_sol, adj_sol_bdy = self._adjoint_solve(adj_inputs[0], compute_bdy) - self.adj_sol = adj_sol - if self.adj_cb is not None: - self.adj_cb(adj_sol) - if self.adj_bdy_cb is not None and compute_bdy: - self.adj_bdy_cb(adj_sol_bdy) - - r = {} - r["form"] = self._create_F_form() - r["adj_sol"] = self.adj_state - r["adj_sol_bdy"] = adj_sol_bdy - return r - - def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, - prepared=None): - if not self.linear and self.func == block_variable.output: - # We are not able to calculate derivatives wrt initial guess. - return None - F_form = prepared["form"] - adj_sol = prepared["adj_sol"] - adj_sol_bdy = prepared["adj_sol_bdy"] - c = block_variable.output - c_rep = block_variable.saved_output - - if isinstance(c, (firedrake.Function, firedrake.Cofunction)): - trial_function = firedrake.TrialFunction(c.function_space()) - elif isinstance(c, firedrake.Constant): - try: - mesh = extract_unique_domain(F_form) - except ValueError: - raise ValueError("Expecting a single mesh") - trial_function = firedrake.TrialFunction( - c._ad_function_space(mesh) - ) - elif isinstance(c, firedrake.DirichletBC): - tmp_bc = c.reconstruct( - g=extract_subfunction(adj_sol_bdy, c.function_space()) - ) - return [tmp_bc] - elif isinstance(c, firedrake.MeshGeometry): - # Using CoordianteDerivative requires us to do action before - # differentiating, might change in the future. - F_form_tmp = firedrake.action(F_form, adj_sol) - X = firedrake.SpatialCoordinate(c_rep) - dFdm = firedrake.derivative(-F_form_tmp, X, firedrake.TestFunction( - c._ad_function_space()) - ) - - dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs) - return dFdm - - # dFdm_cache works with original variables, not block saved outputs. - if c in self._dFdm_cache: - dFdm = self._dFdm_cache[c] - else: - dFdm = -firedrake.derivative(self.lhs, c, trial_function) - dFdm = firedrake.adjoint(dFdm) - self._dFdm_cache[c] = dFdm - - # Replace the form coefficients with checkpointed values. - replace_map = self._replace_map(dFdm) - replace_map[self.func] = self.get_outputs()[0].saved_output - dFdm = replace(dFdm, replace_map) - - if isinstance(dFdm, firedrake.Argument): - # Corner case. Should be fixed more permanently upstream in UFL. - # See: https://github.com/FEniCS/ufl/issues/395 - dFdm = ufl.Action(dFdm, adj_sol) - else: - dFdm = dFdm * adj_sol - dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs) - - return dFdm - - class ProjectBlock(CachedSolverBlock): pass @@ -1100,6 +369,11 @@ class SupermeshProjectBlock(Block): Step 2. solve linear system. """ + + pop_kwargs_keys = ["adj_cb", "adj_bdy_cb", "adj2_cb", "adj2_bdy_cb", + "forward_args", "forward_kwargs", "adj_args", + "adj_kwargs"] + def __init__(self, source, target_space, target, bcs=[], **kwargs): super(SupermeshProjectBlock, self).__init__( ad_block_tag=kwargs.pop("ad_block_tag", None) diff --git a/firedrake/adjoint_utils/projection.py b/firedrake/adjoint_utils/projection.py index c739083b22..08944b53c2 100644 --- a/firedrake/adjoint_utils/projection.py +++ b/firedrake/adjoint_utils/projection.py @@ -38,7 +38,7 @@ def wrapper(self, **kwargs): V = self.target.function_space() if annotate: bcs = kwargs.get("bcs", []) - sb_kwargs = ProjectBlock.pop_kwargs(kwargs) + sb_kwargs = SupermeshProjectBlock.pop_kwargs(kwargs) if self._target_is_function: # block should be created before project because output # might also be an input that needs checkpointing diff --git a/firedrake/adjoint_utils/solving.py b/firedrake/adjoint_utils/solving.py index c7301288f4..d35abcf2fe 100644 --- a/firedrake/adjoint_utils/solving.py +++ b/firedrake/adjoint_utils/solving.py @@ -1,5 +1,5 @@ from pyadjoint.tape import get_working_tape -from firedrake.adjoint_utils.blocks import CachedSolverBlock, GenericSolveBlock, ProjectBlock +from firedrake.adjoint_utils.blocks import CachedSolverBlock, ProjectBlock def get_solve_blocks(): @@ -11,6 +11,6 @@ def get_solve_blocks(): return [ block for block in get_working_tape().get_blocks() - if issubclass(type(block), (CachedSolverBlock, GenericSolveBlock)) + if issubclass(type(block), CachedSolverBlock) and not issubclass(type(block), ProjectBlock) ] diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py index 9e3aedbcb7..b1147f4bf6 100644 --- a/firedrake/adjoint_utils/variational_solver.py +++ b/firedrake/adjoint_utils/variational_solver.py @@ -1,8 +1,7 @@ import copy from functools import wraps, cached_property from pyadjoint.tape import get_working_tape, stop_annotating, annotate_tape, no_annotations -from firedrake.adjoint_utils.blocks import ( - NonlinearVariationalSolveBlock, CachedSolverBlock) +from firedrake.adjoint_utils.blocks import CachedSolverBlock from firedrake.adjoint_utils.blocks.solving import solve_init_params from firedrake.ufl_expr import derivative, adjoint, action from ufl import replace, Action @@ -525,66 +524,6 @@ def wrapper(self, **kwargs): return wrapper - @staticmethod - def _ad_annotate_solve_old(solve): - @wraps(solve) - def wrapper(self, **kwargs): - """To disable the annotation, just pass :py:data:`annotate=False` to this routine, and it acts exactly like the - Firedrake solve call. 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).""" - from firedrake import LinearVariationalSolver - annotate = annotate_tape(kwargs) - if annotate: - bounds = kwargs.pop("bounds", None) - if bounds is not None: - raise ValueError( - "MissingMathsError: we do not know how to differentiate through a variational inequality") - - tape = get_working_tape() - problem = self._ad_problem - sb_kwargs = NonlinearVariationalSolveBlock.pop_kwargs(kwargs) - sb_kwargs.update(kwargs) - - block = NonlinearVariationalSolveBlock(problem._ad_F == 0, - problem._ad_u, - problem._ad_bcs, - adj_cache=self._ad_adj_cache, - problem_J=problem._ad_J, - solver_kwargs=self._ad_kwargs, - ad_block_tag=self.ad_block_tag, - **sb_kwargs) - - # Forward variational solver. - if not self._ad_solvers["forward_nlvs"]: - self._ad_solvers["forward_nlvs"] = type(self)( - self._ad_problem_clone(self._ad_problem, block.get_dependencies()), - **self._ad_kwargs - ) - - # Adjoint variational solver. - if not self._ad_solvers["adjoint_lvs"]: - with stop_annotating(): - self._ad_solvers["adjoint_lvs"] = LinearVariationalSolver( - self._ad_adj_lvs_problem(block, problem._ad_adj_F), - *block.adj_args, **block.adj_kwargs) - if self._ad_problem._constant_jacobian: - self._ad_solvers["update_adjoint"] = False - - block._ad_solvers = self._ad_solvers - - tape.add_block(block) - - with stop_annotating(): - out = solve(self, **kwargs) - - if annotate: - block.add_output(self._ad_problem._ad_u.create_block_variable()) - - return out - - return wrapper - @no_annotations def _ad_problem_clone(self, problem, dependencies): """Replaces every coefficient in the residual, jacobian and Dirichlet boundary values From ecd1467bbcb749c7cb6666a8a970c2ce02d04b14 Mon Sep 17 00:00:00 2001 From: Angus Gibson Date: Tue, 20 Oct 2026 11:07:21 +1100 Subject: [PATCH 6/6] CachedSolverBlock -> NonlinearVariationalSolveBlock Now that the original namespace is clear, we can use a more descriptive name for the block! --- firedrake/adjoint_utils/blocks/__init__.py | 2 +- firedrake/adjoint_utils/blocks/solving.py | 4 ++-- firedrake/adjoint_utils/solving.py | 4 ++-- firedrake/adjoint_utils/variational_solver.py | 4 ++-- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/firedrake/adjoint_utils/blocks/__init__.py b/firedrake/adjoint_utils/blocks/__init__.py index 1e27c51a80..9701c9545d 100644 --- a/firedrake/adjoint_utils/blocks/__init__.py +++ b/firedrake/adjoint_utils/blocks/__init__.py @@ -1,6 +1,6 @@ from firedrake.adjoint_utils.blocks.assembly import AssembleBlock # noqa F401 from firedrake.adjoint_utils.blocks.solving import ( # noqa F401 - CachedSolverBlock, ProjectBlock, SupermeshProjectBlock + NonlinearVariationalSolveBlock, ProjectBlock, SupermeshProjectBlock ) from firedrake.adjoint_utils.blocks.function import ( # noqa F401 FunctionAssignBlock, FunctionMergeBlock, SubfunctionBlock diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py index 1867ff34d5..b5683e57ab 100644 --- a/firedrake/adjoint_utils/blocks/solving.py +++ b/firedrake/adjoint_utils/blocks/solving.py @@ -33,7 +33,7 @@ class SolverType(Enum): HESSIAN = SolverType.HESSIAN -class CachedSolverBlock(Block): +class NonlinearVariationalSolveBlock(Block): def __init__( self, forward_cache, tangent_cache, adjoint_cache, hessian_cache, ad_block_tag=None ): @@ -342,7 +342,7 @@ def solve_init_params(self, args, kwargs, varform): self.assemble_kwargs["appctx"] = kwargs["appctx"] -class ProjectBlock(CachedSolverBlock): +class ProjectBlock(NonlinearVariationalSolveBlock): pass diff --git a/firedrake/adjoint_utils/solving.py b/firedrake/adjoint_utils/solving.py index d35abcf2fe..5c055aa4f9 100644 --- a/firedrake/adjoint_utils/solving.py +++ b/firedrake/adjoint_utils/solving.py @@ -1,5 +1,5 @@ from pyadjoint.tape import get_working_tape -from firedrake.adjoint_utils.blocks import CachedSolverBlock, ProjectBlock +from firedrake.adjoint_utils.blocks import NonlinearVariationalSolveBlock, ProjectBlock def get_solve_blocks(): @@ -11,6 +11,6 @@ def get_solve_blocks(): return [ block for block in get_working_tape().get_blocks() - if issubclass(type(block), CachedSolverBlock) + if issubclass(type(block), NonlinearVariationalSolveBlock) and not issubclass(type(block), ProjectBlock) ] diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py index b1147f4bf6..184da642de 100644 --- a/firedrake/adjoint_utils/variational_solver.py +++ b/firedrake/adjoint_utils/variational_solver.py @@ -1,7 +1,7 @@ import copy from functools import wraps, cached_property from pyadjoint.tape import get_working_tape, stop_annotating, annotate_tape, no_annotations -from firedrake.adjoint_utils.blocks import CachedSolverBlock +from firedrake.adjoint_utils.blocks import NonlinearVariationalSolveBlock from firedrake.adjoint_utils.blocks.solving import solve_init_params from firedrake.ufl_expr import derivative, adjoint, action from ufl import replace, Action @@ -484,7 +484,7 @@ def _ad_hessian_cache(self): @staticmethod def _ad_block_type(): - return CachedSolverBlock + return NonlinearVariationalSolveBlock @staticmethod def _ad_annotate_solve(solve):