From 82d35cb19259b73eb9ad032aab7c41f150e9fb45 Mon Sep 17 00:00:00 2001 From: logan-nc Date: Fri, 11 Sep 2026 21:55:41 -0400 Subject: [PATCH 1/4] KineticForces - PERF - Bracket bounce-point solves with the exact B extrema MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _find_bounce_points_and_grid called Roots.find_zeros — a generic adaptive scan of the whole theta interval — once per pitch angle, so the scan repeated for all 128 lambda points on every flux surface, bounce harmonic and species. The 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 and supplies a guaranteed bracket. Enumerate those stationary points once per surface by walking the cubic's cells through the public coeffs/CellPoly interface and solving the quadratic derivative on each, then per lambda test the cached endpoint B values and Brent-solve only the intervals that actually bracket. This keeps the completeness of the adaptive scan — including root pairs that straddle an extremum too closely for a node-by-node sign scan to separate — while moving the per-cell work above the lambda loop. Degenerate lambdas, where v_par is tangent to zero at an extremum so no interval brackets strictly, fall back to find_zeros and keep their existing behaviour; the lambda grid excludes both trapped-passing endpoints, so these are rare. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_0186CLfjeJeSusNTva5RMfGU --- src/KineticForces/BounceAveraging.jl | 130 ++++++++++++++++++++++++++- 1 file changed, 126 insertions(+), 4 deletions(-) diff --git a/src/KineticForces/BounceAveraging.jl b/src/KineticForces/BounceAveraging.jl index f1bb1fc32..4eac54de6 100644 --- a/src/KineticForces/BounceAveraging.jl +++ b/src/KineticForces/BounceAveraging.jl @@ -10,6 +10,12 @@ Reference: [Logan et al., Phys. Plasmas 20, 122507 (2013)] Ports Fortran torque.F90 lines 530-816. """ +# Guard against a non-terminating spline cell walk; θ grids are O(10²) cells. +const MAX_SPLINE_CELLS = 100_000 +# Stationary points closer than this in θ are the same root seen from both sides +# of a cell boundary. +const EXTREMUM_MERGE_TOL = 1e-12 + # ============================================================================ # BounceData struct # ============================================================================ @@ -208,6 +214,12 @@ function compute_bounce_data( # Per-surface scratch, reused across all λ. scr = BounceScratch(ntheta, mpert) + # Stationary points of B(θ) bound the intervals on which B is monotone. They + # depend only on the surface, so enumerate them once here rather than + # rescanning [0,1] for every λ. + ext = _b_field_extrema(B_vpar) + bpts_buf = Float64[] + # Trapped-passing boundary and λ range lmdatpb = bo / bmax lmdamax = bo / bmin @@ -241,7 +253,7 @@ function compute_bounce_data( # Find bounce points and build θ sub-grid _, _, tdt_pts, tdt_wts = _find_bounce_points_and_grid( - lmda, bo, sigma, B_vpar, theta_bmax, psi, ntheta) + lmda, bo, sigma, B_vpar, theta_bmax, psi, ntheta, ext, bpts_buf) # Bounce integrals over θ (Fortran lines 674-735) wbbar, wdbar, dJdJ_val, wmats_lmda = _bounce_integrate( @@ -417,6 +429,108 @@ consistent with the bounce-point roots as in Fortran's `vspl`. 1.0 - (lmda / bo) * B_vpar(mod(θ, 1.0)) +""" + BFieldExtrema + +Stationary points of the periodic cubic `B_vpar` on θ ∈ [0,1), with `B` cached at +each one. Consecutive points bound intervals on which `B` is monotone, so every +interval holds at most one bounce point and supplies a guaranteed bracket. +Built once per flux surface in `compute_bounce_data` and reused for every λ. + +## Fields +- `theta::Vector{Float64}`: stationary points, sorted ascending, in [0,1) +- `bval::Vector{Float64}`: `B_vpar` evaluated at each stationary point +""" +struct BFieldExtrema + theta::Vector{Float64} + bval::Vector{Float64} +end + +""" +Enumerate the stationary points of the cubic `B_vpar` exactly, by walking its cells +and solving the quadratic dS/dθ = 0 on each. Uses only the public `coeffs`/`CellPoly` +interface, so it holds for whatever θ grid the surface interpolant was built on. +""" +function _b_field_extrema(B_vpar) + theta = Float64[] + x = 0.0 + ncell = 0 + while true + cell = coeffs(B_vpar, x) + h = cell.xR - cell.xL + _, c, b, a = cell.p + # S(u) = d + c·u + b·u² + a·u³ ⇒ S'(u) = c + 2b·u + 3a·u², u = θ − xL ∈ [0, h) + qa, qb, qc = 3a, 2b, c + if abs(qa) <= eps(Float64) * max(abs(qb), abs(qc), 1.0) + if qb != 0 + u = -qc / qb + (0.0 <= u < h) && push!(theta, cell.xL + u) + end + else + disc = qb^2 - 4 * qa * qc + if disc >= 0 + sq = sqrt(disc) + for u in ((-qb - sq) / (2qa), (-qb + sq) / (2qa)) + (0.0 <= u < h) && push!(theta, cell.xL + u) + end + end + end + cell.xR >= 1.0 && break + x = cell.xR + ncell += 1 + ncell > MAX_SPLINE_CELLS && error("ERROR: _b_field_extrema - cell walk did not reach θ=1") + end + + sort!(theta) + # A root sitting on a cell boundary can be reported by both neighbouring cells. + if length(theta) > 1 + keep = 1 + for i in 2:length(theta) + if theta[i] - theta[keep] > EXTREMUM_MERGE_TOL + keep += 1 + theta[keep] = theta[i] + end + end + resize!(theta, keep) + end + + return BFieldExtrema(theta, [B_vpar(t) for t in theta]) +end + +""" +Bounce points of `v_par(θ) = 1 − (λ/bo)·B(θ)` for a trapped particle, found by +bracketed solves on the monotone intervals between stationary points of `B`. +`B = bo/λ` at a bounce point, so an interval brackets one iff `bo/λ` lies strictly +between its endpoint `B` values — a scalar test against the cached values, with no +spline evaluation per λ. Returns roots sorted descending, matching the order the +deepest-well and marginally-trapped logic downstream assumes. An empty or +single-element result signals a degenerate λ and sends the caller to the fallback. +""" +function _bounce_points_from_extrema!(bpts::Vector{Float64}, ext::BFieldExtrema, + lmda::Float64, bo::Float64, B_vpar) + empty!(bpts) + k = length(ext.theta) + k < 2 && return bpts + + btarget = bo / lmda + for i in 1:k + j = i == k ? 1 : i + 1 + fa = ext.bval[i] - btarget + fb = ext.bval[j] - btarget + fa * fb < 0 || continue + # The last interval wraps through the θ = 0/1 seam; _vpar_from_spline takes + # mod(θ, 1), so B is continuous across it for the periodic fit. + ta = ext.theta[i] + tb = i == k ? ext.theta[1] + 1.0 : ext.theta[j] + root = Roots.find_zero(θ -> _vpar_from_spline(B_vpar, lmda, bo, θ), (ta, tb), Roots.Brent()) + push!(bpts, mod(root, 1.0)) + end + + sort!(bpts; rev=true) + return bpts +end + + """ Find bounce points for trapped/passing particles and build θ sub-grid. Returns (t1, t2, theta_points, theta_weights). @@ -424,14 +538,22 @@ Returns (t1, t2, theta_points, theta_weights). function _find_bounce_points_and_grid( lmda::Float64, bo::Float64, sigma::Int, B_vpar, theta_bmax::Float64, psi::Float64, - ntheta::Int + ntheta::Int, ext::BFieldExtrema, bpts_buf::Vector{Float64} ) if sigma == 0 # trapped # Bounce points: all roots of v_par(θ) = 1 − (λ/bo)·B_vpar(θ) in (0,1), # sorted descending — the same order as Fortran spline_roots, which the # marginally-trapped and deepest-well wrap logic below assume. - vpar_fn = θ -> _vpar_from_spline(B_vpar, lmda, bo, θ) - bpts = sort!(Roots.find_zeros(vpar_fn, 0.0, 1.0); rev=true) + bpts = _bounce_points_from_extrema!(bpts_buf, ext, lmda, bo, B_vpar) + + if length(bpts) < 2 + # Degenerate λ: v_par is tangent to zero at an extremum of B (so no + # interval brackets strictly), or B has fewer than two stationary + # points. Rare — the λ grid excludes both trapped-passing endpoints — + # so fall back to the adaptive whole-interval scan and keep its behaviour. + vpar_fn = θ -> _vpar_from_spline(B_vpar, lmda, bo, θ) + bpts = sort!(Roots.find_zeros(vpar_fn, 0.0, 1.0); rev=true) + end nbpts = length(bpts) if nbpts < 1 From 12f2757fdfc6b791d7e13cdf68f754544f78c751 Mon Sep 17 00:00:00 2001 From: logan-nc Date: Fri, 11 Sep 2026 21:59:22 -0400 Subject: [PATCH 2/4] KineticForces - TEST - Cover the bracketed bounce-point enumeration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Assert the stationary-point enumeration is exact (dB/dtheta vanishes at each point, the count matches a dense derivative sign scan, and a period holds an even number) and that the bracketed solver reproduces Roots.find_zeros across the whole trapped lambda range — including lambdas pushed to within 1e-12 of both trapped-passing boundaries, where the two roots straddling an extremum are arbitrarily close and a node-by-node sign scan would merge them. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_0186CLfjeJeSusNTva5RMfGU --- test/runtests_kinetic.jl | 77 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/test/runtests_kinetic.jl b/test/runtests_kinetic.jl index 29955c2c7..b3b0be91e 100644 --- a/test/runtests_kinetic.jl +++ b/test/runtests_kinetic.jl @@ -132,6 +132,83 @@ end end + # ========================================================================= + # Bounce-point root finding + # ========================================================================= + @testset "bounce-point enumeration" begin + # Multi-well periodic B(θ): three distinct wells of unequal depth, so the + # deepest-well selection and the wrap through θ = 0/1 are both exercised. + xs = collect(range(0.0, 1.0; length=257)) + B_vals = @. 2.0 + 0.4 * cos(2pi * xs) + 0.12 * cos(4pi * xs + 0.7) + 0.05 * sin(6pi * xs) + B_vals[end] = B_vals[1] + B_vpar = cubic_interp(xs, B_vals; bc=PeriodicBC()) + bo = 2.0 + + ext = KF._b_field_extrema(B_vpar) + + @testset "stationary points are exact" begin + # A periodic function has an even number of stationary points per period. + @test iseven(length(ext.theta)) + @test length(ext.theta) >= 2 + @test issorted(ext.theta) + @test all(0.0 .<= ext.theta .< 1.0) + # dB/dθ vanishes at each, and the cached B values match the spline. + dB = deriv1(B_vpar) + scale = maximum(abs, B_vals) + for (θ, b) in zip(ext.theta, ext.bval) + @test abs(dB(θ)) < 1e-8 * scale + @test b ≈ B_vpar(θ) rtol=1e-14 + end + # Every stationary point of a dense derivative sign scan is accounted for. + dense = range(0.0, 1.0; length=20001) + dv = [dB(x) for x in dense] + nsign = count(i -> dv[i] * dv[i+1] < 0, 1:length(dv)-1) + @test length(ext.theta) == nsign + end + + @testset "agrees with the adaptive scan across λ" begin + # The bracketed solver must reproduce Roots.find_zeros over the whole + # trapped range, including λ pushed against both boundaries where the + # two roots straddling an extremum are arbitrarily close together. + bmax = maximum(ext.bval) + bmin = minimum(ext.bval) + lmdatpb = bo / bmax + lmdamax = bo / bmin + lambdas = collect(range(lmdatpb * (1 + 1e-6), lmdamax * (1 - 1e-6); length=200)) + for eps in (1e-3, 1e-6, 1e-9, 1e-12) + push!(lambdas, lmdatpb * (1 + eps)) + push!(lambdas, lmdamax * (1 - eps)) + end + + buf = Float64[] + for lmda in lambdas + new_roots = copy(KF._bounce_points_from_extrema!(buf, ext, lmda, bo, B_vpar)) + ref_roots = sort!(KF.Roots.find_zeros(θ -> KF._vpar_from_spline(B_vpar, lmda, bo, θ), 0.0, 1.0); rev=true) + # Fewer than two roots is the documented degenerate signal, where + # the caller falls back to the adaptive scan. + length(new_roots) < 2 && continue + @test length(new_roots) == length(ref_roots) + if length(new_roots) == length(ref_roots) + @test all(abs.(new_roots .- ref_roots) .< 1e-9) + end + # Roots are returned in the descending order downstream assumes. + @test issorted(new_roots; rev=true) + # v_par vanishes at each root. + for θ in new_roots + @test abs(KF._vpar_from_spline(B_vpar, lmda, bo, θ)) < 1e-10 + end + end + end + + @testset "degenerate λ signals the fallback" begin + # λ below the trapped-passing boundary puts bo/λ above every B, so no + # interval brackets and the caller is told to fall back. + buf = Float64[] + bmax = maximum(ext.bval) + @test length(KF._bounce_points_from_extrema!(buf, ext, 0.5 * bo / bmax, bo, B_vpar)) < 2 + end + end + # ========================================================================= # Resonance root solver # ========================================================================= From 5d650390ad36193814c478243d6f7ad76abeb764 Mon Sep 17 00:00:00 2001 From: logan-nc Date: Fri, 11 Sep 2026 22:17:20 -0400 Subject: [PATCH 3/4] KineticForces - TEST - Assert the straddling pair against geometry, not another solver MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The near-boundary lambdas asserted equality with Roots.find_zeros, which encodes that solver's resolution limit as the expected answer: as lambda approaches the trapped-passing boundary the two roots straddling the extremum collapse together, and whether the adaptive scan still separates them depends on its own step heuristics and on floating-point rounding that differs between Julia 1.11 and 1.12. The bracketed solver would then fail CI for being more accurate. Assert the geometry instead — exactly two roots, v_par vanishing at each, and the arc enclosing the extremum closing monotonically as lambda tightens, ending far below the theta-grid spacing. The arc is measured on the circle: this fixture's B maximum sits at theta ~ 0.999, so the enclosing arc runs through the theta = 0/1 seam. Equality with find_zeros is kept for the 200-point bulk sweep, where it is a fair comparison. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_0186CLfjeJeSusNTva5RMfGU --- test/runtests_kinetic.jl | 66 +++++++++++++++++++++++++++++----------- 1 file changed, 49 insertions(+), 17 deletions(-) diff --git a/test/runtests_kinetic.jl b/test/runtests_kinetic.jl index b3b0be91e..8a96266c8 100644 --- a/test/runtests_kinetic.jl +++ b/test/runtests_kinetic.jl @@ -166,27 +166,21 @@ @test length(ext.theta) == nsign end - @testset "agrees with the adaptive scan across λ" begin - # The bracketed solver must reproduce Roots.find_zeros over the whole - # trapped range, including λ pushed against both boundaries where the - # two roots straddling an extremum are arbitrarily close together. - bmax = maximum(ext.bval) - bmin = minimum(ext.bval) - lmdatpb = bo / bmax - lmdamax = bo / bmin - lambdas = collect(range(lmdatpb * (1 + 1e-6), lmdamax * (1 - 1e-6); length=200)) - for eps in (1e-3, 1e-6, 1e-9, 1e-12) - push!(lambdas, lmdatpb * (1 + eps)) - push!(lambdas, lmdamax * (1 - eps)) - end + bmax = maximum(ext.bval) + bmin = minimum(ext.bval) + lmdatpb = bo / bmax + lmdamax = bo / bmin + @testset "agrees with the adaptive scan across λ" begin + # Over the bulk of the trapped range the bracketed solver must reproduce + # Roots.find_zeros exactly. λ crowded against the boundaries is asserted + # separately below: there the two roots straddling an extremum become so + # close that whether the adaptive scan still resolves them depends on its + # own step heuristics, which is not a contract worth pinning a test to. buf = Float64[] - for lmda in lambdas + for lmda in range(lmdatpb * (1 + 1e-6), lmdamax * (1 - 1e-6); length=200) new_roots = copy(KF._bounce_points_from_extrema!(buf, ext, lmda, bo, B_vpar)) ref_roots = sort!(KF.Roots.find_zeros(θ -> KF._vpar_from_spline(B_vpar, lmda, bo, θ), 0.0, 1.0); rev=true) - # Fewer than two roots is the documented degenerate signal, where - # the caller falls back to the adaptive scan. - length(new_roots) < 2 && continue @test length(new_roots) == length(ref_roots) if length(new_roots) == length(ref_roots) @test all(abs.(new_roots .- ref_roots) .< 1e-9) @@ -200,6 +194,44 @@ end end + @testset "resolves the pair straddling an extremum" begin + # Approaching a boundary, the two bounce points collapse onto the extremum + # of B from either side. Assert that directly rather than against another + # solver: exactly two roots, one on each side of the extremum, both closing + # on it as λ tightens, with v_par vanishing at each. + buf = Float64[] + θmax = ext.theta[argmax(ext.bval)] + θmin = ext.theta[argmin(ext.bval)] + + # Width of the arc between the two roots that encloses θx. θ is periodic, + # so the enclosing arc may run through the θ = 0/1 seam — the B maximum of + # this fixture sits at θ ≈ 0.999, which exercises exactly that. + function enclosing_arc(roots, θx) + lo, hi = minimum(roots), maximum(roots) + return lo <= θx <= hi ? hi - lo : 1.0 - (hi - lo) + end + + prev = Dict(:max => Inf, :min => Inf) + for eps in (1e-4, 1e-6, 1e-8, 1e-10, 1e-12) + for (key, lmda, θx) in ((:max, lmdatpb * (1 + eps), θmax), (:min, lmdamax * (1 - eps), θmin)) + roots = copy(KF._bounce_points_from_extrema!(buf, ext, lmda, bo, B_vpar)) + @test length(roots) == 2 + length(roots) == 2 || continue + for θ in roots + @test abs(KF._vpar_from_spline(B_vpar, lmda, bo, θ)) < 1e-10 + end + # The pair hugs the extremum and closes on it as λ tightens. + arc = enclosing_arc(roots, θx) + @test arc < prev[key] + prev[key] = arc + end + end + # Having closed monotonically from 1e-4 to 1e-12, the pair is now far + # tighter than the θ-grid spacing that a node-by-node scan would resolve. + @test prev[:max] < 0.5 * (xs[2] - xs[1]) + @test prev[:min] < 0.5 * (xs[2] - xs[1]) + end + @testset "degenerate λ signals the fallback" begin # λ below the trapped-passing boundary puts bo/λ above every B, so no # interval brackets and the caller is told to fall back. From 01ea6e8a6af5db5199fa1ace51cfeff16a535693 Mon Sep 17 00:00:00 2001 From: logan-nc Date: Sat, 12 Sep 2026 08:31:10 -0400 Subject: [PATCH 4/4] KineticForces - PERF - Solve the bounce points directly instead of iterating The bracketed Brent solve was iterating to find something the spline already writes down. Inside a cell the interpolant is a cubic with known coefficients, so the crossing B(theta) = bo/lambda is a closed-form root, not a search. Decompose B once per surface into its cells, keeping each cell's coefficients and its endpoint value alongside the stationary points already being enumerated. Locating the crossing then needs only the cell, and the cell need not be searched for blind: lambda ascends through compute_bounce_data, so bo/lambda falls monotonically and each interval's crossing walks steadily along the cells. Each lambda resumes from the previous one's cell -- the same hint idiom FastInterpolations uses for its own searches -- which lands in the same cell for the median lambda and steps one or two otherwise, against about eight probes for a bisection that starts from nothing 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. The cubic is solved by the trigonometric form, which stays conditioned where Cardano's radicals cancel, with one Newton step to recover the digits lost when two roots nearly coincide, and an explicit repeated-root branch for the discriminant cancelling to zero. Roots now satisfy v_par = 0 to 4.4e-16 rather than to a solver tolerance, and agree with Roots.find_zeros to 1e-13 over 3000 lambda on each of five surface geometries. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_0186CLfjeJeSusNTva5RMfGU --- src/KineticForces/BounceAveraging.jl | 297 +++++++++++++++++++++++---- test/runtests_kinetic.jl | 73 +++++-- 2 files changed, 310 insertions(+), 60 deletions(-) diff --git a/src/KineticForces/BounceAveraging.jl b/src/KineticForces/BounceAveraging.jl index 4eac54de6..c62ae8518 100644 --- a/src/KineticForces/BounceAveraging.jl +++ b/src/KineticForces/BounceAveraging.jl @@ -15,6 +15,8 @@ const MAX_SPLINE_CELLS = 100_000 # Stationary points closer than this in θ are the same root seen from both sides # of a cell boundary. const EXTREMUM_MERGE_TOL = 1e-12 +# Cells a hinted cell search may step before falling back to bisection. +const HINT_STEP_BUDGET = 8 # ============================================================================ # BounceData struct @@ -214,11 +216,13 @@ function compute_bounce_data( # Per-surface scratch, reused across all λ. scr = BounceScratch(ntheta, mpert) - # Stationary points of B(θ) bound the intervals on which B is monotone. They - # depend only on the surface, so enumerate them once here rather than - # rescanning [0,1] for every λ. - ext = _b_field_extrema(B_vpar) + # B(θ) is decomposed once per surface into its cells and stationary points; every + # λ then reuses it instead of rescanning [0,1]. + bf = _surface_b_field(B_vpar) bpts_buf = Float64[] + # One resumable cell hint per monotone interval; λ ascends, so each interval's + # crossing advances steadily and the next λ starts where the last one finished. + hints = ones(Int, length(bf.theta) + 1) # Trapped-passing boundary and λ range lmdatpb = bo / bmax @@ -253,7 +257,7 @@ function compute_bounce_data( # Find bounce points and build θ sub-grid _, _, tdt_pts, tdt_wts = _find_bounce_points_and_grid( - lmda, bo, sigma, B_vpar, theta_bmax, psi, ntheta, ext, bpts_buf) + lmda, bo, sigma, B_vpar, theta_bmax, psi, ntheta, bf, bpts_buf, hints) # Bounce integrals over θ (Fortran lines 674-735) wbbar, wdbar, dJdJ_val, wmats_lmda = _bounce_integrate( @@ -430,36 +434,52 @@ consistent with the bounce-point roots as in Fortran's `vspl`. """ - BFieldExtrema + SurfaceBField -Stationary points of the periodic cubic `B_vpar` on θ ∈ [0,1), with `B` cached at -each one. Consecutive points bound intervals on which `B` is monotone, so every -interval holds at most one bounce point and supplies a guaranteed bracket. -Built once per flux surface in `compute_bounce_data` and reused for every λ. +The periodic cubic B(θ) of one flux surface, decomposed once and reused for every λ. +Holds the per-cell polynomial coefficients, B at every knot, and the stationary points +of B. Consecutive stationary points bound intervals on which B is monotone, so each +holds at most one bounce point; within such an interval the cached knot values locate +the cell by bisection and the cell's cubic is then solved in closed form. ## Fields -- `theta::Vector{Float64}`: stationary points, sorted ascending, in [0,1) -- `bval::Vector{Float64}`: `B_vpar` evaluated at each stationary point +- `knot::Vector{Float64}`: cell boundaries, ascending, `knot[1] = 0`, `knot[end] = 1` +- `bknot::Vector{Float64}`: B at each knot +- `poly::Vector{NTuple{4,Float64}}`: per-cell `(d, c, b, a)` of `S(u) = d + cu + bu² + au³`, + in the cell-local coordinate `u = θ − knot[i]` +- `theta::Vector{Float64}`: stationary points of B, ascending, in [0,1) +- `bval::Vector{Float64}`: B at each stationary point """ -struct BFieldExtrema +struct SurfaceBField + knot::Vector{Float64} + bknot::Vector{Float64} + poly::Vector{NTuple{4,Float64}} theta::Vector{Float64} bval::Vector{Float64} end """ -Enumerate the stationary points of the cubic `B_vpar` exactly, by walking its cells -and solving the quadratic dS/dθ = 0 on each. Uses only the public `coeffs`/`CellPoly` -interface, so it holds for whatever θ grid the surface interpolant was built on. +Decompose the cubic `B_vpar` by walking its cells once, recording each cell's +polynomial and endpoint value and solving the quadratic dS/dθ = 0 on each to get the +exact stationary points. Uses only the public `coeffs`/`CellPoly` interface, so it +holds for whatever θ grid the surface interpolant was built on. """ -function _b_field_extrema(B_vpar) +function _surface_b_field(B_vpar) + knot = Float64[0.0] + bknot = Float64[] + poly = NTuple{4,Float64}[] theta = Float64[] + x = 0.0 - ncell = 0 while true cell = coeffs(B_vpar, x) h = cell.xR - cell.xL - _, c, b, a = cell.p - # S(u) = d + c·u + b·u² + a·u³ ⇒ S'(u) = c + 2b·u + 3a·u², u = θ − xL ∈ [0, h) + d, c, b, a = cell.p + push!(poly, (d, c, b, a)) + push!(bknot, d) # S(0) = B at the cell's left knot + push!(knot, cell.xR) + + # S'(u) = c + 2b·u + 3a·u², u ∈ [0, h) qa, qb, qc = 3a, 2b, c if abs(qa) <= eps(Float64) * max(abs(qb), abs(qc), 1.0) if qb != 0 @@ -475,14 +495,15 @@ function _b_field_extrema(B_vpar) end end end + cell.xR >= 1.0 && break x = cell.xR - ncell += 1 - ncell > MAX_SPLINE_CELLS && error("ERROR: _b_field_extrema - cell walk did not reach θ=1") + length(poly) > MAX_SPLINE_CELLS && error("ERROR: _surface_b_field - cell walk did not reach θ=1") end + push!(bknot, evalpoly(knot[end] - knot[end-1], poly[end])) # B at θ = 1 sort!(theta) - # A root sitting on a cell boundary can be reported by both neighbouring cells. + # A stationary point on a cell boundary is reported by both neighbouring cells. if length(theta) > 1 keep = 1 for i in 2:length(theta) @@ -494,35 +515,223 @@ function _b_field_extrema(B_vpar) resize!(theta, keep) end - return BFieldExtrema(theta, [B_vpar(t) for t in theta]) + bval = [_b_at(knot, poly, t) for t in theta] + return SurfaceBField(knot, bknot, poly, theta, bval) +end + +"""Cell index holding θ ∈ [0,1], from the ascending knot vector.""" +@inline _cell_index(knot::Vector{Float64}, θ::Float64) = + clamp(searchsortedlast(knot, θ), 1, length(knot) - 1) + +"""Evaluate B at θ from the cached cell polynomials.""" +@inline function _b_at(knot::Vector{Float64}, poly::Vector{NTuple{4,Float64}}, θ::Float64) + i = _cell_index(knot, θ) + return evalpoly(θ - knot[i], poly[i]) +end + +""" +Real roots of `a·u³ + b·u² + c·u + d = 0`, returned as `(count, r1, r2, r3)` with +unused slots `NaN`. Degenerate leading coefficients fall through to the quadratic and +linear cases; three distinct real roots use the trigonometric form, which stays well +conditioned where Cardano's radicals cancel. +""" +function _real_cubic_roots(a::Float64, b::Float64, c::Float64, d::Float64) + scale = max(abs(b), abs(c), abs(d), 1.0) + if abs(a) <= eps(Float64) * scale + if abs(b) <= eps(Float64) * max(abs(c), abs(d), 1.0) + c == 0 && return (0, NaN, NaN, NaN) + return (1, -d / c, NaN, NaN) + end + disc = c * c - 4 * b * d + disc < 0 && return (0, NaN, NaN, NaN) + sq = sqrt(disc) + # Cancellation-free quadratic roots (Numerical Recipes §5.6). + q = c == 0 ? -0.5 * sq : -0.5 * (c + copysign(sq, c)) + r1 = q / b + r2 = q == 0 ? r1 : d / q + return (2, r1, r2, NaN) + end + + B, C, D = b / a, c / a, d / a + shift = B / 3 + p = C - B * B / 3 + q = 2 * B^3 / 27 - B * C / 3 + D + disc = (q / 2)^2 + (p / 3)^3 + + if p == 0 && q == 0 + return (1, -shift, NaN, NaN) # triple root + elseif abs(disc) <= 8 * eps(Float64) * max((q / 2)^2, abs(p / 3)^3) + # Repeated root. The discriminant cancels to ~0 here, so the branches below + # would lose it: disc > 0 by a rounding step reports only the simple root. + t2 = -3q / (2p) + return (3, 3q / p - shift, t2 - shift, t2 - shift) + elseif disc > 0 + s = sqrt(disc) + t = cbrt(-q / 2 + s) + cbrt(-q / 2 - s) + return (1, t - shift, NaN, NaN) + else + # Three real roots: t_k = 2r·cos(φ − 2πk/3), r = √(−p/3), φ = acos(−q/2r³)/3. + r = sqrt(-p / 3) + φ = acos(clamp(-q / (2 * r^3), -1.0, 1.0)) / 3 + return (3, 2r * cos(φ) - shift, 2r * cos(φ - 2π / 3) - shift, 2r * cos(φ - 4π / 3) - shift) + end +end + +""" +Solve `B(θ) = btarget` inside cell `ic`, restricted to `θ ∈ [θlo, θhi]`. The cubic is +solved in closed form and polished with one Newton step, which recovers the digits the +closed form loses when two of its roots are nearly coincident. Returns `NaN` if no root +lies in the restricted range. +""" +function _cell_level_root(bf::SurfaceBField, ic::Int, btarget::Float64, θlo::Float64, θhi::Float64) + d, c, b, a = bf.poly[ic] + x0 = bf.knot[ic] + ulo, uhi = θlo - x0, θhi - x0 + # Admit roots a rounding step outside the cell: θlo/θhi are stationary points and + # cell edges, and the root can sit exactly on one. + pad = 8 * eps(Float64) * max(abs(ulo), abs(uhi), bf.knot[ic+1] - x0) + + n, r1, r2, r3 = _real_cubic_roots(a, b, c, d - btarget) + best, bestres = NaN, Inf + for k in 1:n + u = k == 1 ? r1 : (k == 2 ? r2 : r3) + (isfinite(u) && ulo - pad <= u <= uhi + pad) || continue + u = clamp(u, ulo, uhi) + # One Newton step on S(u) − btarget, skipped at a stationary point. + deriv = c + u * (2b + u * 3a) + if deriv != 0 + un = u - (evalpoly(u, (d - btarget, c, b, a))) / deriv + (ulo - pad <= un <= uhi + pad) && (u = clamp(un, ulo, uhi)) + end + res = abs(evalpoly(u, (d - btarget, c, b, a))) + if res < bestres + best, bestres = u, res + end + end + return isnan(best) ? NaN : x0 + best +end + +""" +Where `btarget` sits relative to cell `ic` of a monotone span: `-1` before it, `0` +inside, `+1` past it. `s` carries the span's direction so one comparison serves both. +The span's first and last cells are entered part-way, at the stationary points bounding +it, so their outer edge value comes from `ba`/`bb` rather than from a knot. +""" +@inline function _cell_position( + bf::SurfaceBField, ic::Int, ia::Int, ib::Int, + ba::Float64, bb::Float64, btarget::Float64, s::Float64 +) + left = ic == ia ? ba : bf.bknot[ic] + right = ic == ib ? bb : bf.bknot[ic+1] + s * btarget < s * left && return -1 + s * btarget > s * right && return 1 + return 0 +end + +""" +Cell holding the `B = btarget` crossing on a monotone span running from cell `ia` to +cell `ib`. + +λ advances monotonically through `compute_bounce_data`, so `btarget = bo/λ` falls +monotonically and each interval's crossing walks steadily along the cells in one +direction. Resuming from the previous λ's cell therefore costs a step or two, the same +hint idiom FastInterpolations uses for its own searches. Bisection stays as the fallback +for the first λ of a surface and for the sweep near a stationary point, where B is flat +and the crossing can cross many cells between consecutive λ. +""" +function _locate_cell( + bf::SurfaceBField, ia::Int, ib::Int, ba::Float64, bb::Float64, + btarget::Float64, hint::Int +) + s = bb >= ba ? 1.0 : -1.0 + ic = clamp(hint, ia, ib) + pos = _cell_position(bf, ic, ia, ib, ba, bb, btarget, s) + steps = 0 + while pos != 0 && steps < HINT_STEP_BUDGET + ic += pos + (ia <= ic <= ib) || return _bisect_cell(bf, ia, ib, btarget, s) + pos = _cell_position(bf, ic, ia, ib, ba, bb, btarget, s) + steps += 1 + end + return pos == 0 ? ic : _bisect_cell(bf, ia, ib, btarget, s) +end + +"""Bisection on the span's monotone knot values, for when the hint does not pay off.""" +function _bisect_cell(bf::SurfaceBField, ia::Int, ib::Int, btarget::Float64, s::Float64) + ia >= ib && return ia + s * bf.bknot[ia+1] > s * btarget && return ia + s * bf.bknot[ib] <= s * btarget && return ib + lo, hi = ia + 1, ib + while hi - lo > 1 + mid = (lo + hi) >>> 1 + s * bf.bknot[mid] <= s * btarget ? (lo = mid) : (hi = mid) + end + return lo +end + +""" +Solve the single `B = btarget` crossing on `[θa, θb]`, a span on which B is monotone and +which does not cross the θ = 0/1 seam. Returns `(θ, cell)`; the cell feeds back as the +next λ's hint. `θ` is `NaN` when no root lies in the span. +""" +function _monotone_segment_root( + bf::SurfaceBField, θa::Float64, θb::Float64, + btarget::Float64, hint::Int +) + ia = _cell_index(bf.knot, θa) + ib = max(_cell_index(bf.knot, θb), ia) + ba = _b_at(bf.knot, bf.poly, θa) + bb = _b_at(bf.knot, bf.poly, θb) + + ic = ia == ib ? ia : _locate_cell(bf, ia, ib, ba, bb, btarget, hint) + root = _cell_level_root(bf, ic, btarget, max(θa, bf.knot[ic]), min(θb, bf.knot[ic+1])) + # A target landing within rounding of a knot can select the neighbouring cell. + isnan(root) && ic > ia && (root = _cell_level_root(bf, ic - 1, btarget, max(θa, bf.knot[ic-1]), min(θb, bf.knot[ic]))) + isnan(root) && ic < ib && (root = _cell_level_root(bf, ic + 1, btarget, max(θa, bf.knot[ic+1]), min(θb, bf.knot[ic+2]))) + return root, ic end """ -Bounce points of `v_par(θ) = 1 − (λ/bo)·B(θ)` for a trapped particle, found by -bracketed solves on the monotone intervals between stationary points of `B`. -`B = bo/λ` at a bounce point, so an interval brackets one iff `bo/λ` lies strictly -between its endpoint `B` values — a scalar test against the cached values, with no -spline evaluation per λ. Returns roots sorted descending, matching the order the -deepest-well and marginally-trapped logic downstream assumes. An empty or -single-element result signals a degenerate λ and sends the caller to the fallback. +Bounce points of `v_par(θ) = 1 − (λ/bo)·B(θ)` for a trapped particle. `B = bo/λ` at a +bounce point, so a monotone interval between stationary points of B contains one iff +`bo/λ` lies strictly between its endpoint B values — a scalar test against cached +values. The crossing is then solved directly from the cell's cubic coefficients: no +iteration, no spline evaluation, and no blind search, since `hints` carries each +interval's cell over from the previous λ. + +Returns roots sorted descending, the order the deepest-well and marginally-trapped logic +downstream assumes; fewer than two roots signals a degenerate λ and sends the caller to +the fallback. `hints` is sized `length(bf.theta) + 1`, the extra slot being the far half +of the interval that wraps through the seam. """ -function _bounce_points_from_extrema!(bpts::Vector{Float64}, ext::BFieldExtrema, - lmda::Float64, bo::Float64, B_vpar) +function _bounce_points_at_lambda!( + bpts::Vector{Float64}, hints::Vector{Int}, + bf::SurfaceBField, lmda::Float64, bo::Float64 +) empty!(bpts) - k = length(ext.theta) + k = length(bf.theta) k < 2 && return bpts btarget = bo / lmda for i in 1:k j = i == k ? 1 : i + 1 - fa = ext.bval[i] - btarget - fb = ext.bval[j] - btarget - fa * fb < 0 || continue - # The last interval wraps through the θ = 0/1 seam; _vpar_from_spline takes - # mod(θ, 1), so B is continuous across it for the periodic fit. - ta = ext.theta[i] - tb = i == k ? ext.theta[1] + 1.0 : ext.theta[j] - root = Roots.find_zero(θ -> _vpar_from_spline(B_vpar, lmda, bo, θ), (ta, tb), Roots.Brent()) + (bf.bval[i] - btarget) * (bf.bval[j] - btarget) < 0 || continue + θa, θb = bf.theta[i], bf.theta[j] + + root = NaN + if i < k + root, hints[i] = _monotone_segment_root(bf, θa, θb, btarget, hints[i]) + else + # The last interval wraps through θ = 0/1. Split it at the seam, where B is + # continuous for the periodic fit, and solve whichever half spans btarget. + if (bf.bval[i] - btarget) * (bf.bknot[end] - btarget) <= 0 + root, hints[k] = _monotone_segment_root(bf, θa, 1.0, btarget, hints[k]) + else + root, hints[k+1] = _monotone_segment_root(bf, 0.0, θb, btarget, hints[k+1]) + end + end + + isnan(root) && return empty!(bpts) push!(bpts, mod(root, 1.0)) end @@ -538,13 +747,13 @@ Returns (t1, t2, theta_points, theta_weights). function _find_bounce_points_and_grid( lmda::Float64, bo::Float64, sigma::Int, B_vpar, theta_bmax::Float64, psi::Float64, - ntheta::Int, ext::BFieldExtrema, bpts_buf::Vector{Float64} + ntheta::Int, bf::SurfaceBField, bpts_buf::Vector{Float64}, hints::Vector{Int} ) if sigma == 0 # trapped # Bounce points: all roots of v_par(θ) = 1 − (λ/bo)·B_vpar(θ) in (0,1), # sorted descending — the same order as Fortran spline_roots, which the # marginally-trapped and deepest-well wrap logic below assume. - bpts = _bounce_points_from_extrema!(bpts_buf, ext, lmda, bo, B_vpar) + bpts = _bounce_points_at_lambda!(bpts_buf, hints, bf, lmda, bo) if length(bpts) < 2 # Degenerate λ: v_par is tangent to zero at an extremum of B (so no diff --git a/test/runtests_kinetic.jl b/test/runtests_kinetic.jl index 8a96266c8..d502385f7 100644 --- a/test/runtests_kinetic.jl +++ b/test/runtests_kinetic.jl @@ -144,18 +144,18 @@ B_vpar = cubic_interp(xs, B_vals; bc=PeriodicBC()) bo = 2.0 - ext = KF._b_field_extrema(B_vpar) + bf = KF._surface_b_field(B_vpar) @testset "stationary points are exact" begin # A periodic function has an even number of stationary points per period. - @test iseven(length(ext.theta)) - @test length(ext.theta) >= 2 - @test issorted(ext.theta) - @test all(0.0 .<= ext.theta .< 1.0) + @test iseven(length(bf.theta)) + @test length(bf.theta) >= 2 + @test issorted(bf.theta) + @test all(0.0 .<= bf.theta .< 1.0) # dB/dθ vanishes at each, and the cached B values match the spline. dB = deriv1(B_vpar) scale = maximum(abs, B_vals) - for (θ, b) in zip(ext.theta, ext.bval) + for (θ, b) in zip(bf.theta, bf.bval) @test abs(dB(θ)) < 1e-8 * scale @test b ≈ B_vpar(θ) rtol=1e-14 end @@ -163,14 +163,35 @@ dense = range(0.0, 1.0; length=20001) dv = [dB(x) for x in dense] nsign = count(i -> dv[i] * dv[i+1] < 0, 1:length(dv)-1) - @test length(ext.theta) == nsign + @test length(bf.theta) == nsign end - bmax = maximum(ext.bval) - bmin = minimum(ext.bval) + bmax = maximum(bf.bval) + bmin = minimum(bf.bval) lmdatpb = bo / bmax lmdamax = bo / bmin + @testset "closed-form cubic roots" begin + # Polynomials with known roots, including the repeated-root case where the + # discriminant cancels to zero and a naive branch reports only the simple root. + for (a, b, c, d, expect) in ((1.0, -6.0, 11.0, -6.0, [1.0, 2.0, 3.0]), # (u-1)(u-2)(u-3) + (1.0, 0.0, -1.0, 0.0, [-1.0, 0.0, 1.0]), # u³ - u + (2.0, -4.0, 2.0, 0.0, [0.0, 1.0]), # 2u(u-1)², double root + (1.0, 0.0, 0.0, -8.0, [2.0]), # u³ - 8, one real root + (0.0, 1.0, -3.0, 2.0, [1.0, 2.0]), # degenerates to a quadratic + (0.0, 0.0, 2.0, -4.0, [2.0])) # degenerates to a linear + n, r1, r2, r3 = KF._real_cubic_roots(a, b, c, d) + got = filter(isfinite, [r1, r2, r3][1:min(n, 3)]) + for e in expect + @test any(abs.(got .- e) .< 1e-9) + end + # Every returned root satisfies the polynomial. + for u in got + @test abs(evalpoly(u, (d, c, b, a))) < 1e-9 * max(1.0, abs(u)^3) + end + end + end + @testset "agrees with the adaptive scan across λ" begin # Over the bulk of the trapped range the bracketed solver must reproduce # Roots.find_zeros exactly. λ crowded against the boundaries is asserted @@ -178,8 +199,9 @@ # close that whether the adaptive scan still resolves them depends on its # own step heuristics, which is not a contract worth pinning a test to. buf = Float64[] + hints = ones(Int, length(bf.theta) + 1) for lmda in range(lmdatpb * (1 + 1e-6), lmdamax * (1 - 1e-6); length=200) - new_roots = copy(KF._bounce_points_from_extrema!(buf, ext, lmda, bo, B_vpar)) + new_roots = copy(KF._bounce_points_at_lambda!(buf, hints, bf, lmda, bo)) ref_roots = sort!(KF.Roots.find_zeros(θ -> KF._vpar_from_spline(B_vpar, lmda, bo, θ), 0.0, 1.0); rev=true) @test length(new_roots) == length(ref_roots) if length(new_roots) == length(ref_roots) @@ -189,7 +211,9 @@ @test issorted(new_roots; rev=true) # v_par vanishes at each root. for θ in new_roots - @test abs(KF._vpar_from_spline(B_vpar, lmda, bo, θ)) < 1e-10 + # Solved directly from the cell's cubic, so v_par vanishes to + # machine precision rather than to a solver tolerance. + @test abs(KF._vpar_from_spline(B_vpar, lmda, bo, θ)) < 1e-13 end end end @@ -200,8 +224,9 @@ # solver: exactly two roots, one on each side of the extremum, both closing # on it as λ tightens, with v_par vanishing at each. buf = Float64[] - θmax = ext.theta[argmax(ext.bval)] - θmin = ext.theta[argmin(ext.bval)] + hints = ones(Int, length(bf.theta) + 1) + θmax = bf.theta[argmax(bf.bval)] + θmin = bf.theta[argmin(bf.bval)] # Width of the arc between the two roots that encloses θx. θ is periodic, # so the enclosing arc may run through the θ = 0/1 seam — the B maximum of @@ -214,7 +239,7 @@ prev = Dict(:max => Inf, :min => Inf) for eps in (1e-4, 1e-6, 1e-8, 1e-10, 1e-12) for (key, lmda, θx) in ((:max, lmdatpb * (1 + eps), θmax), (:min, lmdamax * (1 - eps), θmin)) - roots = copy(KF._bounce_points_from_extrema!(buf, ext, lmda, bo, B_vpar)) + roots = copy(KF._bounce_points_at_lambda!(buf, hints, bf, lmda, bo)) @test length(roots) == 2 length(roots) == 2 || continue for θ in roots @@ -232,12 +257,28 @@ @test prev[:min] < 0.5 * (xs[2] - xs[1]) end + @testset "cell hints never change the answer" begin + # The hint only says where to start looking. A stale or absurd one must cost + # a few extra steps, or a fall back to bisection, and nothing else. + buf = Float64[] + nh = length(bf.theta) + 1 + swept = ones(Int, nh) + for lmda in range(lmdatpb * (1 + 1e-6), lmdamax * (1 - 1e-6); length=120) + want = copy(KF._bounce_points_at_lambda!(buf, swept, bf, lmda, bo)) + for bogus in (ones(Int, nh), fill(length(bf.poly), nh), fill(-5, nh), fill(10^6, nh)) + got = copy(KF._bounce_points_at_lambda!(buf, copy(bogus), bf, lmda, bo)) + @test got == want + end + end + end + @testset "degenerate λ signals the fallback" begin # λ below the trapped-passing boundary puts bo/λ above every B, so no # interval brackets and the caller is told to fall back. buf = Float64[] - bmax = maximum(ext.bval) - @test length(KF._bounce_points_from_extrema!(buf, ext, 0.5 * bo / bmax, bo, B_vpar)) < 2 + hints = ones(Int, length(bf.theta) + 1) + bmax = maximum(bf.bval) + @test length(KF._bounce_points_at_lambda!(buf, hints, bf, 0.5 * bo / bmax, bo)) < 2 end end