Repository navigation
feat(domain): add a Kolmogorov-flow solver-in-the-loop benchmark to ns-grid - #274
Closed
dionhaefner wants to merge 17 commits into
Closed
dionhaefner wants to merge 17 commits into
dionhaefner wants to merge 17 commits into
Conversation
… experiments PhiFlow, Warp-NS and Exponax diverge in the configurations that forward/kolmogorov documents, so they are ineligible for the solver-in-the-loop experiments for the same reasons. The plots now label solvers like the other plots.
At a fixed step of 0.01 the central difference's truncation error grows
with the unroll, as the training loss curves more sharply along the
gradient. On XLB at K = 16 the check failed at 0.10 although the gradient
agrees with the finite difference to 5e-4 at a step of 1e-3. Keep the step
with the smallest error over {1e-3, 3e-3, 1e-2}, as the gradient suite does.
The FD check differentiated the training loss over the longest trained unroll, K = 4 in CI. It now has its own unroll, `fd_unroll` (16), so the gradient through a long chain of solver calls and native states is checked at the cost of one gradient and a few loss evaluations, without training at that length.
…ity and cost Remove the on-demand campaign configuration and the `on_demand` experiment flag that only it used. How much the solver gradient helps training depends on the network, learning rate and budget more than on the solver, so it is not something the benchmark scores. The CI configuration keeps its gates and its four arms, which give the cost of a training update with and without differentiating through the solver. The plots of rollout errors and gap closed are dropped: at 32 updates they do not measure how much the corrector helps. The cost plot now shows seconds per update per solver and arm.
# Conflicts: # mosaic/benchmarks/problems/navier_stokes_grid/exclusions.py
dionhaefner
added a commit
that referenced
this pull request
Oct 29, 2026
#### Relevant issue or PR Found while adding the Kolmogorov solver-in-the-loop experiment (#274), where PICT fails the finite-difference check of the training gradient. #### Description of changes **Deviation from PICT.** PICT's linear solves stop once the RMS of their residual falls below an absolute tolerance, 1e-5 in float32 by default. The adjoint solves of its backward pass use the same tolerance. When the cotangent is small, the adjoint residual starts close to or below that tolerance, so the adjoint solves stop early. PICT's VJP is then not linear in the cotangent, and the gradient depends on how the loss is scaled. Single calls with O(1) cotangents, which is what the gradient suite checks, are unaffected. The problem shows up when a scalar loss sits at the end of a chain of calls. On the Kolmogorov solver-in-the-loop loss (16 chained calls, cotangents of about 1e-4 per entry), the gradient with respect to the state after interval k disagrees with finite differences by: | k (calls left to the loss) | 15 (1) | 13 (3) | 8 (8) | 0 (16) | |---|---|---|---|---| | relative error, as is | 0.04% | 3.7% | 10% | 73% | | loss × 1e4 | — | 0.014% | 0.023% | 0.015% | The tesseract's VJP now scales the cotangents to a largest entry of 1 before PICT's backward pass and scales the gradients back. This is exact for a linear VJP and makes the adjoint tolerance relative to the cotangent. The code comment and the tesseract's description (rendered on the docs solvers page) both state the deviation. An alternative is a relative convergence criterion in PISOtorch. PICT's criteria are all absolute (`ABS_MAX`, `ABS_MEAN`, `ABS_SUM`, `NORM2`, `NORM2_NORMALIZED`), so that would mean patching PICT itself. #### Testing done - New `tests/test_pict_vjp.py`: the VJP for a cotangent scaled by 1e-5, divided by 1e-5, matches the VJP for the unscaled cotangent to within 1e-3. The test fails on main's image (4.3%) and passes with this change. `tests/test_pict_cylinder.py` passes too. - A/B of PICT's gradient and optimization suites on ns-grid and ns-3d-grid, main's image against this branch's (local A100 run, same harness): | | main | this PR | |---|---|---| | ns-grid `fd_check`, rel. error at ε = 1 / 0.1 / 0.01 | 1.1e-3 / 1.4e-3 / 1.1e-3 | 2.7e-4 / 3.2e-4 / 7.2e-4 | | ns-3d-grid `fd_check`, same | 1.1e-4 / 5.0e-4 / 4.6e-3 | 1.4e-4 / 3.7e-4 / 4.5e-3 | | ns-grid `analytic_check` error, ν = 1e-3 … 5e-2 | 0.0043 … 0.0198 | 0.0044 … 0.0198 | | ns-3d-grid recovery, final IC error (3 trials) | 0.0013, 0.0015, 0.0021 | 0.0022, 0.0028, 0.0040 | | ns-3d-grid recovery, final loss | 3.5e-10 … 6.4e-10 | 5.7e-10 … 8.5e-10 | In the other gradient experiments (horizon sweeps, Jacobian spectra), gradient norms and condition numbers change by 6% at most, and every status stays the same. PICT is excluded from `drag_opt`, so recovery is the only optimization it runs. Recovery gets slightly worse: both images drive the loss to the same float32 floor, but the IC error ends about twice as large with this change. A repeat of both arms reproduced every trial to two digits, so it is not run-to-run noise. Both are far below the failure threshold of 2.0 (from an initial error of 1.0). I don't have an explanation that I've verified; the loss is very flat there (the Jacobian's condition number is about 7e7), so where L-BFGS stops in the directions the loss barely sees could depend on details of the gradient.
Contributor
Author
|
Closing in favour of #316, which carries this experiment, a longer stability variant and the study's scripts and results for reproduction. Training through the solver did not separate the solvers beyond forward cost. The PICT adjoint defect it found is caught as well by |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Relevant issue or PR
Stacked on #273 (native solver state), #272 (Kolmogorov forward checks), #268 and #267. Supersedes #116, #121 and #211, which stay open until a maintainer closes them. PICT passes the gradient gate only with #294.
Description of changes
optimization/kolmogorov_silchecks whether a neural corrector can be trained through each differentiable ns-grid solver, and what that costs. A small periodic CNN corrects the solver's coarse run (64², dt = 0.01) of forced 2D Kolmogorov flow every 0.1 time units. It is trained against that solver's own 256² reference in three arms with identical windows, updates and initialisation:The corrector outputs a stream function, so its correction is exactly divergence-free. Its readout starts at zero, so update 0 is the uncorrected solver.
Scoring is feasibility and cost only.
sil_status.feasible()requires:The plot shows the wall time per training update for each arm. The corrector's errors are recorded but not scored: 32 updates are too few to say how much it helps. An earlier version of this PR also registered a 256-update campaign for that. It is gone, because whether the solver gradient helps depends more on the ML setup than on the solver.
Results (local A100 run, not CI)
forward/kolmogorov). OpenFOAM is not differentiable.mainon 2026-10-28, on an otherwise idle GPU.Testing done
tests/test_sil.pyuses closed-form transitions in place of a solver. It covers:tests/test_sil_plots.py: the cost plot renders from a run on the closed-form transitions. The A100 run above renders it too.mosaic run -p ns-grid -e optimization/kolmogorov_sil -s 'jax-cfd,INS.jl,XLB,PICT,Exponax'.