Repository navigation
Diagnose electron energization and validate acceleration objectives - #43
rogeriojorge wants to merge 25 commits into
Conversation
…e training window
Codecov Report❌ Patch coverage is
... and 1 file with indirect coverage changes 🚀 New features to boost your workflow:
|
|
Superseded by the cleanup and independent-agent handoff. The prior 17-figure gallery and detailed exploratory evidence remain available at historical commit adaf542. They were removed from the code-review diff to keep this PR concise. Those historical figures do not establish the final nonthermal-acceleration or same-tail-objective AD/FD claim; follow the current handoff for validated results, failed checks and remaining work. |
Independent-agent handoff — goals, code, evidence, and active physics workGoal: a small, accurate, fast SPECTRAX differentiability API and one physically meaningful plasma-control example, with two publication figures: (1) validated particle energization/acceleration, and (2) AD versus finite differences on that same optimization problem. Prefer a genuinely 2D reconnection/acceleration showcase. A simplified wave–particle case is an acceptable fallback if the 2D physics cannot yet be validated. Do not declare the physics goal achieved merely because an optimizer improves a tail integral. Repositories, branches, and review scope
Keep both SPECTRAX PRs open: only other collaborators may merge them. SOLVAX's admin merge was explicitly authorized and is complete. Numerical reports still use released SOLVAX 0.20.0; do not silently relabel them as measurements of merged SOLVAX main. PR43 includes the API changes from PR42; it is not an independent competing API. The cleanup reduces PR43 from 177 files / 18,129 added lines to 12 files / 1,050 added lines at 6062da5. No generated JSON/NPZ, figure binaries, benchmark campaign or collection of planning documents remains in that review diff. Preserve this discipline: one example, one example-local observable helper, two focused test files, one short user document, and the generic solver change. Evidence belongs on the separate evidence branch or external artifacts, with immutable links. Do not merge the evidence branch into the code branch. Where the implementation is
SOLVAX segmented replay retains O(S sqrt(N)) states by default. A fixed Diffrax checkpoint budget retains O(K S) states, with reverse AD only on that path. Both differentiate the executed discrete RK4 algorithm, not a continuous adjoint; workspace and phase-space DOFs still consume memory. This does not implement the block-local implicit reverse sweep from Differentiate the Solver, Not the Equation. SPECTRAX's explicit global spectral RHS has different structure. Do not add an unrelated implicit solver solely to claim that paper's algorithm. The SOLVAX change removes per-step bounds guards in full checkpoint segments, giving the partial segment its own static loop: 11 added/9 removed numerical lines. It preserves primal/JVP/VJP behavior and memory bounds. Recorded RTX A4000 warm-gradient improvement is 3.40%, with unchanged 66 MiB allocator peak; compilation increased approximately 2 s in the paired case, amortized after roughly 210 calls. It is a modest repeated-execution improvement, not a memory reduction or universally faster cold solve. See #103 for tests and exact measurements. Recorded findings from source review and parallel validationThese findings must survive the cleanup; the historical archive contains their provenance.
Historical detailed sources: local-tail results, acceleration results, methods assessment, 49 archived follow-up reports. Their checksums were audited before removal from the code diff. These are recorded engineering/scientific review findings, not a claim of external collaborator approval. New physics probes: promising directions, not publication-ready results
Preferred 2D route and explicit completion criteriaThe user's preferred final figure remains 2D reconnection or particle acceleration, rather than a purely 1D demonstration. A 2D electrostatic wave–particle extension could be considered only if it contains meaningful transverse dynamics; a y-invariant replication of a 1D calculation is not a 2D physics result. Do not inflate DOFs just to improve an AD/FD ratio. For reconnection, the archived baseline handoff proposes a periodic double-Harris finite-guide-field starting point. No reconnection initialization or simulation has been validated in this work. Verify periodic smoothness, charge/Gauss constraints, normalized Ampère balance, pressure balance, positive distribution and unseeded kinetic relaxation before seeding reconnection. A force-balanced profile is not automatically a kinetic equilibrium; neither is a force-free field with arbitrary shifted Maxwellians. Check nonrelativistic validity, resolved electron scales, energy/work budgets, topology change, local distributions and guide-field-dependent mechanisms. Reproduce matched physical parameters before borrowing a PIC paper's acceleration interpretation. 2D cannot establish the stronger 3D acceleration reported in the literature. Before calling the final result publication-ready:
Literature and what it contributes
Start from GitHub, without machine-specific pathsgh repo clone uwplasma/SPECTRAX
cd SPECTRAX
git fetch origin agent/acceleration-validation agent/checkpointed-plasma-control agent/particle-acceleration-evidence
git switch agent/acceleration-validation
python -m pip install -e .
PYTHONPATH=. JAX_PLATFORMS=cpu python -m pytest tests/test_autodiff.py tests/test_particle_acceleration.py -q
PYTHONPATH=. JAX_PLATFORMS=cpu python Examples/2D_particle_acceleration.py --objective excess --grid 12 --hermite 3 --steps 4 --window 0 .2 --intervals 2 --quadrature 16 --evaluate-only --fd-check --output results/cpu-smoke
# On a single available NVIDIA GPU; reserve/serialize access with other users.
CUDA_VISIBLE_DEVICES=0 XLA_PYTHON_CLIENT_PREALLOCATE=false JAX_PLATFORMS=cuda PYTHONPATH=. python Examples/2D_particle_acceleration.py --objective excess --method ad --output results/ad
# Only after validation, run the matched --method fd case and --validate/--plot modes.The target is one NVIDIA GPU (measure device/version/precision; recent measurements used an RTX A4000), with small matched CPU checks. Avoid multiple local simulations and GPU contention. At the cleaned checkpoint, the 12 focused CPU tests passed. Python 3.11/3.12 CI and patch-coverage checks are green for both cleaned PR43 (6062da5) and cleaned PR42 (3dd7edc). The code cleanup is complete at this checkpoint; the final acceleration physics and same-objective full AD/FD publication comparison are not yet complete. |
Revised plan after an independent review and local tests (2026-09-13)I read every commit and the diff at Short version. The gradient capability already exists in 1. What the local tests showed (CPU, float64/complex128, M3 Max; all on
|
| integrator | gradient vs. centred FD | grad / forward time |
|---|---|---|
| Dopri8, adaptive PID (the TOML default) | 2.1e-9 | 4.6× |
| Dopri8, fixed step | 1.4e-11 | 4.4× |
| Tsit5, fixed step | 1.4e-11 | 4.3× |
No custom VJP/JVP is needed anywhere: real FFTs, Hermite shifts, the 2/3 mask and the collision matrix all have built-in derivatives, and the reverse pass through time is Diffrax's default RecursiveCheckpointAdjoint (binomial/Revolve-type schedule). Forward mode (ForwardMode) agrees with reverse mode to 1e-9. The progress-meter callback does not block differentiation.
B. What bounds gradient memory. Compiled reverse-pass workspace from jit(value_and_grad).lower().compile().memory_analysis() (state 38 KB unless stated):
- Time steps N (fixed step, K = 8 checkpoints): workspace 1.91 MB for every N from 50 to 1600; K = 32: 4.04 MB, also flat. Full tape (K = N): 6.1 → 92.6 MB for N = 50 → 800, i.e. linear. Time ratio with K = 32: 4.9× (N = 50) → 10.6× (N = 1600), the expected O(N log N) recomputation.
- Checkpoint budget K (N = 200): 1.4 MB (K=1) → 7.7 MB (K=64), ≈ 2.6 state copies per checkpoint; time 110× (K=1), 12× (K=4), 9× (K=8), 6.6× (K=16), 5× (K=64) the forward solve. K ≈ 16–64 is the sweet spot; Diffrax's default is log2(max_steps).
- Resolution S (K = 8, N = 50, 8²/3³ … 32²/6³, state 0.04 → 3.8 MB): forward workspace ≈ 16–20 S; gradient workspace = 2.4–3.0 × forward workspace at every size. That is the correct claim for the paper: the gradient costs a constant factor of the forward solve, independent of resolution and of N ("bounded", not "constant").
max_steps=1000000(hard-coded today) costs ≈160 B per potential step in the reverse pass: 2.8 MB (200) → 163 MB (10⁶), independent of S and N. Fixed-step runs should pass the actual step count (PR Make solver step limit configurable #16 already proposes exposing it).
C. Showcase prototype (16², 4³, T = 200, adaptive Dopri8 at 1e-7, nu = 1, 206 steps; controls = amplitudes and phases of the 4 lowest stream-function modes of B⊥ at fixed in-plane magnetic energy, electrons carrying the consistent current; P = 8):
| objective | AD gradient vs. FD | AD grad / forward | FD (2P solves) / AD | 10 L-BFGS-B iterations |
|---|---|---|---|---|
| in-plane magnetic energy at T (minimise = fastest conversion) | 1.6e-8 | 7.1× | 3.2× | 1.081 → 0.753 (30 % more conversion) |
| smooth peak of Jz at T (maximise = strongest current sheet) | 2.3e-8 | 4.8× | 3.2× | 0.0288 → 0.0406 (+41 %) |
Both objectives are plain functions of the simulation output (Fk for energies, plasma_current + inverse FFT for Jz). The FD/AD ratio grows linearly with P; at P = 8 it is already 3.2×.
D. Why the answer is the discrete checkpointed adjoint and not the alternatives ("differentiate the solver, not the equation" made concrete):
- Continuous adjoint (
BacksolveAdjoint, the "differentiate the equation" route): cannot even be constructed for SPECTRAX's complex state in diffrax 0.7.2 (dtype mismatch in the augmented backward state), and is inexact by construction. - Algebraically reversible RK4 (McCallum & Foster 2024, O(1) memory in N, prototyped with a custom VJP on the SPECTRAX RHS, 8²/4³): exact for collisionless runs (gradient error 1e-17 … 1e-13 up to N = 1600, T = 32, λ = 0.999). With collisions the backward reconstruction amplifies round-off by e^(ν T): NaN at ν = 1, T = 32 and a 1e106 reconstruction error at ν = 3, T = 8. The Orszag–Tang case runs at ν = 1 to T = 500, so this route is out; it becomes a one-paragraph negative result in the paper.
- Block-implicit reverse sweep (Shu et al.): applies to block-coordinate implicit solvers; SPECTRAX's explicit global spectral RHS has no such block structure. The transferable principle, differentiating the executed discrete algorithm, is exactly what the checkpointed adjoint does.
- Note for anyone testing collisionality gradients: with
Nn = Nm = Np = 3the collision operator is identically zero (its normalisation divides by N − 3), sod/dnuis exactly 0 there. Use ≥ 4 modes.
E. Cost calibration (forward, adaptive Dopri8 at 1e-7): 16²/4³ 0.010 s per step, 32²/6³ 0.10 s, 64²/6³ 0.40 s on the laptop CPU; the controller takes only ≈1.4 steps per time unit at these sizes with relative energy error 3–5e-11. The full TOML case (128²/6³, T = 500) is a single-GPU job; everything below it runs locally.
2. Critical reading of the current branch
simulation_final(136 lines: fixed RK4 + SOLVAX two-level replay + a Diffrax RK4 shim) re-implements whatsimulation+adjoint=already provides, with a lower-order fixed-step integrator (the study needed 10 RK4 steps per time unit; adaptive Dopri8 needs 1.4). Its one genuine extra, the streamedintegrand, is DiffraxSaveAt(fn=...)in one line if a time-window objective is ever needed.- SOLVAX's
checkpointed_fori_loopis a √N two-level schedule; Diffrax's binomial schedule is better on both axes and is already a dependency. SOLVAX #103 (3.4 % faster replay) is fine but irrelevant to SPECTRAX gradients. Nothing in SOLVAX is needed here, and test D shows the only candidate numerical method worth adding to SOLVAX for this purpose (reversible stepping) does not survive collisions. - The physics (local flow-relative electron tail, matched-Gaussian excess, 20 archived probes) never closed: the optimised excess is still a deficit and the fallback fails its own negativity gate. That work stays archived at
adaf542for a possible future paper; the SPECTRAX methods paper needs a showcase that closes. - The archived AD/FD ratios (3.01×, 60.4×) were measured on different objectives and a different solver path; none transfers. They are replaced by one measured curve in P on the final example.
- The attached
2D_Orszag_Tang.pyruns on theDiscontinuous-Galerkinbranch; its driver has the samediffeqsolve(..., max_steps=1000000, progress_meter=TqdmProgressMeter())structure, so the same three keywords apply there verbatim.
3. Next steps (in order; executing now on branch agent/autodiff-plan from main, draft PR #44)
- Library change (≤ 15 lines), superseding the code of Add checkpointed plasma gradients and constrained optimization examples #42/Diagnose electron energization and validate acceleration objectives #43 and absorbing Make simulation progress meter configurable #11/Make solver step limit configurable #16/Add differentiable simulation example #36:
simulation(..., max_steps=1_000_000, adjoint=RecursiveCheckpointAdjoint(), progress_meter=TqdmProgressMeter())pass-through, and real-valued kinetic energies indiagnostics(they are currently complex with zero imaginary part, which breaksfloat()and optimisers). No new solver, no SOLVAX dependency, no custom VJP. Done locally. - Tests (
tests/test_autodiff.py, 4 tests, 28 s on CPU): reverse = forward = finite differences for the adaptive Dopri8 path; gradient invariant to the checkpoint budget on a fixed-step path; the compiled reverse-pass workspace does not grow with the step count (the memory claim, enforced in CI); a Jz-functional gradient throughplasma_current. Passing locally; flake8 clean. - One example,
Examples/2D_Orszag_Tang_optimization.py(~250 lines):optimize(L-BFGS-B onconversionorpeak_current, JSON report + physics figure: Jz maps initial/final/optimised, energy traces, objective history),benchmark(AD vs. FD gradient cost vs. P; reverse-pass workspace vs. N for tape/K=8/K=32; workspace vs. state size; FD-step error curve; JSON + methods figure),plot(figures from JSON only, PNG + PDF). Observables usediagnostics/plasma_current, so the example survives PR Restore to complex FFT #41 (complex FFT) and the DG branch. Written; smoke-testing now. - Publication runs: locally 32²/4³, T = 200, 8–16 controls, and the benchmark curves; then one RTX A4000 (office GPU 0 is free) for 64²/6³ and 128²/6³ at T = 500 with the TOML parameters, including peak device memory. Raw JSON and figures go in the PR comment and an evidence branch, never in the code diff.
- Docs: one short section in
README.mdwith the four-line usage pattern and the measured bounds. - Draft PR Differentiable simulation: Diffrax adjoint pass-through, tests, and an Orszag-Tang inverse-design example #44 is open from
agent/autodiff-planwith the diff and links to this comment; figures follow in a comment there. Then Add checkpointed plasma gradients and constrained optimization examples #42 can be closed as superseded and Diagnose electron energization and validate acceleration objectives #43 kept for the physics discussion and archive links. - Paper: methods paragraph and captions written from the JSON of steps 3–4; keep the Joglekar–Thomas and Skene–Burns positioning from the handoff, and add the reversible/continuous-adjoint negative result from test D as the justification for the checkpointed discrete adjoint.
|
Superseded by #44 for the gradient capability: #44 differentiates the existing |
Add a compact checkpointed gradient API and one constrained electron-tail control example, with matched AD/FD optimization and streamed physics validation/figure generation.
The review diff has been reduced from 177 files / 18,129 added lines to 12 files / 1,050 added lines at
6062da5. It contains one example, one example-local helper, two focused test files, one short usage document, and the solver/dependency changes. Generated results and figures are outside this code diff.Start with the independent-agent handoff for branches, immutable code/results, SOLVAX work, scientific findings, literature, exact reproduction commands and remaining acceptance gates. New probes live on
agent/particle-acceleration-evidence; do not merge that artifact branch here. Historical reports and figures remain atadaf542.Twelve focused CPU tests passed for the cleaned code. The final nonthermal-acceleration claim and full same-tail-objective AD/FD publication comparison remain incomplete. Current OT excess optimization still yields a deficit; the simplified wave–particle pilot has a positive signal but fails coarse negativity checks. A resolved 2D reconnection/acceleration example remains preferred; the simplified case is a fallback pending validation.
This PR includes the API from #42. Both SPECTRAX PRs remain open for other collaborators to merge. SOLVAX #103 is merged; recorded numerical studies still use released SOLVAX 0.20.0.