Merge remote-tracking branch 'origin/develop' into performance/consistent-surface-theta-parametrization

Author: logan-nc

src/KineticForces/BounceAveraging.jl

@@ -147,7 +147,7 @@ end
 
 """
     compute_bounce_data(psi, n, l, q, bo, bmax, bmin, theta_bmax,
-                        tspl, B_extrap, mfac, chi1, ro, dbob_m_f, divx_m_f,
+                        tspl, B_vpar, mfac, chi1, ro, dbob_m_f, divx_m_f,
                         divxfac, wdfac, mass, chrg, T_s, method;
                         nlmda=128, ntheta=128,
                         smat=nothing, tmat=nothing, xmat=nothing,
@@ -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_vpar`: 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]
@@ -190,7 +190,7 @@ function compute_bounce_data(
     psi::Float64, n::Int, l::Int, q::Float64,
     bo::Float64, bmax::Float64, bmin::Float64,
     theta_bmax::Float64,
-    tspl, B_extrap, mfac::Vector{Int}, chi1::Float64, ro::Float64,
+    tspl, B_vpar, mfac::Vector{Int}, chi1::Float64, ro::Float64,
     dbob_m_f::Vector{ComplexF64}, divx_m_f::Vector{ComplexF64},
     divxfac::Float64, wdfac::Float64,
     mass::Float64, chrg::Float64,
@@ -241,12 +241,12 @@ function compute_bounce_data(
 
         # Find bounce points and build θ sub-grid
         _, _, tdt_pts, tdt_wts = _find_bounce_points_and_grid(
-            lmda, bo, sigma, B_extrap, theta_bmax, psi, ntheta)
+            lmda, bo, sigma, B_vpar, theta_bmax, psi, ntheta)
 
         # Bounce integrals over θ (Fortran lines 674-735)
         wbbar, wdbar, dJdJ_val, wmats_lmda = _bounce_integrate(
             tdt_pts, tdt_wts, lmda, lnq, sigma, n, q, bo,
-            tspl, B_extrap, chi1, ro, mfac, dbob_m_f, divx_m_f, divxfac, wdfac,
+            tspl, B_vpar, chi1, ro, mfac, dbob_m_f, divx_m_f, divxfac, wdfac,
             do_matrices, mpert, smat, tmat, xmat, ymat, zmat, scr)
 
         # Physical frequencies (Fortran lines 744-745)
@@ -409,12 +409,12 @@ end
 
 
 """
-Parallel-velocity factor `v_par = 1 − (λ/bo)·B(θ)` from the endpoint-fit cubic of B
-(`B_extrap`, built where the surface interpolants are constructed), keeping v_par
+Parallel-velocity factor `v_par = 1 − (λ/bo)·B(θ)` from the periodic cubic of B
+(`B_vpar`, 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))
+@inline _vpar_from_spline(B_vpar, lmda::Float64, bo::Float64, θ::Float64) =
+    1.0 - (lmda / bo) * B_vpar(mod(θ, 1.0))
 
 
 """
@@ -423,14 +423,14 @@ Returns (t1, t2, theta_points, theta_weights).
 """
 function _find_bounce_points_and_grid(
     lmda::Float64, bo::Float64, sigma::Int,
-    B_extrap, theta_bmax::Float64, psi::Float64,
+    B_vpar, theta_bmax::Float64, psi::Float64,
     ntheta::Int
 )
     if sigma == 0  # trapped
-        # Bounce points: all roots of v_par(θ) = 1 − (λ/bo)·B_extrap(θ) in (0,1),
+        # 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_extrap(B_extrap, lmda, bo, θ)
+        vpar_fn = θ -> _vpar_from_spline(B_vpar, lmda, bo, θ)
         bpts = sort!(Roots.find_zeros(vpar_fn, 0.0, 1.0); rev=true)
 
         nbpts = length(bpts)
@@ -443,7 +443,7 @@ function _find_bounce_points_and_grid(
             t1 = bpts[1]
             t2 = bpts[1] + 1.0
         else
-            t1, t2 = _find_deepest_well(bpts, B_extrap, lmda, bo)
+            t1, t2 = _find_deepest_well(bpts, B_vpar, lmda, bo)
         end
 
         # Power-law grid refined near bounce points
@@ -463,7 +463,7 @@ end
 Find the deepest potential well (largest midpoint v_par) among bounce-point pairs,
 handling pairs that wrap through θ = 0/1.
 """
-function _find_deepest_well(bpts::Vector{Float64}, B_extrap, lmda::Float64, bo::Float64)
+function _find_deepest_well(bpts::Vector{Float64}, B_vpar, lmda::Float64, bo::Float64)
     nbpts = length(bpts)
     best_vpar = 0.0
     best_t1 = 0.0
@@ -477,7 +477,7 @@ function _find_deepest_well(bpts::Vector{Float64}, B_extrap, lmda::Float64, bo::
         else
             θmid = 0.5 * (bpts[i] + bpts[j])
         end
-        vpar_mid = _vpar_from_extrap(B_extrap, lmda, bo, θmid)
+        vpar_mid = _vpar_from_spline(B_vpar, lmda, bo, θmid)
         if vpar_mid > best_vpar
             best_t1 = bpts[i]
             best_t2 = bpts[j]
@@ -506,7 +506,7 @@ Ports Fortran torque.F90 lines 674-793.
 function _bounce_integrate(
     tdt_pts::Vector{Float64}, tdt_wts::Vector{Float64},
     lmda::Float64, lnq::Float64, sigma::Int, n::Int, q::Float64, bo::Float64,
-    tspl, B_extrap, chi1::Float64, ro::Float64,
+    tspl, B_vpar, chi1::Float64, ro::Float64,
     mfac::Vector{Int}, dbob_m_f::Vector{ComplexF64}, divx_m_f::Vector{ComplexF64},
     divxfac::Float64, wdfac::Float64,
     do_matrices::Bool, mpert::Int,
@@ -550,9 +550,9 @@ 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)
+        vpar = 1.0 - (lmda / bo) * B_vpar(θmod)
 
         if vpar <= 0
             # Negative v_par near a bounce point: same fill rules as the Fortran bounce loop.
@@ -632,9 +632,13 @@ function _bounce_integrate(
         return 0.0, 0.0, 0.0, nothing
     end
 
-    # Bounce-averaged frequencies
+    # Bounce-averaged frequencies. wbbar already carries one factor of ro that its own
+    # normalization bhat = sqrt(2T/m)/ro cancels; reusing it inside wdbar imports that ro
+    # a third time while dhat = (T/q)/(bo·ro²) removes only the two written explicitly, so
+    # the drift prefactor takes ro, not ro². (Otherwise ω_D = wdbar·dhat carries a surplus
+    # length: 4π·(I₂/I₁)·(T/q) is already V/Wb = 1/s, so the extra ro leaves m/s.)
     wbbar = ro * twopi / ((2 - sigma) * total_wb)
-    wdbar = ro^2 * bo * wdfac * wbbar * 2 * (2 - sigma) * total_wd
+    wdbar = ro * bo * wdfac * wbbar * 2 * (2 - sigma) * total_wd
 
     # Phase factor pl_i = exp(-2πi·lnq·fsi_wb(θ_i)/((2-σ)·total_wb)), using the
     # cumulative spline integral of the bounce action.

src/KineticForces/CalculatedKineticMatrices.jl

@@ -136,7 +136,9 @@ function compute_calculated_kinetic_matrices(
                         z_s, m_s, kf_ctrl.wdfac, kf_ctrl.divxfac,
                         el_s, equil, intr_t, prof_s;
                         nutype=kf_ctrl.nutype, f0type=kf_ctrl.f0type, nufac=kf_ctrl.nufac,
-                        atol_xlmda=kf_ctrl.atol_xlmda, rtol_xlmda=kf_ctrl.rtol_xlmda
+                        atol_xlmda=kf_ctrl.atol_xlmda, rtol_xlmda=kf_ctrl.rtol_xlmda,
+                        atol_x=kf_ctrl.atol_x, rtol_x=kf_ctrl.rtol_x,
+                        nested_tolerance_margin=kf_ctrl.nested_tolerance_margin
                     )
                     full_w .+= block_w
                     full_t .+= block_t

src/KineticForces/Compute.jl

@@ -137,7 +137,9 @@ function integrate_psi_quadgk(
                 tpsi!(thread_tpsi[tid], psi, n, l, zi, mi, wdfac, divxfac,
                     electron, method, equil, thread_intrs[tid], kinetic_profiles;
                     op_wmats=w,
-                    atol_xlmda=ctrl.atol_xlmda, rtol_xlmda=ctrl.rtol_xlmda)
+                    atol_xlmda=ctrl.atol_xlmda, rtol_xlmda=ctrl.rtol_xlmda,
+                    atol_x=ctrl.atol_x, rtol_x=ctrl.rtol_x,
+                    nested_tolerance_margin=ctrl.nested_tolerance_margin)
                 harm_vals[ell_idx] = thread_tpsi[tid][]
                 is_matrix_method && (harm_elems[ell_idx] .= w)
             end

src/KineticForces/KineticForcesStructs.jl

@@ -121,11 +121,20 @@ builds a second control for its differing tolerance).
     nn::Int = 1                     # Toroidal mode number
     nl::Int = 1                     # Bounce harmonic number
 
-    # Tolerances.
-    # *_xlmda: shared tolerances for inner λ (pitch) and x (energy) integrations
-    # *_psi:   tolerances for outer ψ quadrature
-    atol_xlmda::Float64 = 1e-8     # Absolute tolerance for inner pitch + energy integrations
-    rtol_xlmda::Float64 = 1e-5     # Relative tolerance for inner pitch + energy integrations
+    # Tolerances, outermost to innermost: ψ quadrature ⊃ λ (pitch) ⊃ x (energy).
+    # Each level must be resolved more tightly than the one enclosing it, or the outer
+    # integrator chases its integrand's own quadrature noise instead of converging.
+    # *_xlmda: tolerances for the λ (pitch) integration
+    # *_x:     tolerances for the x (energy) integration nested inside it; NaN ⇒ derive as
+    #          nested_tolerance_margin × the pitch tolerances
+    # *_psi:   tolerances for the outer ψ quadrature
+    atol_xlmda::Float64 = 1e-8     # Absolute tolerance for the inner pitch integration
+    rtol_xlmda::Float64 = 1e-5     # Relative tolerance for the inner pitch integration
+    atol_x::Float64 = NaN          # Absolute tolerance for the energy integration (NaN ⇒ derived)
+    rtol_x::Float64 = NaN          # Relative tolerance for the energy integration (NaN ⇒ derived)
+    # The pitch integrand IS the energy integral, so the energy level is resolved this much
+    # tighter than the pitch level by default.
+    nested_tolerance_margin::Float64 = 1e-2   # Factor relating derived energy tolerances to the pitch ones
     # rtol_psi is the primary convergence knob: ~2 significant figures matches the validity
     # of the NTV model approximations. Do not set it tighter than the noise floor of the
     # inner integrals (keep rtol_psi ≳ 10 × rtol_xlmda).

src/KineticForces/Torque.jl

@@ -35,7 +35,14 @@ function tpsi!(tpsi_var::Ref{ComplexF64}, psi::Float64, n::Int, l::Int,
               op_wmats::Union{Nothing,Array{ComplexF64,3}}=nothing,
               rex_override::Union{Nothing,Float64}=nothing,
               imx_override::Union{Nothing,Float64}=nothing,
-              atol_xlmda::Float64=1e-9, rtol_xlmda::Float64=1e-6)
+              atol_xlmda::Float64=1e-9, rtol_xlmda::Float64=1e-6,
+              atol_x::Float64=NaN, rtol_x::Float64=NaN,
+              nested_tolerance_margin::Float64=1e-2)
+
+    # The pitch integrand is itself the energy integral, so the energy level is resolved
+    # tighter than the pitch level (explicit atol_x/rtol_x override the derived values).
+    atol_energy = isnan(atol_x) ? nested_tolerance_margin * atol_xlmda : atol_x
+    rtol_energy = isnan(rtol_x) ? nested_tolerance_margin * rtol_xlmda : rtol_x
 
     # Enforce bounds
     if psi > 1
@@ -94,10 +101,12 @@ function tpsi!(tpsi_var::Ref{ComplexF64}, psi::Float64, n::Int, l::Int,
 
     # Create periodic interpolant for poloidal quantities
     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())
+    # v_par and the bounce points use a separate scalar cubic of B (Fortran's vspl).
+    # 1−(λ/bo)B is periodic on the closed θ interval, so the fit must be periodic too:
+    # a non-periodic endpoint fit is only C⁰ at the θ=0/1 seam and manufactures false
+    # near-seam extrema and root pairs there. The fit of 1−(λ/bo)B equals 1−(λ/bo)
+    # times the fit of B, so one B_vpar per surface serves every λ.
+    B_vpar = cubic_interp(xs, B_vals; bc=PeriodicBC())
 
     bmax = maximum(B_vals)
     ibmax = argmax(B_vals)
@@ -228,10 +237,10 @@ function tpsi!(tpsi_var::Ref{ComplexF64}, psi::Float64, n::Int, l::Int,
                                    method, op_wmats;
                                    chi1=intr.chi1, ro=intr.ro, mfac=intr.mfac,
                                    mpert=intr.mpert, theta_bmax=theta_bmax,
-                                   B_extrap=B_extrap,
+                                   B_vpar=B_vpar,
                                    smat=smat_f, tmat=tmat_f, xmat=xmat_f,
                                    ymat=ymat_f, zmat=zmat_f,
-                                   energy_atol=atol_xlmda, energy_rtol=rtol_xlmda,
+                                   energy_atol=atol_energy, energy_rtol=rtol_energy,
                                    pitch_atol=atol_xlmda, pitch_rtol=rtol_xlmda,
                                    rex_override=rex_override, imx_override=imx_override)
     end
@@ -403,7 +412,7 @@ function calculate_gar(psi, n, l, q, epsr, wdian, wdiat, welec, nuk, bo,
                        bmax, bmin, n_s::Float64, T_s::Float64, mass, chrg, tspl,
                        dbob_m_f, divx_m_f, divxfac, wdfac, method, op_wmats;
                        chi1::Float64, ro::Float64, mfac::Vector{Int}, mpert::Int,
-                       theta_bmax::Float64, B_extrap,
+                       theta_bmax::Float64, B_vpar,
                        smat=nothing, tmat=nothing, xmat=nothing,
                        ymat=nothing, zmat=nothing,
                        nlmda::Int=128, ntheta::Int=128,
@@ -417,7 +426,7 @@ function calculate_gar(psi, n, l, q, epsr, wdian, wdiat, welec, nuk, bo,
     # Bounce-averaged quantities per pitch angle
     bounce = compute_bounce_data(
         psi, n, l, q, bo, bmax, bmin, theta_bmax,
-        tspl, B_extrap, mfac, chi1, ro, dbob_m_f, divx_m_f, divxfac, wdfac,
+        tspl, B_vpar, mfac, chi1, ro, dbob_m_f, divx_m_f, divxfac, wdfac,
         mass, chrg, T_s, method;
         nlmda, ntheta, smat, tmat, xmat, ymat, zmat)
 
@@ -617,8 +626,9 @@ 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())
+    # Periodic cubic of B for v_par and bounce points (Fortran vspl equivalent); 1−(λ/bo)B
+    # is periodic on the closed θ interval, so a non-periodic fit would break at the seam.
+    B_vpar = cubic_interp(xs, B_vals; bc=PeriodicBC())
 
     bmax = maximum(B_vals)
     ibmax = argmax(B_vals)
@@ -687,7 +697,7 @@ function _setup_surface_state(
 
     return (;
         chrg, mass,
-        tspl, B_extrap, bmax, bmin, theta_bmax,
+        tspl, B_vpar, bmax, bmin, theta_bmax,
         q, n_s, T_s, welec,
         wdian, wdiat, wtran, wgyro, nuk,
         epsr,
@@ -752,7 +762,7 @@ function kinetic_energy_matrices_for_euler_lagrange!(
 
     bounce = compute_bounce_data(
         psi, n, l, state.q, bo, state.bmax, state.bmin, state.theta_bmax,
-        state.tspl, state.B_extrap, mfac, chi1, ro, dbob_m_f, divx_m_f, 1.0, wdfac,
+        state.tspl, state.B_vpar, mfac, chi1, ro, dbob_m_f, divx_m_f, 1.0, wdfac,
         state.mass, state.chrg, state.T_s, "fwmm";
         nlmda, ntheta, smat=smat_f, tmat=tmat_f, xmat=xmat_f, ymat=ymat_f, zmat=zmat_f)
 
@@ -894,7 +904,9 @@ function compute_kinetic_matrices_at_psi!(
     electron::Bool, equil, intr::KineticForcesInternal,
     kinetic_profiles::Equilibrium.KineticProfileSplines;
     nutype::String="harmonic", f0type::String="maxwellian", nufac::Float64=1.0,
-    atol_xlmda::Float64=1e-9, rtol_xlmda::Float64=1e-6)
+    atol_xlmda::Float64=1e-9, rtol_xlmda::Float64=1e-6,
+    atol_x::Float64=NaN, rtol_x::Float64=NaN,
+    nested_tolerance_margin::Float64=1e-2)
 
     # Bypass ψ > 1 (no kinetic contribution outside plasma)
     if psi > 1
@@ -903,13 +915,17 @@ function compute_kinetic_matrices_at_psi!(
         return nothing
     end
 
+    # See tpsi!: the energy integral is the pitch integrand, so it is resolved tighter.
+    atol_energy = isnan(atol_x) ? nested_tolerance_margin * atol_xlmda : atol_x
+    rtol_energy = isnan(rtol_x) ? nested_tolerance_margin * rtol_xlmda : rtol_x
+
     state = _setup_surface_state(psi, zi, mi, electron,
                                   equil, intr, kinetic_profiles)
 
     kinetic_energy_matrices_for_euler_lagrange!(
         kwmat, ktmat, state, psi, n, l, wdfac, intr;
         nutype, f0type, nufac,
-        energy_atol=atol_xlmda, energy_rtol=rtol_xlmda,
+        energy_atol=atol_energy, energy_rtol=rtol_energy,
         pitch_atol=atol_xlmda, pitch_rtol=rtol_xlmda)
 
     return nothing