[dionhaefner]

Relevant issue or PR

Stacked on #272 (Kolmogorov forward checks), which is stacked on #268 and #267. Re-implements the contract of #121 on current main; #121 itself is superseded and not merged.

Description of changes

The solver-in-the-loop benchmark calls a solver once per correction interval and hands the corrected state back through the canonical cell-centred velocity. For a solver on a staggered grid or a lattice, every such round trip loses information. A corrector trained across them would partly learn to undo our interface rather than the solver's error.

Solvers can now continue from their own state:

• Schema. The canonical input gains state (an opaque, solver-specific array) and return_state; the output gains state. A solver accepts them only if its InputSchema sets supports_native_state, and only on a periodic box. Continuing from state, a solver takes on only the change from the velocity it returned with that state to v0. An unchanged v0 therefore continues the run exactly, and a corrector's change still reaches the solver.
• grid_layout.py (shared): the face↔centre averages and their inverse, the lift. The average damps mode k by cos(πk/N), which vanishes at Nyquist, so the lift caps its gain at 32. That drops the Nyquist mode, and only it on grids up to N = 100. This cap is interface numerics we own, and every staggered solver uses the same one.
• JAX-CFD and INS.jl: the state is the high-face velocity plus the returned cell-centred velocity, (nx, ny, 1, 4). A correction is lifted onto the faces and projected on its own before it is added, so a zero correction adds exactly zero. INS.jl runs the same hand-written RK4 loop in apply and VJP whenever a state is involved, and its VJP takes a cotangent on the returned state.
• XLB: the state is the post-collision populations, the returned velocity and the number of LBM sub-steps, which fixes the lattice units, (2q + 3, nx, ny, 1). A correction moves only the equilibrium momentum: f' = f + f_eq(ρ, u + δ) − f_eq(ρ, u), keeping density and the non-equilibrium moments. XLB runs in float64, so the populations are stored as float32 high and low parts to keep that precision across a restart.
• All three continue in 2D only, and the forcing and drag are inputs on every call, so a continued run carries them unchanged.

Closure check

I ran a chain of one-interval calls against one call over the same time. The setup is the Kolmogorov coarse configuration (64², dt = 0.01, 10 steps per interval), from the three test states of the ground truth. The table gives the relative L2 distance (median / max):

Solver	Restart from the canonical velocity: 2 intervals	16	40	With native state: 2	16	40
JAX-CFD	1.0% / 1.0%	13% / 14%	55% / 64%	0	0	0
INS.jl	1.0% / 1.0%	13% / 14%	55% / 63%	0	0	0
XLB	0.53% / 0.55%	2.3% / 2.7%	9.7% / 11%	1e-6	1e-6	3e-6 / 5e-6
PICT	0	0	0	—	—	—
Exponax	1e-6	7e-6	5e-5	—	—	—
Warp-NS	0	diverges	diverges	—	—	—
PhiFlow	1.1% / 1.1%	diverges	diverges	—	—	—

The thresholds carried over from #116 are a median of 0.5% and a 95th percentile of 1%.

• Pass with native state: JAX-CFD, INS.jl and XLB.
• Pass as they are: PICT, Exponax and Warp-NS. PICT is collocated and its PISO step carries nothing else, so the canonical restart is already exact. Exponax is collocated too and its ETDRK2 step carries no history, so it restarts exactly up to round-off. Warp-NS diverges on this flow, which forward/kolmogorov documents.
• PhiFlow fails the check but diverges on this flow at the benchmark timestep (explicit Euler with central advection), so it is not eligible for the solver-in-the-loop benchmark. I did not add native state to a solver that cannot run the flow.

Testing done

• tests/test_grid_layout.py: the averages follow the high-face convention, and the lift inverts them below Nyquist, drops Nyquist, and works on JAX arrays.
• tests/test_schemas.py: state and return_state are rejected by solvers that do not opt in and outside periodic boxes.
• tests/test_ns_grid_native_state.py (Docker): for JAX-CFD, INS.jl and XLB, four chained calls with state match one call to 1e-5. A VJP through v0 and the state, with cotangents on the result and the returned state, matches a central difference to 2%. All pass on images built from this branch merged with main (2026-10-28).
• Calls without a state are bitwise identical to main for JAX-CFD, XLB and INS.jl. This was an A/B between images before main was merged in, on multimode at N = 32, TGV at N = 16 and 3D TGV at N = 16, forward and VJP with respect to v0 and viscosity.
• The closure table is from the same merged images.

[dionhaefner]

Closing in favour of #316, which carries this change together with the solver-in-the-loop study for reproduction. The benchmark side of the Kolmogorov work is #272 and #315.