diff --git a/src/KineticForces/BounceAveraging.jl b/src/KineticForces/BounceAveraging.jl index f1bb1fc32..c62ae8518 100644 --- a/src/KineticForces/BounceAveraging.jl +++ b/src/KineticForces/BounceAveraging.jl @@ -10,6 +10,14 @@ 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 +# Cells a hinted cell search may step before falling back to bisection. +const HINT_STEP_BUDGET = 8 + # ============================================================================ # BounceData struct # ============================================================================ @@ -208,6 +216,14 @@ function compute_bounce_data( # Per-surface scratch, reused across all λ. scr = BounceScratch(ntheta, mpert) + # 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 lmdamax = bo / bmin @@ -241,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) + 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( @@ -417,6 +433,313 @@ consistent with the bounce-point roots as in Fortran's `vspl`. 1.0 - (lmda / bo) * B_vpar(mod(θ, 1.0)) +""" + SurfaceBField + +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 +- `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 SurfaceBField + knot::Vector{Float64} + bknot::Vector{Float64} + poly::Vector{NTuple{4,Float64}} + theta::Vector{Float64} + bval::Vector{Float64} +end + +""" +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 _surface_b_field(B_vpar) + knot = Float64[0.0] + bknot = Float64[] + poly = NTuple{4,Float64}[] + theta = Float64[] + + x = 0.0 + while true + cell = coeffs(B_vpar, x) + h = cell.xR - cell.xL + 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 + 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 + 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 stationary point on a cell boundary is 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 + + 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. `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_at_lambda!( + bpts::Vector{Float64}, hints::Vector{Int}, + bf::SurfaceBField, lmda::Float64, bo::Float64 +) + empty!(bpts) + k = length(bf.theta) + k < 2 && return bpts + + btarget = bo / lmda + for i in 1:k + j = i == k ? 1 : i + 1 + (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 + + 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 +747,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, 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. - 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_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 + # 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 diff --git a/test/runtests_kinetic.jl b/test/runtests_kinetic.jl index 29955c2c7..d502385f7 100644 --- a/test/runtests_kinetic.jl +++ b/test/runtests_kinetic.jl @@ -132,6 +132,156 @@ 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 + + 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(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(bf.theta, bf.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(bf.theta) == nsign + end + + 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 + # 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[] + 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_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) + @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 + # 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 + + @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[] + 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 + # 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_at_lambda!(buf, hints, bf, lmda, bo)) + @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 "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[] + 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 + # ========================================================================= # Resonance root solver # =========================================================================