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/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/adjoint_utils/projection.py b/firedrake/adjoint_utils/projection.py index e8114b4145..c4f3a05b90 100644 --- a/firedrake/adjoint_utils/projection.py +++ b/firedrake/adjoint_utils/projection.py @@ -5,41 +5,31 @@ 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()) 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. diff --git a/firedrake/projection.py b/firedrake/projection.py index 5a2d83f371..10cc32eb95 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], @@ -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, ) 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)