LocalStability - PERF! - Loosen ballooning ODE abstol and fix pitch-quadrature type instability

[adrianaghiozzi]

Release note

• Audience: developers
• Numerical impact: ballooning Delta' profile checksum in diiid_n1 moves ~1e-5 relative (harness @ 11666e33e, develop @ cbe890c0c). All stability-classification-adjacent quantities (Mercier D_I, resistive interchange D_R, singular surfaces, gal α per surface) are bit-for-bit identical.
• Migration: none

Loosens the abstol on the ballooning marginal-stability ODE solve from machine epsilon (1e-16) to match reltol (1e-8), removing needless step-size control on a non-stiff DP5 solver (~8% speedup on Force-Free/LocalStability-heavy cases). Also fixes a real type instability in the pitch-angle quadrature (PitchIntegration.jl), adds an explicit equilibrium type annotation in two KineticForces/Torque.jl call sites, and refactors a common subexpression in Kernel2D.jl's Green's function.

Regression report

Full 13-case suite, develop @ cbe890c0c vs local (this branch, merged with develop) @ 11666e33e. One expected diff, no regressions:

ggj_ray_q500i             — 4 unchanged, bit-for-bit
solovev_kinetic_nuzero    — 14 unchanged (O(1e-14) noise)
diiid_slayer_n1           — 15 unchanged (SLAYER root-find noise within existing tolerance)
gal_resistive_pe          — 8/8 N/A on BOTH refs (pre-existing coverage gap, not from this branch — see note below)
solovev_kinetic_calculated— 14 unchanged (O(1e-13-1e-14) or exact)
gal_resistive_diiid       — 10 unchanged, ALL bit-for-bit incl. "gal α per surface", "||gal Δ' matrix||"
solovev_n1                — 22 unchanged, q/pressure/ca_left checksums identical
ggj_reference             — 4 unchanged, bit-for-bit
solovev_multi_n           — 15 unchanged, q profile checksum identical
diiid_n1                  — 46 unchanged, 1 CHANGED, 1 untracked runtime
  CHANGED: ballooning Delta' profile (checksum)
    9a1108562647... (develop) -> d9f849242bb0... (local)
    Expected effect of the abstol loosening above; matches prior validated runs
    at 69ad76e4f (pre-refactor) and d266e7976 (post-refactor merge).
  Everything downstream unaffected: Mercier D_I profile, resistive interchange
  D_R profile, q profile, pressure profile all identical; island half-widths
  (1.4e-16), Chirikov parameter (1.1e-14), NTV torque/energy all O(1e-14-1e-19).
solovev_kinetic_ntv       — 6 unchanged, bit-for-bit
diiid_n1_riccati          — 17 unchanged (Riccati Δ' diffs 1e-14/1e-12, noise level)
solovev_kinetic_multiion  — 6 unchanged, NTV per-species torques bit-for-bit to 1e-18-1e-19
efit_fixedbdy_separatrix  — 5 unchanged, bit-for-bit

No quantity literally named alpha_critical is tracked; the closest proxies (gal α per surface, Mercier D_I / resistive D_R checksums, singular-surface locations) are bit-for-bit identical throughout, consistent with the two prior validations of this same change.

Notes for reviewers

On the ballooning abstol change (flagged in review): this line was carried onto this branch by an earlier develop merge, not introduced fresh here — it originates from commit 69ad76e4f (FORCEFREESTATES - OPTIMIZATION - Loosen ballooning ODE abstol to match reltol), which ran an 11-case regression harness against merge-base a0cad260 and validated locstab/alpha_critical-adjacent quantities bit-for-bit identical, with only the raw ballooning_Delta_prime diagnostic drifting ~1e-5. It was re-validated again after the LocalStability module extraction in d266e7976, and now a third time above after merging in 155 more develop commits (including the KineticForces tolerance fixes in #390/#391/#392) — no new diffs from any of that.

gal_resistive_pe coverage gap (unrelated to this PR): this case reports N/A for all 8 tracked quantities (rational-q values, island half-widths, Chirikov parameter, resonant/penetrated b-field, PE Δ' per surface) on both develop and local. Since it's symmetric it isn't a regression from this branch, but it means the case is currently providing zero regression coverage — worth a separate ticket to find out why the DRIVEN/RPEC coil-driven singular-coupling output isn't populating (possibly an HDF5 group rename or extraction-path mismatch between the harness case TOML and current output schema).

On the profiling suggestions in the original PR description: those findings (raw content below) predate #390/#391/#392, and #392 in particular changed the energy-integration tolerances inside the same tpsi! chain the ~84% figure was measured against, so that share needs re-measuring before anything gets scoped off it. The find_zeros → find_sign_change_roots swap is not a safe drop-in substitution for the reason already given below (find_sign_change_roots only catches sign changes between adjacent grid nodes); a follow-up would instead enumerate the exact stationary points of the B spline once per surface and do bracketed Brent solves only on the monotone intervals that bracket bo/λ — same completeness guarantee as the adaptive scan, without the per-λ whole-interval walk. That's out of scope for this PR and should be its own branch off develop.

Original profiling report (superseded context above)

This PR introduces some modest performance improvements by explicitly typing a function argument that was being inferred in Torque.jl, refactoring the green function in Kernel2D and fixing a type instability in the integrate_pitch_gar_quadgk/integrate_pitch_gar_quadgk_wt functions.

The source of much of the compute time for several of the examples goes back to the quadgk step, and the tpsi! function in Torque.jl. Claude has some suggestions about how this tpsi! function could be improved, but it included some caveats that seem worthy of human-review/discussion before implementing and therefore are not pushed onto this branch. Claude's ideas below:

Profiling results — diiid_n1 full pipeline
Ran Julia's statistical profiler across the entire main() call (equilibrium → stability → PE → KineticForces → ballooning). Headline finding:

tpsi! (KineticForces/Torque.jl) and its call chain account for ~84% of all profiled samples — despite KineticForces looking like a modest post-processing stage in the wall-clock logs (~55s of ~133s), it dominates actual CPU-thread time by a wide margin once you account for the equilibrium/stability stage being comparatively cheap per-thread.

#1 hotspot: Roots.find_zeros in bounce-point finding — ~26% of total pipeline time
src/KineticForces/BounceAveraging.jl:434, inside _find_bounce_points_and_grid:

bpts = sort!(Roots.find_zeros(vpar_fn, 0.0, 1.0); rev=true)

This is a generic adaptive full-interval root scanner, called once per trapped-particle λ point — per ψ, per bounce harmonic, per species — so it fires an enormous number of times. It alone accounts for ~9,400 of the ~36,300 total samples.

The interesting part: tpsi! (the caller's caller) already computes theta_bmin/theta_bmax a few lines earlier using a lighter-weight utility that already exists in the codebase — find_sign_change_roots(f, grid) in src/KineticForces/Utils.jl:19, explicitly documented as "Single source of truth for the scan-then-Brent idiom" and already used elsewhere in the same file. _find_bounce_points_and_grid doesn't use it, reaching for the heavier Roots.find_zeros instead.

Why I didn't just swap it: find_zeros's adaptive scan is more robust against closely-spaced multiple roots (the code already handles nbpts > 2 for multi-well B(θ) profiles in _find_deepest_well), whereas find_sign_change_roots only catches sign changes between adjacent grid nodes — too coarse a grid could silently miss a well. This is a physics-sensitive change (bounce-point location feeds the NTV torque/kinetic-energy integrals directly), so it needs a grid-resolution choice validated against find_zeros' output before trusting it, plus a regression-harness pass afterward.

Smaller, secondary costs (for context, not obviously worth chasing)

• QuadGK adaptive quadrature (the pitch-angle integral itself) — substantial but that's inherent to the physics, not obviously wasteful.
• FastInterpolations spline construction/evaluation inside tpsi! — only ~4% of total, much smaller than I initially suspected from reading the code.

---

Set an assignee and at least one human reviewer. Per this project's policy, this PR requires third-party human reviewer approval before merge — the clean regression run above does not substitute for review.

[logan-nc]

@adrianaghiozzi I assume this is ready to review? If not, please convert to draft

[adrianaghiozzi]

@logan-nc Yes this is now fully ready! I just pulled in develop and ran the regression harness again - the output of the regression harness is shown in the commit message of d266e79

[jhalpern30]

Looks great! Only one minor comment is about comments lol, then its ready to merge

Also - I see your notes about benchmarking the quadgk stuff - can you also add some comment about what kind of speedup the vacuum kernel optimizations obtained just for posterity?

[jhalpern30]

Claude is rather inconsistent with its comments - sometimes, it will add large, multi-line comments that saturate the repo, and other times it will remove helpful ones

Can you add back in some of these comments? For sections that you optimized and made individual terms out of date, you can just ask Claude to go through the pdf and revise whatever is necessary. This is just helpful for linking the math here directly to Chance's derivation

[adrianaghiozzi]

Yep, just added it back in e952629 !

[adrianaghiozzi]

@jhalpern30 Adding for posterity this local test of the Vacuum/Kernel2D green function showing the speedup I observed for a test before and after the changes implemented here:

[matt-pharr]

@logan-nc What are your thoughts on the 0-finding? You are the one who knows those physics the best

[logan-nc]

FYI reviewing and thinking about this is on my todo list but it's just a long list, so please be patient

[ebursch]

This also looks good to me! Per Pharr's comment, I also defer to Nik on the 0 finding.

[logan-nc]

Reviewed this while scoping the follow-up performance work it recommends. Three of the four changes look good; one needs a decision before merge.

Please justify or drop the Ballooning tolerance change

src/LocalStability/Ballooning.jl:1183 and :1204 loosen both ODE integrations:

   reltol=TOLERANCE,
-  abstol=TOLERANCE^2,
+  abstol=TOLERANCE,

With TOLERANCE = 1e-8 (:1153) that is abstol 1e-16 → 1e-8, eight orders of magnitude, on both the left and right DP5() solves in integrate_ballooning_ode. It is in a different module from the rest of the PR and isn't mentioned in the description or in any of the three stated changes, so I wanted to flag it rather than assume it was intentional.

If it was deliberate and speeds up the ballooning solve, that's worth having — but it needs a regression-harness pass showing what moves, and the PR should be tagged ! and say so in the release note. If it was incidental, reverting those two lines makes the rest of the PR cleanly inert.

The other three changes look correct and worth keeping as-is:

• equil::Equilibrium.PlasmaEquilibrium on tpsi! (Torque.jl:33) and compute_kinetic_matrices_at_psi! (:895) — two sites, the description mentions one
• the Kernel2D.jl green common-subexpression refactor
• the integrate_pitch_gar_quadgk / _wt fix in PitchIntegration.jl, where a runtime-varying tuple arity was being splatted into quadgk! — a real type instability

CI blockers

Two conventions checks will fail as-is:

1. Title. python3 ci/conventions/check_subject.py --title "..." rejects the current title (needs Area - TAG - Imperative summary). Something like KineticForces - PERF - Fix type instabilities in the pitch-angle quadrature if the Ballooning change is dropped; if it stays, it needs ! and probably a broader Area.
2. Release note. The body has no ## Release note block from the PR template, which is required and must carry the regression-harness result plus the commit it was run at.

On the profiling suggestions in the description

Thanks for writing those up — they're being acted on. Two notes for whoever picks that thread up:

• The profile predates KineticForces - BUGFIX! - Fit the parallel-velocity spline periodically #390/KineticForces - BUGFIX! - Drop the spurious major radius from the omega_D prefactor #391/KineticForces - BUGFIX! - Give the energy integration its own tolerances #392 (merged 2026-09-04), and KineticForces - BUGFIX! - Give the energy integration its own tolerances #392 changed the energy-integration tolerances inside the same tpsi! chain the ~84% figure is measured against, so the shares need re-measuring before anything is scoped on them.
• The find_zeros → find_sign_change_roots swap is the right instinct but not safe as a direct substitution, for exactly the reason given in the description: find_sign_change_roots only catches sign changes between adjacent grid nodes. The planned approach instead enumerates the exact stationary points of the B spline once per surface and does bracketed Brent solves only on the monotone intervals that actually bracket bo/λ — same completeness guarantee as the adaptive scan, without the per-λ whole-interval walk. That will be a separate PR from develop; this one shouldn't grow to hold it.

🤖 Generated with Claude Code

https://claude.ai/code/session_0186CLfjeJeSusNTva5RMfGU

[adrianaghiozzi]

@logan-nc Me and Claude updated the title and description of this PR to clearly show the regression harness was run and to justify the tolerance changes. Should be ready to go now!

[logan-nc]

@jmlmir369 please be aware of the changed tolerance behavior. I cannot think of any reason the abstol should be squared and/or be below machine precision, so I went ahead and merged. If you want tighter abstol than reltol for reasons based on your experience with running this part of the code please open a separate issue/PR to implement it in a more transparent way with annotations explaining the logic and (hopefully) less computational cost. Thanks!!

[logan-nc]

It's also worth pinging @nelsonand here → TLDR is that the ballooning calcs are much faster now.

[jmlmir369]

Thanks for the update. I reviewed the change, and it
looks reasonable to me. I do not see any major concerns. Thank you.