Fix checkpointed recompute aliasing for state-transfer Function targets

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.
 

src/dolfinx_adjoint/blocks/function_assigner.py

@@ -232,9 +232,9 @@ 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, _Function):
+            output = output._ad_new_like()
         if isinstance(prepared, dolfinx.fem.Function):
             output.x.array[:] = prepared.x.array[:]
         elif isinstance(prepared, (float, int)):

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]

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,52 @@ 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)
+
+        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)

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."""
@@ -253,22 +333,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)