Skip to content

The discrete adjoint: built into the solvers, driven from the run transcript - #744

Open
lmoresi wants to merge 1 commit into
docs/timestepping-patternfrom
feature/discrete-adjoint
Open

lmoresi wants to merge 1 commit into
docs/timestepping-patternfrom
feature/discrete-adjoint

Conversation

@lmoresi

@lmoresi lmoresi commented Sep 14, 2026

Copy link
Copy Markdown
Member

What this is

The discrete adjoint, built into the solvers and driven from the run transcript. Stacked on #716, which supplies the step boundary, the snapshots and the record it walks; the diff here is exactly the adjoint additions.

An implicit step is a residual R(u; m) = 0 and the residual is SymPy, so every first derivative is available symbolically: ∂R/∂u is the Jacobian the SNES already assembles, ∂R/∂m and the coupling to every field the residual reads are differentiated, not approximated. The adjoint of one solve is a transpose against that Jacobian; the adjoint of a run is the chain of them in reverse order of the record.

In every solver

b  = -solver.dual_of(T.sym[0] - T_target.sym[0])   # -dJ/dT for J = ½∫(T - T*)²
mu, reason = solver.adjoint_solve(b, target=mu_var) # Kᵀ mu = b
dJ_dkappa  = solver.sensitivity(mu_var, kappa)      # ∫ (∂F/∂κ)·mu, symbolic ∂F/∂κ
  • adjoint_solve transposes against the assembled Jacobian via KSP.solveTranspose. Dirichlet conditions come out homogenised for free: the global vector holds only unconstrained dofs. Stokes does it on the composite (u, p) system.
  • dual_of assembles ∫ e φⱼ on the solver's own constrained space as a generic solver's residual at zero — no linear solve.
  • sensitivity differentiates the residual through the constitutive model's own symbol: the residual holds \upkappa whose value is the user's kappa, and uw.function.derivative substitutes constants' values before differentiating, so both naive routes give zero. _peel_except expands every named expression except the one being differentiated against.
  • If the forward solve used the Picard tangent, the adjoint assembles the consistent tangent itself (a kernel rebuild; the DM and KSP are kept) and puts the Picard kernel back for the next forward. Picard iterations do not spoil the adjoint — the converged state is the same and ∂R/∂u is a function of the state alone; only the compiled kernel was the wrong matrix. Under "continuation" the adjoint assembles at α = 1.

Over a whole run

back   = uw.adjoint.TranscriptAdjoint(model, model.save_state())
result = back.gradient(misfit, parameters=[eta0], fields=[beta])
result["parameters"][eta0]                     # dJ/dη₀
uw.adjoint.inner(beta, result["fields"][beta], dbeta0_dc)   # dJ/dc through β₀

No problem-specific wiring. For each solve, in reverse order of the record: restore the step's snapshot, replay the solves before it, replay it, put each history's input back in the slot the post-solve hook shifted; transpose-solve with the accumulated dual on its unknown; then read the residual to see what the solve depended on. Every other field in F0/F1 gets (∂R/∂f)ᵀμ as a dual on its own space — through its value and, for a field read via its gradient (a Crank–Nicolson step reads κ∇T_old), through the gradient part of the load, assembled as ∫ g₀φⱼ + g₁·∇φⱼ with no integration by parts. A history slot's dual goes to the field it tracks at the previous level; every parameter gets μᵀ∂R/∂m, plus the explicit ∂J/∂m.

Every operator in the transcript carries a verdict, written when it ran: adjoint: {supported, reason}. A rotated constraint (no transpose path through its own Krylov loop), an unconverged solve, a semi-Lagrangian history (interpolation at departure points not materialised), a swarm step whose particle set changed — each refuses with the reason, and transcript_adjoint_segments reads the verdicts back as the strong/weak-constraint partition of the assimilation window.

The default tangent is Newton

consistent_jacobian = True is now the default; Picard is the opt-in, by name, where a hard-yield solve needs it as an entry requirement. The residual is symbolic so the tangent is exact and cheap; Picard converged linearly and left the SNES holding a Jacobian that was not ∂R/∂u for everything downstream. The design note had already named this the intended default and marked it "benchmark before flipping"; the full suite was that benchmark. What it turned up: the sharp-yield cold solve that Picard could not do converges under Newton (test_1057 now sets Picard explicitly on its cold half); and the cold-start "warm-up" on the saddle-point path degraded a linear ksponly solve under Eisenstat–Walker (12% vs 2% on the spherical-shell Nitsche response) — now skipped for ksponly, and honestly described: it is an nrichardson residual sweep, not a Picard step, and picard=-1 switches it off.

Verified

Every case against central finite differences:

case agreement
Poisson, dJ/dκ 1e-4
one SUPG step (non-symmetric K, where a wrong transpose shows) 1e-3
Stokes, constant viscosity 1e-4
Stokes, η(ε̇), consistent tangent / continuation / Picard forward 1e-3
two-solver, two-step run: dJ/dη₀ and the dual on β₀ (θ = 1 and θ = 0.5) 1e-7 measured, serial and np = 2
the sinking-blob example, blob centre as control, through the library Taylor ratio 1.00000 / 0.99997

The example (docs/examples/adjoint/sinker_transcript) keeps its hand-rolled adjoint beside the one-call version, and both Taylor tests give the same gradient to six digits.

Adversarial review

Three reviewers ran probes against the branch (posted on #716, fixed in fba44975 and later): two MPI deadlocks (a rank-local gate before a collective solve; rank-local size checks), a ghost double-count in the dual pairing, the missing explicit ∂J/∂m, sixteen scratch variables leaked per call, a second adjoint_solve refusing after the Picard rebuild, and the warm-up's misdescription. All fixed here except dropping the warm-up, which is a benchmark item in the planning file.

Not in this PR

Semi-Lagrangian schemes still refuse (their two transport operators are not materialised); second-order terms are noted in the design note and not built.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V

@lmoresi

lmoresi commented Sep 15, 2026

Copy link
Copy Markdown
Member Author

Adversarial review — adjoint machinery and the Newton default (2026-09-14)

Two independent reviewers running probes. Findings that survived, and what was done. The fixes were made on the build branch before the split; that history is on archive/timestepping-pattern-history (fixes in fba44975, further in dd928797), and this PR's single commit carries them.

Adjoint machinery (adjoint_solve / dual_of / sensitivity, uw.adjoint.TranscriptAdjoint)

Sev Finding Evidence Status
HIGH _nonzero was a rank-local max gating a collective solve; a misfit supported on one rank's cells deadlocked gradient() np=2, misfit ½|v|²·[y>0.85]: rank 0 solves, rank 1 skips; alive at 100 s vs 10 s serial fixed: reduced over ranks, empty-array safe
HIGH the field-gradient test's dual @ direction was rank-local and counted ghosts np=2: −0.00602 / −0.00614 vs fd −0.010876; owned-dof dot = fd to 7.7e-10 fixed: uw.adjoint.inner() over owned dofs; test passes at np=2
MED gradient() omitted the explicit ∂J/∂m misfit ½η₀|v|²: −0.004824 vs fd −0.002404 fixed, with a test
MED 16 MeshVariables leaked per gradient() call; call time 1.3 → 14.5 s by the sixth vars 5→101 over 6 calls, RSS 335→656 MB fixed: scratch pool; test asserts no growth
MED a second adjoint_solve after the Picard-path rebuild raised "no forward solve" is_setup=False after _restore_tangent fixed: the consistent kernel stays installed until the next forward build
LOW dual-space scratch name collided when the BC set changed without changing its count _dual_SNES_Poisson_29_1 twice fixed: hashed signature
LOW size-mismatch raises were rank-local, before collectives by inspection fixed: decided with an allreduce
LOW AdvDiffusion at θ=0.5 refused ("reads a derivative of psi_star_…") — names the slot, and test_0020 uses θ=1 without saying why _reads on the default solver refusal now names the tracked field and says θ=1 avoids it; the gradient-read term (integration by parts with its boundary part) is an open item

Attacks that failed: token prefix collisions (symbols print braced); a Projection with no BCs (rel 2.3e-9); Vector_Projection through mesh.vector.jacobian (2.4e-9); a parameter two levels deep (2.6e-8); serial gradient() bit-repeatable ×4; state after gradient() is level 0 and the next forward reproduces J to 1e-15; Picard→consistent rebuild cycles identical over 3 rounds.

Newton default (fa83b19d)

Sev Finding Evidence Status
MED the ksponly guard read snes.getType() — the previous solve's type — not the declared option set ksponly after a solve, cold solve: picard = 1 ran fixed: reads petsc_options
MED "a cold start takes one Picard step" is true only on the saddle-point standard path — not scalar/vector solvers, not the rotated/fault-contact path Poisson: no warm-up; rotated: log begins at Newton wording fixed (docstring, CLAUDE.md)
MED the warm-up is an nrichardson residual sweep, not a Picard step: no linear solve, no frozen tangent; 1–12% residual reduction on linear Stokes; the sharp-yield case converges under pure Newton without it (7 its) snes_monitor; 1057 fixture on the rotated path comment corrected; picard=-1 switches it off; dropping it is a planning benchmark item
LOW Stokes_Constrained 2-D annulus with two add_constraint_bc(degree=2) crashes in _scatter_global_to_fields (pre-existing) broadcast (12484,) into (6396,) planning item; own issue

Attacks that failed: linear Stokes newtonls+EW new vs old default identical to 1e-9 on box and annulus; rotated + sharp yield cold: Newton 7 its where Picard did not converge in 50; Poisson κ(T) cold: Newton 8–10 its vs Picard 24–28, no window where the old default converged and the new one does not; continuation unchanged; the homotopy march forces its own tangent.

Also from self-review: the Stokes adjoint verdict called the two-assembly nonlinearity probe on every recorded solve; now the symbolic test, cached until the residual can change.

@lmoresi
lmoresi force-pushed the feature/discrete-adjoint branch from dd7596d to 2848d87 Compare September 15, 2026 00:08
The exact content of the split: adjoint_solve / dual_of / sensitivity on
every solver (composite on Stokes; the consistent tangent assembled for the
adjoint whichever tangent the forward used); uw.adjoint.TranscriptAdjoint,
a reverse driver over the transcript with no problem-specific wiring, with
fields read through their value and their gradient assembled as FEM loads;
uw.adjoint.inner over owned dofs; an adjoint verdict on every recorded
operator and transcript_adjoint_segments; the consistent Newton tangent as
the default with Picard opt-in; the sinking-blob example with its Taylor
test through the library (1.00000); tests 0018-0020, serial and np=2.

The history of how each piece was built and reviewed is on the
docs/timestepping-pattern branch before its split commit (baae081,
258ee6e, fa83b19, 790a97c, 567c8c9, fba4497, dd92879).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V
@lmoresi
lmoresi force-pushed the feature/discrete-adjoint branch from 2848d87 to 6c3a976 Compare September 15, 2026 10:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant