DirichletBC using data from other spaces + adjoint + multigrid support - #5127
DirichletBC using data from other spaces + adjoint + multigrid support#5127pbrubeck wants to merge 22 commits into
Conversation
Not for this PR but I think that this needs cleaning up. This dmhooks business is getting quite out of hand. I plan on doing this at some point.
Seems reasonable. This needs to be written down somewhere though. Both in a docstring and the manual. |
connorjward
left a comment
There was a problem hiding this comment.
Adjoint test please. Control must be a Function in an expression in a BC.
Then definitely qualifies for release.
connorjward
left a comment
There was a problem hiding this comment.
Adjoint test please. Control must be a Function in an expression in a BC.
Then definitely qualifies for release.
DirichletBC data on a different FunctionSpace was interpolated via a raw Interpolator.assemble() call, bypassing pyadjoint annotation. Route it through the top-level assemble() instead, stashing the Interpolate object so BaseFormAssembler reuses its cached Interpolator. DirichletBCBlock's adjoint also assumed the BC data lived on the same space as the BC itself; add the cross-space case, transforming the adjoint value through the adjoint of the (linear) interpolation operator. Add a regression test: a time-dependent heat equation with BC data interpolated from a different space at each step, verified with a Taylor test. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
9d87113 to
784bf85
Compare
784bf85 to
a5eec62
Compare
a5eec62 to
9b4d4cc
Compare
DirichletBC data on a different FunctionSpace was interpolated via a raw Interpolator.assemble() call, bypassing pyadjoint annotation. Route it through the top-level assemble() instead, stashing the Interpolate object so BaseFormAssembler reuses its cached Interpolator. DirichletBCBlock's adjoint also assumed the BC data lived on the same space as the BC itself; add the cross-space case, transforming the adjoint value through the adjoint of the (linear) interpolation operator. Fix two further gaps found while making a reused-DirichletBC variant of the new test actually differentiable, rather than merely documenting it as broken: - DirichletBC.function_arg's setter now keeps the FloatingType's _ad_args in sync with the assigned value, so reusing a BC via set_value()/mutating its data re-tapes against the current value instead of the one captured at __init__. - NonlinearVariationalSolveBlock._forward_solve now applies its own recovered, per-step bcs to the cached forward_nlvs solver before solving, instead of solving with whatever the shared DirichletBC objects' live state currently happens to hold. Add a regression test, parametrized over constructing a fresh DirichletBC each timestep vs. reusing one via NonlinearVariationalSolver, both verified with a Taylor test. Document the underlying patterns in AGENTS.md in generalized form, so they transfer to future pyadjoint (or similar record/replay) bugs beyond this specific one. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
9b4d4cc to
2ecca4f
Compare
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Sort from cheapest diagnostic wins to deepest architectural gotchas, and cut each lesson down to headline + one anchoring example. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Make DirichletBCBlock depend on the Function the BC actually applies (bc._function_arg) rather than the raw user-supplied value. The interpolation (or, for elements without interpolation support, the projection) of the boundary value is re-taped through the annotated assemble/project when the BC is added as a dependency of another block, so its own AssembleBlock/ProjectBlock supplies the adjoint, TLM, and recompute. This deletes the hand-rolled interpolation pullback and the dead constant-handling machinery from DirichletBCBlock. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ples Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
DirichletBC's projection fallback goes back to one cached Projector per boundary value: rebuilding a solver stack on every function_arg access crashes when the accessor fires inside PCSetUp_MG's coarsening callback (gmg on extruded hierarchies, where NCE/NCF cannot interpolate). The projection stays on the tape through a new ProjectorMixin that annotates ProjectorBase.project, taping one ProjectBlock per call against the current source value. Also locate the DirichletBCBlocks in test_bdy_control by type rather than by tape position, since the taped boundary-value interpolation now precedes them. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude Code auto-loads CLAUDE.md from the repository root; the @AGENTS.md import inlines the agent instructions so they reach every session without relying on convention or memory. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
| # recompute, so its problem's bcs must be swapped for this block's | ||
| # own checkpointed ones before solving, or a reused solver would | ||
| # keep solving with whichever BC data was set most recently. | ||
| self._ad_solvers["forward_nlvs"]._problem.bcs = bcs |
There was a problem hiding this comment.
I think the right approach here is that the "forward_nlvs" should have a copy of the bcs which should depend on copies of any coefficients in the BC expressions, and the coefficient values should be overwritten at this point.
There was a problem hiding this comment.
I had the following question in the meeting:
Consider a solve like:
f = Function(V)
solve(a == f * v * dx, bcs=[DirichletBC(V, f, ...)])If we copy and replace the function f for the form and bc separately then we will get something like:
solve(a == g * v * dx, bcs=[DirichletBC(V, h, ...)])Is this something to be concerned about? Is it even legal?
There was a problem hiding this comment.
I'll add such case (f in both RHS and BC data, used as Control) as a Taylor test.
| def rhs(self): | ||
| pass | ||
|
|
||
| @ProjectorMixin._ad_annotate_project |
There was a problem hiding this comment.
Do we not need to annotate init? I thought that that's the pattern for this sort of thing.
Description
DirichletBCnow accepts boundary data as aFunctionon a differentFunctionSpace, not onlycomposite expressions involving one — e.g.
DirichletBC(V1, Function(V2), sub)now works, matchingthe composite case (
DirichletBC(V1, 2*Function(V2), sub)) that already did.This is needed so
DirichletBCs can be refined symbolically rather than numerically in an adaptivelyrefined
MeshHierarchy(#5213): symbolic and numeric coarsening agree when spaces are nested, but notunder refinement, where numeric coarsening would keep applying the coarse BC throughout the adaptive
loop and accumulate large errors. Since the symbolic boundary data can live on a different space than
the BC itself,
BCBase.coefficients()(andEquationBC/EquationBCSplitequivalents) is added so DMhooks can be attached to every coefficient's own function space and the right
TransferManagerfound.Differentiability under
pyadjointcomes from composition with the taped operations rather than newadjoint code:
assemble(), so it is recorded as its own block on the tape; likewise the fallback for elementswithout interpolation support (e.g. HCT) is routed through the now annotated
Projector.DirichletBCBlockdepends on theFunctionthe BC actually applies (the interpolation output),not the raw user-supplied value, and the interpolation is re-taped whenever the BC is added as a
dependency of another block (
_ad_will_add_as_dependency), so the dependency always points at thecurrent boundary data. The adjoint, TLM, and Hessian then flow through the interpolation's or
projection's own block, which removed
DirichletBCBlock's hand-rolled interpolation pullback andits dead constant-handling machinery.
DirichletBC(set_valueor in-place mutation of the data in a time loop, with a cachedNonlinearVariationalSolver) is differentiable too, not just a freshly-constructed BC each step:each solve records its own interpolation + BC chain, and
NonlinearVariationalSolveBlockresyncsits cached solver's BCs per recompute instead of reusing whatever the shared, live BC currently
holds.
Added regression tests with Taylor tests: the time-dependent cross-space BC parametrized over both
configurations above, and a projected (HCT) boundary value. The
AGENTS.mdpyadjoint notes arerewritten as the composition-first design principles that came out of this work.