KineticForces: preserve periodic vpar spline and enumerate bounce roots exactly

src/KineticForces/BounceAveraging.jl

@@ -168,8 +168,8 @@ Ports Fortran torque.F90 lines 530-816 (GAR branch).
 - `bmax, bmin`: Max/min of B(θ) at this ψ
 - `theta_bmax`: θ location of Bmax (nodal knot; the passing-transit start)
 - `tspl`: Periodic poloidal interpolant: tspl(θ) → [B, dB/dψ, dB/dθ, J, dJ/dψ]
-- `B_extrap`: Endpoint-fit (non-periodic) cubic of B(θ) used for v_par and the
-  bounce-point roots (the Fortran `vspl` equivalent)
+- `B_extrap`: Periodic cubic of B(θ) used for v_par and the bounce-point roots
+  (the Fortran `vspl` equivalent)
 - `mfac`: Poloidal mode numbers [mlow:mhigh]
 - `chi1`: 2π·ψ₀ flux normalization
 - `ro`: Major radius [m]
@@ -409,14 +409,106 @@ end
 
 
 """
-Parallel-velocity factor `v_par = 1 − (λ/bo)·B(θ)` from the endpoint-fit cubic of B
+Fit the periodic magnetic-field spline used by both the parallel-velocity
+factor and its bounce-point roots.
+
+The sampled field closes at θ=0/1. A wrapped non-periodic endpoint fit is only
+C⁰ at that seam and can create false near-seam extrema and root pairs.
+"""
+@inline _fit_vpar_B_spline(xs, B_vals) =
+    cubic_interp(xs, B_vals; bc=PeriodicBC())
+
+
+"""
+Parallel-velocity factor `v_par = 1 − (λ/bo)·B(θ)` from the periodic cubic of B
 (`B_extrap`, built where the surface interpolants are constructed), keeping v_par
 consistent with the bounce-point roots as in Fortran's `vspl`.
 """
 @inline _vpar_from_extrap(B_extrap, lmda::Float64, bo::Float64, θ::Float64) =
     1.0 - (lmda / bo) * B_extrap(mod(θ, 1.0))
 
 
+"""
+Return every distinct bounce root on one closed poloidal turn in the descending
+order expected by the well-selection logic.
+
+Each knot interval is a cubic polynomial. Splitting it at its analytically
+computed stationary points makes every resulting interval monotone, so close
+root pairs cannot be skipped by a coarse heuristic scan. The inclusive
+endpoints represent one physical point, so a seam root is retained only at
+θ=0.
+"""
+@inline function _find_bounce_roots(B_extrap, lmda::Float64, bo::Float64)
+    vpar_fn = θ -> _vpar_from_extrap(B_extrap, lmda, bo, θ)
+    roots = Float64[]
+    xs = B_extrap.cache.x
+    ys = B_extrap.y
+    zs = B_extrap.z
+    value_tol = 64eps(Float64) * max(1.0, abs(lmda / bo) * maximum(abs, ys))
+    root_tol = 64eps(Float64)
+
+    function push_distinct!(root)
+        if isempty(roots) || abs(root - last(roots)) > root_tol
+            push!(roots, root)
+        end
+    end
+
+    function scan_monotone_interval(left, right)
+        fleft = vpar_fn(left)
+        fright = vpar_fn(right)
+        if abs(fleft) <= value_tol
+            push_distinct!(left)
+        elseif abs(fright) > value_tol && signbit(fleft) != signbit(fright)
+            push_distinct!(Roots.find_zero(vpar_fn, (left, right), Roots.Brent()))
+        end
+    end
+
+    for i in firstindex(xs):(lastindex(xs)-1)
+        left = xs[i]
+        right = xs[i+1]
+        h = right - left
+        zleft = zs[i]
+        zright = zs[i+1]
+        c1 = (ys[i+1] - ys[i]) / h - h * (2zleft + zright) / 6
+        c2 = zleft / 2
+        c3 = (zright - zleft) / (6h)
+
+        stationary = Float64[]
+        qa = 3c3
+        qb = 2c2
+        qc = c1
+        if iszero(qa)
+            if !iszero(qb)
+                push!(stationary, -qc / qb)
+            end
+        else
+            discriminant = muladd(-4qa, qc, qb^2)
+            if discriminant >= 0
+                sqrt_discriminant = sqrt(discriminant)
+                q = -0.5 * (qb + copysign(sqrt_discriminant, qb))
+                if iszero(q)
+                    push!(stationary, -qb / (2qa))
+                else
+                    push!(stationary, q / qa, qc / q)
+                end
+            end
+        end
+        filter!(u -> 0 < u < h, stationary)
+        sort!(unique!(stationary))
+
+        subleft = left
+        for u in stationary
+            subright = left + u
+            scan_monotone_interval(subleft, subright)
+            subleft = subright
+        end
+        scan_monotone_interval(subleft, right)
+    end
+
+    return sort!(roots; rev=true)
+end
+
+
 """
 Find bounce points for trapped/passing particles and build θ sub-grid.
 Returns (t1, t2, theta_points, theta_weights).
@@ -430,8 +522,7 @@ function _find_bounce_points_and_grid(
         # Bounce points: all roots of v_par(θ) = 1 − (λ/bo)·B_extrap(θ) 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_extrap(B_extrap, lmda, bo, θ)
-        bpts = sort!(Roots.find_zeros(vpar_fn, 0.0, 1.0); rev=true)
+        bpts = _find_bounce_roots(B_extrap, lmda, bo)
 
         nbpts = length(bpts)
         if nbpts < 1
@@ -550,7 +641,7 @@ function _bounce_integrate(
         jac = tspl_f[4]
         djdpsi = tspl_f[5]
 
-        # v_par from the endpoint-fit cubic (consistent with the bounce points);
+        # v_par from the periodic cubic (consistent with the bounce points);
         # the periodic tspl B_val remains the numerator field in the integrands.
         vpar = 1.0 - (lmda / bo) * B_extrap(θmod)
 

src/KineticForces/Torque.jl

@@ -92,12 +92,12 @@ function tpsi!(tpsi_var::Ref{ComplexF64}, psi::Float64, n::Int, l::Int,
         djdpsi_vals[i+1] = equil.rzphi_jac(pt; deriv=DerivOp(1, 0), hint=hJ) / intr.chi1^2
     end
 
-    # Create periodic interpolant for poloidal quantities
+    # Create periodic interpolants for poloidal quantities. v_par and its
+    # bounce-point roots must use the same periodic fit of B: wrapping a
+    # non-periodic endpoint fit preserves B(0)=B(1), but leaves derivative
+    # jumps at the physical seam.
     tspl = cubic_interp(xs, Series(hcat(B_vals, dBdpsi_vals, dBdtheta_vals, jac_vals, djdpsi_vals)); bc=PeriodicBC())
-    # v_par and the bounce points use a separate endpoint-fit (non-periodic) cubic
-    # of B, like Fortran's vspl. The endpoint fit of 1−(λ/bo)B equals 1−(λ/bo)
-    # times the fit of B, so one B_extrap per surface serves every λ.
-    B_extrap = cubic_interp(xs, B_vals; bc=CubicFit())
+    B_extrap = _fit_vpar_B_spline(xs, B_vals)
 
     bmax = maximum(B_vals)
     ibmax = argmax(B_vals)
@@ -615,8 +615,7 @@ function _setup_surface_state(
     end
 
     tspl = cubic_interp(xs, Series(hcat(B_vals, dBdpsi_vals, dBdtheta_vals, jac_vals, djdpsi_vals)); bc=PeriodicBC())
-    # Endpoint-fit (non-periodic) cubic of B for v_par and bounce points (Fortran vspl equivalent).
-    B_extrap = cubic_interp(xs, B_vals; bc=CubicFit())
+    B_extrap = _fit_vpar_B_spline(xs, B_vals)
 
     bmax = maximum(B_vals)
     ibmax = argmax(B_vals)
@@ -912,5 +911,3 @@ function compute_kinetic_matrices_at_psi!(
 
     return nothing
 end
-
-

test/runtests_kinetic.jl

@@ -3,6 +3,63 @@
 
     KF = GeneralizedPerturbedEquilibrium.KineticForces
 
+    # =========================================================================
+    # Periodic parallel-velocity spline
+    # =========================================================================
+    @testset "periodic parallel-velocity spline" begin
+        xs = collect(range(0.0, 1.0; length=13))
+        B_exact(θ) = 2.0 + 0.25cos(2π * θ) + 0.1sin(4π * θ)
+        B_vals = B_exact.(xs)
+        B_vpar = KF._fit_vpar_B_spline(xs, B_vals)
+
+        # The physical θ=0/1 seam is C². An endpoint-fit cubic has independent
+        # one-sided derivatives here, even when it is wrapped before evaluation.
+        @test B_vpar(0.0) ≈ B_vpar(1.0) atol=10eps()
+        @test B_vpar(0.0; deriv=DerivOp(1)) ≈
+              B_vpar(1.0; deriv=DerivOp(1)) atol=100eps()
+        @test B_vpar(0.0; deriv=DerivOp(2)) ≈
+              B_vpar(1.0; deriv=DerivOp(2)) atol=1000eps()
+
+        # This level lies above the analytic field maximum. A wrapped endpoint
+        # fit creates a false near-seam maximum and two spurious bounce roots.
+        artifact_level = 2.3038
+        dense_Bmax = maximum(B_exact, range(0.0, 1.0; length=100_001))
+        @test dense_Bmax < artifact_level
+        @test isempty(KF._find_bounce_roots(B_vpar, 1 / artifact_level, 1.0))
+
+        # Retain complete descending root semantics for a physical level that
+        # crosses the field twice.
+        roots = KF._find_bounce_roots(B_vpar, 1 / 2.2, 1.0)
+        @test length(roots) == 2
+        @test issorted(roots; rev=true)
+        @test all(abs(KF._vpar_from_extrap(B_vpar, 1 / 2.2, 1.0, θ)) < 1e-12
+                  for θ in roots)
+
+        # An isolated knot maximum has one crossing on either side by the
+        # intermediate value theorem, even when the pair is much narrower than
+        # a global heuristic root scan.
+        narrow_xs = collect(range(0.0, 1.0; length=257))
+        narrow_B = fill(2.0, length(narrow_xs))
+        narrow_B[101] = 2.1
+        narrow_spline = KF._fit_vpar_B_spline(narrow_xs, narrow_B)
+        narrow_level = 2.099
+        close_roots = KF._find_bounce_roots(narrow_spline, 1 / narrow_level, 1.0)
+        @test length(close_roots) == 2
+        @test close_roots[2] < narrow_xs[101] < close_roots[1]
+        @test all(abs(KF._vpar_from_extrap(
+                      narrow_spline, 1 / narrow_level, 1.0, θ
+                  )) < 1e-12 for θ in close_roots)
+
+        # Inclusive endpoints are the same physical point. A seam crossing
+        # therefore contributes one root, not separate roots at θ=0 and θ=1.
+        seam_roots = KF._find_bounce_roots(B_vpar, 1.0, B_vpar(0.0))
+        @test length(seam_roots) == 2
+        @test seam_roots[end] == 0.0
+        @test all(θ -> θ < 1.0, seam_roots)
+        @test all(abs(KF._vpar_from_extrap(B_vpar, 1.0, B_vpar(0.0), θ)) < 1e-12
+                  for θ in seam_roots)
+    end
+
     # =========================================================================
     # powspace grid generation
     # =========================================================================