feat(domain): benchmark all differentiable solvers in the loop

[andrinr]

TL;DR — which solvers are good for solver-in-the-loop training, ranked
fairly: six differentiable NS solvers, one corrector, one protocol.

• Ranked result (below): PICT 1.87×, Warp-NS 1.36×, JAX-CFD 0.71×,
  XLB 0.61×, INS.jl 0.54×, PhiFlow not admitted. Each solver is scored on
  the fraction of its own coarse-to-fine gap its corrector closes, so no
  solver is penalised for its raw forward accuracy — the same fairness 3D
  initial-condition recovery gets by having each solver invert its own
  forward map.
• Only the two velocity-complete solvers beat their own uncorrected
  rollout. The three admitted native-checkpoint solvers do not, at a
  converged budget, against their own target.
• The solver VJP is worth the most exactly where the corrector fails.
  Lift is 0.98–0.99 for the two solvers whose correctors work and 1.13–1.40
  for the three that lose. No cell here shows the VJP helping a corrector
  that is itself useful.
• The ranking is unchanged when every solver is instead trained against one
  shared external target, so it is not an artifact of the target choice.
• Draft, stacked on fix(solver): make recurrent state canonical and opt-in #121 and merged up to its current head.

One 2D solver-in-the-loop neural-correction benchmark, shared by all six
differentiable Navier–Stokes grid solvers: JAX-CFD, INS.jl, PhiFlow, PICT,
Warp-NS, XLB. Same corrector, same data, same optimizer, same protocol —
only the differentiable solver transition and its VJP change between cells.

It reports two things that are usually conflated, and keeps them apart:

• Correction gain — does the trained corrector beat the uncorrected
  solver? Reference-dependent, reported per target.
• Solver-VJP lift — within one solver, does backpropagating through eight
  autoregressive solver–corrector transitions beat the identical forward
  recurrence with the state stopped at every solver output? This is the
  primary differentiability result: a recurrent credit-assignment audit.

Stacked on #121 (opt-in canonical recurrent state); draft until that lands.

The ranked result

Which solver should you pick for solver-in-the-loop training? Each solver's
corrector is trained against that solver's own spatially and temporally
refined trajectory, so every solver faces the same task — remove your own
discretization error using your own gradients — and none is penalised for
where its raw forward happens to sit relative to some external reference. This
is the fairness that 3D initial-condition recovery gets by having each solver
invert its own forward map. The ranked quantity is the fraction of a solver's
own coarse-to-fine gap that its corrector closes; raw errors are not
comparable across solvers here, and are shown only for context.

Three model seeds each, 3000 updates, mean objective:

Rank	Solver	Recurrence	Own gap	Corrected	Gap closure	VJP lift
1	PICT	velocity-complete	0.04093	0.01668	1.871× (sd 0.087)	0.992 (sd 0.047)
2	Warp-NS	velocity-complete	0.02780	0.01708	1.355× (sd 0.081)	0.981 (sd 0.079)
3	JAX-CFD	native checkpoint	0.05555	0.15729	0.705× (sd 0.002)	1.305 (sd 0.020)
4	XLB	native checkpoint	0.02685	0.11463	0.611× (sd 0.019)	1.128 (sd 0.016)
5	INS.jl	native checkpoint	0.04130	0.16347	0.535× (sd 0.018)	1.399 (sd 0.026)
—	PhiFlow	native checkpoint	—	—	not admitted	—

Two things fall out, and they are the same two the shared-target study found:

Only velocity-complete solvers produce a corrector worth running. PICT and
Warp-NS close their own refinement gap by 1.87× and 1.36×. The three admitted
checkpoint solvers make their own rollout worse than not correcting at all,
at a converged budget, against their own target — so this is not an artifact of
chasing someone else's discretization.

The solver VJP is worth the most exactly where the corrector fails. Lift is
0.98–0.99 for the two solvers whose correctors work, and 1.13–1.40 for the
three that lose. Backpropagating through the solver reliably separates two bad
correctors and does nothing for a good one.

The ranking does not depend on the target. Repeating the comparison with
every solver trained against one shared 128² pseudo-spectral reference gives
PICT 1.95×, Warp-NS 1.37×, JAX-CFD 0.87×, INS.jl 0.79×, XLB 0.79× — the same
top two, the same split, and the same inverse relation to lift. Since the two
designs fail in opposite ways (a shared target confounds correctability with
forward accuracy; per-solver targets differ in difficulty), agreement between
them is the strongest evidence in this PR.

PhiFlow is the one gap: its native state closes the shared-target task but its
refined-reference closure audit spikes, so it is refused admission to the fair
comparison and can only be ranked in the shared-target view. An admission
failure is reported as a result, as elsewhere in this benchmark.

Paired differentiability result

Differentiating through the solver pays, and that is the robust part of the
story. Under both independent references, four of six solvers show a
resolved paired benefit from full temporal credit over local one-step
supervision. Nothing about the forward recurrence or the supervision changes
between the two arms — only whether gradients cross solver steps.

Solver	VJP lift, spectral target [95% CI]	VJP lift, finite-volume target [95% CI]	Verdict
JAX-CFD	+39.0% [+35.6%, +42.2%]	+49.8% [+35.4%, +77.5%]	beneficial under both
INS.jl	+38.8% [+27.8%, +50.5%]	+38.2% [+35.8%, +41.3%]	beneficial under both
PhiFlow	+37.0% [+21.0%, +48.1%]	+45.1% [+38.2%, +54.3%]	beneficial under both
XLB	+20.5% [+13.1%, +30.2%]	+16.2% [+7.2%, +24.6%]	beneficial under both
PICT	−4.2% [−7.8%, −1.0%]	+4.3% [+2.1%, +6.7%]	flips sign with the target
Warp-NS	−0.3% [−2.8%, +2.0%]	+4.9% [−2.0%, +10.9%]	inconclusive under both

A better gradient does not buy you a better corrector. At this training
budget only PICT's corrector beats its own uncorrected solver under both
references; the other five are worse than doing nothing, and the three with
the largest VJP lift are among the worst absolute performers. Temporal credit
and corrector usefulness are separate axes, and a benchmark that reports only
one of them will mislead. The 200-update budget is part of why the other five
lose here — see the follow-up
for what a 15× budget changes, and what it does not.

Solver	Native final error (spectral / FV)	Corrected final error (spectral / FV)	Correction gain (spectral / FV)
JAX-CFD	0.06842 / 0.06755	0.1526 / 0.1321	0.885× / 0.967×
INS.jl	0.05117 / 0.05108	0.1492 / 0.1526	0.737× / 0.701×
PhiFlow	0.04461 / 0.04461	0.2921 / 0.1743	0.686× / 0.677×
XLB	0.03587 / 0.03489	0.1117 / 0.1255	0.752× / 0.650×
PICT	0.05016 / 0.05039	0.04448 / 0.03957	1.107× / 1.212×
Warp-NS	0.03568 / 0.03608	0.03809 / 0.03705	0.946× / 0.973×

The solver's error is a property of the solver; the corrector's error is a
property of the target you picked. Repeating the whole study against an
independently discretized reference — dealiased pseudo-spectral versus
conservative finite volume, whose held-out fields agree to a median 0.16% —
moves native solver error by at most 2.7% but moves trained corrected error by
up to 40.3% (PhiFlow), and flips PICT's paired VJP verdict from significantly
harmful to significantly beneficial. Very similar targets do not imply the
same learned endpoint. If you benchmark a hybrid solver against one reference,
you are reporting a property of that reference too.

Protocol

Paired training and evaluation — click to expand

Every cell shares a 32² candidate grid, viscosity, physical time step, 16
training ICs, eight held-out ICs, three model seeds, 200 optimizer updates,
recurrent training to t=1.92, and evaluation to t=2.88. The learned
component is the same zero-initialized Equinox periodic residual CNN (56,098
parameters).

Training uses truncated backpropagation through eight consecutive correction
intervals. Each interval advances four native solver steps (4 × 0.02 = 0.08),
applies the corrector, and feeds both the corrected velocity and the native
solver checkpoint into the next interval. There is no teacher forcing inside
the sampled window. The loss averages all eight corrected states and adds a
0.1-weight terminal loss on the last provisional solver state. Held-out
evaluation is fully free-running for 36 solver–corrector intervals — after the
initial condition no reference field re-enters the rollout.

The two paired arms are:

1. full temporal credit — gradients flow through eight solver–corrector
   transitions, including solver velocity and native state;
2. local control — identical forward recurrence and supervision, but
   recurrent velocity and native state are stopped at every solver output.

The VJP lift is the geometric paired error_local / error_full ratio.
Correction gain is a separate quantity: error_uncorrected / error_corrected.

Reference design and convergence audits — click to expand

Reference	Production target	Convergence audit
dealiased pseudo-spectral	128², dt/4	256², dt/8
conservative finite volume	MUSCL/minmod + Rusanov + five-point diffusion, 128², dt/4	256², dt/8

Both start from the same Fourier-prolonged continuum initial condition and
compare seeds 0 and 100 at frames 1, 8, 24, 36. Training is rejected if the p95
128²-versus-256² discrepancy exceeds 0.5%.

Audit	Median / p95 / max relative discrepancy	Gate
pseudo-spectral 128² vs 256²	p95 9.89e-8; max 1.01e-7	pass
finite-volume 128² vs 256²	0.121% / 0.160% / 0.168%	pass
spectral vs finite-volume, all held-out frames	0.162% / 0.207% / 0.224%	diagnostic
spectral vs finite-volume, final frame	0.196% / 0.222% / 0.224%	diagnostic

The finite-volume method uses an FFT only for the periodic discrete Poisson
solve; transport and diffusion use conservative grid stencils, so this is not a
second spectral integrator.

Recurrent-state admission — click to expand

Native recurrence is detected only when state is differentiable in both the
Tesseract input and output schemas. JAX-CFD, INS.jl, PhiFlow and XLB carry a
native checkpoint through training, evaluation, plotting, finite-difference
validation and differentiation; PICT and Warp-NS are velocity-complete. All six
pass the same nonlinear and Taylor–Green recurrence gate.

Solver	Recurrent representation	Nonlinear closure p95	Admitted
JAX-CFD	native checkpoint	0	yes
INS.jl	native checkpoint	2.92e-7	yes
PhiFlow	native checkpoint	4.72e-7	yes
PICT	velocity-complete	0	yes
Warp-NS	velocity-complete	0	yes
XLB	native checkpoint	1.75e-3	yes

Controls: analytic Taylor–Green and method-local reference — click to expand

Analytic Taylor–Green. All six solvers pass and train against the analytic
solution. JAX-CFD, INS.jl, PhiFlow, PICT and Warp-NS improve absolute error;
XLB is close to neutral. The paired VJP effect is small near this error floor,
as expected.

Method-local refined reference. Five solvers also complete a task against
their own 64², half-step target. Absolute errors are not ranked across solvers
because the targets differ. PICT improves absolute error by 1.119×; the paired
VJP benefit is resolved for JAX-CFD (37.0%), INS.jl (29.8%), XLB (13.3%) and
PICT (4.60%), and inconclusive for Warp-NS. PhiFlow is rejected before training
in this control only: its native state closes the common-reference task, but
the chaotic refined-reference audit spikes to 0.0101/0.0252 maximum
coarse/fine closure residuals, 22.8%/70.7% of the refinement signal.

Full fields, physics diagnostics and GIFs for every task are in the
artifact index.

Limits

• The paired VJP comparison is within solver. It isolates temporal
  differentiation; it is not a cross-solver gradient-quality ranking.
• A positive lift says the eight-step training graph beats local supervision.
  It does not say the corrector beats the uncorrected solver: across six
  solvers the two properties are anti-correlated (see the follow-up above).
• Absolute errors remain properties of the chosen target. The two-reference
  sweep tests whether conclusions survive one independent discretization
  change; it does not establish a universally neutral reference.
• Reference agreement is reported for this smooth low-wave-number IC
  distribution, viscosity, resolution and horizon only.

Follow-up: is the sub-one correction gain a budget artifact?

Partly, and the answer separates the solvers along the recurrent-state
interface. A budget ladder (200 → 1000 → 4000 updates) and a paired objective
probe showed both published choices understate the corrector: the plain mean
objective beats the solver_terminal one at equal budget, and correction gain
keeps climbing with updates. Repeating the whole nonlinear task at 3000 updates
with mean, three model seeds per solver:

Solver	Recurrence	Native	Corrected	Correction gain	VJP lift
PICT	velocity-complete	0.05438	0.02061	1.949× (sd 0.027)	0.978 (sd 0.004)
Warp-NS	velocity-complete	0.03683	0.02233	1.368× (sd 0.077)	0.964 (sd 0.055)
JAX-CFD	native checkpoint	0.07187	0.16858	0.872× (sd 0.016)	1.259 (sd 0.018)
INS.jl	native checkpoint	0.05528	0.14706	0.794× (sd 0.035)	1.506 (sd 0.021)
PhiFlow	native checkpoint	0.04659	0.14142	0.743× (sd 0.012)	1.550 (sd 0.007)
XLB	native checkpoint	0.03643	0.10944	0.791× (sd 0.027)	1.067 (sd 0.046)

The two axes oppose each other across the six solvers (Spearman 0.77 between
better gain and lower lift, XLB the exception): every solver with a resolved
VJP lift trails its uncorrected solver, and neither solver that beats its
uncorrected solver shows one. Per-solver, against a zero log-lift:

Solver	Lift	t	Verdict	Correction gain
PhiFlow	1.550	86.9	VJP helps	0.743×
INS.jl	1.506	34.0	VJP helps	0.794×
JAX-CFD	1.259	22.2	VJP helps	0.872×
XLB	1.067	2.4	not resolved	0.791×
Warp-NS	0.964	−1.2	not resolved	1.368×
PICT	0.978	−10.5	VJP slightly hurts	1.949×

No cell in this benchmark shows the solver VJP helping a corrector that is
itself useful. Where the VJP clearly helps, it separates two correctors that
both lose to doing nothing; where the corrector wins, the VJP is worth nothing
or slightly negative. The fair per-solver ranking above reproduces this
independently.

One caveat on the metric itself: the paired control stops the recurrent
velocity and the native checkpoint. A checkpoint solver therefore loses
strictly more gradient path to the control than a velocity-complete solver,
which has no checkpoint to stop. Part of the measured lift may reflect how much
the control removes rather than what differentiating through the solver is
worth, so lift magnitudes are not directly comparable across the two interface
classes.

Two limits on how far that should be read:

• Budget does not rescue the checkpoint group. A single 4000-update JAX-CFD
  cell scored 1.059×, which looked like a crossover. Replicating it across
  three seeds gives 0.949× (log sd 0.096) — the seed spread at 4000 updates is
  six times wider than at 3000, and the apparent jump was seed noise, not
  budget. Checkpoint correctors do not cross 1.0 out to 4000 updates.
• The alignment is across solvers, not a demonstrated mechanism. A
  within-solver probe that forces a checkpoint solver onto the velocity-complete
  path reproduces both halves of the split, but it fails this benchmark's own
  recurrence-closure gate (long_closure_error_p95 0.52 and 0.26 against a 0.01
  tolerance) and is therefore diagnostic only. Two solvers per class is a small
  sample: the split is a robust observation, not an established causal claim
  about recurrent-state interfaces.

One gradient-quality result fell out of the same runs: PhiFlow's training
gradient went non-finite at update 2021 in one of three seeds and the harness
stopped that cell, consistent with the existing repository finding that
projection-based schemes accumulate modes that corrupt long gradients. That
seed is excluded from the PhiFlow row above rather than averaged in.

Implementation

• Detect native recurrence only when state is differentiable in both schemas.
• Thread solver state through corrector training, evaluation,
  finite-difference checks, long-closure checks and self-reference generation.
• Stop both velocity and native state in the paired control.
• Add a conservative finite-volume reference, exact Fourier prolongation for
  resolution audits, and a gated 128²-versus-256² space-time convergence check.
• Register the sensitivity study as the opt-in
  optimization/solver_in_loop_reference_sensitivity experiment.
• Existing drag-optimization registrations are unchanged.

Validation and offline execution — click to expand

All experiments ran offline on Kander through the shared Slurm/Pyxis solver
tooling. No hosted benchmark was triggered; this draft carries
benchmark:none.

• reference convergence audit: 1697345; all-held-out reference-disagreement
  audit: 1697383;
• pseudo-spectral production: 1697346–1697351; finite-volume production:
  1697352, 1697371–1697375; merge/render: 1697376;
• nonlinear, Taylor–Green and method-local controls: see the artifact index;
• numerical-source validation at 21db5ff: Ruff and format passed; 522 tests
  passed, three skipped (1697344);
• final plot-only layout change: Ruff/format and the offline pre-merge renderer
  passed; hosted Python checks are reported on this PR head.

[PasteurBot]

📊 View the full benchmark results

No benchmarks ran for this PR, so there is no status report.

---