Modest performance improvements + some suggestions for future profiling from Claude - #380
Modest performance improvements + some suggestions for future profiling from Claude#380adrianaghiozzi wants to merge 7 commits into
Conversation
…h reltol integrate_ballooning_ode set abstol=reltol^2 (1e-16, at machine epsilon) against a non-stiff DP5 solver, forcing needlessly tight step control. Profiling (warmed, no JIT noise) showed ballooning boundary search as the dominant real cost on gal-resistive-PE-style cases -- ~55% of pipeline runtime, driven by many repeated ODE shoots per flux surface (linear pre-scan + bisection to locate the marginal-stability crossing). Setting abstol=reltol=1e-8 (the standard pairing) gives a measured ~8% speedup on Force-Free States for DIIID-like_gal_resistive_pe_example, confirming reltol was already the binding constraint for most of the integration. Validated bit-for-bit identical locstab/alpha_critical (the value that actually drives stability classification) before/after; the raw ballooning_Delta_prime diagnostic drifts ~1e-5 relative, which flips a few profile checksums in the regression harness but is far below any threshold that matters physically. Full 11-case regression suite against merge-base a0cad26 is otherwise clean.
|
@adrianaghiozzi I assume this is ready to review? If not, please convert to draft |
Pulls in 171 develop commits, including the ForceFreeStates -> standalone LocalStability module extraction (Ballooning.jl moved to src/LocalStability/Ballooning.jl). Clean merge, no conflicts; Git's rename detection correctly carried this branch's ballooning abstol fix (69ad76e) over to the new path. Regression harness (develop @ 9491f89 vs merged local, 13 cases, full suite): clean across the board. 12/13 cases fully unchanged (energies, ODE step counts, equilibrium/stability quantities, all profile checksums identical). The only diff anywhere is the "ballooning Delta' profile" checksum in diiid_n1 (1 changed / 46 unchanged) -- the expected, already-validated fingerprint of this branch's own abstol=reltol fix (~1e-5 relative drift in the raw diagnostic; locstab/alpha_critical, the value that actually drives stability classification, is bit-for-bit identical). No numeric regressions introduced by the merge itself.
There was a problem hiding this comment.
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?
|
|
||
| ρ2 = x_minus2 + ζ2 | ||
|
|
||
| # Distance parameter ℛ [Chance Phys. Plasmas 1997 2161 eq. 41] |
There was a problem hiding this comment.
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
|
@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: |
|
@logan-nc What are your thoughts on the 0-finding? You are the one who knows those physics the best |
|
FYI reviewing and thinking about this is on my todo list but it's just a long list, so please be patient |
5 participants
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)
That's the report — find_zeros → find_sign_change_roots is the one concrete, high-value, evidence-backed lead. Let me know if you want to move forward on it, want me to first check how many wells B(θ) actually has across this equilibrium's ψ range (to size a safe grid), or want to look at something else.