Repository navigation
Differentiable simulation: Diffrax adjoint pass-through, tests, and an Orszag-Tang inverse-design example - #44
rogeriojorge wants to merge 9 commits into
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests.
... and 1 file with indirect coverage changes 🚀 New features to boost your workflow:
|
Phase 1a: evidence for the first CPU showcase runsThe files are on the evidence branch at Current-sheet intensification (
Magnetic-energy conversion (
Gradient cost and memory. The AD gradient costs about 5 forward solves at every P, while FD is 1.6× (P=4) to 25× (P=64) slower. The reverse-pass workspace is flat in N with checkpointing and linear with a full tape. The gradient/forward workspace ratio stays at 4.2–4.5× from 16²×4³ to 64²×6³.
Two figure fixes follow in the next commit. The objective panels will plot accepted L-BFGS iterates instead of line-search trial points (the spikes in the |
Phase 2a gate: passedThe initial and optimised controls from the 32²×4³ runs were frozen and re-evaluated at higher resolution: forward solves only, adaptive Dopri8, T=200. Raw JSON is on the evidence branch at Relative change of the objective, optimised vs initial controls:
Gate criteria:
At 64²×8³ it is 1e-12 even at 1e-7, because the stiffer system forces smaller steps. The objective values change by at most 0.05 % between tolerances 1e-7 and 1e-9. Phase 3 therefore runs at tolerance 1e-9. Physics read-out (consistent across resolutions):
|
Phase 2c: phase-only controls (
|
| baseline | 16 controls (amplitudes + phases) | 8 phase-only controls | |
|---|---|---|---|
conversion: in-plane magnetic energy converted by T |
6.2 % | 26.1 % | 23.9 % |
conversion: ion kinetic-energy gain |
2.2e-5 | 8.7e-4 | 7.6e-4 |
peak_current: objective (smooth peak of Jz at T) |
0.0387 | 0.0616 | 0.0614 |
peak_current: Jz peak at t=0 |
0.0373 | 0.0374 | 0.0493 |
peak_current: dynamical intensification Jz(T)/Jz(0) |
1.04× | 1.65× | 1.25× |
Conversion. Linear evolution on the uniform background cannot change any mode's energy, so the phase-only gain is pure nonlinear coupling. It reaches 89 % of the full-control improvement. This is the cleanest statement of the showcase: with the same initial energy in every Fourier mode, the arrangement of the modes alone controls how much magnetic energy the vortex converts, and a gradient finds it.
Peak current. The p-norm objective depends on phases already at t=0. Phase-only control reaches the same objective partly by building a larger initial peak (0.049 against 0.037), and its dynamical intensification (1.25×) is below the full-control run (1.65×). For current-sheet formation, the full-control result, with an unchanged initial peak, remains the showcase. The phase-only run is reported as a caveat.
Provenance note: the phase-only conversion report records commit 8380edc, but that process imported the example at 620fd14. Provenance was captured when the report was written, not at the start of the run; this is fixed on the PR branch. The optimisation path is numerically identical between the two commits, so the result is unaffected. Its figure was re-plotted from the JSON to pick up the title fix.
Phase 3a/3b:
|
| initial controls | optimised controls | |
|---|---|---|
| objective: in-plane magnetic energy at T / initial | 0.940 | 0.740 |
| magnetic energy converted by T | 6.0 % | 26.0 % |
| electron kinetic-energy gain | 9.4e-5 | 2.2e-4 |
| ion kinetic-energy gain | 4.2e-5 | 9.2e-4 |
| Dopri8 steps | 277 | 290 |
Cost. 30 L-BFGS-B iterations (the cap) used 34 gradient evaluations at 77 s each, for 47.5 min in total including 95 s of compilation. The optimum matches the 32²×4³ CPU run (0.938 → 0.739), so the low-resolution search already found the physical optimum.
Validation (3b). Frozen optimised controls, tolerance 1e-9:
| resolution | initial → optimised | change | energy error at T |
|---|---|---|---|
| 64²×8³ | 0.9409 → 0.7409 | −21.3 % | 7e-11 |
| 128²×6³ | 0.9396 → 0.7400 | −21.2 % | 1.0–1.4e-8 |
Scope change for the rest of Phase 3.
- The job was not stalled; the peak-current optimisation is running normally (76.8 s per gradient evaluation).
- Because the low-resolution optimum already transfers, the remaining GPU work was trimmed. The GPU benchmark now stops at P=64 controls and N=800 steps instead of P=128 and N=1600, saving roughly 25 min of finite-difference and 1600-step-tape work. The CPU benchmark already shows both scaling trends.
- Peak-current validation and the per-process peak-memory sweep are unchanged.
Phase 3 complete, and the publication figures (4a)All numbers are from one NVIDIA RTX A4000 in float64 at commit Figure 1 (main): magnetic-energy conversion by design
Caption. Gradient-based inverse design of the 2D Orszag–Tang vortex.
Figure 2 (methods): cost and memory of gradients
Caption. Cost and memory of reverse-mode gradients through SPECTRAX, 32²×4³, T=100.
Peak device memory, one process per row (adaptive Dopri8, T=20):
With 8 checkpoints a gradient needs a fixed 2.8–2.9× the memory of a forward solve at every resolution. Its compiled workspace stays at 54 MiB for 50 to 800 steps, whereas a full tape grows from 230 MiB to 3.6 GiB. More checkpoints trade memory for less recomputation. The earlier 128²×6³ out-of-memory error came from a GPU shared with another job, not from the method. Supplementary figure S1: current-sheet intensification
Caption. The layout and setup match Figure 1, with the objective changed to the smooth peak (p=8 norm) of J_z at t ω_pe=200.
Validation of both optima (3b), tolerance 1e-9
The draft methods paragraph and captions are kept with the evidence under |
Phase 5: time-window objectives, parameter sensitivities, and the DG driver5c: sensitivity of the optimised designs to physical parametersThe new
Logarithmic sensitivity d ln|J| / d ln p of each objective at the reference controls and at its own optimum:
Reading the table.
5b: time-window objectives (
|
12201a6 to
14fa5c3
Compare
|
Rebased as a stack so the PRs merge in sequence without conflicts: merge (with a merge commit) #48, #46, #50, #51, #9, #12, #13, #34, #52, #53, then #18, #21, #44. Until the earlier ones land this diff also shows their commits. The conflicts with #13/#50 (same |
…ling An under-resolved run (too few Hermite modes or too weak hypercollisions for the filamentation reaching the cutoff, or a species whose Hermite width is narrower than it becomes) makes the adaptive step collapse; the solve then crawls for hours towards the hard-coded 1e6 steps. dtmin (off by default) stops it with an error at the collapse; max_steps is now a parameter.
simulation(throw=False) returns and puts diffrax's sol.result in output['solver_result'].
…s of a failed solve A short physical transient can need steps far below t_max/max_steps and still finish within the budget, so no floor is inferred: dtmin is now off unless set. Step exhaustion (max_steps_reached) and the floor (dt_min_reached) are reported separately, and a failed solve with throw=False returns valid_times and num_valid_times so the inf-filled saves after the failure are never read as data.
Expose max_steps, adjoint and progress_meter on simulation, keep the kinetic energies real, add gradient tests against forward mode and finite differences, and one Orszag-Tang inverse-design example.
…ifference plateau The optimisation example kept its own real-FFT weights (1 for k=0, 2 otherwise), a second copy of the parity-dependent Parseval rule; it now takes the grid mean of Bx^2 + By^2 after irfftn, which needs no weights. The new test checks that centred differences of a real scalar objective converge to the AD directional derivative as eps^2 and agree with it over at least three consecutive step sizes.
49c4c69 to
3ad4a9d
Compare
1 participant
Target: main. Head:
3ad4a9d. Old head:49c4c69, a merge of #21/#18/#13 and other branches carrying about 950 lines of unrelated diff.Stacked on #50 (
2744422), which providesmax_steps,throwand the validity metadata. Until #50 merges, this diff shows #50's four commits. Review only the top two:c4038a7: the original Differentiable simulation: Diffrax adjoint pass-through, tests, and an Orszag-Tang inverse-design example #44 commit;3ad4a9d: the cleanup.This PR does not depend on #13, #18 or #21.
Change
simulationacceptsadjointandprogress_meter, both static.adjointdefaults toRecursiveCheckpointAdjoint(), so behaviour is unchanged.RecursiveCheckpointAdjoint(checkpoints=k)bounds reverse-mode memory.ForwardMode()works withjax.jacfwd.kinetic_energy_speciesis returned real, so scalar objectives are real.Examples/2D_Orszag_Tang_optimization.py: the controls are stream-function phases and the objective is magnetic-energy conversion. It checks the reverse-mode gradient against centred differences, then runs L-BFGS-B.w = 1for k=0 andw = 2otherwise, a second, parity-dependent copy of the Parseval rule that Fix odd-grid RFFT energy weights; add a Hermite boundary diagnostic #9 fixes in the library.Bx² + By²afterirfftn, which needs no weights for anyNx.AD contract
ImplicitMidpointin reverse mode. Its Newton loop is alax.while_loop.Tests (local, macOS CPU, Python 3.12, jax 0.11.2, diffrax 0.7.2, at
3ad4a9d)python -m pytest -q -n 2 tests: 18 passed. The autodiff file alone takes 17 s on 2 workers.JAX_ENABLE_X64=1 python -m pytest -q -n 2 tests/test_autodiff.py: 6 passed.tests/test_autodiff.pyuses a tiny two-species Orszag–Tang case (8², 4³ Hermite):Example:
python Examples/2D_Orszag_Tang_optimization.py --t-max 5 --iterations 2GitHub CI: see the checks on this head. The local results above are not CI results.
Owner decision
Review the final API and gradient contract once:
adjointandprogress_meteras staticsimulationarguments;