diff --git a/src/ForceFreeStates/EulerLagrange.jl b/src/ForceFreeStates/EulerLagrange.jl index 69ea26cdc..f563dca45 100644 --- a/src/ForceFreeStates/EulerLagrange.jl +++ b/src/ForceFreeStates/EulerLagrange.jl @@ -174,7 +174,6 @@ function serial_eulerlagrange_integration(ctrl::ForceFreeStatesControl, equil::E # Initialization odet = OdeState(intr.numpert_total, ctrl.numsteps_init, ctrl.numunorms_init, intr.msing) - odet.du_store_populated = true if ctrl.sing_start <= 0 initialize_el_at_axis!(odet, ctrl, ffit, equil.profiles, intr) elseif ctrl.sing_start <= intr.msing @@ -643,9 +642,8 @@ function cross_ideal_singular_surf!( # so the result is in a different convention. The canonical Δ' is the STRIDE BVP matrix # (compute_delta_prime_matrix!) populated by the parallel FM path. - sing_der!(du1, odet.u, params, odet.psifac) - # Store values after crossing step and advance + odet.q = equil.profiles.q_spline(odet.psifac; hint=odet.spline_hint) store_ode_data!(odet, odet.psifac, odet.u) end @@ -686,10 +684,8 @@ function cross_kinetic_singular_surf!( sing_der!(du2, odet.u, params, odet.psifac) odet.u .+= (du1 .+ du2) .* dpsi - # re-evaluate at the final post-crossing u so the stored derivatives match u_store - sing_der!(du1, odet.u, params, odet.psifac) - # Store crossing step + odet.q = equil.profiles.q_spline(odet.psifac; hint=odet.spline_hint) store_ode_data!(odet, odet.psifac, odet.u) end @@ -741,7 +737,6 @@ function integrate_el_region!( q_range = abs(q_end - q_start) steps_in_segment = Ref(0) - du_buffer = zeros(ComplexF64, intr.numpert_total, intr.numpert_total, 2) function segment_callback!(integrator) ctrl, _, _, intr, odet, chunk = integrator.p @@ -760,7 +755,8 @@ function integrate_el_region!( in_edge_scan = ctrl.psiedge < intr.psilim && integrator.t >= ctrl.psiedge if near_start || near_end || (odet.total_steps % ctrl.save_interval == 0) || in_edge_scan - sing_der!(du_buffer, integrator.u, integrator.p, integrator.t) + # q at the accepted point, not the last internal Runge-Kutta stage + odet.q = equil.profiles.q_spline(integrator.t; hint=odet.spline_hint) store_ode_data!(odet, integrator.t, integrator.u) end end @@ -773,7 +769,7 @@ function integrate_el_region!( # Guarantees the pre-crossing (or pre-edge) state is always stored in u_store, # regardless of where the last accepted step landed relative to the near_end band. if odet.step == 1 || odet.psi_store[odet.step-1] != sol.t[end] - sing_der!(du_buffer, sol.u[end], (ctrl, equil, ffit, intr, odet, chunk), sol.t[end]) + odet.q = equil.profiles.q_spline(sol.t[end]; hint=odet.spline_hint) store_ode_data!(odet, sol.t[end], sol.u[end]) end @@ -1007,12 +1003,14 @@ function transform_u!(odet::OdeState, intr::ForceFreeStatesInternal) odet.u_store[:, :, 1, istep] .= gauss_buffer mul!(gauss_buffer, odet.u_store[:, :, 2, istep], transforms[:, :, ifix]) odet.u_store[:, :, 2, istep] .= gauss_buffer - mul!(gauss_buffer, odet.du_store[:, :, 1, istep], transforms[:, :, ifix]) - odet.du_store[:, :, 1, istep] .= gauss_buffer - mul!(gauss_buffer, odet.du_store[:, :, 2, istep], transforms[:, :, ifix]) - odet.du_store[:, :, 2, istep] .= gauss_buffer - mul!(gauss_buffer, odet.xi_s_store[:, :, istep], transforms[:, :, ifix]) - odet.xi_s_store[:, :, istep] .= gauss_buffer + # Derivative stores are empty unless a path filled them analytically (galerkin); + # materialized ones are computed after this transform and need no fixup. + if !isempty(odet.du_store) + mul!(gauss_buffer, odet.du_store[:, :, istep], transforms[:, :, ifix]) + odet.du_store[:, :, istep] .= gauss_buffer + mul!(gauss_buffer, odet.xi_s_store[:, :, istep], transforms[:, :, ifix]) + odet.xi_s_store[:, :, istep] .= gauss_buffer + end end jfix = kfix + 1 end diff --git a/src/ForceFreeStates/ForceFreeStatesStructs.jl b/src/ForceFreeStates/ForceFreeStatesStructs.jl index 7c2a70d87..d9c4770d2 100644 --- a/src/ForceFreeStates/ForceFreeStatesStructs.jl +++ b/src/ForceFreeStates/ForceFreeStatesStructs.jl @@ -505,11 +505,22 @@ and a small set of temporary matrices and factors used to compute singular-layer - `u_store::Array{ComplexF64,4}` - Stored solution arrays at each saved step with shape `(numpert_total, numpert_total, 2, numsteps_init)` (complex solution state used by the solver). - - `du_store::Array{ComplexF64,4}` - du₁/dψ and du₂/dψ at the accepted point of each saved step, same shape as `u_store`. + - `du_store::Array{ComplexF64,3}` - dΞ_ψ/dψ (the u₁ block only) at each saved step, shape + `(numpert_total, numpert_total, step)`. Empty until `materialize_derivative_stores!` fills it, + except on the galerkin-matched path which supplies the analytic derivative at construction. + du₂/dψ is never stored densely — its only consumer evaluates it on demand at bracket nodes. - - `xi_s_store::Array{ComplexF64,3}` - Clebsch displacement Ξ_s at each saved step, eq. 18 of Glasser 2016, shape `(numpert_total, numpert_total, numsteps_init)`. + - `xi_s_store::Array{ComplexF64,3}` - Clebsch displacement Ξ_s at each saved step, eq. 18 of Glasser 2016, + shape `(numpert_total, numpert_total, step)`. Empty until materialized, same as `du_store`. - - `du_store_populated::Bool` - True once the serial EL path has filled `du_store`/`xi_s_store`; other paths leave it false. + - `u_store_el_basis::Bool` - True when `u_store` holds the Euler-Lagrange state `(u₁, u₂)`, so the + derivative kernel can be re-applied to it. False on the sparse parallel path, whose stored columns + are chunk-endpoint Riccati matrices; `materialize_derivative_stores!` refuses to run there. + + - `du_store_populated::Bool` - True once `du_store`/`xi_s_store` hold valid data in the final + (post-transform, post-normalization) basis. Set by `materialize_derivative_stores!` or by the + galerkin-matched constructor; stays false where the stores cannot be materialized, e.g. the + sparse parallel path whose solution is in the Riccati basis. - `crit_store::Vector{Float64}` - Stored crit parameter values (smallest eigenvalue of W⁻ꜝ) (length `numsteps_init`). @@ -527,10 +538,6 @@ and a small set of temporary matrices and factors used to compute singular-layer - `u::Array{ComplexF64,3}` - Current working solution arrays with shape `(numpert_total, numpert_total, 2)`. - - `du::Array{ComplexF64,3}` - du/dψ from the latest `sing_der!` call, shape `(numpert_total, numpert_total, 2)`. - - - `xi_s::Matrix{ComplexF64}` - Ξ_s from the latest `sing_der!` call, shape `(numpert_total, numpert_total)`. - - `ising_start::Int` - Index of the starting singular surface to be crossed during integration. # Initialization parameters @@ -578,8 +585,9 @@ and a small set of temporary matrices and factors used to compute singular-layer psi_store::Vector{Float64} = Vector{Float64}(undef, numsteps_init) q_store::Vector{Float64} = Vector{Float64}(undef, numsteps_init) u_store::Array{ComplexF64,4} = Array{ComplexF64}(undef, numpert_total, numpert_total, 2, numsteps_init) - du_store::Array{ComplexF64,4} = Array{ComplexF64}(undef, numpert_total, numpert_total, 2, numsteps_init) - xi_s_store::Array{ComplexF64,3} = Array{ComplexF64}(undef, numpert_total, numpert_total, numsteps_init) + du_store::Array{ComplexF64,3} = Array{ComplexF64}(undef, numpert_total, numpert_total, 0) + xi_s_store::Array{ComplexF64,3} = Array{ComplexF64}(undef, numpert_total, numpert_total, 0) + u_store_el_basis::Bool = true du_store_populated::Bool = false crit_store::Vector{Float64} = Vector{Float64}(undef, numsteps_init) ca_r::Array{ComplexF64,4} = Array{ComplexF64}(undef, numpert_total, numpert_total, 2, msing) @@ -592,8 +600,6 @@ and a small set of temporary matrices and factors used to compute singular-layer psifac::Float64 = 0.0 q::Float64 = 0.0 u::Array{ComplexF64,3} = zeros(ComplexF64, numpert_total, numpert_total, 2) - du::Array{ComplexF64,3} = zeros(ComplexF64, numpert_total, numpert_total, 2) - xi_s::Matrix{ComplexF64} = zeros(ComplexF64, numpert_total, numpert_total) ising_start::Int = 0 psimax::Float64 = 0.0 needs_crossing::Bool = false diff --git a/src/ForceFreeStates/Free.jl b/src/ForceFreeStates/Free.jl index 51efa7610..adf545b1c 100644 --- a/src/ForceFreeStates/Free.jl +++ b/src/ForceFreeStates/Free.jl @@ -28,9 +28,10 @@ end """ normalize_eigenfunctions!(odet::OdeState, wt::AbstractMatrix{ComplexF64}, psio::Float64) -> OdeState -Rescale the stored EL solution vectors in `odet.u_store`, `du_store`, and `xi_s_store` so the -edge displacement matches the free-boundary eigenvectors `wt` (scaled by `2π·psio·1e-3`). -Modifies `odet` in place. Call after `free_run` when downstream code consumes the stored ξ profiles. +Rescale the stored EL solution vectors in `odet.u_store` (and `du_store`/`xi_s_store` when a +path has already filled them) so the edge displacement matches the free-boundary eigenvectors +`wt` (scaled by `2π·psio·1e-3`). Modifies `odet` in place. Call after `free_run` when +downstream code consumes the stored ξ profiles. """ @with_pool pool function normalize_eigenfunctions!(odet::OdeState, wt::AbstractMatrix{ComplexF64}, psio::Float64) N = size(wt, 1) @@ -41,12 +42,14 @@ Modifies `odet` in place. Call after `free_run` when downstream code consumes th odet.u_store[:, :, 1, istep] .= tmp_mat mul!(tmp_mat, odet.u_store[:, :, 2, istep], coeffs) odet.u_store[:, :, 2, istep] .= tmp_mat - mul!(tmp_mat, odet.du_store[:, :, 1, istep], coeffs) - odet.du_store[:, :, 1, istep] .= tmp_mat - mul!(tmp_mat, odet.du_store[:, :, 2, istep], coeffs) - odet.du_store[:, :, 2, istep] .= tmp_mat - mul!(tmp_mat, odet.xi_s_store[:, :, istep], coeffs) - odet.xi_s_store[:, :, istep] .= tmp_mat + # Only the galerkin path carries derivative stores this early; materialized ones + # are built from the normalized u_store afterwards. + if !isempty(odet.du_store) + mul!(tmp_mat, odet.du_store[:, :, istep], coeffs) + odet.du_store[:, :, istep] .= tmp_mat + mul!(tmp_mat, odet.xi_s_store[:, :, istep], coeffs) + odet.xi_s_store[:, :, istep] .= tmp_mat + end end return odet end diff --git a/src/ForceFreeStates/Galerkin/GalerkinMatch.jl b/src/ForceFreeStates/Galerkin/GalerkinMatch.jl index 743757f13..cecbb6330 100644 --- a/src/ForceFreeStates/Galerkin/GalerkinMatch.jl +++ b/src/ForceFreeStates/Galerkin/GalerkinMatch.jl @@ -264,33 +264,21 @@ function gal_matched_odestate(gal_result::GalerkinResult, ffit::FourFitVars, int ngrid_f = length(psi_f) u_store = zeros(ComplexF64, mpert, mpert, 2, ngrid_f) - du_store = zeros(ComplexF64, mpert, mpert, 2, ngrid_f) + du_store = zeros(ComplexF64, mpert, mpert, ngrid_f) xi_s_store = zeros(ComplexF64, mpert, mpert, ngrid_f) - amat = Matrix{ComplexF64}(undef, mpert, mpert) - bmat = Matrix{ComplexF64}(undef, mpert, mpert) - cmat = Matrix{ComplexF64}(undef, mpert, mpert) - tmp = Matrix{ComplexF64}(undef, mpert, mpert) hint = Ref(1) for ip in 1:ngrid_f ξ = @view xi_f[:, ip, :] ξ′ = @view dxi_f[:, ip, :] @views u_store[:, :, 1, ip] .= ξ - @views du_store[:, :, 1, ip] .= ξ′ - # ξ_s = −A⁻¹(B·ξ′ + C·ξ), exactly as the ideal path of sing_der! (Sing.jl:1015-1049) - ffit.amats(vec(amat), psi_f[ip]; hint=hint) - ffit.bmats(vec(bmat), psi_f[ip]; hint=hint) - ffit.cmats(vec(cmat), psi_f[ip]; hint=hint) - LinearAlgebra.LAPACK.potrf!('U', amat) - LinearAlgebra.LAPACK.potrs!('U', amat, bmat) # bmat ← A⁻¹ B - LinearAlgebra.LAPACK.potrs!('U', amat, cmat) # cmat ← A⁻¹ C - xs = @view xi_s_store[:, :, ip] - mul!(tmp, bmat, ξ′) - xs .= .-tmp - mul!(tmp, cmat, ξ) - xs .-= tmp + @views du_store[:, :, ip] .= ξ′ + # ξ_s = −A⁻¹(B·ξ′ + C·ξ), the same node quantity the Euler-Lagrange path computes. + @views compute_node_xi_s!(xi_s_store[:, :, ip], ξ′, ξ, ffit, psi_f[ip]; hint=hint) end + # Derivatives here are the analytic galerkin ξ′, not recomputable from the ODE kernel. return OdeState(; numpert_total=mpert, numunorms_init=1, msing=gal_result.msing, numsteps_init=ngrid_f, - step=ngrid_f, total_steps=ngrid_f, psi_store=psi_f, q_store=q_f, u_store=u_store, du_store=du_store, xi_s_store=xi_s_store) + step=ngrid_f, total_steps=ngrid_f, psi_store=psi_f, q_store=q_f, u_store=u_store, du_store=du_store, + xi_s_store=xi_s_store, du_store_populated=true) end diff --git a/src/ForceFreeStates/Riccati.jl b/src/ForceFreeStates/Riccati.jl index 66d1361e7..9e86aadcf 100644 --- a/src/ForceFreeStates/Riccati.jl +++ b/src/ForceFreeStates/Riccati.jl @@ -1009,12 +1009,6 @@ See: Glasser (2018) Phys. Plasmas 25, 032507 — Eq. 19 (dual Riccati form) # dS = w†·v - S·Ḡ·S [Glasser 2018 eq. 19, dual Riccati] mul!(dS, adjoint(w), v) # dS = w†·v - # Store du1/dψ = Q·v as a diagnostic before v is reused - # Q·v = diag(singfac_vec)·v = Ξ'_Ψ (displacement gradient, with U₂ = I) - @. odet.du[:, :, 1] = singfac_vec * v - @view(odet.du[:, :, 2]) .= 0 - odet.xi_s .= 0 - # Subtract S·Ḡ·S (reuse v and tmp to avoid extra allocation) mul!(tmp, gmat, S) # tmp = Ḡ·S mul!(v, S, tmp) # v = S·Ḡ·S @@ -1794,9 +1788,11 @@ function _assemble_propagators_serially!(odet::OdeState, propagators::Vector{Chu riccati_cross_ideal_singular_surf!(odet, ctrl, equil, ffit, intr, chunk.ising) last_crossing_step = odet.step - 1 else - # Save non-crossing end-of-chunk state. du_store is not meaningful here — when - # ctrl.populate_dense_xi=true the entire odet is replaced by a serial-EL pass - # at the end of parallel_eulerlagrange_integration. + # Save non-crossing end-of-chunk state. These columns are FM/Riccati chunk + # endpoints, not the Euler-Lagrange state — when ctrl.populate_dense_xi=true the + # entire odet is replaced by a serial-EL pass at the end of + # parallel_eulerlagrange_integration. + odet.u_store_el_basis = false if odet.step >= size(odet.u_store, 4) resize_storage!(odet) end diff --git a/src/ForceFreeStates/Sing.jl b/src/ForceFreeStates/Sing.jl index 194cf58a9..13d400128 100644 --- a/src/ForceFreeStates/Sing.jl +++ b/src/ForceFreeStates/Sing.jl @@ -1149,14 +1149,43 @@ more simplistic code with similar performance. - `u::Array{ComplexF64,3}`: Current state array, shape (mpert, mpert, 2) - `params::Tuple{ForceFreeStatesControl, PlasmaEquilibrium, FourFitVars, ForceFreeStatesInternal, OdeState, IntegrationChunk}`: Tuple of relevant structs - `psieval::Float64`: Current psi value at which to evaluate the derivative + +The unpacked-argument method carries the arithmetic; this tuple method is the thin adapter the +integrator calls. Ξ_s is *not* computed here — it is a save-point quantity, obtained from +[`compute_node_xi_s!`](@ref) only where it is actually consumed. """ -@with_pool pool function sing_der!(du::Array{ComplexF64,3}, u::Array{ComplexF64,3}, +function sing_der!(du::Array{ComplexF64,3}, u::Array{ComplexF64,3}, params::Tuple{ForceFreeStatesControl,Equilibrium.PlasmaEquilibrium, FourFitVars,ForceFreeStatesInternal,OdeState,IntegrationChunk}, psieval::Float64) - - # Unpack structs ctrl, equil, ffit, intr, odet, _ = params + return sing_der!(du, u, ctrl, equil, ffit, intr, odet, psieval) +end + +""" + sing_der!(du, u, ctrl, equil, ffit, intr, odet, psieval) + +Unpacked-argument form of the Euler-Lagrange derivative, using `odet`'s spline hints and +recording q at `psieval` in `odet.q`. Not safe to call concurrently on a shared `odet`; +multi-threaded callers should use [`el_derivatives!`](@ref) with their own hints. +""" +function sing_der!(du::Array{ComplexF64,3}, u::Array{ComplexF64,3}, + ctrl::ForceFreeStatesControl, equil::Equilibrium.PlasmaEquilibrium, ffit::FourFitVars, + intr::ForceFreeStatesInternal, odet::OdeState, psieval::Float64) + odet.q = el_derivatives!(du, u, ctrl.kinetic_factor > 0, equil, ffit, intr, psieval, odet.spline_hint, odet.ffit_hint) + return nothing +end + +""" + el_derivatives!(du, u, kinetic, equil, ffit, intr, psieval, spline_hint, ffit_hint) -> q + +Euler-Lagrange (or, when `kinetic` is true, FKG) derivative kernel: writes du₁/dψ and du₂/dψ at +`psieval` into `du` and returns q there. Holds no state of its own — the two hints are the +caller's interval-search accelerators, so concurrent callers just pass their own. +""" +@with_pool pool function el_derivatives!(du::Array{ComplexF64,3}, u::Array{ComplexF64,3}, + kinetic::Bool, equil::Equilibrium.PlasmaEquilibrium, ffit::FourFitVars, + intr::ForceFreeStatesInternal, psieval::Float64, spline_hint::Base.RefValue{Int}, ffit_hint::Base.RefValue{Int}) # Allocate temporary arrays from the pool Npert = intr.numpert_total @@ -1164,13 +1193,10 @@ more simplistic code with similar performance. singfac_vec = acquire!(pool, Float64, Npert) singfac_mat = reshape(singfac_vec, intr.mpert, intr.npert) - amat = acquire!(pool, ComplexF64, Npert, Npert) - bmat = similar!(pool, amat) - cmat = similar!(pool, amat) - fmat_lower = similar!(pool, amat) - kmat = similar!(pool, amat) - gmat = similar!(pool, amat) - tmp_mat = similar!(pool, amat) + fmat_lower = acquire!(pool, ComplexF64, Npert, Npert) + kmat = similar!(pool, fmat_lower) + gmat = similar!(pool, fmat_lower) + tmp_mat = similar!(pool, fmat_lower) fill!(tmp_mat, zero(ComplexF64)) u1 = @view(u[:, :, 1]) @@ -1179,45 +1205,33 @@ more simplistic code with similar performance. du2 = @view(du[:, :, 2]) # Compute singfac = 1 / (m - nq) - # Use shared hint for O(1) interval lookup during sequential ODE integration - odet.q = equil.profiles.q_spline(psieval; hint=odet.spline_hint) - singfac_mat .= 1.0 ./ ((intr.mlow:intr.mhigh) .- odet.q .* (intr.nlow:intr.nhigh)') + # Use caller-supplied hint for O(1) interval lookup during sequential ODE integration + q = equil.profiles.q_spline(psieval; hint=spline_hint) + singfac_mat .= 1.0 ./ ((intr.mlow:intr.mhigh) .- q .* (intr.nlow:intr.nhigh)') - if ctrl.kinetic_factor > 0 + if kinetic # ---- Kinetic path with pre-computed FKG matrices ---- - # Load pre-computed kinetic matrices from splines - # amat/bmat/cmat here are the kinetic-modified A_kin/B_kin/C_kin - # Use odet.ffit_hint (per-thread) instead of ffit._hint (shared, racy in parallel BVP) - ffit.amats(vec(amat), psieval; hint=odet.ffit_hint) - ffit.bmats(vec(bmat), psieval; hint=odet.ffit_hint) - ffit.cmats(vec(cmat), psieval; hint=odet.ffit_hint) - + # Use the caller's hint, not ffit._hint (shared, racy in the parallel BVP) # Load FKG sub-matrices (note: reusing fmat_lower/kmat/gmat as workspace) - f0mat = similar!(pool, amat) - pmat_kin = similar!(pool, amat) - paat_kin = similar!(pool, amat) - kkmat_kin = similar!(pool, amat) - kkaat_kin = similar!(pool, amat) - r1mat_kin = similar!(pool, amat) - r2mat_kin = similar!(pool, amat) - r3mat_kin = similar!(pool, amat) - gaat_kin = similar!(pool, amat) - - ffit.f0mats(vec(f0mat), psieval; hint=odet.ffit_hint) - ffit.pmats(vec(pmat_kin), psieval; hint=odet.ffit_hint) - ffit.paats(vec(paat_kin), psieval; hint=odet.ffit_hint) - ffit.kkmats(vec(kkmat_kin), psieval; hint=odet.ffit_hint) - ffit.kkaats(vec(kkaat_kin), psieval; hint=odet.ffit_hint) - ffit.r1mats(vec(r1mat_kin), psieval; hint=odet.ffit_hint) - ffit.r2mats(vec(r2mat_kin), psieval; hint=odet.ffit_hint) - ffit.r3mats(vec(r3mat_kin), psieval; hint=odet.ffit_hint) - ffit.gaats(vec(gaat_kin), psieval; hint=odet.ffit_hint) - - # A⁻¹B, A⁻¹C via LU (A is non-Hermitian with kinetic contributions) - # Direct LAPACK to avoid the ipiv allocation that lu!/ldiv! would do in this hot loop - _, ipiv, _ = LAPACK.getrf!(amat) - LAPACK.getrs!('N', amat, ipiv, bmat) - LAPACK.getrs!('N', amat, ipiv, cmat) + f0mat = similar!(pool, fmat_lower) + pmat_kin = similar!(pool, fmat_lower) + paat_kin = similar!(pool, fmat_lower) + kkmat_kin = similar!(pool, fmat_lower) + kkaat_kin = similar!(pool, fmat_lower) + r1mat_kin = similar!(pool, fmat_lower) + r2mat_kin = similar!(pool, fmat_lower) + r3mat_kin = similar!(pool, fmat_lower) + gaat_kin = similar!(pool, fmat_lower) + + ffit.f0mats(vec(f0mat), psieval; hint=ffit_hint) + ffit.pmats(vec(pmat_kin), psieval; hint=ffit_hint) + ffit.paats(vec(paat_kin), psieval; hint=ffit_hint) + ffit.kkmats(vec(kkmat_kin), psieval; hint=ffit_hint) + ffit.kkaats(vec(kkaat_kin), psieval; hint=ffit_hint) + ffit.r1mats(vec(r1mat_kin), psieval; hint=ffit_hint) + ffit.r2mats(vec(r2mat_kin), psieval; hint=ffit_hint) + ffit.r3mats(vec(r3mat_kin), psieval; hint=ffit_hint) + ffit.gaats(vec(gaat_kin), psieval; hint=ffit_hint) # Build singfac-dependent F̄, K̄, K̄†, Ḡ† matrices (Logan 2015 Appendix C, Eqs C.5-C.11): # F̄(i,j) = q1*f0*q2 - q1*P - P†'*q2 + R1 @@ -1226,10 +1240,10 @@ more simplistic code with similar performance. # where q1 = (m₁ - n*q), q2 = (m₂ - n*q) — direct singfac, NOT 1/(m-nq) as in ideal path singfac_direct = acquire!(pool, Float64, Npert) singfac_direct_mat = reshape(singfac_direct, intr.mpert, intr.npert) - singfac_direct_mat .= (intr.mlow:intr.mhigh) .- odet.q .* (intr.nlow:intr.nhigh)' + singfac_direct_mat .= (intr.mlow:intr.mhigh) .- q .* (intr.nlow:intr.nhigh)' # Build F, K, K† with singfac (using fmat_lower, kmat, gmat as workspace for F, K, K†) - kaat_kin = similar!(pool, amat) # K† matrix + kaat_kin = similar!(pool, fmat_lower) # K† matrix for j in 1:Npert q2 = singfac_direct[j] for i in 1:Npert @@ -1260,18 +1274,10 @@ more simplistic code with similar performance. else # ---- Ideal path ---- - # Evaluate matrix splines at the current psi (odet.ffit_hint is per-thread) - ffit.amats(vec(amat), psieval; hint=odet.ffit_hint) - ffit.bmats(vec(bmat), psieval; hint=odet.ffit_hint) - ffit.cmats(vec(cmat), psieval; hint=odet.ffit_hint) - ffit.fmats_lower(vec(fmat_lower), psieval; hint=odet.ffit_hint) - ffit.kmats(vec(kmat), psieval; hint=odet.ffit_hint) - ffit.gmats(vec(gmat), psieval; hint=odet.ffit_hint) - - # Solve bmat = A⁻¹ * bmat, cmat = A⁻¹ * cmat in-place via Cholesky - LAPACK.potrf!('U', amat) - LAPACK.potrs!('U', amat, bmat) - LAPACK.potrs!('U', amat, cmat) + # Evaluate matrix splines at the current psi (hint is the caller's, never shared) + ffit.fmats_lower(vec(fmat_lower), psieval; hint=ffit_hint) + ffit.kmats(vec(kmat), psieval; hint=ffit_hint) + ffit.gmats(vec(gmat), psieval; hint=ffit_hint) # See equations 22-24 in Glasser 2016 DCON paper for derivation # du[1] = - F̄⁻¹ * K̄ * u[1] + F̄⁻¹ * Q⁻¹ * u[2] @@ -1288,13 +1294,49 @@ more simplistic code with similar performance. # du[1] = - Q⁻¹ * F̄⁻¹ * K̄ * u[1] + Q⁻¹ * F̄⁻¹ * Q⁻¹ * u[2] du1 .*= singfac_vec end + return q +end + +""" + compute_node_xi_s!(xi_s, du1, u1, ffit, psieval; kinetic=false, hint=Ref(1)) + +Evaluate Ξ_s = -A⁻¹(B·Ξ′_ψ + C·Ξ_ψ) [Glasser Phys. Plasmas 2016 112506 eq. 18] at `psieval`, +writing into `xi_s`. `du1` and `u1` are the Ξ′_ψ and Ξ_ψ blocks at the same ψ, i.e. slices of a +`el_derivatives!` result and its input state. + +Split out of the derivative kernel because Ξ_s is needed only at saved nodes, not at every +Runge-Kutta stage. Ideal runs factor the Hermitian A by Cholesky; with `kinetic=true` A picks up +non-Hermitian contributions and needs an LU. +""" +@with_pool pool function compute_node_xi_s!(xi_s::AbstractMatrix{ComplexF64}, du1::AbstractMatrix{ComplexF64}, + u1::AbstractMatrix{ComplexF64}, ffit::FourFitVars, psieval::Float64; kinetic::Bool=false, hint::Base.RefValue{Int}=Ref(1)) + + Npert = size(u1, 1) + amat = acquire!(pool, ComplexF64, Npert, Npert) + bmat = similar!(pool, amat) + cmat = similar!(pool, amat) + tmp_mat = similar!(pool, amat) + + ffit.amats(vec(amat), psieval; hint=hint) + ffit.bmats(vec(bmat), psieval; hint=hint) + ffit.cmats(vec(cmat), psieval; hint=hint) + + # Solve bmat = A⁻¹ * bmat, cmat = A⁻¹ * cmat in-place + if kinetic + _, ipiv, _ = LAPACK.getrf!(amat) + LAPACK.getrs!('N', amat, ipiv, bmat) + LAPACK.getrs!('N', amat, ipiv, cmat) + else + LAPACK.potrf!('U', amat) + LAPACK.potrs!('U', amat, bmat) + LAPACK.potrs!('U', amat, cmat) + end - odet.du .= du - # Ξ_s = - A⁻¹(B * Ξ'_Ψ + C * Ξ_Ψ), eq. 18 of Glasser 2016 mul!(tmp_mat, bmat, du1) - odet.xi_s .= .-tmp_mat + xi_s .= .-tmp_mat mul!(tmp_mat, cmat, u1) - odet.xi_s .-= tmp_mat + xi_s .-= tmp_mat + return xi_s end """ diff --git a/src/ForceFreeStates/Utils.jl b/src/ForceFreeStates/Utils.jl index f3153b0a3..71830e027 100644 --- a/src/ForceFreeStates/Utils.jl +++ b/src/ForceFreeStates/Utils.jl @@ -2,8 +2,9 @@ resize_storage!(odet::OdeState) Resize storage arrays in `odet` when the current step exceeds allocated size. -Doubles the size of the storage arrays for `u_store`, `du_store`, `xi_s_store`, -`psi_store`, and `q_store`, and copies over existing data to the new arrays. +Doubles the size of the storage arrays for `u_store`, `psi_store`, and `q_store`, and +copies over existing data to the new arrays. The derivative stores are not grown here: +they are filled at exact size by `materialize_derivative_stores!` after integration. """ function resize_storage!(odet::OdeState) oldlen = size(odet.u_store, 4) @@ -11,22 +12,16 @@ function resize_storage!(odet::OdeState) # Allocate new arrays u_new = Array{ComplexF64,4}(undef, odet.numpert_total, odet.numpert_total, 2, newlen) - du_new = Array{ComplexF64,4}(undef, odet.numpert_total, odet.numpert_total, 2, newlen) - xi_s_new = Array{ComplexF64,3}(undef, odet.numpert_total, odet.numpert_total, newlen) psi_new = Vector{Float64}(undef, newlen) q_new = Vector{Float64}(undef, newlen) # Copy old data u_new[:, :, :, 1:odet.step] = odet.u_store[:, :, :, 1:odet.step] - du_new[:, :, :, 1:odet.step] = odet.du_store[:, :, :, 1:odet.step] - xi_s_new[:, :, 1:odet.step] = odet.xi_s_store[:, :, 1:odet.step] psi_new[1:odet.step] = odet.psi_store[1:odet.step] q_new[1:odet.step] = odet.q_store[1:odet.step] # Replace old arrays odet.u_store = u_new - odet.du_store = du_new - odet.xi_s_store = xi_s_new odet.psi_store = psi_new odet.q_store = q_new end @@ -35,24 +30,23 @@ end trim_storage!(odet::OdeState) Trim storage arrays in `odet` to the actual number of steps taken. -Resizes `u_store`, `du_store`, `xi_s_store`, `psi_store`, and `q_store` to the -current step count, removing any unused allocated space. +Resizes `u_store`, `psi_store`, and `q_store` to the current step count, removing any +unused allocated space. The derivative stores are left alone — they are either empty +(serial/Riccati, filled later by `materialize_derivative_stores!`) or already exact-size +(galerkin-matched states). """ function trim_storage!(odet::OdeState) resize!(odet.psi_store, odet.step) resize!(odet.q_store, odet.step) odet.u_store = odet.u_store[:, :, :, 1:odet.step] - odet.du_store = odet.du_store[:, :, :, 1:odet.step] - odet.xi_s_store = odet.xi_s_store[:, :, 1:odet.step] end """ store_ode_data!(odet::OdeState, psi::Float64, u) -Save the current integration state at `psi`: `u` plus `odet.du`, `odet.xi_s`, and `odet.q` -from the latest `sing_der!` call. Callers must evaluate `sing_der!` at exactly `(psi, u)` -first, so the stored derivatives belong to the accepted point rather than the last -internal solver stage. +Save the current integration state at `psi`: `u` plus `odet.q`, which callers set for the +accepted point. Derivatives are not stored — `materialize_derivative_stores!` recomputes +them from `(psi_store, u_store)` after integration, where they are actually consumed. """ function store_ode_data!(odet::OdeState, psi::Float64, u) if odet.step >= size(odet.u_store, 4) @@ -61,7 +55,51 @@ function store_ode_data!(odet::OdeState, psi::Float64, u) odet.psi_store[odet.step] = psi odet.q_store[odet.step] = odet.q @views odet.u_store[:, :, :, odet.step] .= u - @views odet.du_store[:, :, :, odet.step] .= odet.du - @views odet.xi_s_store[:, :, odet.step] .= odet.xi_s odet.step += 1 end + +""" + materialize_derivative_stores!(odet, equil, ffit, intr) -> Bool + +Fill `odet.du_store` (dΞ_ψ/dψ) and `odet.xi_s_store` from the stored solution, returning +whether they hold valid data afterwards. Idempotent: a no-op when `du_store_populated` is +already true, so the galerkin-matched path keeps its analytic derivatives. + +Derivatives are recomputed rather than accumulated during integration because they are +needed only at saved nodes, not at every Runge-Kutta stage. Recomputing after the Gaussian +fixup transforms and the free-boundary normalization is exact rather than merely close: the +Euler-Lagrange system is linear in `u`, and both operations right-multiply the solution by a +mixing matrix `T`, so `du(ψ, u·T) = du(ψ, u)·T`. + +Returns `false` without allocating when there is nothing to work from — no `ffit`, no stored +steps, or a solution held in a basis the Euler-Lagrange kernel does not apply to (the sparse +parallel path, which stores chunk-endpoint Riccati matrices). +""" +function materialize_derivative_stores!( + odet::OdeState, + equil::Equilibrium.PlasmaEquilibrium, + ffit::Union{FourFitVars,Nothing}, + intr::ForceFreeStatesInternal +) + odet.du_store_populated && return true + (isnothing(ffit) || odet.step == 0 || isempty(odet.u_store) || !odet.u_store_el_basis) && return false + + kinetic = ffit.kinetic_populated + nstep = min(odet.step, size(odet.u_store, 4)) + npert = odet.numpert_total + odet.du_store = Array{ComplexF64}(undef, npert, npert, nstep) + odet.xi_s_store = Array{ComplexF64}(undef, npert, npert, nstep) + + du = zeros(ComplexF64, npert, npert, 2) + u = zeros(ComplexF64, npert, npert, 2) + @views for istep in 1:nstep + u .= odet.u_store[:, :, :, istep] + odet.q = el_derivatives!(du, u, kinetic, equil, ffit, intr, odet.psi_store[istep], odet.spline_hint, odet.ffit_hint) + odet.du_store[:, :, istep] .= du[:, :, 1] + compute_node_xi_s!(odet.xi_s_store[:, :, istep], du[:, :, 1], u[:, :, 1], ffit, odet.psi_store[istep]; + kinetic=kinetic, hint=odet.ffit_hint) + end + + odet.du_store_populated = true + return true +end diff --git a/src/GeneralizedPerturbedEquilibrium.jl b/src/GeneralizedPerturbedEquilibrium.jl index 9e6cd0aa2..27279f51c 100755 --- a/src/GeneralizedPerturbedEquilibrium.jl +++ b/src/GeneralizedPerturbedEquilibrium.jl @@ -712,6 +712,10 @@ function write_outputs_to_HDF5( ballooning_boundary=(psi=Float64[], alpha=Float64[], alpha_critical=Float64[]) ) + # Idempotent: already done if a PerturbedEquilibrium stage ran. Leaves the stores empty + # (and the datasets below empty) on paths whose solution basis cannot supply them. + ForceFreeStates.materialize_derivative_stores!(odet, equil, ffit, intr) + h5open(joinpath(intr.dir_path, ctrl.HDF5_filename), "w") do out_h5 # Store git version for reproducibility @@ -807,7 +811,7 @@ function write_outputs_to_HDF5( out_h5["integration/q"] = odet.q_store out_h5["integration/xi_psi"] = odet.u_store[:, :, 1, :] out_h5["integration/u2"] = odet.u_store[:, :, 2, :] # TODO: what to name this? These are the "conjugate momenta" of u1 - out_h5["integration/dxi_psi"] = odet.du_store[:, :, 1, :] + out_h5["integration/dxi_psi"] = odet.du_store out_h5["integration/xi_s"] = odet.xi_s_store out_h5["integration/crit"] = odet.crit_store diff --git a/src/PerturbedEquilibrium/FieldReconstruction.jl b/src/PerturbedEquilibrium/FieldReconstruction.jl index 4f7719ec8..e491477c4 100644 --- a/src/PerturbedEquilibrium/FieldReconstruction.jl +++ b/src/PerturbedEquilibrium/FieldReconstruction.jl @@ -6,7 +6,7 @@ field perturbations in mode space, following the GPEC gpeq module approach. Displacement components from ODE integration (u_store/du_store/xi_s_store): - ξ_ψ: radial displacement (u_store[:,:,1,:]) -- dξ_ψ/dψ: radial derivative (du_store[:,:,1,:]) +- dξ_ψ/dψ: radial derivative (du_store) - ξ_s: toroidal displacement (xi_s_store, Glasser 2016 eq. 18) Contravariant perturbed field from ideal MHD (matches Fortran gpeq_sol): @@ -219,7 +219,7 @@ alpha = flux_matrix \\ response_vector Then sum eigenmode contributions at each radial point (matches Fortran gpeq_sol): xi_psi[ipsi, :] = u_store[:, :, 1, ipsi] * alpha # Ξ_ψ -xi_psi1[ipsi, :] = du_store[:, :, 1, ipsi] * alpha # dΞ_ψ/dψ +xi_psi1[ipsi, :] = du_store[:, :, ipsi] * alpha # dΞ_ψ/dψ xi_s[ipsi, :] = xi_s_store[:, :, ipsi] * alpha # Ξ_s (toroidal, Glasser 2016 eq. 18) # Returns @@ -251,9 +251,9 @@ function sum_eigenmode_contributions( mul!(view(xi_psi_modes, ipsi, :), @view(ForceFreeStates_results.u_store[:, :, 1, ipsi]), alpha) - # du_store[:,:,1] = dΞ_ψ/dψ (radial derivative) + # du_store = dΞ_ψ/dψ (radial derivative) mul!(view(xi_psi1_modes, ipsi, :), - @view(ForceFreeStates_results.du_store[:, :, 1, ipsi]), + @view(ForceFreeStates_results.du_store[:, :, ipsi]), alpha) # xi_s_store = Ξ_s = -A⁻¹(B·Ξ'_ψ + C·Ξ_ψ) (toroidal displacement, Glasser 2016 eq. 18) mul!(view(xi_s_modes, ipsi, :), diff --git a/src/PerturbedEquilibrium/PerturbedEquilibrium.jl b/src/PerturbedEquilibrium/PerturbedEquilibrium.jl index 5de559c58..5e9429d5b 100644 --- a/src/PerturbedEquilibrium/PerturbedEquilibrium.jl +++ b/src/PerturbedEquilibrium/PerturbedEquilibrium.jl @@ -83,6 +83,10 @@ function compute_perturbed_equilibrium( # Step 0: Initialize mode arrays for convenient indexing initialize_mode_arrays!(intr, ffs_intr) + # Ξ′ and Ξ_s are recomputed from the stored solution here rather than carried through + # integration; downstream response and coupling code reads them from the stores. + ForceFreeStates.materialize_derivative_stores!(ForceFreeStates_results, equil, ffit, ffs_intr) + # Load forcing data. On the gpec.h5 replay path the caller preloads # `intr.forcing_modes` from the snapshot, so skip re-reading the original file. if isempty(intr.forcing_modes) diff --git a/src/PerturbedEquilibrium/SingularCoupling.jl b/src/PerturbedEquilibrium/SingularCoupling.jl index 94f6aed96..417474dc7 100644 --- a/src/PerturbedEquilibrium/SingularCoupling.jl +++ b/src/PerturbedEquilibrium/SingularCoupling.jl @@ -63,13 +63,18 @@ Least accurate method, kept for solution paths outside the serial EL integrator (gal-matched, Riccati) whose stored derivatives cover only Ξ′. """ function _chord_solution_at(psi::Float64, resnum::Int, odet::OdeState, nstep::Int) + isempty(odet.du_store) && error( + "_chord_solution_at: no derivative store. The solution is in a basis " * + "the Euler-Lagrange kernel cannot be re-applied to (sparse parallel path); " * + "set populate_dense_xi = true for PerturbedEquilibrium runs." + ) il, ir, _ = _psi_bracket(odet.psi_store, psi, nstep) psi_a, psi_b = odet.psi_store[il], odet.psi_store[ir] u_a = odet.u_store[resnum, :, 1, il] u_b = odet.u_store[resnum, :, 1, ir] - du_a = odet.du_store[resnum, :, 1, il] - du_b = odet.du_store[resnum, :, 1, ir] + du_a = odet.du_store[resnum, :, il] + du_b = odet.du_store[resnum, :, ir] u_e = _hermite_cubic_val(u_a, u_b, du_a, du_b, psi_a, psi_b, psi) du_e = (u_b .- u_a) ./ (psi_b - psi_a) @@ -93,8 +98,8 @@ function _gal_solution_at(psi::Float64, resnum::Int, odet::OdeState, nstep::Int) u_a = odet.u_store[resnum, :, 1, il] u_b = odet.u_store[resnum, :, 1, ir] - du_a = odet.du_store[resnum, :, 1, il] - du_b = odet.du_store[resnum, :, 1, ir] + du_a = odet.du_store[resnum, :, il] + du_b = odet.du_store[resnum, :, ir] u_e = _hermite_cubic_val(u_a, u_b, du_a, du_b, psi_a, psi_b, psi) du_e = _hermite_cubic_deriv(u_a, u_b, du_a, du_b, psi_a, psi_b, psi) @@ -125,8 +130,8 @@ function _solution_at( u_a = odet.u_store[resnum, :, 1, il] u_b = odet.u_store[resnum, :, 1, ir] - du_a = odet.du_store[resnum, :, 1, il] - du_b = odet.du_store[resnum, :, 1, ir] + du_a = odet.du_store[resnum, :, il] + du_b = odet.du_store[resnum, :, ir] u_e = _hermite_cubic_val(u_a, u_b, du_a, du_b, psi_a, psi_b, psi) @@ -146,7 +151,7 @@ function _solution_at( j == k && continue w *= (psi - odet.psi_store[idxs[j]]) / (odet.psi_store[idxs[k]] - odet.psi_store[idxs[j]]) end - du_e .+= (w * singfac(odet.psi_store[idxs[k]])) .* @view(odet.du_store[resnum, :, 1, idxs[k]]) + du_e .+= (w * singfac(odet.psi_store[idxs[k]])) .* @view(odet.du_store[resnum, :, idxs[k]]) end du_e ./= singfac(psi) @@ -160,6 +165,10 @@ Evaluate the `resnum` row of Ξ_ψ and Ξ′_ψ at `psi` from the stored ODE sol ideal Euler-Lagrange relation Ξ′ = Q⁻¹·F̄⁻¹·(Q⁻¹·u₂ − K̄·u₁) [Glasser 2016 eqs. 22-24], with u₁, u₂ Hermite-interpolated to `psi`. Only valid for ideal runs where `ffit.fmats_lower` and `kmats` generated the solution. + +The Hermite slopes need du₂ as well as du₁, and du₂ is not stored: both are evaluated here +from the derivative kernel at the two bracketing nodes, which is where the handful of +resonant evaluation points actually need them. """ function _el_solution_at( psi::Float64, @@ -178,12 +187,19 @@ function _el_solution_at( u1_b = @view odet.u_store[:, :, 1, ir] u2_a = @view odet.u_store[:, :, 2, il] u2_b = @view odet.u_store[:, :, 2, ir] - du1_a = @view odet.du_store[:, :, 1, il] - du1_b = @view odet.du_store[:, :, 1, ir] - du2_a = @view odet.du_store[:, :, 2, il] - du2_b = @view odet.du_store[:, :, 2, ir] + # Own hints: this runs inside a threaded loop over rational surfaces. hint = Ref(1) + q_hint = Ref(1) + du_a = zeros(ComplexF64, npert, npert, 2) + du_b = zeros(ComplexF64, npert, npert, 2) + ForceFreeStates.el_derivatives!(du_a, odet.u_store[:, :, :, il], false, equil, ffit, ffs_intr, psi_a, q_hint, hint) + ForceFreeStates.el_derivatives!(du_b, odet.u_store[:, :, :, ir], false, equil, ffit, ffs_intr, psi_b, q_hint, hint) + du1_a = @view du_a[:, :, 1] + du1_b = @view du_b[:, :, 1] + du2_a = @view du_a[:, :, 2] + du2_b = @view du_b[:, :, 2] + kmat = Matrix{ComplexF64}(undef, npert, npert) u1_e = _hermite_cubic_val(u1_a, u1_b, du1_a, du1_b, psi_a, psi_b, psi) @@ -416,7 +432,7 @@ function compute_singular_coupling_metrics!( # Inner-layer (cusp-free) penetrated field: bpen[s, j] is linear in the same identity-at-edge # coil-drive columns as the OdeState solutions, so it contracts with C_coeffs exactly like - # the outer solution values above (xsp = dot(u, ck)); /area matches the area-weighted + # the outer solution values above (xsp = transpose(u) * ck); /area matches the area-weighted # convention of the pointwise row. if have_inner_bpen && s <= size(intr.inner_bpen, 1) pen_row = (transpose(C_coeffs) * @view(intr.inner_bpen[s, :])) ./ area diff --git a/test/runtests_eulerlagrange.jl b/test/runtests_eulerlagrange.jl index f782d314d..7adb75203 100644 --- a/test/runtests_eulerlagrange.jl +++ b/test/runtests_eulerlagrange.jl @@ -1,3 +1,5 @@ +using TOML + # TODO: this helper may belong in a shared test-utilities file rather than here. # TODO: come up with a Gaussian reduction test that doesn't rely on external data. @@ -35,8 +37,6 @@ end odet.psi_store[i] = Float64(i) odet.q_store[i] = Float64(i * 2) odet.u_store[:, :, :, i] .= ComplexF64(i) - odet.du_store[:, :, :, i] .= ComplexF64(i + 0.5) - odet.xi_s_store[:, :, i] .= ComplexF64(i + 0.25) end # Resize storage @@ -46,16 +46,16 @@ end @test length(odet.psi_store) == 2 * numsteps_init @test length(odet.q_store) == 2 * numsteps_init @test size(odet.u_store, 4) == 2 * numsteps_init - @test size(odet.du_store, 4) == 2 * numsteps_init - @test size(odet.xi_s_store, 3) == 2 * numsteps_init + + # Derivative stores are materialized after integration, so growth never touches them + @test isempty(odet.du_store) + @test isempty(odet.xi_s_store) # Check data is preserved @test all(odet.psi_store[1:odet.step] .== Float64.(1:odet.step)) @test all(odet.q_store[1:odet.step] .== Float64.(2:2:(2*odet.step))) for i in 1:odet.step @test all(odet.u_store[:, :, :, i] .== ComplexF64(i)) - @test all(odet.du_store[:, :, :, i] .== ComplexF64(i + 0.5)) - @test all(odet.xi_s_store[:, :, i] .== ComplexF64(i + 0.25)) end # Check that you can resize again @@ -63,8 +63,6 @@ end @test length(odet.psi_store) == 4 * numsteps_init @test length(odet.q_store) == 4 * numsteps_init @test size(odet.u_store, 4) == 4 * numsteps_init - @test size(odet.du_store, 4) == 4 * numsteps_init - @test size(odet.xi_s_store, 3) == 4 * numsteps_init end @testset "trim_storage!" begin @@ -79,8 +77,6 @@ end odet.psi_store[i] = Float64(i) odet.q_store[i] = Float64(i * 2) odet.u_store[:, :, :, i] .= ComplexF64(i) - odet.du_store[:, :, :, i] .= ComplexF64(i + 0.5) - odet.xi_s_store[:, :, i] .= ComplexF64(i + 0.25) end # Trim storage @@ -90,8 +86,6 @@ end @test length(odet.psi_store) == odet.step @test length(odet.q_store) == odet.step @test size(odet.u_store, 4) == odet.step - @test size(odet.du_store, 4) == odet.step - @test size(odet.xi_s_store, 3) == odet.step # Check all data is preserved @test all(odet.psi_store .== Float64.(1:odet.step)) @@ -120,13 +114,10 @@ end # Initialize index (sorted by unorm) odet.index[:, 1] = [1, 2] - # Set up some u_store, du_store, and xi_s_store data + # Set up some u_store data; derivative stores stay empty on this path for i in 1:odet.step odet.u_store[:, :, 1, i] .= ComplexF64(i) odet.u_store[:, :, 2, i] .= ComplexF64(i + 0.1) - odet.du_store[:, :, 1, i] .= ComplexF64(i + 0.2) - odet.du_store[:, :, 2, i] .= ComplexF64(i + 0.3) - odet.xi_s_store[:, :, i] .= ComplexF64(i + 0.4) end u_orig = copy(odet.u_store) @@ -141,6 +132,47 @@ end # transform_u! doesn't resize arrays - it only applies transformations in-place # The storage arrays retain their original allocated size @test size(odet.u_store) == size(u_orig) + + # Empty derivative stores must be skipped, not indexed into + @test isempty(odet.du_store) + @test isempty(odet.xi_s_store) + end + + @testset "transform_u! with pre-filled derivative stores" begin + # The galerkin-matched path supplies analytic derivatives before the fixup transforms, + # so those arrays must be mixed by the same fixfac matrices as u_store. + mpert = 2 + intr = GeneralizedPerturbedEquilibrium.ForceFreeStates.ForceFreeStatesInternal(; mpert=mpert, numpert_total=mpert) + odet = GeneralizedPerturbedEquilibrium.ForceFreeStates.OdeState(mpert, 10, 5, 2) + + odet.ifix = 1 + odet.step = 5 + odet.sing_flag[1] = false + odet.fixstep[1] = 3 + odet.zeroed_idx[1] = Int[] + odet.fixfac[1, 1, 1] = 1.0 + odet.fixfac[1, 2, 1] = 0.5 + odet.fixfac[2, 1, 1] = 0.0 + odet.fixfac[2, 2, 1] = 1.0 + odet.index[:, 1] = [1, 2] + + odet.du_store = zeros(ComplexF64, mpert, mpert, odet.step) + odet.xi_s_store = zeros(ComplexF64, mpert, mpert, odet.step) + for i in 1:odet.step + odet.u_store[:, :, 1, i] .= ComplexF64(i) + odet.u_store[:, :, 2, i] .= ComplexF64(i + 0.1) + odet.du_store[:, :, i] .= ComplexF64(i + 0.2) + odet.xi_s_store[:, :, i] .= ComplexF64(i + 0.4) + end + du_orig = copy(odet.du_store) + xi_s_orig = copy(odet.xi_s_store) + + GeneralizedPerturbedEquilibrium.ForceFreeStates.transform_u!(odet, intr) + + @test size(odet.du_store) == size(du_orig) + @test size(odet.xi_s_store) == size(xi_s_orig) + @test !all(odet.du_store .== du_orig) + @test !all(odet.xi_s_store .== xi_s_orig) end @testset "apply_gaussian_reduction!" begin @@ -282,11 +314,12 @@ end # Check array dimensions @test size(odet.u) == (numpert_total, numpert_total, 2) - @test size(odet.du) == (numpert_total, numpert_total, 2) - @test size(odet.xi_s) == (numpert_total, numpert_total) @test size(odet.u_store) == (numpert_total, numpert_total, 2, numsteps_init) - @test size(odet.du_store) == (numpert_total, numpert_total, 2, numsteps_init) - @test size(odet.xi_s_store) == (numpert_total, numpert_total, numsteps_init) + # Derivative stores start empty and are sized when materialized + @test size(odet.du_store) == (numpert_total, numpert_total, 0) + @test size(odet.xi_s_store) == (numpert_total, numpert_total, 0) + @test odet.du_store_populated == false + @test odet.u_store_el_basis == true @test length(odet.psi_store) == numsteps_init @test length(odet.q_store) == numsteps_init @test size(odet.ca_r) == (numpert_total, numpert_total, 2, msing) @@ -390,3 +423,97 @@ end @test ode.step == 1 end end + +@testset "materialize_derivative_stores!" begin + FFS = GeneralizedPerturbedEquilibrium.ForceFreeStates + + # Integrate a small ideal case and hand back everything the materializer needs. + function setup_solovev_run() + example_dir = joinpath(@__DIR__, "test_data", "regression_solovev_ideal_example") + inputs = TOML.parsefile(joinpath(example_dir, "gpec.toml")) + inputs["ForceFreeStates"]["verbose"] = false + inputs["ForceFreeStates"]["use_parallel"] = false + inputs["ForceFreeStates"]["write_outputs_to_HDF5"] = false + intr = FFS.ForceFreeStatesInternal(; dir_path=example_dir) + ctrl = FFS.ForceFreeStatesControl(; (Symbol(k) => v for (k, v) in inputs["ForceFreeStates"])...) + eq_config = GeneralizedPerturbedEquilibrium.Equilibrium.EquilibriumConfig(inputs["Equilibrium"], example_dir) + sol_cfg = haskey(inputs, "SOL_INPUT") ? GeneralizedPerturbedEquilibrium.Equilibrium.SolovevConfig(inputs["SOL_INPUT"]) : nothing + equil = GeneralizedPerturbedEquilibrium.Equilibrium.setup_equilibrium(eq_config, sol_cfg) + intr.wall_settings = GeneralizedPerturbedEquilibrium.Vacuum.WallShapeSettings(; (Symbol(k) => v for (k, v) in inputs["Wall"])...) + FFS.sing_lim!(intr, ctrl, equil) + intr.nlow = ctrl.nn_low + intr.nhigh = ctrl.nn_high + intr.npert = 1 + FFS.sing_find!(intr, equil) + intr.mlow = min(intr.nlow * equil.params.qmin, 0) - 4 - ctrl.delta_mlow + intr.mhigh = trunc(Int, intr.nhigh * equil.params.qmax) + ctrl.delta_mhigh + intr.mpert = intr.mhigh - intr.mlow + 1 + intr.numpert_total = intr.mpert * intr.npert + metric = FFS.make_metric(equil, intr.mpert) + ffit = FFS.make_matrix(equil, intr, metric) + odet, _, _, _ = FFS.eulerlagrange_integration(ctrl, equil, ffit, intr) + return odet, ctrl, equil, ffit, intr + end + + odet, ctrl, equil, ffit, intr = setup_solovev_run() + # Untouched copy of the solution, for the column-transform check further down. + odet_pristine = deepcopy(odet) + + @testset "fills the stores once" begin + @test isempty(odet.du_store) + @test !odet.du_store_populated + @test FFS.materialize_derivative_stores!(odet, equil, ffit, intr) + @test odet.du_store_populated + @test size(odet.du_store) == (intr.numpert_total, intr.numpert_total, odet.step) + @test size(odet.xi_s_store) == (intr.numpert_total, intr.numpert_total, odet.step) + @test all(isfinite, abs.(odet.du_store)) + @test all(isfinite, abs.(odet.xi_s_store)) + + # Idempotent: a second call must not overwrite what is already there. + du_first = copy(odet.du_store) + @test FFS.materialize_derivative_stores!(odet, equil, ffit, intr) + @test odet.du_store == du_first + end + + @testset "agrees with a direct kernel evaluation" begin + npert = intr.numpert_total + du = zeros(ComplexF64, npert, npert, 2) + xi_s = zeros(ComplexF64, npert, npert) + for istep in (1, odet.step ÷ 2, odet.step) + psi = odet.psi_store[istep] + u = odet.u_store[:, :, :, istep] + FFS.el_derivatives!(du, u, false, equil, ffit, intr, psi, Ref(1), Ref(1)) + FFS.compute_node_xi_s!(xi_s, @view(du[:, :, 1]), @view(u[:, :, 1]), ffit, psi) + @test odet.du_store[:, :, istep] == du[:, :, 1] + @test odet.xi_s_store[:, :, istep] == xi_s + end + end + + @testset "commutes with a column transform" begin + # The design relies on du(psi, u*T) == du(psi, u)*T, which is what makes it exact to + # materialize after the Gaussian fixups and free-boundary normalization rather than + # transforming stored derivatives alongside u_store. + npert = intr.numpert_total + T = Matrix{ComplexF64}(I, npert, npert) .+ 0.25 .* ComplexF64.(reshape(sin.(1:npert^2), npert, npert)) + odet_t = deepcopy(odet_pristine) + for istep in 1:odet_t.step + odet_t.u_store[:, :, 1, istep] = odet_t.u_store[:, :, 1, istep] * T + odet_t.u_store[:, :, 2, istep] = odet_t.u_store[:, :, 2, istep] * T + end + @test FFS.materialize_derivative_stores!(odet_t, equil, ffit, intr) + + for istep in (1, odet.step ÷ 2, odet.step) + @test isapprox(odet_t.du_store[:, :, istep], odet.du_store[:, :, istep] * T; rtol=1e-10) + @test isapprox(odet_t.xi_s_store[:, :, istep], odet.xi_s_store[:, :, istep] * T; rtol=1e-10) + end + end + + @testset "refuses a solution outside the Euler-Lagrange basis" begin + odet.du_store_populated = false + odet.du_store = Array{ComplexF64}(undef, intr.numpert_total, intr.numpert_total, 0) + odet.u_store_el_basis = false + @test !FFS.materialize_derivative_stores!(odet, equil, ffit, intr) + @test isempty(odet.du_store) + @test !odet.du_store_populated + end +end diff --git a/test/runtests_parallel_integration.jl b/test/runtests_parallel_integration.jl index 7f6c57394..666191d17 100644 --- a/test/runtests_parallel_integration.jl +++ b/test/runtests_parallel_integration.jl @@ -428,6 +428,8 @@ using TOML metric = GeneralizedPerturbedEquilibrium.ForceFreeStates.make_metric(equil, intr.mpert) ffit = GeneralizedPerturbedEquilibrium.ForceFreeStates.make_matrix(equil, intr, metric) odet, _, _, _ = GeneralizedPerturbedEquilibrium.ForceFreeStates.eulerlagrange_integration(ctrl, equil, ffit, intr) + # Derivatives are recomputed on demand; materialize so the stores can be compared. + GeneralizedPerturbedEquilibrium.ForceFreeStates.materialize_derivative_stores!(odet, equil, ffit, intr) return odet end @@ -443,6 +445,8 @@ using TOML @test length(odet_a.q_store) == length(odet_b.q_store) @test size(odet_a.u_store) == size(odet_b.u_store) @test size(odet_a.du_store) == size(odet_b.du_store) + @test size(odet_a.xi_s_store) == size(odet_b.xi_s_store) + @test odet_a.du_store_populated == odet_b.du_store_populated @test maximum(abs.(odet_a.psi_store .- odet_b.psi_store)) == 0.0 @test maximum(abs.(odet_a.q_store .- odet_b.q_store)) == 0.0 @test maximum(abs.(odet_a.u_store .- odet_b.u_store)) == 0.0 @@ -476,10 +480,13 @@ using TOML odet_std = run_and_capture(ex, false) odet_sparse = run_and_capture(ex, true; populate_dense_xi=false) @test odet_sparse.step < odet_std.step - # du_store entries inside FM chunks are left at the @kwdef - # `undef` initial value when populate_dense_xi=false; ensure the - # array IS smaller (sparse). @test length(odet_sparse.psi_store) < length(odet_std.psi_store) + # The sparse solution is in the Riccati basis, so the derivative stores cannot be + # materialized from it and stay empty rather than holding unusable values. + @test !odet_sparse.u_store_el_basis + @test !odet_sparse.du_store_populated + @test isempty(odet_sparse.du_store) + @test isempty(odet_sparse.xi_s_store) end end