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 95583f5..d5c9e51 100644 --- a/src/dolfinx_adjoint/blocks/function_assigner.py +++ b/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)): 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/tests/test_assign.py b/tests/test_assign.py index b1fce07..d52f354 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,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) diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index 311e1a2..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.""" @@ -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)