[dionhaefner]

Relevant issue or PR

For reproduction, not for merge. It collects the solver-in-the-loop study of the Kolmogorov flow in one place, so its results can be rerun, for instance for a paper. It replaces #273 (native solver state) and #274 (optimization/kolmogorov_sil), and is stacked on #272.

Description of changes

The study trained a learned corrector through each differentiable ns-grid solver. It did not become a benchmark. Every solver that runs the turbulent flow passed every training check, so training did not separate the solvers beyond the forward cost that forward/kolmogorov already measures. The one defect it found, PICT's adjoint tolerance (#294), is caught as well by gradient/chain_fd_check in #272, which needs neither a corrector nor native state.

Two commits:

1. Native solver state for JAX-CFD, INS.jl and XLB. Optional state and return_state inputs let a chain of calls continue a staggered or lattice solver exactly, instead of restarting it from the cell-centred velocity.
2. The study:
   • optimization/kolmogorov_sil: the CI-size feasibility and cost check from feat(domain): add a Kolmogorov-flow solver-in-the-loop benchmark to ns-grid #274.
   • optimization/kolmogorov_sil_stability: longer training on three seeds, with 200-interval rollouts of every trained corrector. It takes about half an hour per solver on an A100, and most of a day for PICT, so it is not meant for CI.
   • studies/kolmogorov-sil/: the commands, scripts and result tables.

python -m mosaic.benchmarks.cli run -o results/stability-XLB -p ns-grid \
  -e optimization/kolmogorov_sil_stability -s XLB --no-plots
python studies/kolmogorov-sil/summarize.py results/sil results/stability-*

Findings

The full tables are in studies/kolmogorov-sil/README.md.

• Feasibility. All five solvers pass the reference audit, closure and the FD gradient gate. Their FD errors are 1e-4 to 2e-3, against a 5% threshold.
• Stability. None of the 45 trainings failed and none of the 270 long rollouts went non-finite.
• Accuracy. Averaged over the three seeds, supervised training gives the lowest test error on every solver except Exponax, where stopping the gradient is lower. On that average, training through the solver is never better than stopping the gradient, though it is on single seeds.
• Cost. At K = 16, one update through the solver costs 1.0 s (JAX-CFD) to 27 s (PICT), against 0.7 s to 4.7 s with the gradient stopped. These runs partly shared the GPU, so the times are rough.
• Coverage. Warp-NS and PhiFlow are excluded. Warp-NS diverges on the flow. PhiFlow's 256² reference run goes non-finite from two of the three test states.

Testing done

• Unit and dummy-corpus tests for the native state, the corrector, the SIL kernel and its plots pass, along with the config tests.
• The results in the README come from an A100. The images were built from the code merged with main on 2026-10-28, with PICT including fix(solver): make PICT's VJP linear in the cotangent #294.
• This branch stays on the code that produced those results and does not follow main or the stack below it. Since then INS.jl has moved to its CUDA backend (perf(solver): run INS.jl on its CUDA backend #307), and the native-state change has not been ported to that backend.