[dionhaefner]

Relevant issue or PR

All Kolmogorov-flow work in one PR: body forcing for the ns-grid solvers, the three checks and the ground truth. It replaces #267 (forcing inputs) and #310. It builds on #268 (Exponax in 2D ns-grid, including its 2D sign fix) and #317 (experiment-generated references), both now on main. #203 (docs redesign) is now on main too, with the Kolmogorov setup on its scoreboard; the protocol text of the three checks is still to be added here.

The solver-in-the-loop study that led here is in draft #316 for reproduction. It is not meant to merge.

Description of changes

Three checks of the forced 2D Kolmogorov flow (force sin 4y, linear drag), the ground truth they score against, its generator, and the inputs and outputs they need. It also fixes PhiFlow's periodic pressure solve, which made PhiFlow fail both turbulent checks and run for 79 minutes on CI.

Forcing inputs. The canonical ns-grid schema gets two optional, non-differentiable inputs, allowed only on 2D periodic boxes:

inputs = {..., "forcing": f, "linear_drag": np.array([0.1], np.float32)}  # f: (N, N, 1, 2), at cell centers

Each solver applies them inside its own time step, in both apply and VJP, and maps the force to its native locations the same way it maps v0. With both left at None, every solver runs the code it ran before.

Solver	Force	Drag
JAX-CFD	forcing callable of semi_implicit_navier_stokes, averaged onto the faces	explicit, in every RK4 stage
PhiFlow	resampled to the faces like v0, in the momentum rate of every RK4 stage	in every RK4 stage
Warp-NS	in tentative_vel_2d_kernel	explicit
PICT	velocity source set every step, in unit-cell units	explicit, from the velocity at the start of the step
INS.jl	projected staggered force added to the right-hand side	in every RK4 stage
XLB	exact-difference method after collision; half-force velocity shift for the output	semi-implicit in the half-force velocity
Exponax	spectral curl of the force as the injection term of its Kolmogorov stepper (its own injection_mode would put sin κy half a cell off)	Exponax's drag; the mean flow decays exactly
OpenFOAM	rejected: icoFoam has no momentum source term	rejected

With a force, XLB sizes its lattice Mach number from a bound that does not depend on the solution, so a flow that the force spins up from rest stays at the target Mach number.

Snapshots. One more optional input and a non-differentiable output, also 2D periodic boxes only:

inputs = {..., "steps": 400, "snapshot_every": 100}
out["snapshots"]  # (4, N, N, 1, 2): the velocity at the cell centers after 100, 200, 300 and 400 steps

Each solver records them inside its own loop, so its staggered or lattice state never round-trips through the cell centers between times: JAX-CFD slices its trajectory, Exponax, PhiFlow and XLB nest their scans, INS.jl records in its RK4 loop, PICT in a post-step hook, Warp-NS in its step loop. The VJPs skip them. OpenFOAM rejects them, as it rejects the forcing.

forward/kolmogorov_laminar. From rest at ν = 0.05 the flow stays a shear flow with an exact solution, u = (A/λ)(1 − e^{−λt}) sin(κy). The sweep is N ∈ {32, 64, 128} to t = 2. The check is max_error(0.05, at=64) plus min_order(0.8). A correct second-order scheme is off by about 1% at N = 64. A convention bug (a missing 2π/L, a flipped sign, a force half a cell off) is off by 20% or more.

forward/kolmogorov. Each solver runs on a ladder of grids, 32² to 256², at the CFL number of 64² with dt = 0.01, from three turbulent states. Each grid runs once to t = 4, returns the velocity at t = 1, 2, 3 and 4 as snapshots, and is scored on 64² against the truth. From the first non-finite snapshot on, the run counts as diverged. The plot shows error against wall time per time unit. The check is min_refinement_order(0.8): the observed order from 64² to 256² at t = 1 must be at least 0.8, and no grid from 64² up may diverge. The bound comes from the scheme's consistency, not from the other solvers.

gradient/chain_fd_check. Training a model through a solver chains many calls, and its loss is normalized by the target, so each call receives a far smaller cotangent than in the single-call gradient checks. A VJP that is only accurate for large cotangents, for example because its adjoint solves stop at an absolute tolerance, passes those checks and fails this one. The check chains 16 calls of 0.1 time units at 64². It starts from a truth state plus a 1% divergence-free offset, and the loss is the mean relative squared distance to the solver's own unperturbed chain. The VJP is compared with central differences over 8 divergence-free directions at ε ∈ {1, 0.1, 0.01}. The gates are min_cosine(0.99), max_rel_err(1e-2) and rel_err_peer_outlier(50). A correct VJP meets float32 differences to 1e-4–2e-3, and a truncated one is off by tens of percent.

Ground truth. The ground truth is an Exponax run at 1024², with 2/3-rule dealiasing, from three states spun up for 100 time units. Its generator runs through #317's reference path. A 512² run differs from the truth by 0.3% at t = 1. The 0.7 MB file is documented in references/README.md.

PhiFlow's periodic pressure solve. The periodic pressure matrix is singular. Unless the solve states its rank deficiency, phiml adds a Lagrange multiplier to it, which makes the system indefinite, and CG is not guaranteed to converge on that. Under jit, phiml returns an unconverged solve without raising. On the Kolmogorov flow at 256², 852 of 6400 solves to t = 4 stopped unconverged, 128 of them flagged as diverged with residuals up to 10⁴ times the right-hand side. Which grid and state then went non-finite varied from run to run. The solve now passes rank_deficiency=0. make_incompressible already balances the divergence to zero mean, so plain CG on the singular system converges, in about 375 iterations instead of 570. The adjoint solve now stops at the relative tolerance only, because its right-hand side scales with the cotangent. With the absolute tolerance, the chain check's small cotangents stopped it after 0–17 iterations. Both are recorded as interventions. The fix also makes PhiFlow's cost experiments faster, 1.3–7.5× in 2D and 1.4–16× in 3D, growing with N. All four of its ns-grid cost anomalies become ok, so their explanation, which attributed the cost to the sparse matrix-vector product, is removed.

Warp-NS is excluded from the turbulent checks: its explicit-Euler, central-difference scheme is linearly unstable unless dt ≲ 4ν/u², a bound that does not shrink with the grid. OpenFOAM is excluded too: icoFoam has no source term, and OpenFOAM has no VJP.

Results (A100)

forward/kolmogorov_laminar, relative error at N = 32 / 64 / 128. All pass.

Solver	32	64	128
Exponax	3.0e-06	3.5e-05	1.3e-05
JAX-CFD, INS.jl, PhiFlow	2.9e-02	7.4e-03	1.8e-03
PICT	2.9e-02	7.1e-03	1.7e-03
Warp-NS	3.0e-02	7.7e-03	2.2e-03
XLB	7.8e-02	1.9e-02	4.8e-03

forward/kolmogorov (median error at t = 1 per grid, observed order 64² → 256²) and gradient/chain_fd_check (median error and cosine at the best ε):

Solver	32²	64²	128²	256²	Order	Chain error	Chain cosine
Exponax	0.168	0.039	0.0046	0.0028	1.89	1.8e-03	1.0000
JAX-CFD	0.226	0.086	0.029	0.0077	1.73	2.4e-04	1.0000
INS.jl	0.226	0.086	0.029	0.0077	1.73	2.1e-04	1.0000
XLB	0.135	0.032	0.0063	0.0032	1.66	9.8e-05	1.0000
PICT	0.258	0.092	0.034	0.012	1.46	2.1e-04	1.0000
PhiFlow	0.206	0.082	0.028	0.0076	1.72	2.5e-04	1.0000

Every solver passes both checks. Before the PhiFlow fix, its 256² run went non-finite at t = 4 from two of the three states on CI, and its chain gradient was non-finite on CI and off by 12% at the best ε on the A100.

forward/kolmogorov took 1,027 s for all solvers together on the A100: PICT 523 s, PhiFlow 349 s, INS.jl 68 s, the others under 40 s. The last CI run took 6,600 s on its GPU, 4,759 s of it PhiFlow's. Locally, PhiFlow took 5,585 s before the fix.

The chained check also catches the defect that #294 fixed. Before that fix, PICT's adjoint solves stopped at an absolute tolerance, and the check gave a median error of 0.89 at a cosine of 0.51.

Testing done

• tests/test_schemas.py: the forcing inputs and snapshot_every are rejected outside 2D periodic boxes, and snapshot_every must divide steps. tests/test_ns_grid_forcing.py (Docker, skipped without images) checks every solver image on four things: the laminar transient within 5% at N = 64, a zero force with zero drag against the unforced output, a forced VJP against central differences within 1%, and the snapshots of a forced run against separate runs that stop at their steps, to 1e-5. OpenFOAM must reject both. All 30 pass.
• tests/test_kolmogorov.py covers: exact resampling of band-limited fields, the laminar solution, the vorticity identities, the forward kernel on a fake solver including a diverged grid and a run that diverges only at its last snapshot, and min_refinement_order. tests/core/test_reference.py checks the layout of the stored truth.
• tests/test_chain_fd_check.py covers: the loss, the divergence-free directions, a correct VJP, and a VJP that drops small cotangents, which the check must flag.
• The unit suite passes on this branch merged with main.
• No regression from the snapshots: with them, every solver's forward/kolmogorov errors match the previous CI run's separate calls to within 1e-4 relative (PhiFlow aside). mosaic run of the ns-grid forward suite, gradient/fd_check, gradient/chain_fd_check and the ns-3d forward suite changes no status against that CI run apart from PhiFlow's two fixes. PhiFlow with the fix also ran every suite of both problems: no other status changed, and best-ε gradient errors stayed where they were (ns-grid fd_check 1.04e-5 against 1.05e-5).
• No regression from the forcing inputs (A/B on 2026-10-26, unforced multimode, 2D TGV and 3D TGV, forward and VJP). JAX-CFD, XLB and INS.jl were bitwise identical to main. Warp-NS, PICT and PhiFlow were within 2e-7 (forward) and 2e-6 (VJP), the same as main against itself.
• Results come from the solver images of this PR's previous CI run (built from 83a4d7e, so with fix(solver): tighten PICT's periodic linear-solve tolerance to 1e-6 #287, fix(solver): use XLA's default GEMM autotuning in jax-cfd, PhiFlow and XLB #291 and perf(solver): run INS.jl on its CUDA backend #307), with this push's solver sources and canonical schema layered on top, run on an A100 on 2026-11-02.

🤖 Generated with Claude Code

[read-the-docs-business]

Documentation build overview

📚 mosaic | 🛠️ Build #4662316 | 📁 Comparing 8f8be3f against latest (459e875)

  🔍 Preview build  

10 files changed · ± 10 modified

± Modified

• docs/getting-started.html
• docs/index.html
• docs/internals.html
• docs/results.html
• docs/results_ns_3d_grid.html
• docs/results_ns_grid.html
• docs/results_structural_mesh.html
• docs/results_thermal_mesh.html
• docs/solvers.html
• docs/tutorial.html

[PasteurBot]

📊 View the full benchmark results

The rendered docs preview has every plot for this run (forward accuracy, gradients, cost, optimization) merged with existing baseline results on main. The summary below reports pass/fail status.

---

Mosaic status

Benchmarks: skipped by the benchmark:none label.