diff --git a/.gitignore b/.gitignore index fcb3cc96c..6fa3a0d46 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,7 @@ !docs/src/assets/**/*.png # committed documentation figures (generated by docs/src/figures/**/make_*.jl) !docs/src/figures/**/*.png +!docs/development/figures/**/*.png *.jld2 .gitattributes Manifest.toml diff --git a/benchmarks/benchmark_integrators.jl b/benchmarks/benchmark_integrators.jl new file mode 100644 index 000000000..23718974f --- /dev/null +++ b/benchmarks/benchmark_integrators.jl @@ -0,0 +1,338 @@ +# Work-precision comparison of OrdinaryDiffEq integrators on the STRIDE-equivalent step of +# julia_GPEC: the fundamental-matrix (FM) propagator chunks of the Euler-Lagrange ODE. +# +# A chunk propagator is the fundamental matrix Φ(ψ₂,ψ₁) of the Euler-Lagrange system over one +# sub-interval of the outer region, obtained by integrating from the two identity-block initial +# conditions (I,0) and (0,I) [Glasser 2018 Phys. Plasmas 25, 032507]. Assembling the chunk +# propagators serially with periodic renormalization is what makes the chunked BVP well +# conditioned, and it is where essentially all of the Riccati solver's ODE time is spent. +# +# Accuracy criterion: each candidate is compared against a Vern9 reference at reltol=1e-13, +# abstol=1e-15 through the relative Frobenius error of the end-state blocks, maxed over the two +# initial conditions. The row `Vern9` at the case's own `eulerlagrange_tolerance` is the +# production setting and is the accuracy floor every candidate must be judged against — an +# integrator is only a drop-in replacement if it is at least as accurate as that row. +# +# Usage: +# julia --project=. -t 1 benchmarks/benchmark_integrators.jl --case [options] +# +# --case case directory containing eq.geqdsk and gpec.toml (required) +# --algs comma-separated OrdinaryDiffEq algorithm names +# --rtols comma-separated relative tolerances (default 1e-6,1e-8,1e-10) +# --reps timed repetitions, minimum is kept (default 3) +# --out CSV output path (default benchmarks/integrator_results/.csv) +# --quick smoke test: 3 chunks, Vern9/Tsit5/VCABM, rtol 1e-8, 1 rep +# --include-outer also time the serial outer-plasma Riccati solve with its callback + +using LinearAlgebra, Printf, TOML +using GeneralizedPerturbedEquilibrium +using OrdinaryDiffEq + +const GPE = GeneralizedPerturbedEquilibrium +const FFS = GeneralizedPerturbedEquilibrium.ForceFreeStates + +# Absolute tolerance for candidate solves; `nothing` keeps the production default (OrdinaryDiffEq 1e-6). +const CANDIDATE_ABSTOL = Ref{Union{Nothing,Float64}}(nothing) + +const DEFAULT_ALGS = ["Tsit5", "BS5", "DP5", "Vern6", "Vern7", "Vern8", "Vern9", "DP8", "TanYam7", "TsitPap8", + "Feagin10", "Feagin12", "Feagin14", "VCAB4", "VCAB5", "VCABM3", "VCABM4", "VCABM5", "VCABM", "AN5"] + +# ---------------------------------------------------------------- CLI + +function parse_args(args) + opts = Dict{String,Any}("case" => nothing, "algs" => DEFAULT_ALGS, "rtols" => [1e-6, 1e-8, 1e-10], + "reps" => 3, "out" => nothing, "quick" => false, "include_outer" => false, "blas" => nothing, "abstol" => nothing) + i = 1 + while i <= length(args) + a = args[i] + if a == "--case" + opts["case"] = args[i+1] + i += 2 + elseif a == "--algs" + opts["algs"] = String.(split(args[i+1], ",")) + i += 2 + elseif a == "--rtols" + opts["rtols"] = parse.(Float64, split(args[i+1], ",")) + i += 2 + elseif a == "--reps" + opts["reps"] = parse(Int, args[i+1]) + i += 2 + elseif a == "--out" + opts["out"] = args[i+1] + i += 2 + elseif a == "--quick" + opts["quick"] = true + i += 1 + elseif a == "--include-outer" + opts["include_outer"] = true + i += 1 + elseif a == "--blas-threads" + opts["blas"] = parse(Int, args[i+1]) + i += 2 + elseif a == "--abstol" + opts["abstol"] = parse(Float64, args[i+1]) + i += 2 + else + error("Unknown argument: $a") + end + end + opts["case"] === nothing && error("--case is required") + if opts["quick"] + opts["algs"] = ["Vern9", "Tsit5", "VCABM"] + opts["rtols"] = [1e-8] + opts["reps"] = 1 + end + return opts +end + +# ---------------------------------------------------------------- case setup + +# Build (ctrl, equil, mats, intr) exactly as the production pipeline does: mode space from +# resolve_mode_space!, the two-pass auto grid re-formed once, matrices from prepare_force_free_states!. +function setup_case(dir::AbstractString) + inputs = TOML.parsefile(joinpath(dir, "gpec.toml")) + inputs["ForceFreeStates"]["verbose"] = false + inputs["ForceFreeStates"]["write_outputs_to_HDF5"] = false + inputs["ForceFreeStates"]["integrator"] = "riccati" + ctrl = FFS.ForceFreeStatesControl(; (Symbol(k) => v for (k, v) in inputs["ForceFreeStates"])...) + eq_config = GPE.Equilibrium.EquilibriumConfig(inputs["Equilibrium"], dir) + intr = FFS.ForceFreeStatesInternal(; dir_path=dir) + intr.wall_settings = GPE.Vacuum.WallShapeSettings(; (Symbol(k) => v for (k, v) in inputs["Wall"])...) + GPE.resolve_mode_space!(intr, ctrl) + equil = GPE.Equilibrium.setup_equilibrium(eq_config, nothing) + equil = GPE.maybe_reform_equilibrium(equil, eq_config, nothing, intr, ctrl, nothing) + metric, mats = GPE.prepare_force_free_states!(intr, ctrl, equil, GPE.KineticForces.KineticForcesControl(), nothing) + return ctrl, equil, mats, intr +end + +# ---------------------------------------------------------------- one chunk + +# Reproduces integrate_propagator_chunk! standalone: two identity-block ICs, hints reset before +# each solve, tspan reversed for backward (direction=-1) chunks. +function solve_chunk(chunk, alg, ctrl, equil, mats, intr, proxy; reltol, abstol=nothing) + N = intr.numpert_total + tspan = chunk.direction == 1 ? (chunk.psi_start, chunk.psi_end) : (chunk.psi_end, chunk.psi_start) + params = (ctrl, equil, mats, intr, proxy, chunk) + blocks = Vector{Array{ComplexF64,3}}(undef, 2) + nf = naccept = nreject = 0 + for (ic, slot) in ((1, 1), (2, 2)) + u0 = zeros(ComplexF64, N, N, 2) + for i in 1:N + u0[i, i, ic] = 1 + end + proxy.spline_hint[] = 1 + proxy.mats_hint[] = 1 + prob = ODEProblem(FFS.sing_der!, u0, tspan, params) + atol = abstol === nothing ? CANDIDATE_ABSTOL[] : abstol + sol = + atol === nothing ? solve(prob, alg; reltol=reltol, save_everystep=false, save_end=true) : + solve(prob, alg; reltol=reltol, abstol=atol, save_everystep=false, save_end=true) + blocks[slot] = sol.u[end] + nf += sol.stats.nf + naccept += sol.stats.naccept + nreject += sol.stats.nreject + end + return blocks, nf, naccept, nreject +end + +relerr(a, b) = norm(a .- b) / norm(b) +chunk_error(blocks, ref) = max(relerr(blocks[1], ref[1]), relerr(blocks[2], ref[2])) + +# ---------------------------------------------------------------- sweep + +function sweep(algname, rtol, chunks, ctrl, equil, mats, intr, refs, reps) + N = intr.numpert_total + proxy = FFS.OdeState(N, 1, 1, 0) + alg = try + getfield(OrdinaryDiffEq, Symbol(algname))() + catch err + @warn "construction failed for $algname: $(sprint(showerror, err))" + return nothing + end + try + solve_chunk(chunks[1], alg, ctrl, equil, mats, intr, proxy; reltol=rtol) # JIT warm-up, untimed + catch err + @warn "first solve failed for $algname @ rtol=$rtol: $(sprint(showerror, err))" + return nothing + end + best = Inf + nf = naccept = nreject = 0 + held = Vector{Vector{Array{ComplexF64,3}}}(undef, length(chunks)) + per_naccept = zeros(Int, length(chunks)) + for _ in 1:reps + nf_r = naccept_r = nreject_r = 0 + t = @elapsed for (i, ch) in enumerate(chunks) + blocks, a, b, c = solve_chunk(ch, alg, ctrl, equil, mats, intr, proxy; reltol=rtol) + held[i] = blocks + per_naccept[i] = b + nf_r += a + naccept_r += b + nreject_r += c + end + if t < best + best = t + nf, naccept, nreject = nf_r, naccept_r, nreject_r + end + end + errs = [chunk_error(held[i], refs[i]) for i in eachindex(chunks)] + return (; wall=best, nf=nf, naccept=naccept, nreject=nreject, errs=errs, per_naccept=per_naccept) +end + +# Serial outer-plasma Riccati solve (axis to the first crossing) with its DiscreteCallback. +# Only the first chunk is run: it is the representative callback-carrying solve and needs no +# singular-surface crossing machinery. +function sweep_outer(algname, rtol, ctrl, equil, mats, intr, reps) + alg = try + getfield(OrdinaryDiffEq, Symbol(algname))() + catch err + @warn "construction failed for $algname (outer): $(sprint(showerror, err))" + return nothing + end + saved = FFS.EL_ODE_ALGORITHM[] + FFS.EL_ODE_ALGORITHM[] = alg + ctrl_local = ctrl + best = Inf + ustate = nothing + try + for r in 0:reps + odet = FFS._initialize_parallel_odet(ctrl_local, equil, mats, intr) + serial = FFS.chunk_el_integration_bounds(odet, ctrl_local, intr; bidirectional=false) + t = @elapsed FFS.riccati_integrate_chunk!(odet, ctrl_local, equil, mats, intr, serial[1]) + r == 0 && continue # warm-up + best = min(best, t) + ustate = copy(odet.u) + end + catch err + @warn "outer solve failed for $algname @ rtol=$rtol: $(sprint(showerror, err))" + FFS.EL_ODE_ALGORITHM[] = saved + return nothing + end + FFS.EL_ODE_ALGORITHM[] = saved + return (; wall=best, u=ustate) +end + +# ---------------------------------------------------------------- driver + +function main(args) + opts = parse_args(args) + casedir = abspath(opts["case"]) + casename = basename(rstrip(casedir, '/')) + outpath = opts["out"] === nothing ? joinpath(@__DIR__, "integrator_results", casename * ".csv") : abspath(opts["out"]) + mkpath(dirname(outpath)) + chunkpath = replace(outpath, r"\.csv$" => "") * "_chunks.csv" + + opts["blas"] !== nothing && BLAS.set_num_threads(opts["blas"]) + CANDIDATE_ABSTOL[] = opts["abstol"] + ctrl, equil, mats, intr = setup_case(casedir) + N = intr.numpert_total + odet = FFS._initialize_parallel_odet(ctrl, equil, mats, intr) + chunks, _, _ = FFS._setup_parallel_chunks_and_proxies(odet, ctrl, intr) + opts["quick"] && (chunks = chunks[1:min(3, length(chunks))]) + + @printf("\ncase %s\n", casename) + @printf("N (numpert_total) %d\n", N) + @printf("BLAS threads %d\n", BLAS.get_num_threads()) + @printf("candidate abstol %s\n", CANDIDATE_ABSTOL[] === nothing ? "default (1e-6)" : string(CANDIDATE_ABSTOL[])) + @printf("msing %d\n", intr.msing) + @printf("nchunks %d\n", length(chunks)) + @printf("psi range [%.6f, %.6f] over %d chunk boundaries\n", + minimum(c.psi_start for c in chunks), maximum(c.psi_end for c in chunks), length(chunks) + 1) + @printf("production rtol %.1e (eulerlagrange_tolerance, Vern9)\n\n", ctrl.eulerlagrange_tolerance) + + println("computing Vern9 reference (reltol=1e-13, abstol=1e-15) ...") + refproxy = FFS.OdeState(N, 1, 1, 0) + refs = [solve_chunk(ch, Vern9(), ctrl, equil, mats, intr, refproxy; reltol=1e-13, abstol=1e-15)[1] for ch in chunks] + println("reference done.\n") + + rows = Any[] + chunkrows = Dict{Tuple{String,Float64},Any}() + # The production setting first: it defines the accuracy floor the candidates are judged against. + prod_rtol = ctrl.eulerlagrange_tolerance + pairs = [("Vern9", prod_rtol)] + for rtol in opts["rtols"], a in opts["algs"] + (a, rtol) in pairs || push!(pairs, (a, rtol)) + end + + for (algname, rtol) in pairs + res = sweep(algname, rtol, chunks, ctrl, equil, mats, intr, refs, opts["reps"]) + if res === nothing + push!(rows, (casename, algname, rtol, "failed", length(chunks), N, intr.msing, NaN, 0, 0, 0, NaN, NaN)) + continue + end + us_per_rhs = res.nf > 0 ? res.wall / res.nf * 1e6 : NaN + push!(rows, (casename, algname, rtol, "ok", length(chunks), N, intr.msing, res.wall, res.nf, + res.naccept, res.nreject, maximum(res.errs), median_of(res.errs))) + chunkrows[(algname, rtol)] = res + @printf(" %-10s rtol=%.0e %7.3f s nf=%-8d err_max=%.3e\n", algname, rtol, res.wall, res.nf, maximum(res.errs)) + end + + if opts["include_outer"] + for rtol in opts["rtols"], algname in opts["algs"] + r = sweep_outer(algname, rtol, ctrl, equil, mats, intr, opts["reps"]) + r === nothing && continue + push!(rows, (casename, algname * "+outer", rtol, "ok", 1, N, intr.msing, r.wall, 0, 0, 0, NaN, NaN)) + end + end + + write_csv(outpath, rows) + write_chunk_csv(chunkpath, casename, chunks, chunkrows, rows) + print_table(rows, prod_rtol) + @printf("\nwrote %s\n %s\n", outpath, chunkpath) + return nothing +end + +function median_of(v) + isempty(v) && return NaN + s = sort(v) + n = length(s) + return isodd(n) ? s[(n+1)÷2] : 0.5 * (s[n÷2] + s[n÷2+1]) +end + +const HEADER = "case,alg,rtol,status,nchunks,N,msing,wall_s_min,nf,naccept,nreject,us_per_rhs,err_max,err_median" + +function write_csv(path, rows) + open(path, "w") do io + println(io, HEADER) + for r in rows + us = r[9] > 0 ? r[8] / r[9] * 1e6 : NaN + @printf(io, "%s,%s,%.3e,%s,%d,%d,%d,%.6f,%d,%d,%d,%.4f,%.6e,%.6e\n", + r[1], r[2], r[3], r[4], r[5], r[6], r[7], r[8], r[9], r[10], r[11], us, r[12], r[13]) + end + end +end + +# Per-chunk detail for the three fastest algorithms, so step concentration is visible. +function write_chunk_csv(path, casename, chunks, chunkrows, rows) + ok = filter(r -> r[4] == "ok" && !occursin("+outer", r[2]), rows) + ranked = sort(unique(r[2] for r in ok); by=a -> minimum(r[8] for r in ok if r[2] == a)) + best3 = ranked[1:min(3, length(ranked))] + open(path, "w") do io + println(io, "case,alg,rtol,chunk,psi_start,psi_end,direction,naccept,err") + for ((algname, rtol), res) in sort(collect(chunkrows); by=kv -> (kv[1][1], kv[1][2])) + algname in best3 || continue + for (i, ch) in enumerate(chunks) + @printf(io, "%s,%s,%.3e,%d,%.8f,%.8f,%d,%d,%.6e\n", + casename, algname, rtol, i, ch.psi_start, ch.psi_end, ch.direction, res.per_naccept[i], res.errs[i]) + end + end + end +end + +function print_table(rows, prod_rtol) + println("\n", "="^108) + @printf("%-12s %-9s %-7s %10s %10s %9s %9s %10s %11s %11s\n", + "alg", "rtol", "status", "wall_s", "nf", "naccept", "nreject", "us/rhs", "err_max", "err_med") + println("="^108) + for rtol in sort(unique(r[3] for r in rows)) + sub = sort(filter(r -> r[3] == rtol, rows); by=r -> (r[4] == "ok" ? 0 : 1, isnan(r[8]) ? Inf : r[8])) + for r in sub + us = r[9] > 0 ? r[8] / r[9] * 1e6 : NaN + mark = (r[2] == "Vern9" && rtol == prod_rtol) ? " <- production" : "" + @printf("%-12s %-9.0e %-7s %10.3f %10d %9d %9d %10.4f %11.3e %11.3e%s\n", + r[2], r[3], r[4], r[8], r[9], r[10], r[11], us, r[12], r[13], mark) + end + println("-"^108) + end +end + +main(ARGS) diff --git a/benchmarks/benchmark_integrators_pipeline.jl b/benchmarks/benchmark_integrators_pipeline.jl new file mode 100644 index 000000000..55086ecf6 --- /dev/null +++ b/benchmarks/benchmark_integrators_pipeline.jl @@ -0,0 +1,223 @@ +#!/usr/bin/env julia +""" +benchmark_integrators_pipeline.jl - Whole-pipeline integrator comparison for the Euler-Lagrange sweep. + +Runs the Force-Free States stage of one case (equilibrium and matrix splines prepared once, untimed) +with every requested OrdinaryDiffEq algorithm and relative tolerance, selected through the +`ode_solver` / `ode_abstol` control fields. Times the Euler-Lagrange integration, the free-boundary energies and +the Δ′ BVP separately, and compares the products every downstream consumer reads — the free-boundary +energy eigenvalues `et` and the singular-surface Δ′ matrix — against two references: + + - `ref`: Vern9 at the case's shipped `eulerlagrange_tolerance` (today's production numbers), and + - `truth`: Vern9 at reltol 1e-13, the best answer this pipeline can produce. + +The acceptance rule is "no loss in accuracy": a candidate is only interesting if its distance from +`truth` is no larger than the production reference's own distance from `truth`. + +# Usage + +```bash +julia --project=. -t 1 benchmarks/benchmark_integrators_pipeline.jl --case [options] + + --case directory holding eq.geqdsk and gpec.toml (required) + --algs a,b,c OrdinaryDiffEq algorithm names (default: the explicit non-stiff set below) + --rtols 1e-8,1e-10 relative tolerances to sweep (default: the case's eulerlagrange_tolerance) + --reps timed repetitions per (alg, rtol); the minimum is reported (default 2) + --mode riccati|forward Euler-Lagrange formalism (default riccati, the STRIDE-equivalent path) + --blas-threads BLAS.set_num_threads(k) before timing (default: leave as started) + --abstol absolute tolerance for every candidate solve (default 1e-6, the OrdinaryDiffEq default) + --out output path (default benchmarks/integrator_results/_pipeline.csv) + --no-truth skip the reltol 1e-13 truth run +``` + +Inputs are read only from the case directory; outputs are written under `benchmarks/`. +""" + +using LinearAlgebra +using Printf +using Statistics +using TOML +using OrdinaryDiffEq +using GeneralizedPerturbedEquilibrium + +const GPE = GeneralizedPerturbedEquilibrium +const FFS = GPE.ForceFreeStates + +const DEFAULT_ALGS = ["Vern9", "Vern8", "Vern7", "Vern6", "Tsit5", "BS5", "DP5", "DP8", "TanYam7", "TsitPap8", + "Feagin10", "Feagin12", "Feagin14", "VCAB4", "VCAB5", "VCABM3", "VCABM4", "VCABM5", "VCABM", "AN5"] +const TRUTH_RTOL = 1e-13 +const TRUTH_ABSTOL = 1e-15 + +function parse_args(args) + opts = Dict{String,Any}("algs" => DEFAULT_ALGS, "rtols" => nothing, "reps" => 2, "mode" => "riccati", + "blas" => nothing, "out" => nothing, "truth" => true, "case" => nothing, "abstol" => 1e-6) + i = 1 + while i <= length(args) + a = args[i] + if a == "--case" + opts["case"] = abspath(args[i+1]) + i += 2 + elseif a == "--algs" + opts["algs"] = split(args[i+1], ",") + i += 2 + elseif a == "--rtols" + opts["rtols"] = parse.(Float64, split(args[i+1], ",")) + i += 2 + elseif a == "--reps" + opts["reps"] = parse(Int, args[i+1]) + i += 2 + elseif a == "--mode" + opts["mode"] = args[i+1] + i += 2 + elseif a == "--blas-threads" + opts["blas"] = parse(Int, args[i+1]) + i += 2 + elseif a == "--out" + opts["out"] = abspath(args[i+1]) + i += 2 + elseif a == "--abstol" + opts["abstol"] = parse(Float64, args[i+1]) + i += 2 + elseif a == "--no-truth" + opts["truth"] = false + i += 1 + else + error("unknown argument $a") + end + end + opts["case"] === nothing && error("--case is required") + opts["mode"] in ("riccati", "forward") || error("--mode must be riccati or forward") + return opts +end + +# Equilibrium, wall and control keywords from the case's gpec.toml; the two-pass auto grid is +# re-formed once here so every timed run integrates against the same equilibrium. +function prepare_case(case::String, mode::String) + inputs = TOML.parsefile(joinpath(case, "gpec.toml")) + eq_config = GPE.Equilibrium.EquilibriumConfig(inputs["Equilibrium"], case) + wall = GPE.Vacuum.WallShapeSettings(; (Symbol(k) => v for (k, v) in inputs["Wall"])...) + kw = Dict{Symbol,Any}(Symbol(k) => v for (k, v) in inputs["ForceFreeStates"]) + kw[:verbose] = false + kw[:write_outputs_to_HDF5] = false + kw[:integrator] = mode + ctrl = FFS.ForceFreeStatesControl(; kw...) + + intr = FFS.ForceFreeStatesInternal(; dir_path=case) + intr.wall_settings = wall + GPE.resolve_mode_space!(intr, ctrl) + equil = GPE.Equilibrium.setup_equilibrium(eq_config, nothing) + equil = GPE.maybe_reform_equilibrium(equil, eq_config, nothing, intr, ctrl, nothing) + kf_ctrl = GPE.KineticForces.KineticForcesControl() + t_prep = @elapsed metric, mats = GPE.prepare_force_free_states!(intr, ctrl, equil, kf_ctrl, nothing) + return (; equil, wall, kw, intr, metric, mats, t_prep) +end + +# One Force-Free States run mirroring run_force_free_states, with the three stages timed separately. +function run_once(prep, alg::AbstractString, rtol::Float64, case::String; abstol::Float64=1e-6) + kw = copy(prep.kw) + kw[:eulerlagrange_tolerance] = rtol + kw[:ode_solver] = String(alg) + kw[:ode_abstol] = abstol + ctrl = FFS.ForceFreeStatesControl(; kw...) + intr = deepcopy(prep.intr) + equil, mats, metric = prep.equil, prep.mats, prep.metric + + t_int = @elapsed odet, fm_propagators, fm_chunks, fm_S_left = FFS.eulerlagrange_integration(ctrl, equil, mats, intr) + free_energies = nothing + t_free = 0.0 + t_dp = 0.0 + if ctrl.vac_flag + t_free = @elapsed begin + free_energies = FFS.free_run(odet, ctrl, equil, mats, intr) + FFS.normalize_eigenfunctions!(odet, free_energies.wt, equil.psio) + end + if intr.msing > 0 && fm_propagators !== nothing + t_dp = @elapsed FFS.compute_delta_prime_matrix!(intr, fm_propagators, fm_chunks; + wv=free_energies.wv, psio=equil.psio, debug=false, S_at_surface_left=fm_S_left, ctrl=ctrl, equil=equil, mats=mats) + end + end + res = FFS.build_result(Symbol(ctrl.integrator), ctrl, equil, intr, metric, mats, odet, free_energies, nothing, nothing) + et = free_energies === nothing ? ComplexF64[] : free_energies.et + dp = res.delta_prime === nothing ? nothing : res.delta_prime.matrix + return (; t_int, t_free, t_dp, total_steps=odet.total_steps, nzero=odet.nzero, et, dp, nchunks=fm_chunks === nothing ? 0 : length(fm_chunks)) +end + +relerr(a, b) = norm(a - b) / max(norm(b), eps()) +et_err(r, ref) = (isempty(r.et) || isempty(ref.et)) ? NaN : relerr(r.et[1:min(5, end)], ref.et[1:min(5, end)]) +dp_err(r, ref) = (r.dp === nothing || ref.dp === nothing) ? NaN : relerr(r.dp, ref.dp) + +function main(args) + opts = parse_args(args) + case = opts["case"] + casename = basename(case) + opts["blas"] !== nothing && BLAS.set_num_threads(opts["blas"]) + out = something(opts["out"], joinpath(@__DIR__, "integrator_results", "$(casename)_pipeline.csv")) + mkpath(dirname(out)) + + @info "Preparing $casename (mode=$(opts["mode"]))" + prep = prepare_case(case, opts["mode"]) + rtol0 = Float64(prep.kw[:eulerlagrange_tolerance]) + rtols = something(opts["rtols"], [rtol0]) + N = prep.intr.numpert_total + @printf("case %s: N=%d msing=%d threads=%d blas=%d prep=%.2fs shipped rtol=%.1e abstol=%.1e\n", + casename, N, prep.intr.msing, Threads.nthreads(), BLAS.get_num_threads(), prep.t_prep, rtol0, opts["abstol"]) + + # JIT warm-up and the two references. + run_once(prep, "Vern9", rtol0, case) + ref = run_once(prep, "Vern9", rtol0, case) + truth = opts["truth"] ? run_once(prep, "Vern9", TRUTH_RTOL, case; abstol=TRUTH_ABSTOL) : ref + @printf("reference Vern9@%.0e: steps=%d t_int=%.2fs et1=%.6e | truth Vern9@%.0e/abstol 1e-15: steps=%d t_int=%.2fs et1=%.6e | ref-vs-truth et=%.2e dp=%.2e\n", + rtol0, ref.total_steps, ref.t_int, real(ref.et[1]), TRUTH_RTOL, truth.total_steps, truth.t_int, real(truth.et[1]), + et_err(ref, truth), dp_err(ref, truth)) + + header = + "case,alg,rtol,abstol,status,mode,nthreads,blas_threads,N,msing,nchunks,total_steps,t_int_s,t_free_s,t_dp_s,t_total_s," * + "et1_re,nzero,err_et_vs_ref,err_dp_vs_ref,err_et_vs_truth,err_dp_vs_truth" + rows = String[] + for rtol in rtols, name in opts["algs"] + if !isdefined(OrdinaryDiffEq, Symbol(name)) + @warn "unknown OrdinaryDiffEq algorithm $name" + push!( + rows, + "$casename,$name,$rtol,$(opts["abstol"]),failed,$(opts["mode"]),$(Threads.nthreads()),$(BLAS.get_num_threads()),$N,$(prep.intr.msing)," * join(fill("", 12), ",") + ) + continue + end + alg = name + best = nothing + status = "ok" + try + run_once(prep, alg, rtol, case; abstol=opts["abstol"]) # warm-up for this algorithm's compiled code + for _ in 1:opts["reps"] + r = run_once(prep, alg, rtol, case; abstol=opts["abstol"]) + best = (best === nothing || r.t_int + r.t_free + r.t_dp < best.t_int + best.t_free + best.t_dp) ? r : best + end + catch err + status = "failed" + @warn "$name @ $rtol failed: $(sprint(showerror, err))" + end + if best === nothing + push!( + rows, + "$casename,$name,$rtol,$(opts["abstol"]),$status,$(opts["mode"]),$(Threads.nthreads()),$(BLAS.get_num_threads()),$N,$(prep.intr.msing)," * join(fill("", 12), ",") + ) + continue + end + t_total = best.t_int + best.t_free + best.t_dp + line = @sprintf("%s,%s,%.1e,%.1e,%s,%s,%d,%d,%d,%d,%d,%d,%.4f,%.4f,%.4f,%.4f,%.10e,%d,%.3e,%.3e,%.3e,%.3e", + casename, name, rtol, opts["abstol"], status, opts["mode"], Threads.nthreads(), BLAS.get_num_threads(), N, prep.intr.msing, best.nchunks, + best.total_steps, best.t_int, best.t_free, best.t_dp, t_total, isempty(best.et) ? NaN : real(best.et[1]), best.nzero, + et_err(best, ref), dp_err(best, ref), et_err(best, truth), dp_err(best, truth)) + push!(rows, line) + @printf("%-9s rtol=%.0e steps=%6d t_int=%7.3fs t_tot=%7.3fs et1=%.6e vs_ref: et=%.1e dp=%.1e vs_truth: et=%.1e dp=%.1e\n", + name, rtol, best.total_steps, best.t_int, t_total, isempty(best.et) ? NaN : real(best.et[1]), + et_err(best, ref), dp_err(best, ref), et_err(best, truth), dp_err(best, truth)) + end + open(out, "w") do io + println(io, header) + foreach(r -> println(io, r), rows) + end + @info "Wrote $out" +end + +main(ARGS) diff --git a/benchmarks/profile_ffs_stages.jl b/benchmarks/profile_ffs_stages.jl new file mode 100644 index 000000000..c20806d4c --- /dev/null +++ b/benchmarks/profile_ffs_stages.jl @@ -0,0 +1,114 @@ +#!/usr/bin/env julia +""" +profile_ffs_stages.jl - Wall-clock breakdown of the Equilibrium and Force-Free States stages of one case. + +Replays the calls `main` makes for one case directory and times each stage separately: equilibrium +setup and re-forming, local stability, matrix splines, the Euler-Lagrange integration, the +free-boundary energies, the Δ′ BVP, result assembly and the HDF5 write. Use it to see where a run's +time actually goes before touching any single kernel. + +# Usage + +```bash +julia --project=. -t 4 benchmarks/profile_ffs_stages.jl --case [--verbose true|false] [--blas-threads k] [--no-hdf5] [--repeat n] +``` + +`--verbose` overrides the case's `verbose` (default: as in the TOML). `--repeat n` replays the whole +sequence n times in one process, so the first pass shows cold (JIT-inclusive) timings and later passes +the warm compute cost. Inputs come from the case directory; the HDF5 file is written to a temporary +directory and deleted. +""" + +using LinearAlgebra +using Printf +using TOML +using GeneralizedPerturbedEquilibrium + +const GPE = GeneralizedPerturbedEquilibrium +const FFS = GPE.ForceFreeStates + +function parse_args(args) + opts = Dict{String,Any}("case" => nothing, "verbose" => nothing, "blas" => nothing, "hdf5" => true, "repeat" => 1) + i = 1 + while i <= length(args) + a = args[i] + if a == "--case" + opts["case"] = abspath(args[i+1]) + i += 2 + elseif a == "--verbose" + opts["verbose"] = parse(Bool, args[i+1]) + i += 2 + elseif a == "--blas-threads" + opts["blas"] = parse(Int, args[i+1]) + i += 2 + elseif a == "--no-hdf5" + opts["hdf5"] = false + i += 1 + elseif a == "--repeat" + opts["repeat"] = parse(Int, args[i+1]) + i += 2 + else + error("unknown argument $a") + end + end + opts["case"] === nothing && error("--case is required") + return opts +end + +function main(args) + opts = parse_args(args) + case = opts["case"] + opts["blas"] !== nothing && BLAS.set_num_threads(opts["blas"]) + inputs = TOML.parsefile(joinpath(case, "gpec.toml")) + ffs_table = inputs["ForceFreeStates"] + opts["verbose"] !== nothing && (ffs_table["verbose"] = opts["verbose"]) + ffs_table["write_outputs_to_HDF5"] = false + for pass in 1:opts["repeat"] + workdir = mktempdir() + ffs_table["HDF5_filename"] = joinpath(workdir, "gpec.h5") + timings = Pair{String,Float64}[] + stage(name, f) = (t = @elapsed(r = f()); push!(timings, name => t); r) + + eq_config = GPE.Equilibrium.EquilibriumConfig(inputs["Equilibrium"], case) + ctrl = FFS.ForceFreeStatesControl(; (Symbol(k) => v for (k, v) in ffs_table)...) + intr = FFS.ForceFreeStatesInternal(; dir_path=case) + intr.wall_settings = GPE.Vacuum.WallShapeSettings(; (Symbol(k) => v for (k, v) in inputs["Wall"])...) + GPE.resolve_mode_space!(intr, ctrl) + + equil = stage("equilibrium setup", () -> GPE.Equilibrium.setup_equilibrium(eq_config, nothing)) + equil = stage("equilibrium re-form (two-pass grid)", () -> GPE.maybe_reform_equilibrium(equil, eq_config, nothing, intr, ctrl, nothing)) + locstab, ballooning_boundary = stage("local stability (Mercier/ballooning)", () -> GPE.run_local_stability(ctrl, equil)) + kf_ctrl = GPE.KineticForces.KineticForcesControl() + metric, mats = stage("metric + F/G/K matrix splines", () -> GPE.prepare_force_free_states!(intr, ctrl, equil, kf_ctrl, nothing)) + odet, fm_propagators, fm_chunks, fm_S_left = stage("Euler-Lagrange integration", () -> FFS.eulerlagrange_integration(ctrl, equil, mats, intr)) + free_energies = nothing + if ctrl.vac_flag + free_energies = stage("free-boundary energies (vacuum + eigen)", () -> FFS.free_run(odet, ctrl, equil, mats, intr)) + stage("normalize eigenfunctions", () -> FFS.normalize_eigenfunctions!(odet, free_energies.wt, equil.psio)) + if ctrl.kinetic_factor == 0 && intr.msing > 0 && fm_propagators !== nothing + stage( + "Δ′ BVP (debug=$(ctrl.verbose))", + () -> FFS.compute_delta_prime_matrix!(intr, fm_propagators, fm_chunks; + wv=free_energies.wv, psio=equil.psio, debug=ctrl.verbose, S_at_surface_left=fm_S_left, ctrl=ctrl, equil=equil, mats=mats) + ) + end + end + result = stage("build_result", () -> FFS.build_result(Symbol(ctrl.integrator), ctrl, equil, intr, metric, mats, odet, free_energies, nothing, nothing)) + if opts["hdf5"] + stage("HDF5 write", () -> GPE.write_outputs_to_HDF5(result; locstab=locstab, ballooning_boundary=ballooning_boundary)) + end + rm(workdir; recursive=true, force=true) + + total = sum(last, timings) + println() + @printf("pass %d/%d (%s) case %s N=%d msing=%d integrator=%s threads=%d blas=%d verbose=%s ODE steps=%d\n", + pass, opts["repeat"], pass == 1 ? "cold, includes JIT" : "warm", basename(case), intr.numpert_total, intr.msing, ctrl.integrator, Threads.nthreads(), + BLAS.get_num_threads(), ctrl.verbose, odet.total_steps) + for (name, t) in timings + @printf(" %-42s %8.2f s %5.1f%%\n", name, t, 100t / total) + end + @printf(" %-42s %8.2f s\n", "total", total) + end +end + +main(ARGS) diff --git a/benchmarks/run_geqdsk_corpus.jl b/benchmarks/run_geqdsk_corpus.jl new file mode 100644 index 000000000..bef6dea49 --- /dev/null +++ b/benchmarks/run_geqdsk_corpus.jl @@ -0,0 +1,260 @@ +#!/usr/bin/env julia +""" +run_geqdsk_corpus.jl - Run the Force-Free States stage on every geqdsk in a list and tabulate the outcome. + +Robustness sweep for the Riccati Euler-Lagrange path: each geqdsk is run with one template `gpec.toml` +(the shipped DIII-D-like deck by default) on a pool of worker processes, with a wall-clock cap per case. +A worker that exceeds the cap is killed and replaced, so one pathological equilibrium cannot stall the +sweep, and the JIT cost is paid once per worker rather than once per case. One CSV row per case: +q-profile shape from the geqdsk, mode space, rational surfaces, step count, stage timings, the lowest +free-boundary eigenvalues and the Δ′ diagonal. + +# Usage + +```bash +julia --project=. benchmarks/run_geqdsk_corpus.jl --list --root --out \\ + [--template ] [--workers 2] [--threads 4] [--timeout 600] [--limit n] +``` + +`--list` holds one geqdsk path per line, relative to `--root`. Outputs are written under `benchmarks/`. +""" + +using Distributed +using Printf +using TOML + +function parse_args(args) + opts = Dict{String,Any}("list" => nothing, "root" => ".", "out" => nothing, "workers" => 2, "threads" => 4, + "timeout" => 600.0, "limit" => typemax(Int), + "template" => joinpath(@__DIR__, "..", "examples", "DIIID-like_ideal_example", "gpec.toml")) + i = 1 + while i <= length(args) + a = args[i] + if a == "--list" + opts["list"] = abspath(args[i+1]) + i += 2 + elseif a == "--root" + opts["root"] = abspath(args[i+1]) + i += 2 + elseif a == "--out" + opts["out"] = abspath(args[i+1]) + i += 2 + elseif a == "--template" + opts["template"] = abspath(args[i+1]) + i += 2 + elseif a == "--workers" + opts["workers"] = parse(Int, args[i+1]) + i += 2 + elseif a == "--threads" + opts["threads"] = parse(Int, args[i+1]) + i += 2 + elseif a == "--timeout" + opts["timeout"] = parse(Float64, args[i+1]) + i += 2 + elseif a == "--limit" + opts["limit"] = parse(Int, args[i+1]) + i += 2 + else + error("unknown argument $a") + end + end + opts["list"] === nothing && error("--list is required") + opts["out"] === nothing && (opts["out"] = joinpath(@__DIR__, "integrator_results", "geqdsk_corpus.csv")) + return opts +end + +# Template deck with the Force-Free States stage only (no PerturbedEquilibrium / KineticForces). +function corpus_template(path) + inputs = TOML.parsefile(path) + keep = Dict{String,Any}(k => inputs[k] for k in ("Equilibrium", "Wall", "ForceFreeStates") if haskey(inputs, k)) + keep["Equilibrium"]["eq_filename"] = "eq.geqdsk" + keep["ForceFreeStates"]["integrator"] = "riccati" + keep["ForceFreeStates"]["verbose"] = false + keep["ForceFreeStates"]["write_outputs_to_HDF5"] = false + keep["ForceFreeStates"]["local_stability_flag"] = false + return keep +end + +const HEADER = "idx,geqdsk,status,elapsed_s,nw,q0,qmin,psi_qmin,qedge,N,msing,surfaces,total_steps," * + "t_equil_s,t_prep_s,t_int_s,t_free_s,t_dp_s,et1,et2,et3,dp_diag,message" + +function main(args) + opts = parse_args(args) + files = filter(!isempty, readlines(opts["list"]))[1:min(end, opts["limit"])] + mkpath(dirname(opts["out"])) + template = corpus_template(opts["template"]) + exeflags = ["--project=$(Base.active_project())", "--threads=$(opts["threads"])"] + + # Worker-side case runner; the imports are evaluated first so the macros below resolve. + worker_imports = quote + using LinearAlgebra, TOML, Printf + using GeneralizedPerturbedEquilibrium + end + worker_setup = quote + const GPE = GeneralizedPerturbedEquilibrium + const FFS = GPE.ForceFreeStates + function read_q_profile(path) + lines = readlines(path) + # nw, nh are the last two integers of the header line; the leading label field is + # not column-aligned across EFIT and TokaMaker writers. + ints = [parse(Int, m.match) for m in eachmatch(r"(?= 2 || error("cannot read nw, nh from geqdsk header: $(lines[1])") + nw, nh = ints[end-1], ints[end] + nums = Float64[] + for l in lines[2:end] + for m in eachmatch(r"[-+]?\d*\.\d+(?:[eE][-+]?\d+)?", l) + push!(nums, parse(Float64, m.match)) + end + end + off = 20 + 4nw + nw * nh + q = abs.(nums[off+1:off+nw]) + imin = argmin(q) + return nw, q[1], q[imin], (imin - 1) / (nw - 1), q[end] + end + function run_case(geqdsk::String, template::Dict{String,Any}) + t0 = time() + nw, q0, qmin, psi_qmin, qedge = read_q_profile(geqdsk) + dir = mktempdir() + cp(geqdsk, joinpath(dir, "eq.geqdsk")) + open(joinpath(dir, "gpec.toml"), "w") do io + TOML.print(io, template) + end + inputs = TOML.parsefile(joinpath(dir, "gpec.toml")) + eq_config = GPE.Equilibrium.EquilibriumConfig(inputs["Equilibrium"], dir) + wall = GPE.Vacuum.WallShapeSettings(; (Symbol(k) => v for (k, v) in inputs["Wall"])...) + ctrl = FFS.ForceFreeStatesControl(; (Symbol(k) => v for (k, v) in inputs["ForceFreeStates"])...) + intr = FFS.ForceFreeStatesInternal(; dir_path=dir) + intr.wall_settings = wall + GPE.resolve_mode_space!(intr, ctrl) + t_equil = @elapsed begin + equil = GPE.Equilibrium.setup_equilibrium(eq_config, nothing) + equil = GPE.maybe_reform_equilibrium(equil, eq_config, nothing, intr, ctrl, nothing) + end + t_prep = @elapsed metric, mats = GPE.prepare_force_free_states!(intr, ctrl, equil, GPE.KineticForces.KineticForcesControl(), nothing) + t_int = @elapsed odet, props, chunks, S_left = FFS.eulerlagrange_integration(ctrl, equil, mats, intr) + t_free = @elapsed begin + free = FFS.free_run(odet, ctrl, equil, mats, intr) + FFS.normalize_eigenfunctions!(odet, free.wt, equil.psio) + end + t_dp = 0.0 + if intr.msing > 0 && props !== nothing + t_dp = @elapsed FFS.compute_delta_prime_matrix!(intr, props, chunks; wv=free.wv, psio=equil.psio, debug=false, + S_at_surface_left=S_left, ctrl=ctrl, equil=equil, mats=mats) + end + res = FFS.build_result(Symbol(ctrl.integrator), ctrl, equil, intr, metric, mats, odet, free, nothing, nothing) + surfaces = join([@sprintf("%d@%.4f", s.m[1], s.psifac) for s in intr.sing], ";") + dp = res.delta_prime === nothing ? "" : join([@sprintf("%.6g%+.6gi", real(z), imag(z)) for z in LinearAlgebra.diag(res.delta_prime.matrix)], ";") + et = real.(free.et) + rm(dir; recursive=true, force=true) + return (status="ok", elapsed=time() - t0, nw, q0, qmin, psi_qmin, qedge, N=intr.numpert_total, msing=intr.msing, surfaces, + total_steps=odet.total_steps, t_equil, t_prep, t_int, t_free, t_dp, + et1=get(et, 1, NaN), et2=get(et, 2, NaN), et3=get(et, 3, NaN), dp, message="") + end + # Errors are turned into a row on the worker: a raised exception would carry method + # instances the driver process cannot deserialize. + function run_case_safe(geqdsk::String, template::Dict{String,Any}) + t0 = time() + try + return run_case(geqdsk, template) + catch err + msg = first(sprint(showerror, err), 400) + return (status="failed", elapsed=time() - t0, nw=0, q0=NaN, qmin=NaN, psi_qmin=NaN, qedge=NaN, N=0, msing=0, surfaces="", + total_steps=0, t_equil=NaN, t_prep=NaN, t_int=NaN, t_free=NaN, t_dp=NaN, et1=NaN, et2=NaN, et3=NaN, dp="", message=msg) + end + end + nothing # the eval's return value travels back to the driver; a function object would not deserialize there + end + + os_pids = Dict{Int,Int}() + + function spawn_worker() + pid = only(addprocs(1; exeflags=exeflags)) + os_pids[pid] = remotecall_fetch(getpid, pid) + remotecall_fetch(Core.eval, pid, Main, worker_imports) + remotecall_fetch(Core.eval, pid, Main, worker_setup) + remotecall_fetch(Core.eval, pid, Main, :(const TEMPLATE = $template)) + return pid + end + + # A worker stuck inside a solve does not answer the cooperative shutdown and `rmprocs` then + # throws; SIGKILL the OS process so one pathological equilibrium cannot end the whole sweep. + function kill_worker(pid) + try + rmprocs(pid; waitfor=5) + catch + ospid = get(os_pids, pid, 0) + if ospid > 0 + try + run(`kill -9 $ospid`) + catch + end + end + try + rmprocs(pid; waitfor=0) + catch + end + end + delete!(os_pids, pid) + return nothing + end + + io = open(opts["out"], "w") + println(io, HEADER) + flush(io) + write_row(idx, rel, r) = begin + println( + io, + join( + [idx, rel, r.status, @sprintf("%.1f", r.elapsed), r.nw, @sprintf("%.4f", r.q0), @sprintf("%.4f", r.qmin), + @sprintf("%.3f", r.psi_qmin), @sprintf("%.3f", r.qedge), r.N, r.msing, r.surfaces, r.total_steps, + @sprintf("%.2f", r.t_equil), @sprintf("%.2f", r.t_prep), @sprintf("%.2f", r.t_int), @sprintf("%.2f", r.t_free), + @sprintf("%.2f", r.t_dp), @sprintf("%.8e", r.et1), @sprintf("%.8e", r.et2), @sprintf("%.8e", r.et3), r.dp, + replace(r.message, "," => ";", "\n" => " ")], ",") + ) + flush(io) + end + blank(status, elapsed, msg) = (status, elapsed, nw=0, q0=NaN, qmin=NaN, psi_qmin=NaN, qedge=NaN, N=0, msing=0, surfaces="", total_steps=0, + t_equil=NaN, t_prep=NaN, t_int=NaN, t_free=NaN, t_dp=NaN, et1=NaN, et2=NaN, et3=NaN, dp="", message=msg) + + queue = Channel{Tuple{Int,String}}(length(files)) + for (i, f) in enumerate(files) + put!(queue, (i, f)) + end + close(queue) + lock = ReentrantLock() + + @sync for w in 1:opts["workers"] + @async begin + pid = spawn_worker() + for (idx, rel) in queue + path = joinpath(opts["root"], rel) + t0 = time() + # The task must never throw: killing a timed-out worker makes its pending + # fetch fail, and an exception here would escape @sync and end the sweep. + task = @async try + remotecall_fetch(Core.eval, pid, Main, :(run_case_safe($path, TEMPLATE))) + catch err + err + end + while !istaskdone(task) && time() - t0 < opts["timeout"] + sleep(1) + end + r = if istaskdone(task) + out = fetch(task) + out isa Exception ? blank("failed", time() - t0, first(sprint(showerror, out), 300)) : out + else + kill_worker(pid) + pid = spawn_worker() + blank("timeout", time() - t0, "exceeded $(opts["timeout"]) s") + end + @lock lock write_row(idx, rel, r) + @printf("%3d/%d %-8s %6.1fs %s\n", idx, length(files), r.status, r.elapsed, rel) + end + kill_worker(pid) + end + end + close(io) + @info "Wrote $(opts["out"])" +end + +main(ARGS) diff --git a/docs/development/figures/integrator-comparison/chunks_wp_nf.png b/docs/development/figures/integrator-comparison/chunks_wp_nf.png new file mode 100644 index 000000000..8834d5e2a Binary files /dev/null and b/docs/development/figures/integrator-comparison/chunks_wp_nf.png differ diff --git a/docs/development/figures/integrator-comparison/pipeline_wp_dp.png b/docs/development/figures/integrator-comparison/pipeline_wp_dp.png new file mode 100644 index 000000000..cc534fa36 Binary files /dev/null and b/docs/development/figures/integrator-comparison/pipeline_wp_dp.png differ diff --git a/docs/development/figures/integrator-comparison/pipeline_wp_et.png b/docs/development/figures/integrator-comparison/pipeline_wp_et.png new file mode 100644 index 000000000..74c468454 Binary files /dev/null and b/docs/development/figures/integrator-comparison/pipeline_wp_et.png differ diff --git a/examples/DIIID-like_ideal_example/gpec.toml b/examples/DIIID-like_ideal_example/gpec.toml index 24b2b7f24..57431a1c7 100644 --- a/examples/DIIID-like_ideal_example/gpec.toml +++ b/examples/DIIID-like_ideal_example/gpec.toml @@ -45,6 +45,8 @@ mthvac = 512 # Number of points used in splines over poloidal kinetic_source = "fixed" # Kinetic matrix source: "fixed" test matrices, or "calculated" from the kinetic NTV model kinetic_factor = 0.0 # Scaling of kinetic matrices (0 = ideal path; >0 enables kinetic mode) eulerlagrange_tolerance = 1e-10 # Relative tolerance for ODE integration of Euler-Lagrange equations +ode_abstol = 1e-8 # Absolute tolerance for the same integration +ode_solver = "Vern7" # OrdinaryDiffEq method for every Euler-Lagrange solve save_interval = 3 # Save every Nth ODE step (1=all). Always saves near rational surfaces. singfac_min = 1e-4 # Fractional distance from rational q at which ideal jump enforced ucrit = 1e4 # Column-norm threshold that triggers solution renormalization diff --git a/src/ForceFreeStates/CoreTypes.jl b/src/ForceFreeStates/CoreTypes.jl index 98d4807e1..f2907bc16 100644 --- a/src/ForceFreeStates/CoreTypes.jl +++ b/src/ForceFreeStates/CoreTypes.jl @@ -130,6 +130,8 @@ gpec.toml. - `nstep::Int` - Maximum number of integration steps (not yet implemented) - `ksing::Int` - Singular surface handling parameter - `eulerlagrange_tolerance::Float64` - Relative tolerance for ODE integration of Euler-Lagrange equations + - `ode_abstol::Float64` - Absolute tolerance for the same integration. Default `1e-8`: the OrdinaryDiffEq default of `1e-6` lets small state entries escape the relative control, so the error stops responding to `eulerlagrange_tolerance` below about `1e-8`. + - `ode_solver::String` - OrdinaryDiffEq explicit Runge-Kutta method used for every Euler-Lagrange solve (propagator chunks, Riccati outer plasma, forward sweep, Δ′ shooting). Default `"Vern7"`: on the DIII-D corpus it reaches the same δW and a closer Δ′ than `"Vern9"` for about half the RHS evaluations, because the ninth-order method rejects most of its steps near the rational surfaces. Any name in `OrdinaryDiffEq` that accepts complex states is allowed (`"Vern6"`, `"Vern9"`, `"DP8"`, ...). - `ucrit::Float64` - Critical value of unorm ratio to trigger solution normalization. In the standard path it triggers Gaussian reduction; in the Riccati path it triggers `renormalize_riccati_inplace!`. Default `1e4` empirically keeps max(|U₁|, |U₂|) in O(1)–O(10⁴) over the integration domain on DIII-D / Solovev sweeps; lower triggers excess renorms without accuracy gain, higher risks overflow before the next renorm. - `numsteps_init::Int` - Initial array size for ODE data storage - `numunorms_init::Int` - Initial array size for solution normalization data @@ -167,6 +169,8 @@ gpec.toml. nstep::Int = typemax(Int) ksing::Int = -1 eulerlagrange_tolerance::Float64 = 1e-8 + ode_abstol::Float64 = 1e-8 + ode_solver::String = "Vern7" ucrit::Float64 = 1e4 numsteps_init::Int = 4000 numunorms_init::Int = 100 diff --git a/src/ForceFreeStates/EulerLagrange.jl b/src/ForceFreeStates/EulerLagrange.jl index 462d7b86f..8693d2c24 100644 --- a/src/ForceFreeStates/EulerLagrange.jl +++ b/src/ForceFreeStates/EulerLagrange.jl @@ -45,75 +45,49 @@ and a small set of temporary matrices and factors used to compute singular-layer - `numpert_total::Int` - Total number of Fourier mode combinations (m × n) used in the calculation. - `numunorms_init::Int` - Initial allocation size for the number of normalization operations recorded. - - `msing::Int` - Number of singular surfaces in the equilibrium (used to size asymptotic coefficient arrays). - - `numsteps_init::Int` - Initial allocation size for the number of integration steps to store. - - `step::Int` - Current integration step index (1-based, like `istep` in the original Fortran). - - `psi_store::Vector{Float64}` - Stored psi values at each saved integration step (length `numsteps_init`). - - `q_store::Vector{Float64}` - Stored q values at each saved integration step (length `numsteps_init`). - - `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,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, step)`. Empty until materialized, same as `du_store`. - - `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`). - - `ca_r::Array{ComplexF64,4}` - Asymptotic coefficients just to the right of each singular surface with shape `(numpert_total, numpert_total, 2, msing)`. - - `ca_l::Array{ComplexF64,4}` - Asymptotic coefficients just to the left of each singular surface with shape `(numpert_total, numpert_total, 2, msing)`. - - `ca_populated::Bool` - True once an ideal singular-surface crossing has filled `ca_l`/`ca_r`; kinetic and galerkin-matched runs never populate them and leave this false, and the HDF5 writer then emits zero-extent `ca_left`/`ca_right` datasets instead of unpopulated arrays. - - `edge_scan::EdgeScanState` - Edge dW scan state and results. Initialized as a disabled sentinel (N_edge=0) and replaced by `findmax_dW_edge!` when a scan runs. - - `psifac::Float64` - Current normalized flux coordinate for the integrator. - - `q::Float64` - Safety factor value at `psifac` (current q during integration). - - `u::Array{ComplexF64,3}` - Current working solution arrays with shape `(numpert_total, numpert_total, 2)`. - - `ising_start::Int` - Index of the starting singular surface to be crossed during integration. - - `psimax::Float64` - Maximum psi value for which the integrator is allowed to run in next integration region. - - `needs_crossing::Bool` - Flag indicating whether a rational surface needs to be crossed after the current integration region. - - `nzero::Int` - Count of detected zero crossings (used for diagnostics). - - `new::Bool` - Flag indicating whether a new `unorm0` should be computed after a fixup. # Initialization parameters - - `unorm::Vector{Float64}` - Current norms of the solution vectors (length `numpert_total`). - - `unorm0::Vector{Float64}` - Reference/initial norms of the solution vectors (length `numpert_total`). # Saved data throughout integration - - `ifix::Int` - Number of normalization operations performed (index into normalization arrays). # Total ODE solver steps taken (all steps, not just saved ones) @@ -122,11 +96,8 @@ and a small set of temporary matrices and factors used to compute singular-layer - `sing_flag::Vector{Bool}` - Boolean flags indicating which stored normalizations correspond to singular solutions # Edge dW scan state and results (disabled sentinel when psiedge >= psilim, i.e. no edge scan) (length `numunorms_init`). - - `zeroed_idx::Vector{Vector{Int}}` - For each ideal rational surface jump, a vector of indices of solutions that were zeroed. # Data for integrator - - `fixfac::Array{ComplexF64,3}` - Fix-up factors for Gaussian reduction with shape `(numpert_total, numpert_total, numunorms_init)`. - - `fixstep::Vector{Int64}` - Step indices (psi step positions) at which normalization/fixups were performed (length `numunorms_init`). """ @kwdef mutable struct OdeState @@ -232,9 +203,9 @@ end # at the interval endpoints. Coefficients are ported from STRIDE's ode_itime cost model # (Fortran reference) and unchanged here. Tune only after re-fitting against a per-chunk # step-count sweep; touching these affects parallel-chunk load balancing. -const ODE_COST_AXIS = (a = 39695.0, b = 212830.0) -const ODE_COST_RAT = (a = 17147.0, b = 470710.0) -const ODE_COST_EDGE = (a = 1646.0, b = 4683.0) +const ODE_COST_AXIS = (a=39695.0, b=212830.0) +const ODE_COST_RAT = (a=17147.0, b=470710.0) +const ODE_COST_EDGE = (a=1646.0, b=4683.0) """ ode_itime_cost(psi1, psi2, intr) -> Float64 @@ -267,7 +238,7 @@ never from `Threads.nthreads()` — so the chunk list, and hence every Riccati o identical whatever thread count `julia -t` provides. Each split finds the equal-cost midpoint ψ_mid via bisection: - ode_itime_cost(psi_start, psi_mid) ≈ ode_itime_cost(psi_start, psi_end) / 2 +ode_itime_cost(psi_start, psi_mid) ≈ ode_itime_cost(psi_start, psi_end) / 2 Sub-chunks inherit `needs_crossing=false` and `ising=0`. Only the LAST sub-chunk of each original chunk retains `needs_crossing=true` and the original `ising`, so the @@ -326,10 +297,10 @@ function balance_integration_chunks(chunks::Vector{IntegrationChunk}, ctrl::Forc psi_mid = (lo + hi) / 2.0 left = IntegrationChunk(; psi_start=chunk.psi_start, psi_end=psi_mid, - needs_crossing=false, ising=0, direction=1) + needs_crossing=false, ising=0, direction=1) right = IntegrationChunk(; psi_start=psi_mid, psi_end=chunk.psi_end, - needs_crossing=chunk.needs_crossing, ising=chunk.ising, - direction=chunk.direction) + needs_crossing=chunk.needs_crossing, ising=chunk.ising, + direction=chunk.direction) splice!(result, best_idx, [left, right]) end @@ -349,17 +320,27 @@ Only the Riccati branch populates `propagators` / `chunks` / `S_left`, which for all three. """ function eulerlagrange_integration(ctrl::ForceFreeStatesControl, equil::Equilibrium.PlasmaEquilibrium, mats::MatrixSplines, intr::ForceFreeStatesInternal) - - if ctrl.integrator == "riccati" - ctrl.kinetic_factor > 0 && error("kinetic runs require integrator=\"forward\"; the Riccati integrator has no kinetic crossing.") - return riccati_eulerlagrange_integration(ctrl, equil, mats, intr) - elseif ctrl.integrator == "forward" - return forward_eulerlagrange_integration(ctrl, equil, mats, intr) - elseif ctrl.integrator == "galerkin" + ctrl.integrator == "galerkin" && error("integrator = \"galerkin\" solves the Euler-Lagrange system variationally, not by ODE integration; " * "it is dispatched to galerkin_solve.") + ctrl.integrator in ("riccati", "forward") || + error("Unknown integrator: $(ctrl.integrator). Expected \"forward\", \"riccati\", or \"galerkin\".") + ctrl.integrator == "riccati" && ctrl.kinetic_factor > 0 && + error("kinetic runs require integrator=\"forward\"; the Riccati integrator has no kinetic crossing.") + + # The RHS works on mpert×mpert blocks, where multithreaded BLAS costs more in synchronization + # than it saves; pin BLAS to one thread for the sweep and restore it afterwards. + blas_threads = BLAS.get_num_threads() + BLAS.set_num_threads(1) + try + if ctrl.integrator == "riccati" + return riccati_eulerlagrange_integration(ctrl, equil, mats, intr) + else + return forward_eulerlagrange_integration(ctrl, equil, mats, intr) + end + finally + BLAS.set_num_threads(blas_threads) end - error("Unknown integrator: $(ctrl.integrator). Expected \"forward\", \"riccati\", or \"galerkin\".") end """ @@ -479,20 +460,20 @@ is identified by dominant |U₁| component, giving the physically correct consta Frobenius solution and avoiding the spurious logarithmic irregularity. """ function compute_axis_init(mats::MatrixSplines, profiles::Equilibrium.ProfileSplines, - intr::ForceFreeStatesInternal, psi_low::Float64) - N = intr.numpert_total + intr::ForceFreeStatesInternal, psi_low::Float64) + N = intr.numpert_total hint = Ref(1) # Evaluate stability matrices at psi_low F_lower = zeros(ComplexF64, N, N) - kmat = zeros(ComplexF64, N, N) - gmat = zeros(ComplexF64, N, N) + kmat = zeros(ComplexF64, N, N) + gmat = zeros(ComplexF64, N, N) mats.ideal.F_spline_lower(vec(F_lower), psi_low; hint=hint) - mats.ideal.K_spline(vec(kmat), psi_low; hint=hint) - mats.ideal.G_spline(vec(gmat), psi_low; hint=hint) + mats.ideal.K_spline(vec(kmat), psi_low; hint=hint) + mats.ideal.G_spline(vec(gmat), psi_low; hint=hint) # singfac[j] = 1 / (m_j − n_j · q) for each mode j - q0 = profiles.q_spline(psi_low; hint=hint) + q0 = profiles.q_spline(psi_low; hint=hint) singfac = vec(1.0 ./ ((intr.mlow:intr.mhigh) .- q0 .* (intr.nlow:intr.nhigh)')) # F̄⁻¹ = (F_lower · F_lower')⁻¹ via the Cholesky factor @@ -506,15 +487,15 @@ function compute_axis_init(mats::MatrixSplines, profiles::Equilibrium.ProfileSpl for j in 1:N sf = singfac[j] fi = Finv[j, j] - k = kmat[j, j] + k = kmat[j, j] kd = conj(k) # K̄†[j,j] - g = gmat[j, j] + g = gmat[j, j] # 2×2 ODE matrix block for mode j [Glasser 2016 Eq. 22-24, diagonal approximation] m11 = -sf * fi * k - m12 = sf^2 * fi - m21 = g - kd * fi * k - m22 = sf * kd * fi + m12 = sf^2 * fi + m21 = g - kd * fi * k + m22 = sf * kd * fi # Frobenius matrix A₀_j = ψ_low · M_j [Glasser 2016 Eq. 51] #! format: off @@ -541,7 +522,7 @@ function compute_axis_init(mats::MatrixSplines, profiles::Equilibrium.ProfileSpl if abs(v2) > Base.sqrt(Base.eps(Float64)) * abs(v1) U1_init[j, j] = v1 / v2 else - U1_init[j, j] = one(ComplexF64) + U1_init[j, j] = one(ComplexF64) U2_init[j, j] = zero(ComplexF64) end end @@ -560,7 +541,7 @@ Formerly `ode_axis_init!`. This now only initializes `psifac`, `ising_start`, an Move ising_start logic to chunk_el_integration_bounds? """ function initialize_el_at_axis!(odet::OdeState, ctrl::ForceFreeStatesControl, mats::MatrixSplines, - profiles::Equilibrium.ProfileSplines, intr::ForceFreeStatesInternal) + profiles::Equilibrium.ProfileSplines, intr::ForceFreeStatesInternal) # Default psifac to minimum equilibrium psi value odet.psifac = profiles.xs[1] @@ -737,7 +718,7 @@ function chunk_el_integration_bounds(odet::OdeState, ctrl::ForceFreeStatesContro psi_end=psi_end, needs_crossing=true, ising=ising_current, - direction = bidirectional ? -1 : 1 + direction=bidirectional ? -1 : 1 )) # After crossing, we jump to the other side of the singular surface @@ -966,7 +947,7 @@ function integrate_el_region!( cb = DiscreteCallback((u, t, integrator) -> true, segment_callback!) prob = ODEProblem(sing_der!, odet.u, (chunk.psi_start, chunk.psi_end), (ctrl, equil, mats, intr, odet, chunk)) - sol = solve(prob, Vern9(); reltol=ctrl.eulerlagrange_tolerance, callback=cb, save_everystep=false, save_end=true) + sol = solve(prob, el_ode_algorithm(ctrl); reltol=ctrl.eulerlagrange_tolerance, abstol=ctrl.ode_abstol, callback=cb, save_everystep=false, save_end=true) # Unconditionally save the final step if the callback did not already capture it. # Guarantees the pre-crossing (or pre-edge) state is always stored in u_store, @@ -1171,7 +1152,7 @@ function transform_u!(odet::OdeState, intr::ForceFreeStatesInternal) temp[ksol, jsol] = odet.fixfac[ksol, jsol, ifix] end end - mul!(gauss_buffer, view(gauss,:,:,ifix), temp) + mul!(gauss_buffer, view(gauss, :, :, ifix), temp) gauss[:, :, ifix] .= gauss_buffer end # Account for zeroed indices at singular surfaces in `ode_ideal_cross` @@ -1188,7 +1169,7 @@ function transform_u!(odet::OdeState, intr::ForceFreeStatesInternal) # and mfix + 1 is the for the region after the last fixup and before the edge transforms[:, :, end] .= identity for ifix in odet.ifix:-1:1 - mul!(view(transforms,:,:,ifix), view(gauss,:,:,ifix), view(transforms,:,:,(ifix+1))) + mul!(view(transforms, :, :, ifix), view(gauss, :, :, ifix), view(transforms, :, :, (ifix + 1))) end # Now that we have the transform matrices, we can apply them to the solution vectors @@ -1448,4 +1429,3 @@ non-Hermitian contributions and needs an LU. xi_s .-= tmp_mat return xi_s end - diff --git a/src/ForceFreeStates/ForceFreeStates.jl b/src/ForceFreeStates/ForceFreeStates.jl index 1dc93e776..4505c5492 100644 --- a/src/ForceFreeStates/ForceFreeStates.jl +++ b/src/ForceFreeStates/ForceFreeStates.jl @@ -6,6 +6,17 @@ using LinearAlgebra.LAPACK using TOML using FFTW using OrdinaryDiffEq + +""" + el_ode_algorithm(ctrl) -> OrdinaryDiffEq algorithm + +The solver named by `ctrl.ode_solver`, instantiated for the Euler-Lagrange `solve` calls. +""" +function el_ode_algorithm(ctrl) + isdefined(OrdinaryDiffEq, Symbol(ctrl.ode_solver)) || + error("ode_solver = \"$(ctrl.ode_solver)\" is not an OrdinaryDiffEq algorithm name (e.g. \"Vern7\", \"Vern9\", \"DP8\")") + return getfield(OrdinaryDiffEq, Symbol(ctrl.ode_solver))() +end using HDF5 using JLD2 using FastInterpolations diff --git a/src/ForceFreeStates/Riccati/DeltaPrimeBVP.jl b/src/ForceFreeStates/Riccati/DeltaPrimeBVP.jl index 0a955fde7..6dc2c3d23 100644 --- a/src/ForceFreeStates/Riccati/DeltaPrimeBVP.jl +++ b/src/ForceFreeStates/Riccati/DeltaPrimeBVP.jl @@ -7,6 +7,7 @@ Compute the inter-surface tearing stability matrix (msing × msing) using the STRIDE global BVP formulation [Glasser 2018 Phys. Plasmas 25, 032501, Sec. III.B]. The BVP encodes the full plasma response with unknowns at each surface boundary: + ``` x_axis (N): free IC parameters at the axis (U₁ = 0 regular solutions) x_left[j] (2N): state at left inner-layer boundary of surface j @@ -19,9 +20,11 @@ The BVP encodes the full plasma response with unknowns at each surface boundary: When `wv` is provided (the vacuum response matrix, singfac-scaled), the edge BC follows the Fortran STRIDE convention: + ``` U₁ = c, U₂ = -wv·ψ₀²·c ``` + which is the free-boundary condition `wp + wv = 0` at the edge. When `wv` is `nothing`, a conducting wall BC (`U₁ = 0`) is used. @@ -38,9 +41,11 @@ keeps the BVP matrix full-rank and well-conditioned. The raw BVP solution is a 2·msing × 2·msing matrix `dp` with left/right sub-indices at each surface. The PEST3-convention Δ' matrix is the linear combination [Chance, PPPL-2527]: + ``` deltap(i,j) = dp(2i,2j) - dp(2i,2j-1) - dp(2i-1,2j) + dp(2i-1,2j-1) ``` + stored in `intr.delta_prime_matrix` (msing × msing). ## Limitations @@ -65,16 +70,28 @@ function compute_delta_prime_matrix!( intr::ForceFreeStatesInternal, propagators::Vector{ChunkPropagator}, chunks::Vector{IntegrationChunk}; - wv::Union{Nothing,Matrix{ComplexF64}} = nothing, - psio::Float64 = 0.0, - debug::Bool = false, - S_at_surface_left::Union{Nothing,Vector{Matrix{ComplexF64}}} = nothing, - ctrl::Union{Nothing,ForceFreeStatesControl} = nothing, - equil::Union{Nothing,Equilibrium.PlasmaEquilibrium} = nothing, - mats::Union{Nothing,MatrixSplines} = nothing + wv::Union{Nothing,Matrix{ComplexF64}}=nothing, + psio::Float64=0.0, + debug::Bool=false, + S_at_surface_left::Union{Nothing,Vector{Matrix{ComplexF64}}}=nothing, + ctrl::Union{Nothing,ForceFreeStatesControl}=nothing, + equil::Union{Nothing,Equilibrium.PlasmaEquilibrium}=nothing, + mats::Union{Nothing,MatrixSplines}=nothing ) intr.msing == 0 && return _has_unsupported_multi_resonance(intr) && return + # Same small-block BLAS pin as eulerlagrange_integration: the shooting solves re-run the RHS. + blas_threads = BLAS.get_num_threads() + BLAS.set_num_threads(1) + try + _compute_delta_prime_matrix!(intr, propagators, chunks, wv, psio, debug, S_at_surface_left, ctrl, equil, mats) + finally + BLAS.set_num_threads(blas_threads) + end + return +end + +function _compute_delta_prime_matrix!(intr, propagators, chunks, wv, psio, debug, S_at_surface_left, ctrl, equil, mats) sing, i_crossings, msing = _select_active_surfaces(intr, chunks) msing == 0 && return @@ -91,8 +108,10 @@ function compute_delta_prime_matrix!( if !use_S_axis for ic in i_crossings chunks[ic].direction == 1 || - error("compute_delta_prime_matrix!: FM-axis fallback (use_S_axis=false) requires forward crossing chunks; " * - "chunk $ic has direction=$(chunks[ic].direction). Either provide S_at_surface_left or use bidirectional=false.") + error( + "compute_delta_prime_matrix!: FM-axis fallback (use_S_axis=false) requires forward crossing chunks; " * + "chunk $ic has direction=$(chunks[ic].direction). Either provide S_at_surface_left or use bidirectional=false." + ) end end @@ -105,15 +124,15 @@ function compute_delta_prime_matrix!( _build_asymptotic_basis_matrices(sing, has_ua, N, msing) debug && _log_bvp_setup(chunks, sing, S_at_surface_left, use_S_axis, has_ua, - Phi_L_mats, Phi_R_mats, Phi_R_halves, ipert_all, wv, psio, N, msing) + Phi_L_mats, Phi_R_mats, Phi_R_halves, ipert_all, wv, psio, N, msing) if use_S_axis uShootR, uShootL, uAxis = _build_S_axis_shooting_propagators( propagators, chunks, i_crossings, sing, msing, N, T_left_mats, T_right_mats, has_ua, ctrl, equil, mats, intr, debug) debug && _log_S_axis_shooting_propagators(uShootR, uShootL, uAxis, - S_at_surface_left, T_left_mats, - ipert_all, has_ua, msing, N) + S_at_surface_left, T_left_mats, + ipert_all, has_ua, msing, N) M, nMat, col_edge = _assemble_bvp_S_axis( uShootR, uShootL, uAxis, ipert_all, msing, N, wv, psio) else @@ -140,13 +159,13 @@ function compute_delta_prime_matrix!( # The raw matrix is consumed by `pest3_decompose` to recover (A', B', Γ', Δ') for the full # det(D' − D(γ)) = 0 eigenvalue problem; see the `delta_prime_raw` docstring in CoreTypes.jl. intr.delta_prime_matrix = deltap - intr.delta_prime_raw = dp_raw_persisted + intr.delta_prime_raw = dp_raw_persisted end # Column index helpers for the BVP matrix. j is the 1-based singular-surface index, # N is numpert_total. Layout: c_axis(N), c_left[1](2N), c_right[1](2N), ..., c_edge(N). -_col_left(j::Int, N::Int) = (N + 4N*(j-1) + 1):(N + 4N*(j-1) + 2N) -_col_right(j::Int, N::Int) = (N + 4N*(j-1) + 2N + 1):(N + 4N*j) +_col_left(j::Int, N::Int) = (N+4N*(j-1)+1):(N+4N*(j-1)+2N) +_col_right(j::Int, N::Int) = (N+4N*(j-1)+2N+1):(N+4N*j) # Multi-resonance surfaces (one q value satisfying multiple (m,n) tuples in a multi-n run) # are not yet handled by the inter-surface BVP. Returns true if any surface has >1 modes; @@ -185,9 +204,9 @@ end # Midpoint splitting halves each inter-surface span's condition number — STRIDE's trick: # cond(full) = 10¹⁵ → cond(half) ≈ 10⁷·⁵, an 8-digit accuracy gain. function _assemble_segment_propagators(propagators::Vector{ChunkPropagator}, - chunks::Vector{IntegrationChunk}, - i_crossings::Vector{Int}, msing::Int, N::Int, - use_S_axis::Bool) + chunks::Vector{IntegrationChunk}, + i_crossings::Vector{Int}, msing::Int, N::Int, + use_S_axis::Bool) Phi_L_mats = [assemble_fm_matrix(propagators, i_crossings[j]:i_crossings[j]) for j in 1:msing] Phi_R_mats = Vector{Matrix{ComplexF64}}(undef, msing + 1) if !use_S_axis @@ -201,11 +220,11 @@ function _assemble_segment_propagators(propagators::Vector{ChunkPropagator}, Phi_R_halves = Vector{Tuple{Matrix{ComplexF64},Matrix{ComplexF64}}}(undef, msing - 1) for j in 1:msing-1 chunk_start = i_crossings[j] + 1 - chunk_end = i_crossings[j+1] - 1 - n_chunks = chunk_end - chunk_start + 1 + chunk_end = i_crossings[j+1] - 1 + n_chunks = chunk_end - chunk_start + 1 if n_chunks >= 2 i_mid = chunk_start + div(n_chunks, 2) - 1 - Phi_left_half = assemble_fm_matrix(propagators, chunk_start:i_mid) + Phi_left_half = assemble_fm_matrix(propagators, chunk_start:i_mid) Phi_right_half = assemble_fm_matrix(propagators, i_mid+1:chunk_end) Phi_R_halves[j] = (Phi_left_half, Phi_right_half) else @@ -220,17 +239,17 @@ end # N+1:2N = small solutions (z^{+α}, bounded). Fortran STRIDE bakes T into the shooting # propagators (uFM_sing_init); we multiply T into the BVP propagator blocks at each surface. function _build_asymptotic_basis_matrices(sing::Vector{SingType}, has_ua::Bool, N::Int, msing::Int) - T_left_mats = Vector{Matrix{ComplexF64}}(undef, msing) + T_left_mats = Vector{Matrix{ComplexF64}}(undef, msing) T_right_mats = Vector{Matrix{ComplexF64}}(undef, msing) - T_left_inv = Vector{Matrix{ComplexF64}}(undef, msing) - T_right_inv = Vector{Matrix{ComplexF64}}(undef, msing) + T_left_inv = Vector{Matrix{ComplexF64}}(undef, msing) + T_right_inv = Vector{Matrix{ComplexF64}}(undef, msing) if has_ua for j in 1:msing sp = sing[j] - T_left_mats[j] = [sp.ua_left[:,:,1]; sp.ua_left[:,:,2]] - T_right_mats[j] = [sp.ua_right[:,:,1]; sp.ua_right[:,:,2]] - T_left_inv[j] = inv(T_left_mats[j]) - T_right_inv[j] = inv(T_right_mats[j]) + T_left_mats[j] = [sp.ua_left[:, :, 1]; sp.ua_left[:, :, 2]] + T_right_mats[j] = [sp.ua_right[:, :, 1]; sp.ua_right[:, :, 2]] + T_left_inv[j] = inv(T_left_mats[j]) + T_right_inv[j] = inv(T_right_mats[j]) end end return T_left_mats, T_right_mats, T_left_inv, T_right_inv @@ -261,7 +280,7 @@ function _build_S_axis_shooting_propagators( end if can_reintegrate && !isempty(shoot_range_R) uShootR[j] = integrate_fm_with_ua_ic(chunks, shoot_range_R, sing[j].ua_right, - ctrl, equil, mats, intr; backward=false, psi_ua=sing[j].psi_ua_right) + ctrl, equil, mats, intr; backward=false, psi_ua=sing[j].psi_ua_right) else T_init = has_ua ? T_right_mats[j] : nothing uShootR[j] = assemble_fm_matrix(propagators, shoot_range_R; T_init=T_init) @@ -278,7 +297,7 @@ function _build_S_axis_shooting_propagators( end if can_reintegrate && !isempty(shoot_range_L) uShootL[j] = integrate_fm_with_ua_ic(chunks, shoot_range_L, sing[j].ua_left, - ctrl, equil, mats, intr; backward=true, psi_ua=sing[j].psi_ua_left) + ctrl, equil, mats, intr; backward=true, psi_ua=sing[j].psi_ua_left) else T_init = has_ua ? T_left_mats[j] : nothing uShootL[j] = assemble_fm_matrix(propagators, shoot_range_L; T_init=T_init) @@ -287,10 +306,10 @@ function _build_S_axis_shooting_propagators( uAxis, i_axis_mid = _build_conditioned_axis_propagator(propagators, i_crossings, N) uShootL[1] = _build_uShootL_first(propagators, chunks, i_crossings, sing, - T_left_mats, has_ua, can_reintegrate, i_axis_mid, - ctrl, equil, mats, intr, N) + T_left_mats, has_ua, can_reintegrate, i_axis_mid, + ctrl, equil, mats, intr, N) if debug - shoot_range_L1 = (i_axis_mid + 1):(i_crossings[1] - 1) + shoot_range_L1 = (i_axis_mid+1):(i_crossings[1]-1) @info " Axis propagator: $(i_axis_mid) chunks, cond=$(@sprintf("%.2e", cond(uAxis)))" @info " uShootL[1]: range=$(shoot_range_L1), cond=$(@sprintf("%.2e", cond(uShootL[1])))" end @@ -303,14 +322,14 @@ end # chunk+1 to chunk(i_crossings[j]-1). The ψ midpoint is used (not the chunk-index midpoint) # because chunks near singularities are packed tighter in ψ — Fortran convention. function _midpoint_shoot_range(chunks::Vector{IntegrationChunk}, i_crossings::Vector{Int}, - j::Int, msing::Int; side::Symbol) + j::Int, msing::Int; side::Symbol) if side === :right - j == msing && return (i_crossings[msing] + 1):length(chunks) + j == msing && return (i_crossings[msing]+1):length(chunks) chunk_start = i_crossings[j] + 1 - chunk_end = i_crossings[j+1] - 1 + chunk_end = i_crossings[j+1] - 1 else # :left, j >= 2 chunk_start = i_crossings[j-1] + 1 - chunk_end = i_crossings[j] - 1 + chunk_end = i_crossings[j] - 1 end psi_mid_target = (chunks[chunk_start].psi_start + chunks[chunk_end].psi_end) / 2 i_mid_inter = chunk_start @@ -321,7 +340,7 @@ function _midpoint_shoot_range(chunks::Vector{IntegrationChunk}, i_crossings::Ve end i_mid_inter = ic end - return side === :right ? (chunk_start:i_mid_inter) : ((i_mid_inter + 1):chunk_end) + return side === :right ? (chunk_start:i_mid_inter) : ((i_mid_inter+1):chunk_end) end # Build a well-conditioned axis propagator by forward-propagating [0; I] through the @@ -329,7 +348,7 @@ end # midpoint is placed one chunk before the first surface so that uShootL[1] covers only the # last chunk, keeping it well-conditioned. function _build_conditioned_axis_propagator(propagators::Vector{ChunkPropagator}, - i_crossings::Vector{Int}, N::Int) + i_crossings::Vector{Int}, N::Int) n_pre_cross = i_crossings[1] - 1 i_axis_mid = max(1, n_pre_cross - 1) uAxis = zeros(ComplexF64, 2N, N) @@ -340,8 +359,8 @@ function _build_conditioned_axis_propagator(propagators::Vector{ChunkPropagator} prop = propagators[ic] upper_old = uAxis[1:N, :] lower_old = uAxis[N+1:2N, :] - uAxis[1:N, :] .= prop.block_upper_ic[:,:,1] * upper_old .+ prop.block_lower_ic[:,:,1] * lower_old - uAxis[N+1:2N, :] .= prop.block_upper_ic[:,:,2] * upper_old .+ prop.block_lower_ic[:,:,2] * lower_old + uAxis[1:N, :] .= prop.block_upper_ic[:, :, 1] * upper_old .+ prop.block_lower_ic[:, :, 1] * lower_old + uAxis[N+1:2N, :] .= prop.block_upper_ic[:, :, 2] * upper_old .+ prop.block_lower_ic[:, :, 2] * lower_old Q, _ = qr(uAxis) uAxis .= Matrix(Q)[:, 1:N] end @@ -355,18 +374,18 @@ end # Falls back to T_left_mats[1] (or identity if no ua) when there's only 1 chunk before the # first crossing. function _build_uShootL_first(propagators::Vector{ChunkPropagator}, - chunks::Vector{IntegrationChunk}, i_crossings::Vector{Int}, - sing::Vector{SingType}, T_left_mats::Vector{Matrix{ComplexF64}}, - has_ua::Bool, can_reintegrate::Bool, i_axis_mid::Int, - ctrl, equil, mats, intr::ForceFreeStatesInternal, N::Int) - shoot_range_L1 = (i_axis_mid + 1):(i_crossings[1] - 1) + chunks::Vector{IntegrationChunk}, i_crossings::Vector{Int}, + sing::Vector{SingType}, T_left_mats::Vector{Matrix{ComplexF64}}, + has_ua::Bool, can_reintegrate::Bool, i_axis_mid::Int, + ctrl, equil, mats, intr::ForceFreeStatesInternal, N::Int) + shoot_range_L1 = (i_axis_mid+1):(i_crossings[1]-1) if can_reintegrate && !isempty(shoot_range_L1) return integrate_fm_with_ua_ic(chunks, shoot_range_L1, sing[1].ua_left, - ctrl, equil, mats, intr; - backward=true, psi_ua=sing[1].psi_ua_left) + ctrl, equil, mats, intr; + backward=true, psi_ua=sing[1].psi_ua_left) elseif !isempty(shoot_range_L1) return assemble_fm_matrix(propagators, shoot_range_L1; - T_init=has_ua ? T_left_mats[1] : nothing) + T_init=has_ua ? T_left_mats[1] : nothing) else return has_ua ? T_left_mats[1] : Matrix{ComplexF64}(I, 2N, 2N) end @@ -377,19 +396,19 @@ end # the catastrophically ill-conditioned axis FM. Fortran-matched structure with # nMat = (2 + 4·msing)·N. Returns (M, nMat, col_edge). function _assemble_bvp_S_axis(uShootR::Vector{Matrix{ComplexF64}}, - uShootL::Vector{Matrix{ComplexF64}}, - uAxis::Matrix{ComplexF64}, ipert_all::Vector{Int}, - msing::Int, N::Int, - wv::Union{Nothing,Matrix{ComplexF64}}, psio::Float64) + uShootL::Vector{Matrix{ComplexF64}}, + uAxis::Matrix{ComplexF64}, ipert_all::Vector{Int}, + msing::Int, N::Int, + wv::Union{Nothing,Matrix{ComplexF64}}, psio::Float64) # STRIDE global BVP block structure [Glasser-Kolemen 2018 PoP 25, 032501 Eq. 37]. nMat = (2 + 4 * msing) * N col_axis = 1:N - col_edge = (nMat - N + 1):nMat + col_edge = (nMat-N+1):nMat M = zeros(ComplexF64, nMat, nMat) # Axis matching: uShootL[1] · c_left[1] = uAxis · c_axis (2N equations) M[1:2N, _col_left(1, N)] .= uShootL[1] - M[1:2N, col_axis] .= -uAxis + M[1:2N, col_axis] .= -uAxis row_offset = 2N for j in 1:msing @@ -398,21 +417,21 @@ function _assemble_bvp_S_axis(uShootR::Vector{Matrix{ComplexF64}}, for i in 1:2N if i != ipert_j && i != ipert_j + N row_offset += 1 - M[row_offset, _col_left(j, N)[i]] = 1 + M[row_offset, _col_left(j, N)[i]] = 1 M[row_offset, _col_right(j, N)[i]] = -1 end end - junc_rows = (row_offset + 1):(row_offset + 2N) + junc_rows = (row_offset+1):(row_offset+2N) if j < msing # Midpoint matching between consecutive surfaces - M[junc_rows, _col_right(j, N)] .= -uShootR[j] - M[junc_rows, _col_left(j+1, N)] .= uShootL[j+1] + M[junc_rows, _col_right(j, N)] .= -uShootR[j] + M[junc_rows, _col_left(j + 1, N)] .= uShootL[j+1] else # Edge junction M[junc_rows, _col_right(msing, N)] .= uShootR[msing] if wv !== nothing - M[junc_rows[1:N], col_edge] .= -I(N) + M[junc_rows[1:N], col_edge] .= -I(N) M[junc_rows[N+1:end], col_edge] .= wv .* psio^2 else M[junc_rows[N+1:end], col_edge] .= -I(N) @@ -425,7 +444,7 @@ function _assemble_bvp_S_axis(uShootR::Vector{Matrix{ComplexF64}}, for j in 1:msing ipert_j = ipert_all[j] row_offset += 1 - M[row_offset, _col_left(j, N)[ipert_j]] = 1 + M[row_offset, _col_left(j, N)[ipert_j]] = 1 row_offset += 1 M[row_offset, _col_right(j, N)[ipert_j]] = 1 end @@ -460,51 +479,51 @@ end # Fallback BVP assembly with FM-based axis BC (used when no Riccati S matrices are available). # Uses the conditioned axis propagator Phi_R[1][:,N+1:2N] in place of S-axis matching. function _assemble_bvp_FM_axis(Phi_L_mats::Vector{Matrix{ComplexF64}}, - Phi_R_mats::Vector{Matrix{ComplexF64}}, ipert_all::Vector{Int}, - msing::Int, N::Int, - T_left_inv::Vector{Matrix{ComplexF64}}, - T_right_inv::Vector{Matrix{ComplexF64}}, has_ua::Bool, - wv::Union{Nothing,Matrix{ComplexF64}}, psio::Float64) + Phi_R_mats::Vector{Matrix{ComplexF64}}, ipert_all::Vector{Int}, + msing::Int, N::Int, + T_left_inv::Vector{Matrix{ComplexF64}}, + T_right_inv::Vector{Matrix{ComplexF64}}, has_ua::Bool, + wv::Union{Nothing,Matrix{ComplexF64}}, psio::Float64) nMat = (2 + 4 * msing) * N col_axis = 1:N - col_edge = (N + 4N*msing + 1):nMat + col_edge = (N+4N*msing+1):nMat M = zeros(ComplexF64, nMat, nMat) M[1:2N, (N+1):(N+2N)] .= Phi_L_mats[1] - M[1:2N, col_axis] .= -view(Phi_R_mats[1], :, N+1:2N) + M[1:2N, col_axis] .= -view(Phi_R_mats[1], :, N+1:2N) - row_drive_base = 2N + (4N-2)*msing + row_drive_base = 2N + (4N - 2) * msing for j in 1:msing ipert_j = ipert_all[j] cl = _col_left(j, N) cr = _col_right(j, N) - row_cont = 2N + (4N-2)*(j-1) + row_cont = 2N + (4N - 2) * (j - 1) for i in 1:2N if i != ipert_j && i != ipert_j + N row_cont += 1 - M[row_cont, cl[i]] = 1 + M[row_cont, cl[i]] = 1 M[row_cont, cr[i]] = -1 end end - junc_rows = (row_cont + 1):(2N + (4N-2)*j) + junc_rows = (row_cont+1):(2N+(4N-2)*j) if j < msing - M[junc_rows, cr] .= Phi_R_mats[j+1] - M[junc_rows, _col_left(j+1, N)] .= -Phi_L_mats[j+1] + M[junc_rows, cr] .= Phi_R_mats[j+1] + M[junc_rows, _col_left(j + 1, N)] .= -Phi_L_mats[j+1] else M[junc_rows, cr] .= Phi_R_mats[msing+1] if wv !== nothing - M[junc_rows[1:N], col_edge] .= -I(N) + M[junc_rows[1:N], col_edge] .= -I(N) M[junc_rows[N+1:end], col_edge] .= wv .* psio^2 else M[junc_rows[N+1:end], col_edge] .= -I(N) end end if has_ua - M[row_drive_base + 2j-1, cl] .= T_left_inv[j][ipert_j, :] - M[row_drive_base + 2j, cr] .= T_right_inv[j][ipert_j, :] + M[row_drive_base+2j-1, cl] .= T_left_inv[j][ipert_j, :] + M[row_drive_base+2j, cr] .= T_right_inv[j][ipert_j, :] else - M[row_drive_base + 2j-1, cl[ipert_j]] = 1 - M[row_drive_base + 2j, cr[ipert_j]] = 1 + M[row_drive_base+2j-1, cl[ipert_j]] = 1 + M[row_drive_base+2j, cr[ipert_j]] = 1 end end return M, nMat, col_edge @@ -515,8 +534,8 @@ end # combination subtracts dp_raw entries up to ~3×10⁴ larger than the result, and Float64 # precision lets the imaginary part drift 2–5× on DIIID-class equilibria. function _solve_bvp_and_combine_pest3(M::Matrix{ComplexF64}, msing::Int, N::Int, nMat::Int, - use_S_axis::Bool, ipert_all::Vector{Int}, col_edge, - ctrl, debug::Bool) + use_S_axis::Bool, ipert_all::Vector{Int}, col_edge, + ctrl, debug::Bool) s2 = 2 * msing Tc = (ctrl === nothing || ctrl.extended_precision_bvp) ? Complex{Double64} : ComplexF64 M_solve = Tc.(M) @@ -533,17 +552,17 @@ function _solve_bvp_and_combine_pest3(M::Matrix{ComplexF64}, msing::Int, N::Int, for jsing in 1:msing, side in 1:2 dRow = 2jsing - (2 - side) fill!(b, 0) - drive_row = use_S_axis ? (nMat - s2 + dRow) : (2N + (4N-2)*msing + dRow) + drive_row = use_S_axis ? (nMat - s2 + dRow) : (2N + (4N - 2) * msing + dRow) b[drive_row] = 1 x = use_lu ? (M_lu \ b) : (M_pinv * b) debug && _log_bvp_solve(x, b, M_solve, jsing, side, dRow, msing, N, - ipert_all, col_edge, use_S_axis) + ipert_all, col_edge, use_S_axis) for ksing in 1:msing ipert_k = ipert_all[ksing] dp_raw[dRow, 2ksing-1] = x[_col_left(ksing, N)[ipert_k+N]] - dp_raw[dRow, 2ksing] = x[_col_right(ksing, N)[ipert_k+N]] + dp_raw[dRow, 2ksing] = x[_col_right(ksing, N)[ipert_k+N]] end end @@ -563,7 +582,7 @@ end # Logging helpers for `compute_delta_prime_matrix!`. Called only when debug=true. function _log_bvp_setup(chunks, sing, S_at_surface_left, use_S_axis, has_ua, - Phi_L_mats, Phi_R_mats, Phi_R_halves, ipert_all, wv, psio, N, msing) + Phi_L_mats, Phi_R_mats, Phi_R_halves, ipert_all, wv, psio, N, msing) @info "Δ' BVP: $(length(chunks)) chunks, $msing surfaces, N=$N" @info "Δ' BVP: Axis BC: $(use_S_axis ? "S-based (Riccati)" : "FM-based (conditioned)")" @info "Δ' BVP: Asymptotic basis: $(has_ua ? "available" : "NOT available (raw basis driving)")" @@ -575,8 +594,8 @@ function _log_bvp_setup(chunks, sing, S_at_surface_left, use_S_axis, has_ua, if has_ua for j in 1:msing sp = sing[j] - T_l = [sp.ua_left[:,:,1]; sp.ua_left[:,:,2]] - T_r = [sp.ua_right[:,:,1]; sp.ua_right[:,:,2]] + T_l = [sp.ua_left[:, :, 1]; sp.ua_left[:, :, 2]] + T_r = [sp.ua_right[:, :, 1]; sp.ua_right[:, :, 2]] @info " Surface $j: cond(T_left)=$(@sprintf("%.2e", cond(T_l))), cond(T_right)=$(@sprintf("%.2e", cond(T_r)))" ipert_j = ipert_all[j] @info " Surface $j ua_left (ipert=$ipert_j, psi_ua_left=$(@sprintf("%.8f", sp.psi_ua_left))):" @@ -603,7 +622,7 @@ function _log_bvp_setup(chunks, sing, S_at_surface_left, use_S_axis, has_ua, end function _log_S_axis_shooting_propagators(uShootR, uShootL, uAxis, S_at_surface_left, - T_left_mats, ipert_all, has_ua, msing, N) + T_left_mats, ipert_all, has_ua, msing, N) @info " Shooting propagators (S-based axis BC, no axis unknowns):" for j in 1:msing shoot_R_str = @sprintf("%.2e", cond(uShootR[j])) @@ -636,7 +655,7 @@ function _log_S_axis_shooting_propagators(uShootR, uShootL, uAxis, S_at_surface_ end function _log_bvp_solve(x, b, M_solve, jsing, side, dRow, msing, N, - ipert_all, col_edge, use_S_axis) + ipert_all, col_edge, use_S_axis) residual = norm(ComplexF64.(M_solve * x - b)) side_str = side == 1 ? "left" : "right" @info " BVP solve: jsing=$jsing side=$side_str (dRow=$dRow): ||Mx-b||=$(@sprintf("%.2e", residual)), ||x||=$(@sprintf("%.2e", Float64(norm(x))))" @@ -644,9 +663,9 @@ function _log_bvp_solve(x, b, M_solve, jsing, side, dRow, msing, N, ipert_ks = ipert_all[ks] cl = _col_left(ks, N) cr = _col_right(ks, N) - xl_big = ComplexF64(x[cl[ipert_ks]]) + xl_big = ComplexF64(x[cl[ipert_ks]]) xl_small = ComplexF64(x[cl[ipert_ks+N]]) - xr_big = ComplexF64(x[cr[ipert_ks]]) + xr_big = ComplexF64(x[cr[ipert_ks]]) xr_small = ComplexF64(x[cr[ipert_ks+N]]) @info " surf $ks: x_left[big]=$(@sprintf("%+.4e%+.4ei", real(xl_big), imag(xl_big))), x_left[small]=$(@sprintf("%+.4e%+.4ei", real(xl_small), imag(xl_small)))" @info " surf $ks: x_right[big]=$(@sprintf("%+.4e%+.4ei", real(xr_big), imag(xr_big))), x_right[small]=$(@sprintf("%+.4e%+.4ei", real(xr_small), imag(xr_small)))" @@ -660,7 +679,7 @@ end function _log_bvp_pest3(dp_raw, deltap, s2, msing, Tc) @info "Δ' BVP: Full dp_raw matrix ($(s2)×$(s2)) [$(Tc)]:" for i in 1:s2 - row_str = join([@sprintf("%+.6e", Float64(real(dp_raw[i,j]))) for j in 1:s2], " ") + row_str = join([@sprintf("%+.6e", Float64(real(dp_raw[i, j]))) for j in 1:s2], " ") @info " dp_raw[$i,:] = $row_str" end @info "Δ' BVP: Raw dp diagonal = $([@sprintf("%.4f%+.4fi", Float64(real(dp_raw[i,i])), Float64(imag(dp_raw[i,i]))) for i in 1:s2])" @@ -720,8 +739,8 @@ function pest3_decompose(dp_raw::AbstractMatrix) for i in 1:m, j in 1:m LL = dp_raw[2i-1, 2j-1] LR = dp_raw[2i-1, 2j] - RL = dp_raw[2i, 2j-1] - RR = dp_raw[2i, 2j] + RL = dp_raw[2i, 2j-1] + RR = dp_raw[2i, 2j] Ap[i, j] = RR + RL + LR + LL Bp[i, j] = RR - RL + LR - LL Gp[i, j] = RR + RL - LR - LL diff --git a/src/ForceFreeStates/Riccati/Propagators.jl b/src/ForceFreeStates/Riccati/Propagators.jl index 285695c5b..d888750b7 100644 --- a/src/ForceFreeStates/Riccati/Propagators.jl +++ b/src/ForceFreeStates/Riccati/Propagators.jl @@ -6,7 +6,7 @@ # length chunks; the absolute floor catches short chunks where 5% of the span would be # smaller than the typical ODE step. const SAVE_NEAR_END_FRAC = 0.05 -const SAVE_NEAR_END_PSI = 1e-4 +const SAVE_NEAR_END_PSI = 1e-4 """ assemble_fm_matrix(propagators, idx_range; condition=false) -> Matrix{ComplexF64} @@ -16,6 +16,7 @@ in order for indices `idx_range`. Returns Φ_end * ... * Φ_start, so that the r maps the IC at the start of `idx_range[1]` to the state at the end of `idx_range[end]`. Each `ChunkPropagator` stores the 2N columns of Φ split into two N×N×2 blocks: + ``` block_upper_ic[:,:,1:2] ↔ Φ[:,1:N] (result from IC=(I,0)) block_lower_ic[:,:,1:2] ↔ Φ[:,N+1:2N] (result from IC=(0,I)) @@ -33,8 +34,8 @@ means only U₂ ICs are needed. Do NOT use for inter-surface segments where both and U₂ components carry physical information. """ function assemble_fm_matrix(propagators::Vector{ChunkPropagator}, idx_range; - condition::Bool=false, - T_init::Union{Nothing,Matrix{ComplexF64}}=nothing) + condition::Bool=false, + T_init::Union{Nothing,Matrix{ComplexF64}}=nothing) # Determine matrix size from T_init if provided (lets us handle empty idx_range and even # an empty propagators list, provided T_init carries the dimension). Otherwise fall back # to the first propagator that actually exists in idx_range, with a final fallback to @@ -123,7 +124,7 @@ end riccati_der!(du, u, params, psieval) Evaluate the explicit dual Riccati ODE right-hand side: - dS/dψ = w†·F̄⁻¹·w - S·Ḡ·S, w = Q - K̄·S +dS/dψ = w†·F̄⁻¹·w - S·Ḡ·S, w = Q - K̄·S where Q = diag(1/(m - n·q)) is the diagonal singular factor matrix. The identity slice u[:,:,2] = I does not evolve (du[:,:,2] = 0). @@ -147,7 +148,7 @@ See: Glasser (2018) Phys. Plasmas 25, 032507 — Eq. 19 (dual Riccati form) _, equil, mats, intr, odet, _ = params Npert = intr.numpert_total - S = @view u[:, :, 1] + S = @view u[:, :, 1] dS = @view du[:, :, 1] @view(du[:, :, 2]) .= 0 # identity does not evolve @@ -162,9 +163,9 @@ See: Glasser (2018) Phys. Plasmas 25, 032507 — Eq. 19 (dual Riccati form) fmat_lower = acquire!(pool, ComplexF64, Npert, Npert) kmat = similar!(pool, fmat_lower) gmat = similar!(pool, fmat_lower) - w = similar!(pool, fmat_lower) # w = Q - K̄·S - v = similar!(pool, fmat_lower) # v = F̄⁻¹·w (then reused for S·Ḡ·S) - tmp = similar!(pool, fmat_lower) # scratch + w = similar!(pool, fmat_lower) # w = Q - K̄·S + v = similar!(pool, fmat_lower) # v = F̄⁻¹·w (then reused for S·Ḡ·S) + tmp = similar!(pool, fmat_lower) # scratch # Evaluate F̄ (Cholesky factor), K̄, Ḡ splines at current ψ mats.ideal.F_spline_lower(vec(fmat_lower), psieval; hint=mats._hint) @@ -255,8 +256,8 @@ function riccati_integrate_chunk!( cb = DiscreteCallback((u, t, integrator) -> true, riccati_integrator_callback!) rtol = ctrl.eulerlagrange_tolerance prob = ODEProblem(sing_der!, odet.u, (chunk.psi_start, chunk.psi_end), - (ctrl, equil, mats, intr, odet, chunk)) - sol = solve(prob, Vern9(); reltol=rtol, callback=cb, save_everystep=false, save_end=true) + (ctrl, equil, mats, intr, odet, chunk)) + sol = solve(prob, el_ode_algorithm(ctrl); reltol=rtol, abstol=ctrl.ode_abstol, callback=cb, save_everystep=false, save_end=true) odet.u .= sol.u[end] odet.psifac = sol.t[end] # Renormalize end state to (S, I) convention for the next chunk. @@ -273,8 +274,8 @@ end renormalize_riccati!(odet, intr) After a singular surface crossing, restore the canonical Riccati storage convention: - u[:,:,1] = S_new = U₁_new · U₂_new⁻¹ - u[:,:,2] = I +u[:,:,1] = S_new = U₁_new · U₂_new⁻¹ +u[:,:,2] = I `riccati_cross_ideal_singular_surf!` leaves u[:,:,1] = U₁_new and u[:,:,2] = U₂_new (not I), so this step is required before continuing the Riccati integration. @@ -298,8 +299,8 @@ end renormalize_riccati_inplace!(u, N) In-place Riccati renormalization on an arbitrary N×N×2 array: - u[:,:,1] = U₁ · U₂⁻¹ (new S) - u[:,:,2] = I +u[:,:,1] = U₁ · U₂⁻¹ (new S) +u[:,:,2] = I Used in `riccati_integrator_callback!` to renormalize the integrator's live state when column norms grow beyond `ctrl.ucrit`, analogous to Gaussian reduction in the @@ -346,8 +347,8 @@ function integrate_propagator_chunk!( # naturally. The resulting propagator maps state at psi_end → psi_start, which is # well-conditioned because exponentially growing solutions (forward) decay backward. tspan = chunk.direction == 1 ? - (chunk.psi_start, chunk.psi_end) : - (chunk.psi_end, chunk.psi_start) + (chunk.psi_start, chunk.psi_end) : + (chunk.psi_end, chunk.psi_start) rtol = ctrl.eulerlagrange_tolerance params = (ctrl, equil, mats, intr, odet_proxy, chunk) @@ -359,7 +360,7 @@ function integrate_propagator_chunk!( odet_proxy.spline_hint[] = 1 odet_proxy.mats_hint[] = 1 prob = ODEProblem(sing_der!, u_upper, tspan, params) - sol = solve(prob, Vern9(); reltol=rtol, save_everystep=false, save_end=true) + sol = solve(prob, el_ode_algorithm(ctrl); reltol=rtol, abstol=ctrl.ode_abstol, save_everystep=false, save_end=true) prop.block_upper_ic .= sol.u[end] odet_proxy.total_steps += sol.stats.naccept # thread-local; summed into odet after the BVP barrier @@ -371,7 +372,7 @@ function integrate_propagator_chunk!( odet_proxy.spline_hint[] = 1 odet_proxy.mats_hint[] = 1 prob = ODEProblem(sing_der!, u_lower, tspan, params) - sol = solve(prob, Vern9(); reltol=rtol, save_everystep=false, save_end=true) + sol = solve(prob, el_ode_algorithm(ctrl); reltol=rtol, abstol=ctrl.ode_abstol, save_everystep=false, save_end=true) prop.block_lower_ic .= sol.u[end] odet_proxy.total_steps += sol.stats.naccept end @@ -401,12 +402,12 @@ function integrate_fm_with_ua_ic( equil::Equilibrium.PlasmaEquilibrium, mats::MatrixSplines, intr::ForceFreeStatesInternal; - backward::Bool = false, - psi_ua::Float64 = NaN + backward::Bool=false, + psi_ua::Float64=NaN ) N = intr.numpert_total psi_start = chunks[first(chunk_range)].psi_start - psi_end = chunks[last(chunk_range)].psi_end + psi_end = chunks[last(chunk_range)].psi_end # Use stored ua ψ location if provided; otherwise fall back to chunk boundary. # The ua is evaluated at the inner-layer boundary (exact ψ from singular crossing), # which may differ slightly from the nearest chunk boundary. @@ -437,9 +438,9 @@ function integrate_fm_with_ua_ic( odet_proxy.spline_hint[] = 1 odet_proxy.mats_hint[] = 1 prob = ODEProblem(sing_der!, u0, tspan, params) - sol = solve(prob, Vern9(); reltol=rtol, abstol=abstol_arr, save_everystep=false, save_end=true) - result[1:N, 1:N] .= sol.u[end][:, :, 1] - result[N+1:2N, 1:N] .= sol.u[end][:, :, 2] + sol = solve(prob, el_ode_algorithm(ctrl); reltol=rtol, abstol=abstol_arr, save_everystep=false, save_end=true) + result[1:N, 1:N] .= sol.u[end][:, :, 1] + result[N+1:2N, 1:N] .= sol.u[end][:, :, 2] # Batch 2: columns N+1:2N of T (small solutions) u0[:, :, 1] .= ua[:, N+1:2N, 1] @@ -450,9 +451,9 @@ function integrate_fm_with_ua_ic( odet_proxy.spline_hint[] = 1 odet_proxy.mats_hint[] = 1 prob = ODEProblem(sing_der!, u0, tspan, params) - sol = solve(prob, Vern9(); reltol=rtol, abstol=abstol_arr, save_everystep=false, save_end=true) - result[1:N, N+1:2N] .= sol.u[end][:, :, 1] - result[N+1:2N, N+1:2N] .= sol.u[end][:, :, 2] + sol = solve(prob, el_ode_algorithm(ctrl); reltol=rtol, abstol=abstol_arr, save_everystep=false, save_end=true) + result[1:N, N+1:2N] .= sol.u[end][:, :, 1] + result[N+1:2N, N+1:2N] .= sol.u[end][:, :, 2] return result end @@ -464,8 +465,8 @@ Apply the chunk propagator `prop` to the current state `odet.u` in-place. The propagator acts as a linear map on the (U₁, U₂) pair: - U₁_new = block_upper_ic[:,:,1] · U₁_prev + block_lower_ic[:,:,1] · U₂_prev - U₂_new = block_upper_ic[:,:,2] · U₁_prev + block_lower_ic[:,:,2] · U₂_prev +U₁_new = block_upper_ic[:,:,1] · U₁_prev + block_lower_ic[:,:,1] · U₂_prev +U₂_new = block_upper_ic[:,:,2] · U₁_prev + block_lower_ic[:,:,2] · U₂_prev This correctly propagates any state (not just the identity), including the (S, I) form produced by Riccati-style crossings. @@ -520,8 +521,8 @@ function apply_propagator_inverse!(odet::OdeState, prop::ChunkPropagator) # Φ_bwd maps state at psi_end → psi_start (well-conditioned). # We want Φ_fwd = Φ_bwd⁻¹ to advance state from psi_start → psi_end. # Solving Φ_bwd · x = [U₁_old; U₂_old] gives x = Φ_bwd⁻¹ · [U₁_old; U₂_old]. - u_old = [odet.u[:,:,1]; odet.u[:,:,2]] # 2N × N + u_old = [odet.u[:, :, 1]; odet.u[:, :, 2]] # 2N × N u_new = Φ \ u_old # LU solve, 2N × N - odet.u[:,:,1] .= u_new[1:N, :] - odet.u[:,:,2] .= u_new[N+1:2N, :] + odet.u[:, :, 1] .= u_new[1:N, :] + odet.u[:, :, 2] .= u_new[N+1:2N, :] end