[dionhaefner]

Relevant issue or PR

#268 (Exponax in 2D ns-grid) builds on this. The Kolmogorov-flow checks in #272 are the first experiments to use the new inputs.

Description of changes

Adds two optional inputs to the canonical ns-grid schema:

• forcing: a constant body force per unit mass at the cell centers, with the same shape as v0;
• linear_drag: the coefficient α of a drag term −αu.

Both are allowed only on 2-D periodic boxes; the canonical schema rejects them with an obstacle, non-periodic BCs, in 3-D, or when the force's shape does not match v0. Neither is differentiable. With both left at None, every solver runs exactly the code it ran before.

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

Each solver applies the force and the drag inside its own time step, in both apply and VJP. The force reaches each solver's native locations through the same map it uses for v0.

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, part of the momentum rate of incompressible_rk4	in every RK4 stage
Warp-NS	in tentative_vel_2d_kernel	explicit
PICT	velocity source, set every step in a PRE callback, 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	Exponax's own drag
OpenFOAM	rejected: icoFoam has no momentum source term	rejected

Solver changes beyond the new inputs:

• Exponax 2D convention (bug fix). Exponax's 2D vorticity steppers evolve q = −ω, while the tesseract passed ω. Its 2D path therefore advected with the wrong sign and returned −S(−u₀). It now converts with the right sign. The 2D path has never run in ns-grid (Exponax is still popped there; fix(domain): run Exponax on the periodic ns-grid experiments #268 adds it), and the 3D path is unchanged.
• Exponax force placement. The canonical layout centers cell j at (j + ½)h; Exponax samples at jh. Its own injection_mode/injection_scale would put sin(κy) half a cell off, a 20% error at 64² in the laminar test. Feeding the force's own curl as the injection term is exact for any force, so the tesseract does that instead of mapping to the single-mode inputs.
• Exponax mean flow now decays under drag (and grows under a mean force) exactly; before, the mean was restored unchanged.
• XLB sub-steps: with a force, the lattice Mach number is sized from a bound that does not depend on the solution, max(max|v0|, ‖F‖∞/(νκ_min² + α)). This keeps a flow that the force spins up from rest at the target Mach number.

Testing done

• tests/test_schemas.py: the new inputs are rejected outside 2-D periodic boxes and accepted inside.

• tests/test_ns_grid_forcing.py (Docker, skipped without images), for every solver image:

  • the laminar Kolmogorov transient from rest against its exact solution u = (A/λ)(1 − e^{−λt}) sin(κy), λ = νκ² + α, within 5% at N = 64;
  • a zero force with zero drag reproduces the unforced output;
  • a forced VJP against a central difference within 1%.

  All pass on images built from this branch, and OpenFOAM's rejection is tested. Measured laminar errors at dt = 0.002:

  Solver	32²	64²	128²
  JAX-CFD	2.9e-2	7.4e-3	1.8e-3
  INS.jl	2.9e-2	7.4e-3	1.8e-3
  PhiFlow	2.9e-2	7.4e-3	1.8e-3
  Warp-NS	3.0e-2	7.7e-3	2.2e-3
  XLB	7.8e-2	1.9e-2	4.8e-3
  Exponax	4e-7	7e-6	1e-5

  PICT passes the 64² test (dt = 0.01).

• After the latest merge with main, where INS.jl runs on its CUDA backend (perf(solver): run INS.jl on its CUDA backend #307), the forcing tests pass on freshly built INS.jl and XLB images.

• No regression (A/B, 2026-10-26, before refactor(solver): use PhiFlow's documented schemes for the periodic box and the channel #288 and perf(solver): run INS.jl on its CUDA backend #307 changed PhiFlow and INS.jl on main). Images from main and from this branch ran the same unforced cases: multimode at N = 32, TGV at N = 16, and 3D TGV at N = 16. Each case compared the forward result and the VJP with respect to v0 and viscosity.

  • JAX-CFD, XLB and INS.jl are bitwise identical.
  • Warp-NS, PICT and PhiFlow differ by at most 2e-7 (forward) and 2e-6 (VJP). That is the same size as main against itself in a fresh container.
  • Exponax's 3D and TGV outputs are unchanged. 2D multimode changes by the sign fix above.

• No benchmark label: the A/B comparison above covers the unforced paths. Whether a CI benchmark pass is worth its cost is the maintainer's call.

[dionhaefner]

Folded into #272. The forcing inputs have no user outside the Kolmogorov-flow checks, so they now land with them. The one general fix this PR carried, Exponax's 2D sign convention, moved to #268.