KineticForces - FEATURE - Per-bounce-harmonic torque density output

docs/development/hdf5-conventions.md

@@ -44,7 +44,7 @@ Top level (11 groups):
 | `LocalStability/` | Mercier `D_I`, resistive interchange `D_R`, `ballooning_Delta_prime` on `psi`; the ballooning α boundary on `ballooning_psi` |
 | `SingularSurfaces/` | Per-rational-surface data: `rational_psi`/`rational_q`/`rational_m`/`rational_n`, GGJ coefficients, `Delta_prime_matrix`/`Delta_prime_raw`/`Delta_coil`/`pest3_A`/`pest3_B`/`pest3_Gamma` (Riccati or Galerkin alike), the inner-layer `Match/` results, `Kinetic/` |
 | `PerturbedEquilibrium/` | `ForcingModes/`, `Response/`, `ResponseMatrices/`, `SingularCoupling/`, `Energies/`, control-surface spectra |
-| `KineticForces/` | `<method>/` (torque/energy profiles, `EnergyIntegrals/`, `KineticMatrices/`); multi-ion runs add `PerSpecies/<species>/<method>/` with the same per-method layout, summing to the top-level total |
+| `KineticForces/` | `<method>/` (torque/energy profiles, `dTdpsi_ell` resolved by bounce harmonic on the `ell` scale, `EnergyIntegrals/`, `KineticMatrices/`); multi-ion runs add `PerSpecies/<species>/<method>/` with the same per-method layout, summing to the top-level total |
 | `ErrorFields/` | `CoilSensitivities/` (per-coil-set control-surface spectra and their rigid shift/tilt derivatives, `DominantMode/` full-window projection); `MonteCarlo/` (intrinsic and corrected `\|δ\|` histograms over the sampled tolerances, per batch and averaged); `Risk/` (threshold density, `P(lock\|δ)`, locking probabilities, `ToleranceScan/`); `NTV/` (correction-coil overlap and NTV torque couplings per kAt) |
 | `Tearing/` | `PerSurface/` (+ `DpMatrix/`), `Roots/`, `LayerWidths/`, `Diagnostics/{ValidRoots,Poles,FilteredRoots}`, `Scan/Surface_<k>/` |
 | `SurfaceGeometries/` | `{Plasma,Wall}/{x,y,z}` point clouds |

src/KineticForces/Compute.jl

@@ -132,6 +132,7 @@ function integrate_psi_quadgk(
 
     logged_psi = Float64[]
     logged_dtdpsi = ComplexF64[]
+    logged_dtdpsi_ell = Vector{Vector{ComplexF64}}()
     logged_elems = is_matrix_method ? Vector{Array{ComplexF64,3}}() : nothing
 
     function psi_batch!(y::AbstractVector{ComplexF64}, x::AbstractVector)
@@ -160,17 +161,20 @@ function integrate_psi_quadgk(
             for ell_idx in 1:nharm
                 total += harm_vals[ell_idx]
             end
-            total *= kinetic_axis_validity_envelope(psi, psi_c)
+            envelope = kinetic_axis_validity_envelope(psi, psi_c)
+            total *= envelope
             y[k] = total
 
             push!(logged_psi, psi)
             push!(logged_dtdpsi, total)
+            # Same envelope as `total` so Σ_ℓ still reproduces it; the broadcast also copies the reused `harm_vals`.
+            push!(logged_dtdpsi_ell, harm_vals .* envelope)
             if is_matrix_method
                 elems_accum = zeros(ComplexF64, mpert, mpert, 6)
                 for ell_idx in 1:nharm
                     elems_accum .+= harm_elems[ell_idx]
                 end
-                elems_accum .*= kinetic_axis_validity_envelope(psi, psi_c)
+                elems_accum .*= envelope
                 push!(logged_elems, elems_accum)
             end
         end
@@ -190,6 +194,11 @@ function integrate_psi_quadgk(
     perm = sortperm(logged_psi)
     sorted_psi = logged_psi[perm]
     sorted_dtdpsi = logged_dtdpsi[perm]
+    # Per-ℓ profile, permuted by the SAME `perm`: QuadGK visits nodes in refinement order, not in
+    # ψ order, so a per-ℓ log left unsorted would misalign with `sorted_psi` without any error.
+    # `reduce(hcat, ...)` builds (nharm, npsi); transpose so ψ is the first axis, as `dtdpsi` is.
+    sorted_dtdpsi_ell = isempty(logged_dtdpsi_ell) ? zeros(ComplexF64, 0, 0) :
+                        permutedims(reduce(hcat, logged_dtdpsi_ell[perm]))
 
     # Cumulative trapezoidal integration for T(ψ) profile
     npts = length(sorted_psi)
@@ -202,7 +211,9 @@ function integrate_psi_quadgk(
         end
     end
 
-    torque_profile = npts > 1 ? (psi=sorted_psi, dtdpsi=sorted_dtdpsi, t_cumulative=t_cumulative) : nothing
+    torque_profile = npts > 1 ?
+                     (psi=sorted_psi, dtdpsi=sorted_dtdpsi, t_cumulative=t_cumulative,
+        dtdpsi_ell=sorted_dtdpsi_ell, ell=collect(-nl:nl)) : nothing
 
     # Trapezoidal integration of kinetic matrices over ψ (matrix methods only)
     matrix_integrated = nothing
@@ -230,6 +241,28 @@ end
 # High-level orchestration
 # ============================================================================
 
+"""
+    _accumulate_interp!(acc, grid, xs, ys)
+
+Add the linear interpolant of `(xs, ys)` onto `grid` into `acc`, contributing zero outside
+`[xs[1], xs[end]]`. `xs` must be sorted ascending. A duplicate node or the right endpoint takes
+`ys[j]` directly rather than dividing by a zero interval. Shared by the total and per-ℓ profiles in
+[`combine_species_states`](@ref) so both use identical interpolation.
+"""
+function _accumulate_interp!(acc, grid, xs::AbstractVector{Float64}, ys::AbstractVector{ComplexF64})
+    for (k, ψ) in enumerate(grid)
+        (ψ < xs[1] || ψ > xs[end]) && continue
+        j = searchsortedlast(xs, ψ)
+        if j == length(xs) || xs[j+1] == xs[j]   # endpoint or duplicate node: no interpolation
+            acc[k] += ys[j]
+        else
+            t = (ψ - xs[j]) / (xs[j+1] - xs[j])
+            acc[k] += (1 - t) * ys[j] + t * ys[j+1]
+        end
+    end
+    return acc
+end
+
 """
     combine_species_states(states) -> KineticForcesState
 
@@ -259,19 +292,30 @@ function combine_species_states(states::AbstractVector{KineticForcesState})
         for r in results
             length(r.psi_grid) >= 2 || continue
             o = sortperm(r.psi_grid)
-            xs = r.psi_grid[o]
-            ys = r.dtdpsi[o]
-            for (k, ψ) in enumerate(grid)
-                (ψ < xs[1] || ψ > xs[end]) && continue
-                j = searchsortedlast(xs, ψ)
-                if j == length(xs) || xs[j+1] == xs[j]   # endpoint or duplicate node: no interpolation
-                    dtdpsi[k] += ys[j]
-                else
-                    t = (ψ - xs[j]) / (xs[j+1] - xs[j])
-                    dtdpsi[k] += (1 - t) * ys[j] + t * ys[j+1]
+            _accumulate_interp!(dtdpsi, grid, r.psi_grid[o], r.dtdpsi[o])
+        end
+
+        # Per-ℓ profile, interpolated onto the same union grid column by column. Only combined
+        # when every species carries it on an identical ℓ range; otherwise left empty.
+        ell = first(results).ell
+        has_ell = !isempty(ell) && all(r -> r.ell == ell && size(r.dtdpsi_ell, 2) == length(ell) &&
+                                            size(r.dtdpsi_ell, 1) == length(r.psi_grid), results)
+        dtdpsi_ell = zeros(ComplexF64, has_ell ? length(grid) : 0, has_ell ? length(ell) : 0)
+        if has_ell
+            for r in results
+                length(r.psi_grid) >= 2 || continue
+                o = sortperm(r.psi_grid)
+                xs = r.psi_grid[o]
+                for j in eachindex(ell)
+                    _accumulate_interp!(view(dtdpsi_ell, :, j), grid, xs, r.dtdpsi_ell[o, j])
                 end
             end
+            # Take the total from the per-ℓ row sum so Σ_ℓ ≡ dtdpsi holds exactly rather than to
+            # interpolation roundoff. Linear interpolation commutes with the sum, so this is the
+            # same quantity the loop above produced.
+            dtdpsi = vec(sum(dtdpsi_ell; dims=2))
         end
+
         tcum = zeros(ComplexF64, length(grid))
         for j in 2:length(grid)
             tcum[j] = tcum[j-1] + 0.5 * (dtdpsi[j] + dtdpsi[j-1]) * (grid[j] - grid[j-1])
@@ -281,6 +325,7 @@ function combine_species_states(states::AbstractVector{KineticForcesState})
             total_torque=sum(r.total_torque for r in results),
             total_energy=sum(r.total_energy for r in results),
             psi_grid=grid, dtdpsi=dtdpsi, t_cumulative=tcum,
+            dtdpsi_ell=dtdpsi_ell, ell=(has_ell ? ell : Int[]),
             psi_nsteps=sum(r.psi_nsteps for r in results))
     end
     combined.completed = true
@@ -338,6 +383,8 @@ function compute_torque_all_methods!(state::KineticForcesState, intr::KineticFor
         psi_grid_out = Float64[]
         dtdpsi_out = ComplexF64[]
         t_cum_out = ComplexF64[]
+        dtdpsi_ell_out = zeros(ComplexF64, 0, 0)
+        ell_out = Int[]
         psi_nsteps_total = 0
         panel_psis_out = Float64[]
         resonance_psis_out = Float64[]
@@ -365,6 +412,11 @@ function compute_torque_all_methods!(state::KineticForcesState, intr::KineticFor
                     psi_grid_out = result.torque_profile.psi
                     dtdpsi_out = result.torque_profile.dtdpsi
                     t_cum_out = result.torque_profile.t_cumulative
+                    # Profiles come from the first n only, while `total_torque` sums over all n.
+                    # So sum(dtdpsi_ell; dims=2) == dtdpsi always, but neither integrates to
+                    # `total_torque` when npert > 1.
+                    dtdpsi_ell_out = result.torque_profile.dtdpsi_ell
+                    ell_out = result.torque_profile.ell
                 end
             end
 
@@ -387,6 +439,8 @@ function compute_torque_all_methods!(state::KineticForcesState, intr::KineticFor
             psi_grid=psi_grid_out,
             dtdpsi=dtdpsi_out,
             t_cumulative=t_cum_out,
+            dtdpsi_ell=dtdpsi_ell_out,
+            ell=ell_out,
             psi_nsteps=psi_nsteps_total,
             panel_psis=panel_psis_out,
             resonance_psis=resonance_psis_out

src/KineticForces/KineticForcesStructs.jl

@@ -519,6 +519,10 @@ Results for one NTV computation method across all flux surfaces.
     psi_grid::Vector{Float64} = Float64[]
     dtdpsi::Vector{ComplexF64} = ComplexF64[]
     t_cumulative::Vector{ComplexF64} = ComplexF64[]
+    # dT/dψ resolved by bounce harmonic ℓ, shape (length(psi_grid), length(ell)) with ell = -nl:nl.
+    # Summing over ℓ reproduces `dtdpsi`. Empty when no ψ profile was produced.
+    dtdpsi_ell::Matrix{ComplexF64} = zeros(ComplexF64, 0, 0)
+    ell::Vector{Int} = Int[]
     psi_nsteps::Int = 0
     # ψ-quadrature panel boundaries and the located kinetic-resonance surfaces (first n)
     panel_psis::Vector{Float64} = Float64[]

src/KineticForces/Output.jl

@@ -67,6 +67,10 @@ function write_to_hdf5!(h5file::HDF5.File, state::KineticForcesState; dVdpsi_spl
             mg["psi"] = result.psi_grid
             mg["dTdpsi"] = result.dtdpsi
             mg["T"] = result.t_cumulative
+            if !isempty(result.dtdpsi_ell)
+                mg["ell"] = result.ell
+                mg["dTdpsi_ell"] = result.dtdpsi_ell
+            end
             if dVdpsi_spline !== nothing
                 mg["dVdpsi"] = [dVdpsi_spline(p) for p in result.psi_grid]
             end
@@ -115,6 +119,8 @@ const KF_METHOD_H5_ANNOTATIONS = [
     "psi" => (; long_name="normalized poloidal flux ψ_N at quadrature evaluation points", scale="psi"),
     "dTdpsi" => (; long_name="complex torque density dT/dψ_N = dT_φ/dψ_N + 2i·n·dδW_k/dψ_N; divide by dVdpsi for dT/dV", units="N*m", dims=("psi",), attach=(1 => "psi",)),
     "T" => (; long_name="cumulative complex torque T(ψ_N) = T_φ + 2i·n·δW_k (trapezoidal)", units="N*m", dims=("psi",), attach=(1 => "psi",)),
+    "ell" => (; long_name="bounce harmonic ℓ ∈ -nl:nl of the per-harmonic torque density", scale="ell"),
+    "dTdpsi_ell" => (; long_name="complex torque density dT/dψ_N resolved by bounce harmonic ℓ; summing over ℓ reproduces dTdpsi", units="N*m", dims=("psi", "ell"), attach=(1 => "psi", 2 => "ell")),
     "dVdpsi" => (; long_name="flux-surface volume derivative dV/dψ_N at quadrature points", units="m^3", dims=("psi",), attach=(1 => "psi",)),
     "EnergyIntegrals/psi" => (; long_name="ψ_N of each energy-integration record"),
     "EnergyIntegrals/lambda" => (; long_name="pitch λ = μB0/E of each record"),

test/runtests_kinetic.jl

@@ -664,6 +664,7 @@
         @test mr.total_torque == 0.0 + 0.0im
         @test mr.total_energy == 0.0
         @test isempty(mr.records)
+        @test isempty(mr.dtdpsi_ell) && isempty(mr.ell)
     end
 
     # =========================================================================

test/runtests_multiion.jl

@@ -182,7 +182,8 @@ using HDF5
             st = KF.KineticForcesState();
             st.method_results["fgar"] = KF.MethodResult(; method="fgar", nn=1,
                 total_torque=ComplexF64(sum(dt)), psi_grid=collect(Float64, psi),
-                dtdpsi=ComplexF64.(dt), t_cumulative=zeros(ComplexF64, length(dt)));
+                dtdpsi=ComplexF64.(dt), t_cumulative=zeros(ComplexF64, length(dt)),
+                dtdpsi_ell=ComplexF64.(hcat(0.2 .* dt, 0.5 .* dt, 0.3 .* dt)), ell=collect(-1:1));
             st
         )
         s1 = mk(0.0:0.1:1.0, fill(1.0, 11))
@@ -197,6 +198,23 @@ using HDF5
         s3 = mk([0.0, 0.5, 0.5, 1.0], fill(1.0, 4))
         c2 = KF.combine_species_states([s1, s3])
         @test all(isfinite, real.(c2.method_results["fgar"].dtdpsi))
+        @test all(isfinite, real.(c2.method_results["fgar"].dtdpsi_ell))
+
+        # Per-ℓ survives the combine on the same union grid, and summing over ℓ reproduces the
+        # total exactly (it is derived as the row sum, so this is not merely interpolation noise).
+        @test r.ell == collect(-1:1)
+        @test size(r.dtdpsi_ell) == (length(r.psi_grid), 3)
+        @test vec(sum(r.dtdpsi_ell; dims=2)) ≈ r.dtdpsi
+        ov = [real(d) for (p, d) in zip(r.psi_grid, r.dtdpsi_ell[:, 2]) if 0.1 <= p <= 0.9]
+        @test all(x -> isapprox(x, 1.5; atol=1e-9), ov)        # 0.5·(1 + 2) on the overlap
+        # A species with no per-ℓ data must leave the combined matrix empty, not error.
+        s4 = KF.KineticForcesState()
+        s4.method_results["fgar"] = KF.MethodResult(; method="fgar", nn=1,
+            psi_grid=[0.0, 0.5, 1.0], dtdpsi=fill(ComplexF64(1.0), 3),
+            t_cumulative=zeros(ComplexF64, 3))
+        c3 = KF.combine_species_states([s1, s4])
+        @test isempty(c3.method_results["fgar"].dtdpsi_ell)
+        @test isempty(c3.method_results["fgar"].ell)
     end
 
     @testset "per-species HDF5 layout" begin
@@ -206,7 +224,8 @@ using HDF5
             st = KF.KineticForcesState();
             st.method_results["fgar"] = KF.MethodResult(; method="fgar", nn=1,
                 total_torque=ComplexF64(tt), psi_grid=collect(0.0:0.25:1.0),
-                dtdpsi=fill(ComplexF64(tt), 5), t_cumulative=zeros(ComplexF64, 5));
+                dtdpsi=fill(ComplexF64(tt), 5), t_cumulative=zeros(ComplexF64, 5),
+                dtdpsi_ell=fill(ComplexF64(tt / 3), 5, 3), ell=collect(-1:1));
             st
         )
         h5 = tempname() * ".h5"
@@ -221,10 +240,18 @@ using HDF5
             @test read(f["KineticForces"]["fgar"]["total_torque"]) == 3.0
             @test read(f["KineticForces"]["PerSpecies"]["ion_z1_m2"]["fgar"]["total_torque"]) == 1.0
             # The metadata pass must reach the per-species method level, not stop at PerSpecies.
-            for d in ("total_torque", "dTdpsi")
+            for d in ("total_torque", "dTdpsi", "dTdpsi_ell")
                 a = HDF5.attributes(f["KineticForces"]["PerSpecies"]["electron"]["fgar"][d])
                 @test haskey(a, "long_name") && haskey(a, "units")
             end
+            # HDF5.jl round-trips the Julia shape, so a read here gives (npsi, nell); a
+            # C-ordered reader (h5py, xarray) sees the transpose, (nell, npsi).
+            g = f["KineticForces"]["fgar"]
+            @test size(read(g["dTdpsi_ell"])) == (5, 3)
+            @test read(g["ell"]) == [-1, 0, 1]
+            # Summing over ℓ must reproduce dTdpsi.
+            @test vec(sum(read(g["dTdpsi_ell"]); dims=2)) ≈ read(g["dTdpsi"])
+            @test haskey(HDF5.attributes(g["dTdpsi_ell"]), "dims")
         end
         rm(h5)
     end