KineticForces - PERF! - Solve the bounce points directly from the spline coefficients

[logan-nc]

Release note

• Audience: developers
• Numerical impact: NTV torque moves by 1.1e-11 relative on diiid_n1; the two kinetic_source = "calculated" cases move at 0.00% relative (1e-7 to 1e-8 absolute energies, 1-4 different adaptive ODE steps out of ~730) (harness @ 2df2724)
• Migration: not required; no interface or configuration changes

Bounce-point location in the NTV kinetic module no longer searches. It solves the crossing directly from the spline's own cubic coefficients, cutting the KineticForces stage of the DIII-D-like example by roughly a quarter. It also resolves bounce-point pairs closer together than the poloidal grid spacing, which the previous adaptive scan could merge into one.

What changed

_find_bounce_points_and_grid called Roots.find_zeros — a generic adaptive scan over all of theta — once per pitch angle, so the scan repeated for all 128 lambda points on every flux surface, bounce harmonic and species.

Bounce points are the level set B(theta) = bo/lambda, and B is monotone between its stationary points, so each such interval holds at most one root. That is enough to make the whole thing direct, with no search of any kind:

1. Decompose B once per surface. One walk over the spline's cells, through the public coeffs/CellPoly interface, records each cell's polynomial and endpoint value and solves the quadratic dS/dtheta = 0 on each to get the exact stationary points.
2. Pick the interval by comparison. An interval contains a root iff bo/lambda lies strictly between its cached endpoint B values — two scalar comparisons, no evaluation.
3. Resume the cell from the previous lambda. Lambda ascends through compute_bounce_data, so bo/lambda falls monotonically and each interval's crossing walks steadily along the cells. Starting where the last lambda finished is the same hint idiom FastInterpolations uses for its own searches: measured on the production grid over 256 cells, the median lambda stays in the same cell and the mean moves 2.84, against ~8 probes for a bisection that starts blind every time. Bisection remains the fallback for the first lambda and for the sweep near a stationary point, where B is flat and the crossing can jump many cells at once.
4. Solve the cell's cubic in closed form. The interpolant already is a cubic with known coefficients, so the crossing is a root formula, not an iteration. The trigonometric form is used because it stays conditioned where Cardano's radicals cancel, with one Newton step to recover digits lost when two roots nearly coincide, and an explicit repeated-root branch for the discriminant cancelling to zero.

Cost accounting

Julia's Profile sampler around the whole main() call on examples/DIIID-like_ideal_example, 16 threads, JIT excluded by a warm-up run. Shares are against attributed CPU work; the raw sample count is dominated by idle worker threads (~210k samples parked in poptask), which would roughly halve every share if left in.

bounce-point cost, share of CPU work	per lambda	per surface	total
develop — adaptive scan of all of theta, every lambda	15.03%	—	15.03%
direct cubic solve, hinted cell	0.062%	0.57%	0.63%

The per-lambda work is 241x cheaper. The per-surface decomposition is the new cost: it also builds the knot, knot-value and coefficient vectors, three 256-element allocations per (psi, ell) call, and compute_bounce_data runs 13 times per surface at nl = 6. That 0.57% is recoverable by decomposing once per psi instead, which needs a signature change in Torque.jl; it is left as a follow-up to keep this change inside one file.

Wall time: 92.5 to 76.1 s overall, KineticForces stage 54.2 to 41.9 s. QuadGK is now the dominant remaining cost at 43.6% and the natural next target. Read the shares within a run rather than absolute counts across runs — the profiles ran under different machine load, which moves absolute counts by several percent either way.

Correctness

Every trapped lambda on 40 flux surfaces of the DIII-D-like equilibrium, checked against Roots.find_zeros on the identical spline, plus lambda injected to within 1e-13 of both trapped/passing boundaries:

check	result
trapped lambda tested	3,040
root-set mismatches on the production lambda grid	0 of 2,560
max abs(v_par) at our roots	4.441e-16
max abs(v_par) at find_zeros' roots	4.441e-16
count mismatches on injected near-tangent lambda	80
— where the new method found more roots	78
— where it deferred to the fallback instead	2
worst root displacement where counts agree	2.5e-10
falls back to find_zeros (production grid: 0)	2

Both solvers zero v_par to the same 4.441e-16, so the direct solve is not more accurate than find_zeros — it is equally accurate without iterating. The 2.5e-10 displacement is conditioning, not error: near a fold abs(dB/dtheta) is ~1e-6, so two theta values that far apart are both roots to the last representable bit.

Of the 80 count disagreements, 78 are ones where the new method is more complete: a part in 1e13 from the boundary the adaptive scan reports a single root where there are two, and the interval method separates them. The other 2 are the closed form declining to place a root in the expected cell at that same extreme and deferring to the fallback — more conservative than a bracketed iteration, which always returns something.

This matters beyond the tangency limit. At lambda within 1e-6 of the boundary the two bounce points straddling the B maximum sit 5.3e-4 apart against a 3.9e-3 theta-grid spacing — both inside a single cell. That is why the tempting swap to the existing find_sign_change_roots helper, which only detects sign changes between adjacent grid nodes, is not a safe substitution: it would find neither root and silently change which well _find_deepest_well selects.

Exact tangency, where no interval brackets strictly, is unreachable by construction because _build_lambda_grid excludes both endpoints in all three method branches. Any lambda yielding fewer than two roots still falls back to find_zeros, preserving existing behaviour.

Regression report

Re-run at 2df272486, after merging develop at 4d47467db (which added the ErrorFields NTV-limited correction case, diiid_error_field, also run here since it exercises the bounce solve). Both refs were built from the pinned CI manifest.

case	path	result	largest movement
diiid_n1	psi quadrature	53 unchanged	torque 5.7e-12, dW 1.3e-12; the other 51 quantities are exactly 0.0
diiid_error_field	psi quadrature	8 unchanged	EFC NTV torques 1.2e-13 to 3.6e-13
solovev_kinetic_ntv	psi quadrature	6 unchanged	7.3e-11
solovev_kinetic_multiion	psi quadrature	6 unchanged	per-species D / T / electron, 3.6e-11 to 2.2e-10
solovev_kinetic_calculated	Mode A	7 changed, 7 unchanged	all 0.00% relative: energies 1e-8 to 1e-7, ODE steps 733 to 730
solovev_kinetic_nuzero	Mode A	6 changed, 8 unchanged	all 0.00% relative: energies 2e-8 to 3e-7, saved ODE steps 598 to 599

Regression Report: diiid_n1
Ref 1: develop  @ 4d47467db (2026-09-29)
Ref 2: 2df272486  @ 2df272486 (2026-09-29)
       env: julia 1.11.6, arm64-apple-darwin24.0.0, manifest f1e0eb38 (pinned), 4 threads/4 BLAS
-------------------------------------------------------------------------------------------------
Quantity                                      develop          2df272486        Diff       Status
-------------------------------------------------------------------------------------------------
NTV torque FGAR [N·m]                         5.335149e-01     5.335149e-01     5.7e-12    OK
NTV kinetic energy dW FGAR [J]                7.994507e-02     7.994507e-02     1.3e-12    OK
total energy Re(et[1])                        8.013525e-01     8.013525e-01     0.0e+00    OK
... 50 further quantities, all 0.0e+00 / identical checksums
-------------------------------------------------------------------------------------------------
Summary: 53 unchanged

diiid_error_field           Summary: 8 unchanged
solovev_kinetic_ntv         Summary: 6 unchanged
solovev_kinetic_multiion    Summary: 6 unchanged
solovev_kinetic_calculated  Summary: 7 changed, 7 unchanged   (all 0.00% relative)
solovev_kinetic_nuzero      Summary: 6 changed, 8 unchanged   (all 0.00% relative)

Notes for reviewers

Why the Mode A cases move. Same mechanism, not a separate effect: root positions shift in their last digits, the kinetic matrices follow, and the adaptive Euler-Lagrange integrator then takes one to three different steps out of ~735, surfacing as 1e-8 energy differences. Nothing moves at a physically meaningful level. The PR is tagged ! because results do move, not because anything meaningful changed.

Where the speed actually came from. An intermediate version of this branch kept a bracketed Brent solve inside the monotone interval; it already had the per-lambda cost down to 0.70%. Replacing it with the direct solve took that to 0.062% but bought almost no wall time (77.6 to 76.1 s), because there was little left to win. The direct solve earns its place on exactness and on removing an iterative solver with a tolerance from a physics kernel, not on speed. Worth knowing before weighing the added cubic-root code.

An honest negative. The three psi-quadrature harness cases got faster, but the two Mode A cases did not. Those runs shared the node with other jobs, so the timings are noisy in both directions, but no Mode A speedup is claimed. The profile exercised the quadrature path only, since the DIII-D example runs kinetic_source = "fixed".

Scope. The diff touches BounceAveraging.jl only. Both compute_bounce_data call sites in Torque.jl are unchanged, so this does not collide with the open _compute_fkg_matrices work, and it serves both the quadrature and Mode A entry points without a signature change. The thread-local deepcopy pool in Compute.jl and the threadid() buffer indexing in ForceFreeStates/Kinetic.jl are deliberately left alone because both sit in files that open work already edits.

Tests. 1844 assertions in runtests_kinetic.jl, including the cubic solver against polynomials with known roots — which is how the repeated-root branch was found necessary, since the discriminant cancels to zero for 2u(u-1)^2 and a naive branch reports only the simple root — and 480 assertions that a deliberately corrupt cell hint (out of range, negative, 1e6) yields byte-identical roots, so the hint can only cost steps, never change an answer.

Prior art. The completeness idea here — that a node-by-node sign scan is not a safe replacement for the adaptive scan — comes from krystophny's #341. This rebuilds it on the supported coeffs/CellPoly API rather than reaching into the interpolant's private cache fields.

Review package: https://claude.ai/code/artifact/edf0c0db-776c-4c50-a283-7413bbc0ba96

Why you. Tagging you as reviewer because the profiling in #380 is what pointed at this hotspot, and because your report is also the reason it was re-measured first: those shares predated #390/#391/#392, and find_zeros turned out to be 15.0% of CPU work rather than the ~26% quoted. The find_sign_change_roots swap your report suggested and then declined is treated at length under Correctness above — declining it was right, and the figure in the review package shows the case that would have broken it.

🤖 Generated with Claude Code

https://claude.ai/code/session_0186CLfjeJeSusNTva5RMfGU

[adrianaghiozzi]

Reviewed this (with Claude Code) — solid change, no correctness issues survived review. Notes below, roughly in order of how much I'd want an answer before approving.

What I checked

• Traced the algorithm by hand: stationary-point decomposition (_surface_b_field), interval bracketing (_bounce_points_at_lambda!), cell location with hint/fallback (_locate_cell/_bisect_cell), and the closed-form cubic solve with Newton polish (_real_cubic_roots/_cell_level_root).
• Verified the periodic-seam handling (the wrap interval at θ=0/1) is internally consistent — the boundary value used for the branch decision (bknot[end]) and the value re-derived inside _monotone_segment_root(bf, θa, 1.0, ...) come from evaluating the same last cell at the same point, so there's no seam mismatch.
• Confirmed the claimed scope: _find_bounce_points_and_grid's signature changed, but it's a private helper called only from compute_bounce_data, whose signature didn't change, and both Torque.jl call sites are untouched.
• Confirmed the title passes ci/conventions/check_subject.py as-is.

Verified with an out-of-tree probe: the same-cell double-extremum path is correct, but untested

_monotone_segment_root has a distinct branch (ia == ib) for when both stationary points bounding a monotone interval land inside the same spline cell — the 257-point, well-separated-wells fixture in the added tests never exercises it. Rather than just flag it, I checked out the branch in a scratch worktree and searched for a periodic B(θ) that actually produces this case:

• Found one on the first random trial (11-knot grid, two harmonics): two stationary points 0.082 apart, both inside cell 5.
• Swept 4,000 λ values through that region, comparing against Roots.find_zeros as ground truth: 0 count mismatches, roots accurate to 8.9e-16.

So the branch is correct — it bounds the cubic solve by the stationary points themselves rather than the cell's knot boundaries, which is what makes it work — but it's presently relying on code-reading rather than test coverage. Given it's the one branch this PR adds that the existing fixture can't reach, worth adding a small test case with two close-together extrema in one cell (happy to hand over the probe script if useful, it's disposable throwaway code, not proposing it as-is for the suite).

Found and confirmed a real misclassification bug in _real_cubic_roots — not blocking, but worth knowing about

The 8 * eps(Float64) * max((q/2)^2, abs(p/3)^3) tolerance in _real_cubic_roots decides whether a near-zero discriminant gets the "repeated root" treatment. I stress-tested it directly with a cubic built from three genuine, distinct real roots (verified: the polynomial evaluates to ~1e-17 at all three, so this isn't a construction artifact), two of them 1e-8 apart:

n=1  x1=0.9099999999999998  x2=NaN  x3=NaN

(true roots were 0.37, 0.37000001, and 0.91 — only the far one came back). At that separation the discriminant lands just outside the tolerance (disc=8.8e-20 vs tol=6.0e-20), so the code takes the single-real-root branch and silently drops two real roots — not imprecisely locates them, drops them entirely.

I then checked whether this reaches a wrong final answer in the actual bounce-point pipeline, and it doesn't: _cell_level_root, restricted to the monotone sub-window that should bracket one of the close roots, gets back NaN (the one candidate _real_cubic_roots returned falls outside that window). That NaN propagates to _bounce_points_at_lambda!, which empties its result on any NaN, and _find_bounce_points_and_grid's existing length(bpts) < 2 check sends it to the Roots.find_zeros fallback. So the pipeline is safety-netted — this degrades to pre-PR performance for that one λ, it does not produce a wrong physics answer.

Given the compound rarity here (needs both two stationary points sharing a cell and a target value producing roots ~1e-8 apart), it's plausible this never fires on real equilibria, and the fallback means it's not a merge-blocker either way. But _real_cubic_roots is generic enough that a future caller using it outside this pipeline's fallback could hit it directly, so flagging it now rather than letting it surface later as a mystery.

---

Nice work tying this back to the #380 thread, and to your credit — no shortcuts against find_sign_change_roots; the completeness argument here holds up.

[adrianaghiozzi]

Claude left a few comments for posterity but the TLDR is that this should be good for merge.

It suggested an optional addition to the test suite - up to you whether that seems worth doing. And then it also pointed out an edge case in which '_real_cubic_roots' would fail, but noted that it's an issue that would likely never arise on a real equilibrium. Claude's analysis about this is documented here in case the issue ever does pop up.

[logan-nc]

Thanks @adrianaghiozzi !

[logan-nc]

Thanks @adrianaghiozzi — both findings were right, and the probe made the first one trivial to reproduce (disc = +8.8e-20 against a true negative; the ratio test it should have used reads 1.0000000000000013).

Follow-up is #483. Beyond the two items here, re-reviewing the merged code with your example in hand turned up the same cancellation at two more sites — the quadratic branch for a degenerate leading coefficient, and the stationary-point search in _surface_b_field, where a dropped near-double stationary point would have lost a root pair with no fallback at all. All three now share one conditioned test, and the ad-hoc tolerance and explicit repeated-root branch are gone.

The same-cell path now has a deterministic fixture rather than a random one (2 + 0.45cos2πθ + 0.20cos6πθ on 11 knots, a pair 0.077 apart sharing a cell). Harness: every row unchanged across the five kinetic cases, as predicted — the misclassification only ever fired on the injected near-tangent λ, where the pipeline had been falling back.

🤖 Generated with Claude Code

https://claude.ai/code/session_0186CLfjeJeSusNTva5RMfGU