Skip to content

Diagnose electron energization and validate acceleration objectives - #43

Closed
rogeriojorge wants to merge 25 commits into
mainfrom
agent/acceleration-validation
Closed

rogeriojorge wants to merge 25 commits into
mainfrom
agent/acceleration-validation

Conversation

@rogeriojorge

@rogeriojorge rogeriojorge commented Sep 12, 2026 •

Copy link
Copy Markdown
Member

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 at adaf542.

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.

@codecov

codecov Bot commented Sep 12, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 92.75362% with 5 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
spectrax/_autodiff.py 92.30% 3 Missing and 2 partials ⚠️
Files with missing lines Coverage Δ
spectrax/__init__.py 100.00% <100.00%> (ø)
spectrax/_simulation.py 100.00% <100.00%> (ø)
spectrax/_autodiff.py 92.30% <92.30%> (ø)

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@rogeriojorge

rogeriojorge commented Sep 12, 2026 •

Copy link
Copy Markdown
Member Author

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.

@rogeriojorge

rogeriojorge commented Sep 13, 2026 •

Copy link
Copy Markdown
Member Author

Independent-agent handoff — goals, code, evidence, and active physics work

Goal: 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

Purpose Branch / immutable checkpoint
Main application/code PR SPECTRAX #43, agent/acceleration-validation; cleaned checkpoint 6062da5
Reusable gradient API SPECTRAX #42, agent/checkpointed-plasma-control; reduced checkpoint 3dd7edc
New experiments, outside the code diff agent/particle-acceleration-evidence, checkpoint 9607fd8
Historical complete study/archive adaf542: raw reports, failed cases, figures and detailed research notes; intentionally removed from the current code-review diff
SOLVAX replay improvement SOLVAX #103, agent/spectrax-replay-performance; merged 917d1a6

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 validation

These findings must survive the cleanup; the historical archive contains their provenance.

Experiment Finding What it does not establish
Original 2D OT kinetic-energy control Eight phases preserve initial spectra/integrated energies; about 19.8% improvement of the small weighted baseline gain Nonthermal acceleration; the original fixed-nu Hermite refinement changes shared-mode damping
Energy accounting Extra electron kinetic gain split approximately 51.3% bulk / 48.7% internal; exact work bookkeeping implemented in the archived study Irreversible heating or an acceleration mechanism from total energy alone
Lab-frame tail control Fine collisionless gain +17.61%; old lab-tail optimum reduces the subsequent local-frame objective by 8.59% A bulk-sensitive tail benefit does not transfer to random-particle energization
Local flow-relative tail control Fine matched-q48 training gain +10.28%, held-out q96 +9.40%; added training gain only 0.002235% of initial electron thermal energy Fine stationarity/global optimum; search hit its 20-iteration cap
Local Gaussian diagnostic At T50, candidate excess −6.92e−6 versus baseline −6.06e−6 No enhanced nonthermal excess at that snapshot, despite increased local tail
Quadrature audit q24 baseline and q48 near-stationary gradient checks failed; higher q96/q128 checks were added and retained Never silently discard failed checks or call every q48 gradient converged
Matched full energy optimization AD 73.48 s versus FD 221.39 s, 3.01× total; same final objective within 2.03e−14 This is not a tail-objective comparison; one paired run, not timing statistics
Short gradient benchmark 128 controls, T2: reverse 0.2818 s versus serial FD 17.014 s, 60.4× Not end-to-end optimization or a benchmark against an independently optimized competing code
Checkpoint memory Replay 260 MiB versus tape 6457 MiB in the spatial case; fixed K16 retained 66 MiB at 256/1024 steps Constant memory as spatial/velocity DOFs increase

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

  1. Direct OT Gaussian-excess optimization: exact source and report. Twenty iterations at 16²/H8/q96 improve the weighted own-initial-subtracted excess from −7.7019e−6 to −3.1731e−6 in 652.32 s total. It remains a deficit, so this is not a nonthermal-tail success. A matched full FD optimization for this objective is not yet complete. CPU/GPU smoke reports agree to about 9e−18 in raw objective; the small full gradient agrees with centered FD.
  2. Simplified electrostatic wave–particle fallback: pilot source, reports and run instructions, based on SPECTRAX's existing Landau example. Eight density harmonics have fixed amplitudes 0.45/m² and variable phases; density stays positive initially and the electric field obeys Gauss's law. Use vth=0.02c, L=4πvth/ωpe, mass ratio 1836, nu=D=0. Positive local Gaussian excess appears by T3–4. At 64 cells/H256/q160, T4 excess is 0.02607 in initial-thermal-energy units, but worst-cell negative mass is 7.85e−5, failing the 1e−8 gate. The 32/H128 case also fails. This is a baseline probe, not a validated optimization or a 2D result.
  3. Compact observable performance prototype: reduced 1D3V helper and tests. Fixed transverse Maxwellians allow longitudinal Gauss–Hermite × perpendicular-energy Gauss–Laguerre quadrature, O(q²) instead of q³. Differentiate local flow and variance; use weighted Hermite recurrence at high H. Sixteen algebra/AD/finiteness tests passed, including H512/q320, H768/q512 and H1024/q640. It is an unintegrated prototype with a different interface, not the 2D example's current helper. Next: use it to make high-H baseline refinement affordable, then integrate only a validated scenario.

Preferred 2D route and explicit completion criteria

The 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:

  • Establish a baseline and frozen-control candidate with independently refined space, Hermite order, timestep, diagnostic sampling and velocity quadrature. Keep the collision operator fixed; nu=0 is the current choice. Require positive density/variance and sampled worst-cell negative mass below 1e−8, plus energy/work errors small relative to the claimed signal.
  • Optimize under a nontrivial fixed-energy/particle constraint. Separate baseline gain, absolute added energy, Gaussian-relative excess, bulk/internal partition, training and held-out results. A tail increase alone can be heating; a positive shape residual alone does not identify a mechanism. If the physics gate fails, record that outcome rather than weakening the label.
  • Run AD and FD on the same final objective, same controls, resolution, quadrature, optimizer/tolerances and budget. Validate FD step size and gradients, count all evaluations, report compilation and full wall time separately, and compare achieved objective values. Do not transplant the old 3.01× or 60.4× ratios. One-pair and serial-FD limitations must remain visible; batch FD where resource-efficient and report the choice.
  • Finish two sober figures: physics/distribution and energy-transfer evidence; matched optimization accuracy/cost with checkpoint memory context. Generate PNG and vector PDF from scripts. Put the final figures and immutable evidence links in a PR comment, not a new directory of generated assets in the code diff.
  • Keep the library API and code diff small. Do not reintroduce the archived research campaign, multiple nearly identical examples, numerous Markdown plans, or exhaustive implementation-mirroring tests.

Literature and what it contributes

Start from GitHub, without machine-specific paths

gh 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.

@rogeriojorge

rogeriojorge commented Sep 13, 2026 •

Copy link
Copy Markdown
Member Author

Continued in #44. Its description is now the handoff and single source of truth (goal, decisions, commands, measured results, gated work plan, agent rules). Read that first.

Revised plan after an independent review and local tests (2026-09-13)

I read every commit and the diff at 6062da5, the handoff comment, the archived study at adaf542, PR #42, PR #36, PR #11, PR #16, SOLVAX #103, and the three cited papers, then ran local tests against SPECTRAX main (2afb1f6) before deciding anything. Scripts and raw JSON for every number below are kept outside the code diff (local spectrax-autodiff-evidence/; they will be attached to the implementation PR).

Short version. The gradient capability already exists in main: jax.grad flows through the unmodified simulation() with Diffrax's discrete checkpointed adjoint, and the reverse pass costs a constant multiple of the forward solve regardless of resolution or step count. The library change is about ten lines, not a parallel RK4 solver; the showcase should be inverse design of the Orszag–Tang vortex against the energy and current-sheet objectives the paper needs; and the memory story should be measured, which is now done.

1. What the local tests showed (CPU, float64/complex128, M3 Max; all on main)

A. jax.grad already flows through simulation(). Tiny 2D Orszag–Tang (8², 3³ Hermite, two species), objective EM + 0.1×kinetic energy at t_max, differentiated with respect to deltaB and nu:

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 = 3 the collision operator is identically zero (its normalisation divides by N − 3), so d/dnu is 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 what simulation + 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 streamed integrand, is Diffrax SaveAt(fn=...) in one line if a time-window objective is ever needed.
  • SOLVAX's checkpointed_fori_loop is 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 adaf542 for 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.py runs on the Discontinuous-Galerkin branch; its driver has the same diffeqsolve(..., 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)

  1. 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 in diagnostics (they are currently complex with zero imaginary part, which breaks float() and optimisers). No new solver, no SOLVAX dependency, no custom VJP. Done locally.
  2. 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 through plasma_current. Passing locally; flake8 clean.
  3. One example, Examples/2D_Orszag_Tang_optimization.py (~250 lines): optimize (L-BFGS-B on conversion or peak_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 use diagnostics/plasma_current, so the example survives PR Restore to complex FFT #41 (complex FFT) and the DG branch. Written; smoke-testing now.
  4. 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.
  5. Docs: one short section in README.md with the four-line usage pattern and the measured bounds.
  6. Draft PR Differentiable simulation: Diffrax adjoint pass-through, tests, and an Orszag-Tang inverse-design example #44 is open from agent/autodiff-plan with 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.
  7. 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.

@rogeriojorge

Copy link
Copy Markdown
Member Author

Superseded by #44 for the gradient capability: #44 differentiates the existing simulation through Diffrax's checkpointed discrete adjoint instead of adding simulation_final, as when #42 was closed. The unvalidated electron-tail study stays available on branch agent/acceleration-validation and in the archived commits linked above. Closing without merging.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant