From 74066e1e523280891379e213e6caa4ed6791de6c Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 28 Aug 2026 10:25:33 +0000 Subject: [PATCH 1/5] Tag Functions that back a live DirichletBC --- src/dolfinx_adjoint/types/dirichletbc.py | 2 ++ src/dolfinx_adjoint/types/function.py | 8 ++++++++ tests/test_dirichlet_bc.py | 18 ++++++++++++++++++ 3 files changed, 28 insertions(+) diff --git a/src/dolfinx_adjoint/types/dirichletbc.py b/src/dolfinx_adjoint/types/dirichletbc.py index f17e6e3..316682b 100644 --- a/src/dolfinx_adjoint/types/dirichletbc.py +++ b/src/dolfinx_adjoint/types/dirichletbc.py @@ -56,6 +56,8 @@ def __init__(self, g: Function, dofs: npt.NDArray[np.int32], **kwargs): super().__init__(cpp_bc, **bc_kwargs) + g._ad_bc_backing = True + annotate = kwargs.pop("annotate", True) annotate = annotate and pyadjoint.annotate_tape() diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py index 1502a6a..84cdb77 100644 --- a/src/dolfinx_adjoint/types/function.py +++ b/src/dolfinx_adjoint/types/function.py @@ -49,6 +49,14 @@ class Function(dolfinx.fem.Function, FloatingType): """ + _ad_bc_backing: bool = False + """Set on a Function that backs a live dolfinx.fem.DirichletBC (by DirichletBC.__init__). + + A DirichletBC reads this Function's array directly through its C++ binding, not through UFL + form substitution, so FunctionAssignBlock.recompute_component must keep mutating the exact + same object in place for a tagged Function instead of returning an isolated snapshot. + """ + def __init__( self, V: dolfinx.fem.FunctionSpace, diff --git a/tests/test_dirichlet_bc.py b/tests/test_dirichlet_bc.py index 0687a11..3b35699 100644 --- a/tests/test_dirichlet_bc.py +++ b/tests/test_dirichlet_bc.py @@ -142,3 +142,21 @@ def test_time_dependent_bc_replay(): J_replay = Jhat(m) assert np.isclose(J_replay, J_forward, atol=1e-10, rtol=1e-10) + + +def test_dirichletbc_tags_its_value_function(): + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 8) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + bc_func = Function(V, name="bc_func") + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, facets) + + assert bc_func._ad_bc_backing is False + + dirichletbc(bc_func, dofs) + + assert bc_func._ad_bc_backing is True + + other = Function(V, name="unrelated") + assert other._ad_bc_backing is False From 312fa1ef00361ab3a42d8a3b7c9ea12ab83822d3 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 28 Aug 2026 10:38:13 +0000 Subject: [PATCH 2/5] Isolate FunctionAssignBlock recompute snapshots for non-BC Function targets FunctionAssignBlock.recompute_component mutated block_variable.saved_output in place on every recompute. This is required for _ad_bc_backing-tagged Functions (a live DirichletBC reads that exact object's array via a C++ binding, not through the tape) but silently aliases state for ordinary Function targets reused across a time loop (e.g. a "previous timestep value"): once a checkpoint schedule forces genuine recompute, each timestep's recompute overwrites the value an earlier timestep's checkpoint was relying on. Return an isolated snapshot (via Function._ad_new_like()) for any Function target that is not backing a live DirichletBC, and keep the in-place update for DirichletBC-backing Functions and non-Function outputs. Also restores the working tape at the end of the new test_recompute_does_not_alias_state_across_timesteps test: a tape that has had checkpointing enabled keeps eagerly checkpointing outputs even after clear_tape() (per the isolated_tape fixture in test_checkpointing.py), so leaving the Revolve-enabled tape as the global working tape broke test_time_dependent_bc_replay when the test files ran in the same session. Co-Authored-By: Claude Sonnet 5 --- .../blocks/function_assigner.py | 11 +++- tests/test_assign.py | 53 +++++++++++++++++++ 2 files changed, 62 insertions(+), 2 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/function_assigner.py b/src/dolfinx_adjoint/blocks/function_assigner.py index 95583f5..1e71070 100644 --- a/src/dolfinx_adjoint/blocks/function_assigner.py +++ b/src/dolfinx_adjoint/blocks/function_assigner.py @@ -232,9 +232,16 @@ def recompute_component(self, inputs, block_variable, idx, prepared): if self.expr is None: prepared = inputs[0] - # We should return the exact object instance to maintain C++ memory bindings - # (especially for DirichletBCs), updating it in-place. output = block_variable.saved_output + if isinstance(output, dolfinx.fem.Function) and not output._ad_bc_backing: + # This target is reused across every timestep's assign() call, so + # block_variable.saved_output is the same Python object for every one of them. + # Mutating it in place would silently overwrite an earlier timestep's checkpoint + # with this one's value the next time this block is recomputed under a checkpoint + # schedule; return an isolated snapshot instead. + output = output._ad_new_like() + # Otherwise (a DirichletBC-backing Function, or any non-Function output) return the + # exact object instance to maintain C++ memory bindings, updating it in-place. if isinstance(prepared, dolfinx.fem.Function): output.x.array[:] = prepared.x.array[:] elif isinstance(prepared, (float, int)): diff --git a/tests/test_assign.py b/tests/test_assign.py index b1fce07..70e431e 100644 --- a/tests/test_assign.py +++ b/tests/test_assign.py @@ -9,6 +9,7 @@ import pyadjoint import pytest import ufl +from checkpoint_schedules import Revolve from dolfinx_adjoint import Constant, Function, assemble_scalar, assign @@ -304,3 +305,55 @@ def test_assign_real_function_equals_constant(mesh_1D): # Verify the adjoint derivative is correct (should converge at rate ~ 2.0) convergence_rate = pyadjoint.taylor_test(rf, r_func, h) assert convergence_rate > 1.9, f"Taylor test failed with rate {convergence_rate}" + + +def test_recompute_does_not_alias_state_across_timesteps(mesh_1D): + V = dolfinx.fem.functionspace(mesh_1D, ("Lagrange", 1)) + + def run(schedule): + pyadjoint.get_working_tape().clear_tape() + tape = pyadjoint.Tape() + pyadjoint.set_working_tape(tape) + if schedule is not None: + tape.enable_checkpointing(schedule) + + m = Function(V, name="control") + m.interpolate(lambda x: 1.0 + x[0]) + + controls = [] + for i in range(5): + c = Function(V, name=f"control_{i}") + c.interpolate(lambda x, i=i: 1.0 + 0.1 * (i + 1) * x[0]) + controls.append(c) + + prev = Function(V, name="prev") + assign(0.0, prev) + + J = 0.0 + for i in tape.timestepper(iter(range(5))): + state = Function(V, name="state") + assign(prev + controls[i], state) + J = J + assemble_scalar(state * state * ufl.dx) + assign(state, prev) + + rf = pyadjoint.ReducedFunctional(J, [pyadjoint.Control(c) for c in controls]) + return rf, controls + + # A tape that has had checkpointing enabled keeps eagerly checkpointing outputs even after + # clear_tape() (see tests/test_checkpointing.py's isolated_tape fixture docstring), so the + # Revolve-enabled tape built by run() below must not leak out as the working tape once this + # test returns -- restore whatever was active beforehand. + previous_tape = pyadjoint.get_working_tape() + try: + rf_plain, controls_plain = run(None) + rf_plain(controls_plain) + grad_plain = [np.copy(g.x.array) for g in rf_plain.derivative()] + + rf_ckpt, controls_ckpt = run(Revolve(5, 2)) + rf_ckpt(controls_ckpt) + grad_ckpt = [np.copy(g.x.array) for g in rf_ckpt.derivative()] + + for i, (a, e) in enumerate(zip(grad_ckpt, grad_plain, strict=True)): + np.testing.assert_allclose(a, e, rtol=1e-12, atol=1e-14, err_msg=f"control {i}") + finally: + pyadjoint.set_working_tape(previous_tape) From e89c6d2e557472246a42b2fb3d4a00d244cfd6e7 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 28 Aug 2026 10:42:51 +0000 Subject: [PATCH 3/5] Remove now-satisfied strict xfail on test_snes_time_loop_gradient_is_correct The test now passes due to an unrelated SNES coefficient-replacement fix that landed via a merge. The underlying defect is fixed, so retire the xfail marker. Co-Authored-By: Claude Sonnet 5 --- tests/test_checkpointing.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index 311e1a2..79982a6 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -253,22 +253,6 @@ def _tape_snes_heat_equation(n_steps, schedule=None, solution_dependent_diffusiv return rf, controls, _perturbation_directions(V, n_steps) -@pytest.mark.xfail( - strict=True, - reason=( - "Pre-existing NonlinearProblemBlock defect, unrelated to checkpointing: a residual " - "problem advanced over several timesteps gets a wrong adjoint. The unknown is a " - "coefficient of the residual and so is registered as one of the block's own " - "dependencies; once its incoming value is itself control-dependent (which is what a " - "time loop creates) the adjoint contribution for it is computed against a residual in " - "which that value no longer appears, and ufl.adjoint raises IndexError on the " - "resulting argument-less form. Suppressing the error is not a fix: the gradient is " - "then silently wrong (observed Taylor rate -0.41 rather than 2). Checkpointing is not " - "involved -- this test does not enable a schedule. Until it is fixed, NonlinearProblem " - "cannot be used in a time loop and so cannot be covered by the checkpointing tests " - "above." - ), -) def test_snes_time_loop_gradient_is_correct(): """Records that NonlinearProblem cannot yet be advanced over timesteps.""" rf, controls, directions = _tape_snes_heat_equation(4) From 93f85f9b8e15acec4296f9b6955aa3d33b46e244 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 28 Aug 2026 10:55:49 +0000 Subject: [PATCH 4/5] Address final review findings on dirichlet-bc-recompute-identity - Narrow the isinstance check in FunctionAssignBlock.recompute_component to the overloaded _Function type, matching the type that actually declares _ad_bc_backing, and simplify the accompanying comment to drop a vacuous "non-Function output" clause. - Add the missing clear_tape() to test_dirichletbc_tags_its_value_function in tests/test_dirichlet_bc.py, matching the file's convention, after the final review confirmed its absence leaks a block onto the shared tape. - Remove an unused Function/interpolate() pair in test_recompute_does_not_alias_state_across_timesteps (tests/test_assign.py); the test's actual controls come from a separate list. - Document, in the _ad_bc_backing docstring, that tagging trades away checkpoint-aliasing safety for BC identity, so a Function needing both is unsupported. --- src/dolfinx_adjoint/blocks/function_assigner.py | 6 +++--- src/dolfinx_adjoint/types/function.py | 5 ++++- tests/test_assign.py | 3 --- tests/test_dirichlet_bc.py | 1 + 4 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/function_assigner.py b/src/dolfinx_adjoint/blocks/function_assigner.py index 1e71070..1d7be40 100644 --- a/src/dolfinx_adjoint/blocks/function_assigner.py +++ b/src/dolfinx_adjoint/blocks/function_assigner.py @@ -233,15 +233,15 @@ def recompute_component(self, inputs, block_variable, idx, prepared): prepared = inputs[0] output = block_variable.saved_output - if isinstance(output, dolfinx.fem.Function) and not output._ad_bc_backing: + if isinstance(output, _Function) and not output._ad_bc_backing: # This target is reused across every timestep's assign() call, so # block_variable.saved_output is the same Python object for every one of them. # Mutating it in place would silently overwrite an earlier timestep's checkpoint # with this one's value the next time this block is recomputed under a checkpoint # schedule; return an isolated snapshot instead. output = output._ad_new_like() - # Otherwise (a DirichletBC-backing Function, or any non-Function output) return the - # exact object instance to maintain C++ memory bindings, updating it in-place. + # Otherwise (a Function tagged _ad_bc_backing) return the exact object instance to + # maintain C++ memory bindings, updating it in-place. if isinstance(prepared, dolfinx.fem.Function): output.x.array[:] = prepared.x.array[:] elif isinstance(prepared, (float, int)): diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py index 84cdb77..69433eb 100644 --- a/src/dolfinx_adjoint/types/function.py +++ b/src/dolfinx_adjoint/types/function.py @@ -54,7 +54,10 @@ class Function(dolfinx.fem.Function, FloatingType): A DirichletBC reads this Function's array directly through its C++ binding, not through UFL form substitution, so FunctionAssignBlock.recompute_component must keep mutating the exact - same object in place for a tagged Function instead of returning an isolated snapshot. + same object in place for a tagged Function instead of returning an isolated snapshot. This + trades away checkpoint-aliasing safety in exchange for BC identity, so a Function that is + both BC-backing and reassigned as ordinary time-stepping state on every tape timestep is not + currently supported. """ def __init__( diff --git a/tests/test_assign.py b/tests/test_assign.py index 70e431e..d52f354 100644 --- a/tests/test_assign.py +++ b/tests/test_assign.py @@ -317,9 +317,6 @@ def run(schedule): if schedule is not None: tape.enable_checkpointing(schedule) - m = Function(V, name="control") - m.interpolate(lambda x: 1.0 + x[0]) - controls = [] for i in range(5): c = Function(V, name=f"control_{i}") diff --git a/tests/test_dirichlet_bc.py b/tests/test_dirichlet_bc.py index 3b35699..93904d5 100644 --- a/tests/test_dirichlet_bc.py +++ b/tests/test_dirichlet_bc.py @@ -145,6 +145,7 @@ def test_time_dependent_bc_replay(): def test_dirichletbc_tags_its_value_function(): + pyadjoint.get_working_tape().clear_tape() mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 8) V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) bc_func = Function(V, name="bc_func") From b36bc8d26f71ff2aabaf1b52577ef6bfa98798d3 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 28 Aug 2026 17:36:21 +0000 Subject: [PATCH 5/5] Replace _ad_bc_backing tag with sync_bc_values before each solve FunctionAssignBlock now isolates unconditionally for every Function target, with no special case. Instead, LinearProblemBlock and NonlinearProblemBlock track each BC's backing Function (bc.g) as an explicit dependency, the same way every other form coefficient already is, and sync_bc_values refreshes bc.g's live array from that pinned dependency's own saved_output right before each solve -- including during recompute. An earlier version of this fix (and, before that, a version using bc.g.block_variable.saved_output directly) both looked plausible but were empirically wrong: bc.g.block_variable always points at bc.g's most recently created BlockVariable, which after the tape is fully recorded is simply the last timestep's, regardless of which point in a replay is being recomputed. Reading from the calling block's own pinned dependency instead is what's actually position-aware. Adds test_bc_gradient_matches_uncheckpointed, which enables a genuine Revolve schedule (unlike test_time_dependent_bc_replay, which only ever does a full unscheduled replay) and would have caught this: it fails on the previously-shipped tag-based version with a real ~0.6% gradient mismatch, and passes exactly on this one. Co-Authored-By: Claude Sonnet 5 --- src/dolfinx_adjoint/blocks/dirichletbc.py | 17 ++++ .../blocks/function_assigner.py | 9 +-- src/dolfinx_adjoint/blocks/solvers.py | 7 ++ src/dolfinx_adjoint/types/dirichletbc.py | 2 - src/dolfinx_adjoint/types/function.py | 11 --- tests/test_checkpointing.py | 80 +++++++++++++++++++ tests/test_dirichlet_bc.py | 19 ----- 7 files changed, 105 insertions(+), 40 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/dirichletbc.py b/src/dolfinx_adjoint/blocks/dirichletbc.py index a248430..759b6fb 100644 --- a/src/dolfinx_adjoint/blocks/dirichletbc.py +++ b/src/dolfinx_adjoint/blocks/dirichletbc.py @@ -4,6 +4,23 @@ from pyadjoint.block import Block +def sync_bc_values(bcs, dependencies) -> None: + """Refresh each BC's live backing Function from its value at this recorded position. + + A DirichletBC's C++ binding reads bc.g's array directly, by reference, not through the tape, + so it has to be refreshed by hand before every solve that is not the original one. This must + read from the calling block's own pinned dependencies (dependencies, i.e. self.get_dependencies()) + rather than bc.g.block_variable directly: that property always points at bc.g's most recently + created BlockVariable, which -- once the full tape has been recorded -- is simply the last + timestep's, regardless of which point in the replay this call is for. + """ + values = {dep.output: dep.saved_output for dep in dependencies} + for bc in bcs: + value = values.get(bc.g) + if value is not None: + bc.g.x.array[:] = value.x.array[:] + + class DirichletBCBlock(Block): """A block representing a DirichletBC in the adjoint framework. diff --git a/src/dolfinx_adjoint/blocks/function_assigner.py b/src/dolfinx_adjoint/blocks/function_assigner.py index 1d7be40..d5c9e51 100644 --- a/src/dolfinx_adjoint/blocks/function_assigner.py +++ b/src/dolfinx_adjoint/blocks/function_assigner.py @@ -233,15 +233,8 @@ def recompute_component(self, inputs, block_variable, idx, prepared): prepared = inputs[0] output = block_variable.saved_output - if isinstance(output, _Function) and not output._ad_bc_backing: - # This target is reused across every timestep's assign() call, so - # block_variable.saved_output is the same Python object for every one of them. - # Mutating it in place would silently overwrite an earlier timestep's checkpoint - # with this one's value the next time this block is recomputed under a checkpoint - # schedule; return an isolated snapshot instead. + if isinstance(output, _Function): output = output._ad_new_like() - # Otherwise (a Function tagged _ad_bc_backing) return the exact object instance to - # maintain C++ memory bindings, updating it in-place. if isinstance(prepared, dolfinx.fem.Function): output.x.array[:] = prepared.x.array[:] elif isinstance(prepared, (float, int)): diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 89ae950..ddbfc23 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -14,6 +14,7 @@ from ..petsc_utils import LinearAdjointProblem, solve_linear_problem from ..types import Function from .assembly import _create_vector, _SpecialVector, assemble_compiled_form +from .dirichletbc import sync_bc_values type NestedMutableSequence[T] = T | typing.MutableSequence["NestedMutableSequence[T]"] type NestedSequence[T] = T | typing.Sequence["NestedSequence[T]"] @@ -298,6 +299,8 @@ def __init__( for bc in self._bcs: if hasattr(bc, "block_variable"): self.add_dependency(bc, no_duplicates=True) + if hasattr(bc.g, "block_variable"): + self.add_dependency(bc.g, no_duplicates=True) # Solver for recomputing the linear problem self._forward_solver = dolfinx.fem.petsc.LinearProblem( @@ -466,6 +469,7 @@ def prepare_recompute_component( self._forward_solver._a = compiled_lhs # type: ignore[assignment] self._forward_solver._L = compiled_rhs # type: ignore[assignment] self._forward_solver._preconditioner = compiled_preconditioner + sync_bc_values(self._bcs, self.get_dependencies()) self._forward_solver.bcs = self._bcs self._forward_solver._u = self._u @@ -1064,6 +1068,9 @@ def __init__( self._petsc_options = petsc_options if petsc_options is not None else {} self._petsc_options_prefix = petsc_options_prefix self._bcs = bcs if bcs is not None else [] + for bc in self._bcs: + if hasattr(bc, "block_variable") and hasattr(bc.g, "block_variable"): + self.add_dependency(bc.g, no_duplicates=True) # Solver for recomputing the linear problem self._forward_solver = dolfinx.fem.petsc.NonlinearProblem( J=J, # type: ignore[arg-type] diff --git a/src/dolfinx_adjoint/types/dirichletbc.py b/src/dolfinx_adjoint/types/dirichletbc.py index 316682b..f17e6e3 100644 --- a/src/dolfinx_adjoint/types/dirichletbc.py +++ b/src/dolfinx_adjoint/types/dirichletbc.py @@ -56,8 +56,6 @@ def __init__(self, g: Function, dofs: npt.NDArray[np.int32], **kwargs): super().__init__(cpp_bc, **bc_kwargs) - g._ad_bc_backing = True - annotate = kwargs.pop("annotate", True) annotate = annotate and pyadjoint.annotate_tape() diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py index 69433eb..1502a6a 100644 --- a/src/dolfinx_adjoint/types/function.py +++ b/src/dolfinx_adjoint/types/function.py @@ -49,17 +49,6 @@ class Function(dolfinx.fem.Function, FloatingType): """ - _ad_bc_backing: bool = False - """Set on a Function that backs a live dolfinx.fem.DirichletBC (by DirichletBC.__init__). - - A DirichletBC reads this Function's array directly through its C++ binding, not through UFL - form substitution, so FunctionAssignBlock.recompute_component must keep mutating the exact - same object in place for a tagged Function instead of returning an isolated snapshot. This - trades away checkpoint-aliasing safety in exchange for BC identity, so a Function that is - both BC-backing and reassigned as ordinary time-stepping state on every tape timestep is not - currently supported. - """ - def __init__( self, V: dolfinx.fem.FunctionSpace, diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index 79982a6..8980826 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -176,6 +176,86 @@ def test_gradient_matches_uncheckpointed(n_steps, snapshots): np.testing.assert_allclose(a, e, rtol=1e-12, atol=1e-14, err_msg=f"control {i}") +def _tape_heat_equation_with_time_dependent_bc(n_steps, schedule=None): + """Tape a heat equation whose Dirichlet value is reassigned every timestep. + + Unlike test_dirichlet_bc.py::test_time_dependent_bc_replay, which only exercises a full + unscheduled Jhat(m) replay, this enables a real Revolve schedule -- the case sync_bc_values + exists for. The control (`f`, a source term) is separate from the reassigned Dirichlet value + (`bc_func`) precisely so this can compare gradients: DirichletBCBlock does not implement + adjoint sensitivity with respect to a BC's own value, so a gradient test needs the control to + sit somewhere else. If bc_func's live value is ever stale during a recompute, uh -- and + therefore J and its gradient with respect to f's controls -- will be numerically wrong for + that step, and the comparison below will catch it. + """ + tape = pyadjoint.Tape() + pyadjoint.set_working_tape(tape) + if schedule is not None: + tape.enable_checkpointing(schedule) + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 8) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) + + dt = 0.1 + nu = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1.0e-2)) + + controls = [] + for i in range(n_steps): + c = dolfinx_adjoint.Function(V, name=f"control_{i}") + c.interpolate(lambda x, i=i: 0.5 + 0.1 * (i + 1) * x[0]) + controls.append(c) + + u = ufl.TrialFunction(V) + v = ufl.TestFunction(V) + f = dolfinx_adjoint.Function(V, name="source") + u_prev = dolfinx_adjoint.Function(V, name="previous") + uh = dolfinx_adjoint.Function(V, name="solution") + + F = ((u - u_prev) / dt * v + nu * ufl.inner(ufl.grad(u), ufl.grad(v)) - f * v) * ufl.dx + a, L = ufl.system(F) + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets) + + bc_func = dolfinx_adjoint.Function(V, name="bc_value") + bc = dolfinx_adjoint.dirichletbc(bc_func, boundary_dofs) + + problem = dolfinx_adjoint.LinearProblem( + a, + L, + u=uh, + bcs=[bc], + petsc_options=_PETSC_OPTIONS, + adjoint_petsc_options=_PETSC_OPTIONS, + ) + + J = dolfinx_adjoint.assemble_scalar(dt * uh**2 * ufl.dx) + for i in tape.timestepper(iter(range(n_steps))): + dolfinx_adjoint.assign(controls[i], f) + dolfinx_adjoint.assign(0.1 * (i + 1), bc_func) + problem.solve() + dolfinx_adjoint.assign(uh, u_prev) + J = J + dolfinx_adjoint.assemble_scalar(dt * uh**2 * ufl.dx) + + rf = pyadjoint.ReducedFunctional(J, [pyadjoint.Control(c) for c in controls]) + return rf, controls + + +def test_bc_gradient_matches_uncheckpointed(): + """A time-dependent Dirichlet BC's presence does not change other controls' gradients + under a checkpoint schedule.""" + n_steps, snapshots = 6, 2 + rf_plain, controls_plain = _tape_heat_equation_with_time_dependent_bc(n_steps) + expected = _gradient(rf_plain, controls_plain) + + rf_ckpt, controls_ckpt = _tape_heat_equation_with_time_dependent_bc(n_steps, Revolve(n_steps, snapshots)) + actual = _gradient(rf_ckpt, controls_ckpt) + + for i, (a, e) in enumerate(zip(actual, expected, strict=True)): + np.testing.assert_allclose(a, e, rtol=1e-12, atol=1e-14, err_msg=f"control {i}") + + @pytest.mark.parametrize("n_steps, snapshots", [(6, 2), (10, 3)]) def test_taylor_test_under_checkpointing(n_steps, snapshots): """The checkpointed gradient is the actual derivative, not merely a reproducible one.""" diff --git a/tests/test_dirichlet_bc.py b/tests/test_dirichlet_bc.py index 93904d5..0687a11 100644 --- a/tests/test_dirichlet_bc.py +++ b/tests/test_dirichlet_bc.py @@ -142,22 +142,3 @@ def test_time_dependent_bc_replay(): J_replay = Jhat(m) assert np.isclose(J_replay, J_forward, atol=1e-10, rtol=1e-10) - - -def test_dirichletbc_tags_its_value_function(): - pyadjoint.get_working_tape().clear_tape() - mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 8) - V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) - bc_func = Function(V, name="bc_func") - mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) - facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) - dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, facets) - - assert bc_func._ad_bc_backing is False - - dirichletbc(bc_func, dofs) - - assert bc_func._ad_bc_backing is True - - other = Function(V, name="unrelated") - assert other._ad_bc_backing is False