From 47ffb439a27d53d4b2954d269b4649ce2b18f45b Mon Sep 17 00:00:00 2001 From: Matthew Pharr Date: Wed, 21 Oct 2026 10:36:24 +0200 Subject: [PATCH 1/3] Equilibrium - REFACTOR - Move construction and post-processing out of Equilibrium.jl into their own files Co-Authored-By: Claude Opus 5.5 --- docs/development/architecture.md | 2 + src/Equilibrium/Construction.jl | 199 +++++++++ src/Equilibrium/Equilibrium.jl | 696 +----------------------------- src/Equilibrium/PostProcessing.jl | 491 +++++++++++++++++++++ 4 files changed, 700 insertions(+), 688 deletions(-) create mode 100644 src/Equilibrium/Construction.jl create mode 100644 src/Equilibrium/PostProcessing.jl diff --git a/docs/development/architecture.md b/docs/development/architecture.md index 059539a8b..5e8719aba 100644 --- a/docs/development/architecture.md +++ b/docs/development/architecture.md @@ -34,6 +34,8 @@ Splines are provided by the external `FastInterpolations` package rather than by - `sol` - Solovev analytical equilibrium - Key files: - `EquilibriumTypes.jl` - Core data structures + - `Construction.jl` - Forming a `PlasmaEquilibrium` from each source type, and re-forming it + - `PostProcessing.jl` - Separatrix, global parameters, q extrema, Grad-Shafranov residual - `ReadEquilibrium.jl` - Parsing equilibrium files - `DirectEquilibrium.jl` - Direct Grad-Shafranov solver - `InverseEquilibrium.jl` - Inverse equilibrium solver diff --git a/src/Equilibrium/Construction.jl b/src/Equilibrium/Construction.jl new file mode 100644 index 000000000..001353a7b --- /dev/null +++ b/src/Equilibrium/Construction.jl @@ -0,0 +1,199 @@ +# Construction.jl +# +# How a PlasmaEquilibrium is formed: the analytic registry, the solver input for each source type, +# the constructors, and re-forming on a new grid. + +""" + AnalyticEqSpec + +Single-source-of-truth entry describing one analytic-equilibrium kind: the `gpec.toml` +section that carries its parameters, the `*Config` type built from that section, and the +run function that turns the config into solver input. `tj_analytic` and `tj_analytic_direct` +share `TJAnalyticConfig` but bind different run functions, so `run_fn` is per-entry rather +than derived from `config_type`. + +## Fields + + - `section::String` — `gpec.toml` section key (e.g. `"SOL_INPUT"`) + - `config_type::DataType` — `*Config` type, constructed via its `(::Dict)` / `(::String)` ctor + - `run_fn::Function` — `(EquilibriumConfig, config) -> eq_input` solver entry point +""" +struct AnalyticEqSpec + section::String + config_type::DataType + run_fn::Function +end + +# Registry of analytic equilibrium kinds. Both `run_input` (fresh runs) and the rerun input +# builder dispatch off this table, so adding a new analytic kind is one new row. +const ANALYTIC_EQ = Dict( + "sol" => AnalyticEqSpec("SOL_INPUT", SolovevConfig, sol_run), + "lar" => AnalyticEqSpec("LAR_INPUT", LargeAspectRatioConfig, lar_run), + "tj_analytic" => AnalyticEqSpec("TJ_ANALYTIC_INPUT", TJAnalyticConfig, tj_analytic_run), + "tj_analytic_direct" => AnalyticEqSpec("TJ_ANALYTIC_INPUT", TJAnalyticConfig, tj_analytic_run_direct) +) + +""" + run_input(eq_config::EquilibriumConfig, source) -> DirectRunInput or InverseRunInput + +Build the solver input for `eq_config` from `source`: `nothing` reads `eq_config.eq_filename`, an +[`AnalyticEquilibriumConfig`](@ref) runs its analytic solver, an `IMASdd.dd` is read by +[`read_imas`](@ref), and a run input rebuilt from a stored ingest is rebound onto `eq_config`. +""" +function run_input(eq_config::EquilibriumConfig, ::Nothing) + eq_type = eq_config.eq_type + eq_type in EFIT_KINDS && return resolve_psihigh!(read_efit(eq_config)) + eq_type in ("chease2", "chease_ascii") && return read_chease_ascii(eq_config) + eq_type in ("chease", "chease_binary") && return read_chease_binary(eq_config) + eq_type == "ifile" && return resolve_psihigh!(read_ifile(eq_config)) + if haskey(ANALYTIC_EQ, eq_type) + spec = ANALYTIC_EQ[eq_type] + error("analytic eq_type=\"$eq_type\" needs its $(spec.config_type) as the source (built from the embedded [$(spec.section)] section)") + end + eq_type == "imas" && error("eq_type=\"imas\" needs an IMASdd.dd as the source") + return error("Equilibrium type $(eq_type) is not implemented") +end + +function run_input(eq_config::EquilibriumConfig, cfg::AnalyticEquilibriumConfig) + spec = get(ANALYTIC_EQ, eq_config.eq_type, nothing) + spec !== nothing && cfg isa spec.config_type || + error("eq_type=\"$(eq_config.eq_type)\" is not an analytic kind solved from a $(typeof(cfg))") + return spec.run_fn(eq_config, cfg) +end + +function run_input(eq_config::EquilibriumConfig, dd::IMASdd.dd) + eq_config.eq_type == "imas" || error("an IMASdd.dd source needs eq_type=\"imas\", not \"$(eq_config.eq_type)\"") + return read_imas(eq_config, dd) +end + +# A run input rebuilt from an ingest is rebound onto the current config, resetting psihigh so a +# pass-1 clamped value does not shadow a psihigh overridden in the rerun TOML. +function run_input(eq_config::EquilibriumConfig, input::DirectRunInput) + input.config = eq_config + input.psihigh_resolved = eq_config.psihigh + return eq_config.eq_type in EFIT_KINDS ? resolve_psihigh!(input) : input +end + +function run_input(eq_config::EquilibriumConfig, input::InverseRunInput) + input.config = eq_config + input.psihigh_resolved = eq_config.psihigh + return resolve_psihigh!(input) +end + +""" + PlasmaEquilibrium(eq_config::EquilibriumConfig, source=nothing; override_psi_nodes=nothing) -> PlasmaEquilibrium + +Form the equilibrium `eq_config` describes from `source` (see [`run_input`](@ref)), run the +solver, and return it with global parameters, q-profile, and GSE diagnostics. Every other +constructor and [`setup_equilibrium`](@ref) end here. + +`override_psi_nodes` bypasses the config-driven radial grid (`grid_type`/`mpsi`) and forms +the equilibrium on the given strictly increasing ψ_N node vector, whose endpoints must match +`psilow`/`psihigh`. Used by the two-pass auto-grid refinement through [`reform`](@ref); standalone +callers with `mpsi=0` get the coarse first-pass grid unless they refine and re-form. +""" +function PlasmaEquilibrium(eq_config::EquilibriumConfig, source=nothing; override_psi_nodes::Union{Nothing,Vector{Float64}}=nothing) + eq_input = run_input(eq_config, source) + + plasma_equilibrium = if eq_config.eq_type == "efit_by_inversion" + equilibrium_solver_by_inversion(eq_input; override_psi_nodes) + elseif eq_config.eq_type == "efit_arclength" + equilibrium_solver(eq_input, arclength_fieldline_int; override_psi_nodes) + else + equilibrium_solver(eq_input; override_psi_nodes) + end + + # Forward the captured ingest so the gpec.h5 writer can snapshot it (nothing for analytic). + plasma_equilibrium.ingest = eq_input.ingest + # Analytic configs are still mutable; copy so later edits by the caller cannot change a re-form. + source isa AnalyticEquilibriumConfig && (plasma_equilibrium.analytic = deepcopy(source)) + + equilibrium_global_parameters!(plasma_equilibrium) + equilibrium_qfind!(plasma_equilibrium) + equilibrium_gse!(plasma_equilibrium) + + return plasma_equilibrium +end + +""" + setup_equilibrium(eq_config::EquilibriumConfig, source=nothing; override_psi_nodes=nothing) + setup_equilibrium(path::String="equil.toml") + +Alias of [`PlasmaEquilibrium(eq_config, source)`](@ref PlasmaEquilibrium); the `path` form reads +the [`EquilibriumConfig`](@ref) from a TOML file. +""" +setup_equilibrium(path::String="equil.toml") = PlasmaEquilibrium(EquilibriumConfig(path)) +setup_equilibrium(eq_config::EquilibriumConfig, source=nothing; kwargs...) = PlasmaEquilibrium(eq_config, source; kwargs...) + +""" + reform(equil::PlasmaEquilibrium, psi_nodes::Vector{Float64}) -> PlasmaEquilibrium + +Re-form `equil` on the ψ_N nodes `psi_nodes` from the source it remembers: the stored +`analytic` config, else its raw `ingest`, so no file is re-read. Only an equilibrium carrying +neither (one built with the positional constructor, or by a reader that captures no ingest) +falls back to re-reading `config.eq_filename`. The node endpoints must match `psilow`/`psihigh`. +""" +function reform(equil::PlasmaEquilibrium, psi_nodes::Vector{Float64}) + source = if equil.analytic !== nothing + equil.analytic + elseif equil.ingest isa DirectIngest + build_direct_from_ingest(equil.config, equil.ingest) + elseif equil.ingest isa InverseIngest + build_inverse_from_ingest(equil.config, equil.ingest) + else + nothing + end + return PlasmaEquilibrium(equil.config, source; override_psi_nodes=psi_nodes) +end + +""" + PlasmaEquilibrium(path::AbstractString; eq_type="efit", kwargs...) -> PlasmaEquilibrium + +Read the equilibrium file at `path` and return the processed equilibrium. Convenience entry +point of the scripting API: `kwargs` are [`EquilibriumConfig`](@ref) fields, so +`PlasmaEquilibrium("g000001.00001"; jac_type="hamada", mpsi=128)` is the whole setup. + +Only file-based equilibria go through this constructor. Analytic and IMAS equilibria have +their own `PlasmaEquilibrium(::AnalyticEquilibriumConfig)` and `PlasmaEquilibrium(::IMASdd.dd)` forms. + +```julia +eq = PlasmaEquilibrium("input.geqdsk"; jac_type="hamada") +``` +""" +function PlasmaEquilibrium(path::AbstractString; eq_type::String="efit", kwargs...) + haskey(ANALYTIC_EQ, eq_type) && + error("$eq_type is an analytic equilibrium: build it with PlasmaEquilibrium($(ANALYTIC_EQ[eq_type].config_type)(...)) instead") + return PlasmaEquilibrium(EquilibriumConfig(; eq_type, eq_filename=abspath(path), kwargs...)) +end + +# Default analytic eq_type for a config: the first registry kind, in sorted order, that takes it. +_analytic_eq_type(cfg::AnalyticEquilibriumConfig) = first(k for k in sort!(collect(keys(ANALYTIC_EQ))) if ANALYTIC_EQ[k].config_type === typeof(cfg)) + +""" + PlasmaEquilibrium(cfg::AnalyticEquilibriumConfig; eq_type=, kwargs...) -> PlasmaEquilibrium + +Solve the analytic equilibrium described by `cfg` and return the processed equilibrium. The +`eq_type` is implied by the config type; `TJAnalyticConfig` defaults to `tj_analytic` and takes +`eq_type="tj_analytic_direct"` for the direct solve. `kwargs` are [`EquilibriumConfig`](@ref) fields. + +```julia +eq = PlasmaEquilibrium(Equilibrium.SolovevConfig(; elongation=1.5); mpsi=0, psi_accuracy=1e-3) +``` +""" +function PlasmaEquilibrium(cfg::AnalyticEquilibriumConfig; eq_type::String=_analytic_eq_type(cfg), kwargs...) + return PlasmaEquilibrium(EquilibriumConfig(; eq_type, kwargs...), cfg) +end + +""" + PlasmaEquilibrium(dd::IMASdd.dd; kwargs...) -> PlasmaEquilibrium + +Read the active equilibrium time slice of the IMAS data dictionary `dd` and return the processed +equilibrium. `kwargs` are [`EquilibriumConfig`](@ref) fields; `imas_cocos` sets the ψ convention. + +```julia +eq = PlasmaEquilibrium(dd; imas_cocos=11, jac_type="hamada") +``` +""" +function PlasmaEquilibrium(dd::IMASdd.dd; kwargs...) + return PlasmaEquilibrium(EquilibriumConfig(; eq_type="imas", kwargs...), dd) +end diff --git a/src/Equilibrium/Equilibrium.jl b/src/Equilibrium/Equilibrium.jl index 76b72d376..55d881b0d 100644 --- a/src/Equilibrium/Equilibrium.jl +++ b/src/Equilibrium/Equilibrium.jl @@ -12,6 +12,12 @@ import StaticArrays: @MMatrix, SVector import ..Utilities import IMASdd +# --- Constants --- +const mu0 = 4π * 1e-7 + +# efit-family equilibrium kinds that share the separatrix-clamp / inversion handling. +const EFIT_KINDS = ("efit", "efit_arclength", "efit_by_inversion") + # --- Internal Module Structure --- include("EquilibriumTypes.jl") include("SolveTolerances.jl") @@ -27,6 +33,8 @@ include("InverseEquilibrium.jl") include("AnalyticEquilibrium.jl") include("GeometryProfiles.jl") include("KineticProfiles.jl") +include("Construction.jl") +include("PostProcessing.jl") # --- Expose types and functions to the user --- export setup_equilibrium, reform, EquilibriumConfig, PlasmaEquilibrium, EquilibriumParameters, AnalyticEquilibriumConfig, @@ -37,692 +45,4 @@ export flux_surface_metric, flux_surface_area export wants_two_pass, refined_psi_grid, merge_mandatory_nodes, bracket_mandatory_nodes, enforce_min_spacing, implied_knot_count export compute_sqrt_jac_delpsi, compute_sqrtamat, rootarea_to_area_weight, area_to_rootarea_weight -# --- Constants --- -const mu0 = 4π * 1e-7 - -# efit-family equilibrium kinds that share the separatrix-clamp / inversion handling. -const EFIT_KINDS = ("efit", "efit_arclength", "efit_by_inversion") - -""" - AnalyticEqSpec - -Single-source-of-truth entry describing one analytic-equilibrium kind: the `gpec.toml` -section that carries its parameters, the `*Config` type built from that section, and the -run function that turns the config into solver input. `tj_analytic` and `tj_analytic_direct` -share `TJAnalyticConfig` but bind different run functions, so `run_fn` is per-entry rather -than derived from `config_type`. - -## Fields - - - `section::String` — `gpec.toml` section key (e.g. `"SOL_INPUT"`) - - `config_type::DataType` — `*Config` type, constructed via its `(::Dict)` / `(::String)` ctor - - `run_fn::Function` — `(EquilibriumConfig, config) -> eq_input` solver entry point -""" -struct AnalyticEqSpec - section::String - config_type::DataType - run_fn::Function -end - -# Registry of analytic equilibrium kinds. Both `run_input` (fresh runs) and the rerun input -# builder dispatch off this table, so adding a new analytic kind is one new row. -const ANALYTIC_EQ = Dict( - "sol" => AnalyticEqSpec("SOL_INPUT", SolovevConfig, sol_run), - "lar" => AnalyticEqSpec("LAR_INPUT", LargeAspectRatioConfig, lar_run), - "tj_analytic" => AnalyticEqSpec("TJ_ANALYTIC_INPUT", TJAnalyticConfig, tj_analytic_run), - "tj_analytic_direct" => AnalyticEqSpec("TJ_ANALYTIC_INPUT", TJAnalyticConfig, tj_analytic_run_direct) -) - -""" - run_input(eq_config::EquilibriumConfig, source) -> DirectRunInput or InverseRunInput - -Build the solver input for `eq_config` from `source`: `nothing` reads `eq_config.eq_filename`, an -[`AnalyticEquilibriumConfig`](@ref) runs its analytic solver, an `IMASdd.dd` is read by -[`read_imas`](@ref), and a run input rebuilt from a stored ingest is rebound onto `eq_config`. -""" -function run_input(eq_config::EquilibriumConfig, ::Nothing) - eq_type = eq_config.eq_type - eq_type in EFIT_KINDS && return resolve_psihigh!(read_efit(eq_config)) - eq_type in ("chease2", "chease_ascii") && return read_chease_ascii(eq_config) - eq_type in ("chease", "chease_binary") && return read_chease_binary(eq_config) - eq_type == "ifile" && return resolve_psihigh!(read_ifile(eq_config)) - if haskey(ANALYTIC_EQ, eq_type) - spec = ANALYTIC_EQ[eq_type] - error("analytic eq_type=\"$eq_type\" needs its $(spec.config_type) as the source (built from the embedded [$(spec.section)] section)") - end - eq_type == "imas" && error("eq_type=\"imas\" needs an IMASdd.dd as the source") - return error("Equilibrium type $(eq_type) is not implemented") -end - -function run_input(eq_config::EquilibriumConfig, cfg::AnalyticEquilibriumConfig) - spec = get(ANALYTIC_EQ, eq_config.eq_type, nothing) - spec !== nothing && cfg isa spec.config_type || - error("eq_type=\"$(eq_config.eq_type)\" is not an analytic kind solved from a $(typeof(cfg))") - return spec.run_fn(eq_config, cfg) -end - -function run_input(eq_config::EquilibriumConfig, dd::IMASdd.dd) - eq_config.eq_type == "imas" || error("an IMASdd.dd source needs eq_type=\"imas\", not \"$(eq_config.eq_type)\"") - return read_imas(eq_config, dd) -end - -# A run input rebuilt from an ingest is rebound onto the current config, resetting psihigh so a -# pass-1 clamped value does not shadow a psihigh overridden in the rerun TOML. -function run_input(eq_config::EquilibriumConfig, input::DirectRunInput) - input.config = eq_config - input.psihigh_resolved = eq_config.psihigh - return eq_config.eq_type in EFIT_KINDS ? resolve_psihigh!(input) : input -end - -function run_input(eq_config::EquilibriumConfig, input::InverseRunInput) - input.config = eq_config - input.psihigh_resolved = eq_config.psihigh - return resolve_psihigh!(input) -end - -""" - PlasmaEquilibrium(eq_config::EquilibriumConfig, source=nothing; override_psi_nodes=nothing) -> PlasmaEquilibrium - -Form the equilibrium `eq_config` describes from `source` (see [`run_input`](@ref)), run the -solver, and return it with global parameters, q-profile, and GSE diagnostics. Every other -constructor and [`setup_equilibrium`](@ref) end here. - -`override_psi_nodes` bypasses the config-driven radial grid (`grid_type`/`mpsi`) and forms -the equilibrium on the given strictly increasing ψ_N node vector, whose endpoints must match -`psilow`/`psihigh`. Used by the two-pass auto-grid refinement through [`reform`](@ref); standalone -callers with `mpsi=0` get the coarse first-pass grid unless they refine and re-form. -""" -function PlasmaEquilibrium(eq_config::EquilibriumConfig, source=nothing; override_psi_nodes::Union{Nothing,Vector{Float64}}=nothing) - eq_input = run_input(eq_config, source) - - plasma_equilibrium = if eq_config.eq_type == "efit_by_inversion" - equilibrium_solver_by_inversion(eq_input; override_psi_nodes) - elseif eq_config.eq_type == "efit_arclength" - equilibrium_solver(eq_input, arclength_fieldline_int; override_psi_nodes) - else - equilibrium_solver(eq_input; override_psi_nodes) - end - - # Forward the captured ingest so the gpec.h5 writer can snapshot it (nothing for analytic). - plasma_equilibrium.ingest = eq_input.ingest - # Analytic configs are still mutable; copy so later edits by the caller cannot change a re-form. - source isa AnalyticEquilibriumConfig && (plasma_equilibrium.analytic = deepcopy(source)) - - equilibrium_global_parameters!(plasma_equilibrium) - equilibrium_qfind!(plasma_equilibrium) - equilibrium_gse!(plasma_equilibrium) - - return plasma_equilibrium -end - -""" - setup_equilibrium(eq_config::EquilibriumConfig, source=nothing; override_psi_nodes=nothing) - setup_equilibrium(path::String="equil.toml") - -Alias of [`PlasmaEquilibrium(eq_config, source)`](@ref PlasmaEquilibrium); the `path` form reads -the [`EquilibriumConfig`](@ref) from a TOML file. -""" -setup_equilibrium(path::String="equil.toml") = PlasmaEquilibrium(EquilibriumConfig(path)) -setup_equilibrium(eq_config::EquilibriumConfig, source=nothing; kwargs...) = PlasmaEquilibrium(eq_config, source; kwargs...) - -""" - reform(equil::PlasmaEquilibrium, psi_nodes::Vector{Float64}) -> PlasmaEquilibrium - -Re-form `equil` on the ψ_N nodes `psi_nodes` from the source it remembers: the stored -`analytic` config, else its raw `ingest`, so no file is re-read. Only an equilibrium carrying -neither (one built with the positional constructor, or by a reader that captures no ingest) -falls back to re-reading `config.eq_filename`. The node endpoints must match `psilow`/`psihigh`. -""" -function reform(equil::PlasmaEquilibrium, psi_nodes::Vector{Float64}) - source = if equil.analytic !== nothing - equil.analytic - elseif equil.ingest isa DirectIngest - build_direct_from_ingest(equil.config, equil.ingest) - elseif equil.ingest isa InverseIngest - build_inverse_from_ingest(equil.config, equil.ingest) - else - nothing - end - return PlasmaEquilibrium(equil.config, source; override_psi_nodes=psi_nodes) -end - -""" - PlasmaEquilibrium(path::AbstractString; eq_type="efit", kwargs...) -> PlasmaEquilibrium - -Read the equilibrium file at `path` and return the processed equilibrium. Convenience entry -point of the scripting API: `kwargs` are [`EquilibriumConfig`](@ref) fields, so -`PlasmaEquilibrium("g000001.00001"; jac_type="hamada", mpsi=128)` is the whole setup. - -Only file-based equilibria go through this constructor. Analytic and IMAS equilibria have -their own `PlasmaEquilibrium(::AnalyticEquilibriumConfig)` and `PlasmaEquilibrium(::IMASdd.dd)` forms. - -```julia -eq = PlasmaEquilibrium("input.geqdsk"; jac_type="hamada") -``` -""" -function PlasmaEquilibrium(path::AbstractString; eq_type::String="efit", kwargs...) - haskey(ANALYTIC_EQ, eq_type) && - error("$eq_type is an analytic equilibrium: build it with PlasmaEquilibrium($(ANALYTIC_EQ[eq_type].config_type)(...)) instead") - return PlasmaEquilibrium(EquilibriumConfig(; eq_type, eq_filename=abspath(path), kwargs...)) -end - -# Default analytic eq_type for a config: the first registry kind, in sorted order, that takes it. -_analytic_eq_type(cfg::AnalyticEquilibriumConfig) = first(k for k in sort!(collect(keys(ANALYTIC_EQ))) if ANALYTIC_EQ[k].config_type === typeof(cfg)) - -""" - PlasmaEquilibrium(cfg::AnalyticEquilibriumConfig; eq_type=, kwargs...) -> PlasmaEquilibrium - -Solve the analytic equilibrium described by `cfg` and return the processed equilibrium. The -`eq_type` is implied by the config type; `TJAnalyticConfig` defaults to `tj_analytic` and takes -`eq_type="tj_analytic_direct"` for the direct solve. `kwargs` are [`EquilibriumConfig`](@ref) fields. - -```julia -eq = PlasmaEquilibrium(Equilibrium.SolovevConfig(; elongation=1.5); mpsi=0, psi_accuracy=1e-3) -``` -""" -function PlasmaEquilibrium(cfg::AnalyticEquilibriumConfig; eq_type::String=_analytic_eq_type(cfg), kwargs...) - return PlasmaEquilibrium(EquilibriumConfig(; eq_type, kwargs...), cfg) -end - -""" - PlasmaEquilibrium(dd::IMASdd.dd; kwargs...) -> PlasmaEquilibrium - -Read the active equilibrium time slice of the IMAS data dictionary `dd` and return the processed -equilibrium. `kwargs` are [`EquilibriumConfig`](@ref) fields; `imas_cocos` sets the ψ convention. - -```julia -eq = PlasmaEquilibrium(dd; imas_cocos=11, jac_type="hamada") -``` -""" -function PlasmaEquilibrium(dd::IMASdd.dd; kwargs...) - return PlasmaEquilibrium(EquilibriumConfig(; eq_type="imas", kwargs...), dd) -end - -""" - equilibrium_separatrix_find!(pe::PlasmaEquilibrium) - -Finds the separatrix locations in the plasma equilibrium (rsep, zsep, rext, zext). -Performs the same function as equil_out_sep_find in the Fortran code. -""" -function equilibrium_separatrix_find!(pe::PlasmaEquilibrium) - mpsi = length(pe.rzphi_xs) - 1 - edge_idx = mpsi + 1 # Edge flux surface index - psi_edge = pe.rzphi_xs[edge_idx] - rsep = zeros(2) - - # Outboard and inboard midplane R via bracketed Brent on θ + η(θ) - η₀ = 0. - # iside=1 → outboard (η₀=0.0, θ near 0), iside=2 → inboard (η₀=0.5, θ near 0.5). - for iside in 1:2 - eta0 = (iside == 1) ? 0.0 : 0.5 - hint2d = (Ref(1), Ref(1)) - theta_lo, theta_hi = (iside == 1) ? (-0.25, 0.25) : (0.25, 0.75) - side_label = (iside == 1) ? "outboard" : "inboard" - theta = try - find_zero( - theta -> theta + pe.rzphi_offset((psi_edge, theta); hint=hint2d) - eta0, - (theta_lo, theta_hi), Roots.Brent(); - atol=1e-12, rtol=1e-12) - catch e - error("Separatrix $side_label midplane root not found in bracket " * - "[$(theta_lo), $(theta_hi)]: $(e.msg)") - end - r2 = pe.rzphi_rsquared((psi_edge, theta)) - offset = pe.rzphi_offset((psi_edge, theta)) - rsep[iside] = pe.ro + sqrt(r2) * cos(2π * (theta + offset)) - end - - # Top and bottom separatrix Z extrema via bracketed Brent on ∂z/∂θ = 0. - # iside=1 → top (θ ∈ [0.0, 0.5]), iside=2 → bottom (θ ∈ [0.5, 1.0]). - # Matches Fortran convention so (r|z)ext[1] / zsep[1] refer to the top extremum - # and (r|z)ext[2] / zsep[2] to the bottom (see equil_out.f::equil_out_sep_find). - # Splines use PeriodicBC + WrapExtrap, so θ outside [0,1] is valid. - zsep = zeros(2) - rext = zeros(2) - zext = zeros(2) - - for iside in 1:2 - hint2d = (Ref(1), Ref(1)) - - # Cache variables populated by z_deriv, read after convergence - rfac = Ref(0.0) - cos_phase = Ref(0.0) - z_val = Ref(0.0) - - # ∂z/∂θ where z(θ) = zo + √r²(θ) · sin(2π(θ + η(θ))) - function z_deriv(theta_inner) - r2 = pe.rzphi_rsquared((psi_edge, theta_inner); hint=hint2d) - r2y = pe.rzphi_rsquared((psi_edge, theta_inner); deriv=DerivOp(0, 1), hint=hint2d) - η = pe.rzphi_offset((psi_edge, theta_inner); hint=hint2d) - η1 = pe.rzphi_offset((psi_edge, theta_inner); deriv=DerivOp(0, 1), hint=hint2d) - rfac_local = sqrt(max(0.0, r2)) - rfac1 = (rfac_local > 0) ? r2y / (2 * rfac_local) : 0.0 - phase1 = 2π * (1 + η1) - sin_phase = sin(2π * (theta_inner + η)) - cos_phase_local = cos(2π * (theta_inner + η)) - - rfac[] = rfac_local - cos_phase[] = cos_phase_local - z_val[] = pe.zo + rfac_local * sin_phase - - return rfac_local * phase1 * cos_phase_local + rfac1 * sin_phase - end - - theta_lo, theta_hi = (iside == 1) ? (0.0, 0.5) : (0.5, 1.0) - side_label = (iside == 1) ? "top" : "bottom" - theta = try - find_zero(z_deriv, (theta_lo, theta_hi), Roots.Brent(); - atol=1e-12, rtol=1e-12) - catch e - error("Separatrix $side_label Z-extremum root not found in bracket " * - "[$(theta_lo), $(theta_hi)]: $(e.msg)") - end - - rext[iside] = pe.ro + rfac[] * cos_phase[] - zsep[iside] = zext[iside] = z_val[] - end - - pe.params.rsep = rsep - pe.params.zsep = zsep - pe.params.rext = rext - pe.params.zext = zext - return (rsep, zsep, rext, zext) -end - -""" - equilibrium_global_parameters!(pe::PlasmaEquilibrium) - -Computes and populates global equilibrium parameters in the `PlasmaEquilibrium` -struct, such as rmean, amean, kappa, bt0, crnt, betat, betan, li1, etc. Performs -the same function as equil_out_global in the Fortran code. -""" -function equilibrium_global_parameters!(pe::PlasmaEquilibrium) - profiles = pe.profiles - mpsi = length(pe.rzphi_xs) - 1 - mtheta = length(pe.rzphi_ys) - 1 - - # Use separatrix geometry - rsep, zsep, rext, _ = equilibrium_separatrix_find!(pe) - - rmean = (rsep[1] + rsep[2]) / 2 - amean = (rsep[1] - rsep[2]) / 2 - aratio = rmean / amean - kappa = (zsep[1] - zsep[2]) / (rsep[1] - rsep[2]) - delta1 = (rmean - rext[1]) / amean - delta2 = (rmean - rext[2]) / amean - dpsi = 1.0 - pe.rzphi_xs[mpsi+1] - psi_edge = profiles.xs[end] - bt0 = (profiles.F_spline.y[end] + profiles.F_deriv(psi_edge; hint=Ref(profiles.npts_minus_1)) * dpsi) / (2π * rmean) - - pe.params.rmean = rmean - pe.params.amean = amean - pe.params.aratio = aratio - pe.params.kappa = kappa - pe.params.delta1 = delta1 - pe.params.delta2 = delta2 - pe.params.bt0 = bt0 - - psio = pe.psio - gs1 = zeros(Float64, mtheta + 1) - gs2 = zeros(Float64, mtheta + 1) - - # Direct array access at edge flux surface grid points using nodal_derivs - for itheta in 0:mtheta - # Function values (partials[1,:,:]) - r2 = pe.rzphi_rsquared.nodal_derivs.partials[1, end, itheta+1] - offset = pe.rzphi_offset.nodal_derivs.partials[1, end, itheta+1] - jac = pe.rzphi_jac.nodal_derivs.partials[1, end, itheta+1] - # Theta derivatives (partials[3,:,:] = ∂f/∂y) - r2_y = pe.rzphi_rsquared.nodal_derivs.partials[3, end, itheta+1] - offset_y = pe.rzphi_offset.nodal_derivs.partials[3, end, itheta+1] - - chi1 = 2π * psio / jac - jacfac = π / jac - rfac = sqrt(r2) - eta = 2π * (pe.rzphi_ys[itheta+1] + offset) - r = pe.ro + rfac * cos(eta) - v21 = jacfac * r2_y / (2π * rfac) - v22 = jacfac * (1 + offset_y) * (2 * rfac) - v33 = jacfac * 2π * (r / π) - dvsq = (v21^2 + v22^2) * (v33 * jac^2)^2 - gs1[itheta+1] = sqrt(dvsq) / (2π * r) - gs2[itheta+1] = chi1 * dvsq / (2π * r)^2 - end - - int1 = sum(gs1) / (mtheta + 1) - int2 = sum(gs2) / (mtheta + 1) - crnt = int2 / (1e6 * mu0) - bp0 = int2 / int1 - bwall = 1e6 * mu0 * crnt / (2π * amean) - - pe.params.crnt = crnt - pe.params.bwall = bwall - - # Flux-surface integrals over the full normalized flux [0, 1]; ExtendExtrap carries the integral past the first/last grid points (integrands are smooth/vanishing at the axis). - xs = profiles.xs - fsi(y) = FastInterpolations.integrate(cubic_interp(xs, y; extrap=ExtendExtrap()), 0.0, 1.0) - P_vals = profiles.P_spline.y - dVdpsi_vals = profiles.dVdpsi_spline.y - - fsi_pdv = fsi(P_vals .* dVdpsi_vals) # ∫ p dV/dψ - fsi_dv = fsi(dVdpsi_vals) # ∫ dV/dψ - fsi_p2dv = fsi(P_vals .^ 2 .* dVdpsi_vals) # ∫ p² dV/dψ - volume = fsi_dv # same integrand as hs col 2 in Fortran - - # Poloidal-field surface integral hs_bp2(ψ) = ψ₀² ∮dθ |∇ψ|² / (R² J). - # This is Fortran equil_out.f's hs%fs(:,3) and is the correct integrand for - # the internal inductance li1/li2/li3 (distinct from the p²·dV integrand - # that drives betaj). - hs_bp2 = zeros(Float64, mpsi + 1) - for ipsi in 0:mpsi - acc = 0.0 - for itheta in 0:mtheta - r2 = pe.rzphi_rsquared.nodal_derivs.partials[1, ipsi+1, itheta+1] - offset = pe.rzphi_offset.nodal_derivs.partials[1, ipsi+1, itheta+1] - jac = pe.rzphi_jac.nodal_derivs.partials[1, ipsi+1, itheta+1] - r2_y = pe.rzphi_rsquared.nodal_derivs.partials[3, ipsi+1, itheta+1] - offset_y = pe.rzphi_offset.nodal_derivs.partials[3, ipsi+1, itheta+1] - - jacfac = π / jac - rfac = sqrt(r2) - eta = 2π * (pe.rzphi_ys[itheta+1] + offset) - r = pe.ro + rfac * cos(eta) - v21 = jacfac * r2_y / (2π * rfac) - v22 = jacfac * (1 + offset_y) * (2 * rfac) - v33 = jacfac * 2π * (r / π) - dvsq = (v21^2 + v22^2) * (v33 * jac^2)^2 - acc += dvsq / (r^2) / jac - end - # Periodic trapezoidal rule on uniform θ grid reduces to a plain mean - # because the first and last grid points coincide — matches the int1/int2 - # pattern used for the edge-surface integrals above. - hs_bp2[ipsi+1] = (acc / (mtheta + 1)) * psio^2 - end - fsi_bp2 = fsi(hs_bp2) - - p0 = P_vals[1] - profiles.P_deriv(profiles.xs[1]; hint=Ref(1)) * profiles.xs[1] # linear extrapolation - betat = 2 * (fsi_pdv / fsi_dv) / bt0^2 - betaj = 2 * sqrt(fsi_p2dv / fsi_dv) / bwall^2 - betan = 100 * amean * bt0 * betat / crnt - betap1 = 2 * (fsi_pdv / fsi_dv) / bp0^2 - betap2 = 4 * fsi_pdv / ((1e6 * mu0 * crnt)^2 * pe.ro) - betap3 = 4 * fsi_pdv / ((1e6 * mu0 * crnt)^2 * rmean) - li1 = fsi_bp2 / fsi_dv / bp0^2 - li2 = 2 * fsi_bp2 / ((1e6 * mu0 * crnt)^2 * pe.ro) - li3 = 2 * fsi_bp2 / ((1e6 * mu0 * crnt)^2 * rmean) - - pe.params.psi0 = psio - pe.params.psi_axis = pe.psio - pe.params.psi_boundary = 1.0 - pe.params.psi_boundary_norm = 1.0 - pe.params.psi_axis_norm = 0.0 - pe.params.psi_norm = 0.0 - pe.params.psi_axis_offset = 0.0 - pe.params.psi_boundary_offset = 0.0 - pe.params.psi_axis_sign = 1 - pe.params.psi_boundary_sign = -1 - pe.params.psi_boundary_zero = false - - pe.params.q0 = profiles.q_spline.y[1] - pe.params.b0 = bt0 - - pe.params.volume = volume - pe.params.betat = betat - pe.params.betan = betan - pe.params.betaj = betaj - pe.params.betap1 = betap1 - pe.params.betap2 = betap2 - pe.params.betap3 = betap3 - pe.params.li1 = li1 - pe.params.li2 = li2 - pe.params.li3 = li3 -end - -""" - equilibrium_qfind!(equil::PlasmaEquilibrium) - -Finds the extrema of the safety factor profile q(ψ) in the plasma equilibrium -and computes derived q-values such as q0, qmin, qmax, qa, and q95. Performs -the same function as equil_out_qfind in the Fortran code. -""" -function equilibrium_qfind!(equil::PlasmaEquilibrium) - - profiles = equil.profiles - xs = profiles.xs - mpsi = length(xs) - 1 - psiexl = Float64[] - qexl = Float64[] - - # Create derivative views for q-spline - q_spline = profiles.q_spline - q_d1 = deriv1(q_spline) - q_d2 = deriv2(q_spline) - q_d3 = deriv3(q_spline) - - # Left endpoint - push!(psiexl, xs[1]) - push!(qexl, q_spline.y[1]) - - # Search for extrema in q(ψ) - for ipsi in 1:mpsi - x0 = xs[ipsi] - x1 = xs[ipsi+1] - xmax = x1 - x0 - - a = q_spline(x0) - b = q_d1(x0) - c = q_d2(x0) - d = q_d3(x0) - - if d != 0.0 - xcrit = -c / d - dx2 = xcrit^2 - 2b / d - if dx2 ≥ 0 - dx = sqrt(dx2) - for delta in (dx, -dx) - x = xcrit - delta - if 0 ≤ x < xmax - ψ = x0 + x - push!(psiexl, ψ) - push!(qexl, q_spline(ψ)) - end - end - end - end - end - - # Right endpoint - push!(psiexl, xs[end]) - push!(qexl, q_spline.y[end]) - - equil.params.qextrema_psi = psiexl - equil.params.qextrema_q = qexl - equil.params.mextrema = length(psiexl) - # Compute derived q-values - q0 = q_spline.y[1] - profiles.q_deriv(xs[1]; hint=Ref(1)) * xs[1] - qmax_edge = q_spline.y[end] - qmin = min(minimum(qexl), q0) - qmax = max(maximum(qexl), qmax_edge) - qa = q_spline.y[end] + profiles.q_deriv(xs[end]; hint=Ref(profiles.npts_minus_1)) * (1.0 - xs[end]) - - q95 = q_spline(0.95) - - # Store derived values - equil.params.q0 = q0 - equil.params.qmin = qmin - equil.params.qmax = qmax - equil.params.qa = qa - equil.params.q95 = q95 -end - -""" - equilibrium_gse!(equil::PlasmaEquilibrium) - -Diagnoses the Grad-Shafranov solution by computing the residual of the -Grad-Shafranov equation across the grid and writing diagnostic data to HDF5 files. -Performs the same function as equil_out_gse in the Fortran code. - -Returns `(; xs, ys, flux_x, flux_y, source, total, error, errori)` on the `rzphi` (ψ, θ) nodes: the two -flux-divergence terms, the source, the residual `total = flux_x - flux_y + source`, `error = |total|` over -the largest term on the grid, and `errori`, the magnitude of the θ-integrated residual on each surface. -""" -function equilibrium_gse!(equil::PlasmaEquilibrium) - - profiles = equil.profiles - mpsi = length(equil.rzphi_xs) - 1 - mtheta = length(equil.rzphi_ys) - 1 - ro, zo = equil.ro, equil.zo - psio = equil.psio - verbose = equil.params.verbose - diagnose_src = equil.params.diagnose_src - diagnose_maxima = equil.params.diagnose_maxima - - if verbose - @info "Diagnosing Grad-Shafranov solution" - end - - # Compute R, Z coordinates using nodal_derivs access - r = zeros(Float64, mpsi + 1, mtheta + 1) - z = zeros(Float64, mpsi + 1, mtheta + 1) - - for ipsi in 1:(mpsi+1) - rfac = @. sqrt(equil.rzphi_rsquared.nodal_derivs.partials[1, ipsi, :]) - angle = @. 2π * (equil.rzphi_ys + equil.rzphi_offset.nodal_derivs.partials[1, ipsi, :]) - r[ipsi, :] .= ro .+ rfac .* cos.(angle) - z[ipsi, :] .= zo .+ rfac .* sin.(angle) - end - - # Compute flux quantities using direct array access at grid points - flux_fs = zeros(Float64, mpsi + 1, mtheta + 1, 2) - flux_fsx = zeros(Float64, mpsi + 1, mtheta + 1, 2) # x-derivatives - flux_fsy = zeros(Float64, mpsi + 1, mtheta + 1, 2) # y-derivatives - for ipsi in 1:(mpsi+1) - for itheta in 1:(mtheta+1) - f1 = equil.rzphi_rsquared.nodal_derivs.partials[1, ipsi, itheta] - f2 = equil.rzphi_offset.nodal_derivs.partials[1, ipsi, itheta] - f4 = equil.rzphi_jac.nodal_derivs.partials[1, ipsi, itheta] - fy1 = equil.rzphi_rsquared.nodal_derivs.partials[3, ipsi, itheta] - fy2 = equil.rzphi_offset.nodal_derivs.partials[3, ipsi, itheta] - fx1 = equil.rzphi_rsquared.nodal_derivs.partials[2, ipsi, itheta] - fx2 = equil.rzphi_offset.nodal_derivs.partials[2, ipsi, itheta] - - flux_fs[ipsi, itheta, 1] = fy1^2 / (4π^2 * f1) + (1 + fy2)^2 * 4 * f1 - flux_fs[ipsi, itheta, 2] = fx1 * fy1 / (4π^2 * f1) + fx2 * (1 + fy2) * 4 * f1 - - flux_fs[ipsi, itheta, 1] *= 2π * psio / f4 - flux_fs[ipsi, itheta, 2] *= 2π * psio / f4 - end - end - # Create flux interpolants for Grad-Shafranov diagnostics - @views flux_fs[:, end, :] .= flux_fs[:, 1, :] - flux_opts = (bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) - flux1 = cubic_interp((equil.rzphi_xs, equil.rzphi_ys), flux_fs[:, :, 1]; flux_opts...) - flux2 = cubic_interp((equil.rzphi_xs, equil.rzphi_ys), flux_fs[:, :, 2]; flux_opts...) - - # Compute flux derivatives at all grid points for diagnostics - hint2d = (Ref(1), Ref(1)) # Shared 2D hint for hot loop optimization - for ipsi in 0:mpsi - for itheta in 0:mtheta - query_point = (equil.rzphi_xs[ipsi+1], equil.rzphi_ys[itheta+1]) - flux_fsx[ipsi+1, itheta+1, 1] = flux1(query_point; deriv=DerivOp(1, 0), hint=hint2d) - flux_fsx[ipsi+1, itheta+1, 2] = flux2(query_point; deriv=DerivOp(1, 0), hint=hint2d) - flux_fsy[ipsi+1, itheta+1, 1] = flux1(query_point; deriv=DerivOp(0, 1), hint=hint2d) - flux_fsy[ipsi+1, itheta+1, 2] = flux2(query_point; deriv=DerivOp(0, 1), hint=hint2d) - end - end - - # Compute source term using direct array access - source = zeros(Float64, mpsi + 1, mtheta + 1) - hint = Ref(1) # Linear search hint for sequential psi access - for ipsi in 1:(mpsi+1) - psi = profiles.xs[ipsi] - s1 = profiles.F_spline.y[ipsi] - s1p = profiles.F_deriv(psi; hint=hint) - s2p = profiles.P_deriv(psi; hint=hint) - for itheta in 1:(mtheta+1) - f4 = equil.rzphi_jac.nodal_derivs.partials[1, ipsi, itheta] - denom = (2π * r[ipsi, itheta])^2 - source[ipsi, itheta] = f4 / (2π * psio * π^2) * (s1 * s1p / denom + s2p) - end - end - - total = flux_fsx[:, :, 1] .- flux_fsy[:, :, 2] .+ source - error = abs.(total) ./ maximum([maximum(abs.(flux_fsx[:, :, 1])), maximum(abs.(flux_fsy[:, :, 2])), maximum(abs.(source))]) - errlog = ifelse.(error .> 0, log10.(error), 0.0) - - if diagnose_maxima - fxmax = maximum(abs.(flux_fsx[:, :, 1])) - fymax = maximum(abs.(flux_fsy[:, :, 2])) - smax = maximum(abs.(source)) - emax = maximum(abs.(error)) - lmax = maximum(errlog) - jmax = ind2sub(size(errlog), argmax(errlog)) - @info "GS residuals: fxmax = $(@sprintf("%.3e", fxmax)), fymax = $(@sprintf("%.3e", fymax)), smax = $(@sprintf("%.3e", smax)), emax = $(@sprintf("%.3e", emax)), lmax = $(@sprintf("%.3f", lmax)), maxloc = $(jmax .- 1)" - end - - # Integrated error criterion - term = zeros(Float64, mpsi + 1, 2) - for ipsi in 1:(mpsi+1) - fs_matrix = zeros(Float64, mtheta + 1, 2) - fs_matrix[:, 1] = flux_fsx[ipsi, :, 1] - fs_matrix[:, 2] = source[ipsi, :] - # Snap the repeated endpoint exactly equal to the start - fs_matrix[end, :] .= fs_matrix[1, :] - - # Compute total integral using FastInterpolations native integration - itp = cubic_interp(equil.rzphi_ys, Series(fs_matrix); bc=PeriodicBC()) - term[ipsi, :] .= FastInterpolations.integrate(itp) - end - - totali = sum(term; dims=2) - errori = abs.(totali) - errlogi = @. ifelse(errori > 0, log10(errori), 0.0) - - if diagnose_src - if verbose - @info "Writing diagnostics to HDF5 files" - end - - # Write contour data - h5open(joinpath(dirname(equil.config.eq_filename), "gsec.h5"), "w") do file - file["mpsi"] = mpsi - file["mtheta"] = mtheta - file["r"] = Float32.(r) - file["z"] = Float32.(z) - file["flux_fsx"] = Float32.(flux_fsx[:, :, 1]) - file["flux_fsy"] = Float32.(flux_fsy[:, :, 2]) - file["source"] = Float32.(source) - file["total"] = Float32.(total) - file["error"] = Float32.(error) - file["errlog"] = Float32.(errlog) - end - - # Write xy plot data - h5open(joinpath(dirname(equil.config.eq_filename), "gse.h5"), "w") do file - gse_data = Array{Float32,3}(undef, mpsi + 1, mtheta + 1, 7) - for ipsi in 0:mpsi - for itheta in 0:mtheta - gse_data[ipsi+1, itheta+1, 1] = Float32(flux.ys[itheta+1]) - gse_data[ipsi+1, itheta+1, 2] = Float32(flux.xs[ipsi+1]) - gse_data[ipsi+1, itheta+1, 3] = Float32(flux_fs[ipsi+1, itheta+1, 1]) - gse_data[ipsi+1, itheta+1, 4] = Float32(flux_fs[ipsi+1, itheta+1, 2]) - gse_data[ipsi+1, itheta+1, 5] = Float32(source[ipsi+1, itheta+1]) - gse_data[ipsi+1, itheta+1, 6] = Float32(total[ipsi+1, itheta+1]) - gse_data[ipsi+1, itheta+1, 7] = Float32(error[ipsi+1, itheta+1]) - end - end - file["gse_data"] = gse_data - end - - # Write integrated error criterion - h5open(joinpath(dirname(equil.config.eq_filename), "gsei.h5"), "w") do file - file["xs"] = Float32.(flux.xs) - file["term"] = Float32.(term) - file["totali"] = Float32.(totali) - file["errori"] = Float32.(errori) - file["errlogi"] = Float32.(errlogi) - end - end - return (; xs=equil.rzphi_xs, ys=equil.rzphi_ys, flux_x=flux_fsx[:, :, 1], flux_y=flux_fsy[:, :, 2], source, total, error, errori=vec(errori)) -end - end # module Equilibrium diff --git a/src/Equilibrium/PostProcessing.jl b/src/Equilibrium/PostProcessing.jl new file mode 100644 index 000000000..665d69e25 --- /dev/null +++ b/src/Equilibrium/PostProcessing.jl @@ -0,0 +1,491 @@ +# PostProcessing.jl +# +# Quantities derived from a formed equilibrium: separatrix, global parameters, q extrema, and the +# Grad-Shafranov residual. + +""" + equilibrium_separatrix_find!(pe::PlasmaEquilibrium) + +Finds the separatrix locations in the plasma equilibrium (rsep, zsep, rext, zext). +Performs the same function as equil_out_sep_find in the Fortran code. +""" +function equilibrium_separatrix_find!(pe::PlasmaEquilibrium) + mpsi = length(pe.rzphi_xs) - 1 + edge_idx = mpsi + 1 # Edge flux surface index + psi_edge = pe.rzphi_xs[edge_idx] + rsep = zeros(2) + + # Outboard and inboard midplane R via bracketed Brent on θ + η(θ) - η₀ = 0. + # iside=1 → outboard (η₀=0.0, θ near 0), iside=2 → inboard (η₀=0.5, θ near 0.5). + for iside in 1:2 + eta0 = (iside == 1) ? 0.0 : 0.5 + hint2d = (Ref(1), Ref(1)) + theta_lo, theta_hi = (iside == 1) ? (-0.25, 0.25) : (0.25, 0.75) + side_label = (iside == 1) ? "outboard" : "inboard" + theta = try + find_zero( + theta -> theta + pe.rzphi_offset((psi_edge, theta); hint=hint2d) - eta0, + (theta_lo, theta_hi), Roots.Brent(); + atol=1e-12, rtol=1e-12) + catch e + error("Separatrix $side_label midplane root not found in bracket " * + "[$(theta_lo), $(theta_hi)]: $(e.msg)") + end + r2 = pe.rzphi_rsquared((psi_edge, theta)) + offset = pe.rzphi_offset((psi_edge, theta)) + rsep[iside] = pe.ro + sqrt(r2) * cos(2π * (theta + offset)) + end + + # Top and bottom separatrix Z extrema via bracketed Brent on ∂z/∂θ = 0. + # iside=1 → top (θ ∈ [0.0, 0.5]), iside=2 → bottom (θ ∈ [0.5, 1.0]). + # Matches Fortran convention so (r|z)ext[1] / zsep[1] refer to the top extremum + # and (r|z)ext[2] / zsep[2] to the bottom (see equil_out.f::equil_out_sep_find). + # Splines use PeriodicBC + WrapExtrap, so θ outside [0,1] is valid. + zsep = zeros(2) + rext = zeros(2) + zext = zeros(2) + + for iside in 1:2 + hint2d = (Ref(1), Ref(1)) + + # Cache variables populated by z_deriv, read after convergence + rfac = Ref(0.0) + cos_phase = Ref(0.0) + z_val = Ref(0.0) + + # ∂z/∂θ where z(θ) = zo + √r²(θ) · sin(2π(θ + η(θ))) + function z_deriv(theta_inner) + r2 = pe.rzphi_rsquared((psi_edge, theta_inner); hint=hint2d) + r2y = pe.rzphi_rsquared((psi_edge, theta_inner); deriv=DerivOp(0, 1), hint=hint2d) + η = pe.rzphi_offset((psi_edge, theta_inner); hint=hint2d) + η1 = pe.rzphi_offset((psi_edge, theta_inner); deriv=DerivOp(0, 1), hint=hint2d) + rfac_local = sqrt(max(0.0, r2)) + rfac1 = (rfac_local > 0) ? r2y / (2 * rfac_local) : 0.0 + phase1 = 2π * (1 + η1) + sin_phase = sin(2π * (theta_inner + η)) + cos_phase_local = cos(2π * (theta_inner + η)) + + rfac[] = rfac_local + cos_phase[] = cos_phase_local + z_val[] = pe.zo + rfac_local * sin_phase + + return rfac_local * phase1 * cos_phase_local + rfac1 * sin_phase + end + + theta_lo, theta_hi = (iside == 1) ? (0.0, 0.5) : (0.5, 1.0) + side_label = (iside == 1) ? "top" : "bottom" + theta = try + find_zero(z_deriv, (theta_lo, theta_hi), Roots.Brent(); + atol=1e-12, rtol=1e-12) + catch e + error("Separatrix $side_label Z-extremum root not found in bracket " * + "[$(theta_lo), $(theta_hi)]: $(e.msg)") + end + + rext[iside] = pe.ro + rfac[] * cos_phase[] + zsep[iside] = zext[iside] = z_val[] + end + + pe.params.rsep = rsep + pe.params.zsep = zsep + pe.params.rext = rext + pe.params.zext = zext + return (rsep, zsep, rext, zext) +end + +""" + equilibrium_global_parameters!(pe::PlasmaEquilibrium) + +Computes and populates global equilibrium parameters in the `PlasmaEquilibrium` +struct, such as rmean, amean, kappa, bt0, crnt, betat, betan, li1, etc. Performs +the same function as equil_out_global in the Fortran code. +""" +function equilibrium_global_parameters!(pe::PlasmaEquilibrium) + profiles = pe.profiles + mpsi = length(pe.rzphi_xs) - 1 + mtheta = length(pe.rzphi_ys) - 1 + + # Use separatrix geometry + rsep, zsep, rext, _ = equilibrium_separatrix_find!(pe) + + rmean = (rsep[1] + rsep[2]) / 2 + amean = (rsep[1] - rsep[2]) / 2 + aratio = rmean / amean + kappa = (zsep[1] - zsep[2]) / (rsep[1] - rsep[2]) + delta1 = (rmean - rext[1]) / amean + delta2 = (rmean - rext[2]) / amean + dpsi = 1.0 - pe.rzphi_xs[mpsi+1] + psi_edge = profiles.xs[end] + bt0 = (profiles.F_spline.y[end] + profiles.F_deriv(psi_edge; hint=Ref(profiles.npts_minus_1)) * dpsi) / (2π * rmean) + + pe.params.rmean = rmean + pe.params.amean = amean + pe.params.aratio = aratio + pe.params.kappa = kappa + pe.params.delta1 = delta1 + pe.params.delta2 = delta2 + pe.params.bt0 = bt0 + + psio = pe.psio + gs1 = zeros(Float64, mtheta + 1) + gs2 = zeros(Float64, mtheta + 1) + + # Direct array access at edge flux surface grid points using nodal_derivs + for itheta in 0:mtheta + # Function values (partials[1,:,:]) + r2 = pe.rzphi_rsquared.nodal_derivs.partials[1, end, itheta+1] + offset = pe.rzphi_offset.nodal_derivs.partials[1, end, itheta+1] + jac = pe.rzphi_jac.nodal_derivs.partials[1, end, itheta+1] + # Theta derivatives (partials[3,:,:] = ∂f/∂y) + r2_y = pe.rzphi_rsquared.nodal_derivs.partials[3, end, itheta+1] + offset_y = pe.rzphi_offset.nodal_derivs.partials[3, end, itheta+1] + + chi1 = 2π * psio / jac + jacfac = π / jac + rfac = sqrt(r2) + eta = 2π * (pe.rzphi_ys[itheta+1] + offset) + r = pe.ro + rfac * cos(eta) + v21 = jacfac * r2_y / (2π * rfac) + v22 = jacfac * (1 + offset_y) * (2 * rfac) + v33 = jacfac * 2π * (r / π) + dvsq = (v21^2 + v22^2) * (v33 * jac^2)^2 + gs1[itheta+1] = sqrt(dvsq) / (2π * r) + gs2[itheta+1] = chi1 * dvsq / (2π * r)^2 + end + + int1 = sum(gs1) / (mtheta + 1) + int2 = sum(gs2) / (mtheta + 1) + crnt = int2 / (1e6 * mu0) + bp0 = int2 / int1 + bwall = 1e6 * mu0 * crnt / (2π * amean) + + pe.params.crnt = crnt + pe.params.bwall = bwall + + # Flux-surface integrals over the full normalized flux [0, 1]; ExtendExtrap carries the integral past the first/last grid points (integrands are smooth/vanishing at the axis). + xs = profiles.xs + fsi(y) = FastInterpolations.integrate(cubic_interp(xs, y; extrap=ExtendExtrap()), 0.0, 1.0) + P_vals = profiles.P_spline.y + dVdpsi_vals = profiles.dVdpsi_spline.y + + fsi_pdv = fsi(P_vals .* dVdpsi_vals) # ∫ p dV/dψ + fsi_dv = fsi(dVdpsi_vals) # ∫ dV/dψ + fsi_p2dv = fsi(P_vals .^ 2 .* dVdpsi_vals) # ∫ p² dV/dψ + volume = fsi_dv # same integrand as hs col 2 in Fortran + + # Poloidal-field surface integral hs_bp2(ψ) = ψ₀² ∮dθ |∇ψ|² / (R² J). + # This is Fortran equil_out.f's hs%fs(:,3) and is the correct integrand for + # the internal inductance li1/li2/li3 (distinct from the p²·dV integrand + # that drives betaj). + hs_bp2 = zeros(Float64, mpsi + 1) + for ipsi in 0:mpsi + acc = 0.0 + for itheta in 0:mtheta + r2 = pe.rzphi_rsquared.nodal_derivs.partials[1, ipsi+1, itheta+1] + offset = pe.rzphi_offset.nodal_derivs.partials[1, ipsi+1, itheta+1] + jac = pe.rzphi_jac.nodal_derivs.partials[1, ipsi+1, itheta+1] + r2_y = pe.rzphi_rsquared.nodal_derivs.partials[3, ipsi+1, itheta+1] + offset_y = pe.rzphi_offset.nodal_derivs.partials[3, ipsi+1, itheta+1] + + jacfac = π / jac + rfac = sqrt(r2) + eta = 2π * (pe.rzphi_ys[itheta+1] + offset) + r = pe.ro + rfac * cos(eta) + v21 = jacfac * r2_y / (2π * rfac) + v22 = jacfac * (1 + offset_y) * (2 * rfac) + v33 = jacfac * 2π * (r / π) + dvsq = (v21^2 + v22^2) * (v33 * jac^2)^2 + acc += dvsq / (r^2) / jac + end + # Periodic trapezoidal rule on uniform θ grid reduces to a plain mean + # because the first and last grid points coincide — matches the int1/int2 + # pattern used for the edge-surface integrals above. + hs_bp2[ipsi+1] = (acc / (mtheta + 1)) * psio^2 + end + fsi_bp2 = fsi(hs_bp2) + + p0 = P_vals[1] - profiles.P_deriv(profiles.xs[1]; hint=Ref(1)) * profiles.xs[1] # linear extrapolation + betat = 2 * (fsi_pdv / fsi_dv) / bt0^2 + betaj = 2 * sqrt(fsi_p2dv / fsi_dv) / bwall^2 + betan = 100 * amean * bt0 * betat / crnt + betap1 = 2 * (fsi_pdv / fsi_dv) / bp0^2 + betap2 = 4 * fsi_pdv / ((1e6 * mu0 * crnt)^2 * pe.ro) + betap3 = 4 * fsi_pdv / ((1e6 * mu0 * crnt)^2 * rmean) + li1 = fsi_bp2 / fsi_dv / bp0^2 + li2 = 2 * fsi_bp2 / ((1e6 * mu0 * crnt)^2 * pe.ro) + li3 = 2 * fsi_bp2 / ((1e6 * mu0 * crnt)^2 * rmean) + + pe.params.psi0 = psio + pe.params.psi_axis = pe.psio + pe.params.psi_boundary = 1.0 + pe.params.psi_boundary_norm = 1.0 + pe.params.psi_axis_norm = 0.0 + pe.params.psi_norm = 0.0 + pe.params.psi_axis_offset = 0.0 + pe.params.psi_boundary_offset = 0.0 + pe.params.psi_axis_sign = 1 + pe.params.psi_boundary_sign = -1 + pe.params.psi_boundary_zero = false + + pe.params.q0 = profiles.q_spline.y[1] + pe.params.b0 = bt0 + + pe.params.volume = volume + pe.params.betat = betat + pe.params.betan = betan + pe.params.betaj = betaj + pe.params.betap1 = betap1 + pe.params.betap2 = betap2 + pe.params.betap3 = betap3 + pe.params.li1 = li1 + pe.params.li2 = li2 + pe.params.li3 = li3 +end + +""" + equilibrium_qfind!(equil::PlasmaEquilibrium) + +Finds the extrema of the safety factor profile q(ψ) in the plasma equilibrium +and computes derived q-values such as q0, qmin, qmax, qa, and q95. Performs +the same function as equil_out_qfind in the Fortran code. +""" +function equilibrium_qfind!(equil::PlasmaEquilibrium) + + profiles = equil.profiles + xs = profiles.xs + mpsi = length(xs) - 1 + psiexl = Float64[] + qexl = Float64[] + + # Create derivative views for q-spline + q_spline = profiles.q_spline + q_d1 = deriv1(q_spline) + q_d2 = deriv2(q_spline) + q_d3 = deriv3(q_spline) + + # Left endpoint + push!(psiexl, xs[1]) + push!(qexl, q_spline.y[1]) + + # Search for extrema in q(ψ) + for ipsi in 1:mpsi + x0 = xs[ipsi] + x1 = xs[ipsi+1] + xmax = x1 - x0 + + a = q_spline(x0) + b = q_d1(x0) + c = q_d2(x0) + d = q_d3(x0) + + if d != 0.0 + xcrit = -c / d + dx2 = xcrit^2 - 2b / d + if dx2 ≥ 0 + dx = sqrt(dx2) + for delta in (dx, -dx) + x = xcrit - delta + if 0 ≤ x < xmax + ψ = x0 + x + push!(psiexl, ψ) + push!(qexl, q_spline(ψ)) + end + end + end + end + end + + # Right endpoint + push!(psiexl, xs[end]) + push!(qexl, q_spline.y[end]) + + equil.params.qextrema_psi = psiexl + equil.params.qextrema_q = qexl + equil.params.mextrema = length(psiexl) + # Compute derived q-values + q0 = q_spline.y[1] - profiles.q_deriv(xs[1]; hint=Ref(1)) * xs[1] + qmax_edge = q_spline.y[end] + qmin = min(minimum(qexl), q0) + qmax = max(maximum(qexl), qmax_edge) + qa = q_spline.y[end] + profiles.q_deriv(xs[end]; hint=Ref(profiles.npts_minus_1)) * (1.0 - xs[end]) + + q95 = q_spline(0.95) + + # Store derived values + equil.params.q0 = q0 + equil.params.qmin = qmin + equil.params.qmax = qmax + equil.params.qa = qa + equil.params.q95 = q95 +end + +""" + equilibrium_gse!(equil::PlasmaEquilibrium) + +Diagnoses the Grad-Shafranov solution by computing the residual of the +Grad-Shafranov equation across the grid and writing diagnostic data to HDF5 files. +Performs the same function as equil_out_gse in the Fortran code. + +Returns `(; xs, ys, flux_x, flux_y, source, total, error, errori)` on the `rzphi` (ψ, θ) nodes: the two +flux-divergence terms, the source, the residual `total = flux_x - flux_y + source`, `error = |total|` over +the largest term on the grid, and `errori`, the magnitude of the θ-integrated residual on each surface. +""" +function equilibrium_gse!(equil::PlasmaEquilibrium) + + profiles = equil.profiles + mpsi = length(equil.rzphi_xs) - 1 + mtheta = length(equil.rzphi_ys) - 1 + ro, zo = equil.ro, equil.zo + psio = equil.psio + verbose = equil.params.verbose + diagnose_src = equil.params.diagnose_src + diagnose_maxima = equil.params.diagnose_maxima + + if verbose + @info "Diagnosing Grad-Shafranov solution" + end + + # Compute R, Z coordinates using nodal_derivs access + r = zeros(Float64, mpsi + 1, mtheta + 1) + z = zeros(Float64, mpsi + 1, mtheta + 1) + + for ipsi in 1:(mpsi+1) + rfac = @. sqrt(equil.rzphi_rsquared.nodal_derivs.partials[1, ipsi, :]) + angle = @. 2π * (equil.rzphi_ys + equil.rzphi_offset.nodal_derivs.partials[1, ipsi, :]) + r[ipsi, :] .= ro .+ rfac .* cos.(angle) + z[ipsi, :] .= zo .+ rfac .* sin.(angle) + end + + # Compute flux quantities using direct array access at grid points + flux_fs = zeros(Float64, mpsi + 1, mtheta + 1, 2) + flux_fsx = zeros(Float64, mpsi + 1, mtheta + 1, 2) # x-derivatives + flux_fsy = zeros(Float64, mpsi + 1, mtheta + 1, 2) # y-derivatives + for ipsi in 1:(mpsi+1) + for itheta in 1:(mtheta+1) + f1 = equil.rzphi_rsquared.nodal_derivs.partials[1, ipsi, itheta] + f2 = equil.rzphi_offset.nodal_derivs.partials[1, ipsi, itheta] + f4 = equil.rzphi_jac.nodal_derivs.partials[1, ipsi, itheta] + fy1 = equil.rzphi_rsquared.nodal_derivs.partials[3, ipsi, itheta] + fy2 = equil.rzphi_offset.nodal_derivs.partials[3, ipsi, itheta] + fx1 = equil.rzphi_rsquared.nodal_derivs.partials[2, ipsi, itheta] + fx2 = equil.rzphi_offset.nodal_derivs.partials[2, ipsi, itheta] + + flux_fs[ipsi, itheta, 1] = fy1^2 / (4π^2 * f1) + (1 + fy2)^2 * 4 * f1 + flux_fs[ipsi, itheta, 2] = fx1 * fy1 / (4π^2 * f1) + fx2 * (1 + fy2) * 4 * f1 + + flux_fs[ipsi, itheta, 1] *= 2π * psio / f4 + flux_fs[ipsi, itheta, 2] *= 2π * psio / f4 + end + end + # Create flux interpolants for Grad-Shafranov diagnostics + @views flux_fs[:, end, :] .= flux_fs[:, 1, :] + flux_opts = (bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) + flux1 = cubic_interp((equil.rzphi_xs, equil.rzphi_ys), flux_fs[:, :, 1]; flux_opts...) + flux2 = cubic_interp((equil.rzphi_xs, equil.rzphi_ys), flux_fs[:, :, 2]; flux_opts...) + + # Compute flux derivatives at all grid points for diagnostics + hint2d = (Ref(1), Ref(1)) # Shared 2D hint for hot loop optimization + for ipsi in 0:mpsi + for itheta in 0:mtheta + query_point = (equil.rzphi_xs[ipsi+1], equil.rzphi_ys[itheta+1]) + flux_fsx[ipsi+1, itheta+1, 1] = flux1(query_point; deriv=DerivOp(1, 0), hint=hint2d) + flux_fsx[ipsi+1, itheta+1, 2] = flux2(query_point; deriv=DerivOp(1, 0), hint=hint2d) + flux_fsy[ipsi+1, itheta+1, 1] = flux1(query_point; deriv=DerivOp(0, 1), hint=hint2d) + flux_fsy[ipsi+1, itheta+1, 2] = flux2(query_point; deriv=DerivOp(0, 1), hint=hint2d) + end + end + + # Compute source term using direct array access + source = zeros(Float64, mpsi + 1, mtheta + 1) + hint = Ref(1) # Linear search hint for sequential psi access + for ipsi in 1:(mpsi+1) + psi = profiles.xs[ipsi] + s1 = profiles.F_spline.y[ipsi] + s1p = profiles.F_deriv(psi; hint=hint) + s2p = profiles.P_deriv(psi; hint=hint) + for itheta in 1:(mtheta+1) + f4 = equil.rzphi_jac.nodal_derivs.partials[1, ipsi, itheta] + denom = (2π * r[ipsi, itheta])^2 + source[ipsi, itheta] = f4 / (2π * psio * π^2) * (s1 * s1p / denom + s2p) + end + end + + total = flux_fsx[:, :, 1] .- flux_fsy[:, :, 2] .+ source + error = abs.(total) ./ maximum([maximum(abs.(flux_fsx[:, :, 1])), maximum(abs.(flux_fsy[:, :, 2])), maximum(abs.(source))]) + errlog = ifelse.(error .> 0, log10.(error), 0.0) + + if diagnose_maxima + fxmax = maximum(abs.(flux_fsx[:, :, 1])) + fymax = maximum(abs.(flux_fsy[:, :, 2])) + smax = maximum(abs.(source)) + emax = maximum(abs.(error)) + lmax = maximum(errlog) + jmax = ind2sub(size(errlog), argmax(errlog)) + @info "GS residuals: fxmax = $(@sprintf("%.3e", fxmax)), fymax = $(@sprintf("%.3e", fymax)), smax = $(@sprintf("%.3e", smax)), emax = $(@sprintf("%.3e", emax)), lmax = $(@sprintf("%.3f", lmax)), maxloc = $(jmax .- 1)" + end + + # Integrated error criterion + term = zeros(Float64, mpsi + 1, 2) + for ipsi in 1:(mpsi+1) + fs_matrix = zeros(Float64, mtheta + 1, 2) + fs_matrix[:, 1] = flux_fsx[ipsi, :, 1] + fs_matrix[:, 2] = source[ipsi, :] + # Snap the repeated endpoint exactly equal to the start + fs_matrix[end, :] .= fs_matrix[1, :] + + # Compute total integral using FastInterpolations native integration + itp = cubic_interp(equil.rzphi_ys, Series(fs_matrix); bc=PeriodicBC()) + term[ipsi, :] .= FastInterpolations.integrate(itp) + end + + totali = sum(term; dims=2) + errori = abs.(totali) + errlogi = @. ifelse(errori > 0, log10(errori), 0.0) + + if diagnose_src + if verbose + @info "Writing diagnostics to HDF5 files" + end + + # Write contour data + h5open(joinpath(dirname(equil.config.eq_filename), "gsec.h5"), "w") do file + file["mpsi"] = mpsi + file["mtheta"] = mtheta + file["r"] = Float32.(r) + file["z"] = Float32.(z) + file["flux_fsx"] = Float32.(flux_fsx[:, :, 1]) + file["flux_fsy"] = Float32.(flux_fsy[:, :, 2]) + file["source"] = Float32.(source) + file["total"] = Float32.(total) + file["error"] = Float32.(error) + file["errlog"] = Float32.(errlog) + end + + # Write xy plot data + h5open(joinpath(dirname(equil.config.eq_filename), "gse.h5"), "w") do file + gse_data = Array{Float32,3}(undef, mpsi + 1, mtheta + 1, 7) + for ipsi in 0:mpsi + for itheta in 0:mtheta + gse_data[ipsi+1, itheta+1, 1] = Float32(flux.ys[itheta+1]) + gse_data[ipsi+1, itheta+1, 2] = Float32(flux.xs[ipsi+1]) + gse_data[ipsi+1, itheta+1, 3] = Float32(flux_fs[ipsi+1, itheta+1, 1]) + gse_data[ipsi+1, itheta+1, 4] = Float32(flux_fs[ipsi+1, itheta+1, 2]) + gse_data[ipsi+1, itheta+1, 5] = Float32(source[ipsi+1, itheta+1]) + gse_data[ipsi+1, itheta+1, 6] = Float32(total[ipsi+1, itheta+1]) + gse_data[ipsi+1, itheta+1, 7] = Float32(error[ipsi+1, itheta+1]) + end + end + file["gse_data"] = gse_data + end + + # Write integrated error criterion + h5open(joinpath(dirname(equil.config.eq_filename), "gsei.h5"), "w") do file + file["xs"] = Float32.(flux.xs) + file["term"] = Float32.(term) + file["totali"] = Float32.(totali) + file["errori"] = Float32.(errori) + file["errlogi"] = Float32.(errlogi) + end + end + return (; xs=equil.rzphi_xs, ys=equil.rzphi_ys, flux_x=flux_fsx[:, :, 1], flux_y=flux_fsy[:, :, 2], source, total, error, errori=vec(errori)) +end From 8060c37688be26af798db6c4d616ac5e3384b089 Mon Sep 17 00:00:00 2001 From: Matthew Pharr Date: Fri, 23 Oct 2026 20:56:29 +0200 Subject: [PATCH 2/3] Equilibrium - REFACTOR - Arrange the Equilibrium module by pipeline stage --- docs/development/architecture.md | 18 +- docs/src/conventions.md | 14 +- .../Direct.jl} | 71 -- .../DirectArcLength.jl} | 0 .../DirectByInversion.jl} | 0 .../Axisymmetric/GradShafranovResidual.jl | 170 ++++ .../Inverse.jl} | 0 .../Separatrix.jl} | 95 +- .../{ => Axisymmetric}/SolveTolerances.jl | 0 src/Equilibrium/Config.jl | 306 ++++++ src/Equilibrium/Construction.jl | 34 +- .../{ => Derived}/CoordinateInvariant.jl | 0 .../{ => Derived}/FluxSurfaceMetrics.jl | 0 .../{ => Derived}/GeometryProfiles.jl | 0 src/Equilibrium/Derived/GlobalParameters.jl | 225 +++++ src/Equilibrium/Equilibrium.jl | 35 +- src/Equilibrium/EquilibriumTypes.jl | 923 ------------------ src/Equilibrium/{ => Grid}/GridRefinement.jl | 0 src/Equilibrium/Grid/PsiGrid.jl | 70 ++ src/Equilibrium/PlasmaEquilibrium.jl | 295 ++++++ src/Equilibrium/PostProcessing.jl | 491 ---------- src/Equilibrium/ReadEquilibrium.jl | 625 ------------ src/Equilibrium/Sources/Analytic/Analytic.jl | 42 + .../Sources/Analytic/LargeAspectRatio.jl | 271 +++++ src/Equilibrium/Sources/Analytic/Solovev.jl | 118 +++ .../Analytic/TJAnalytic.jl} | 331 +------ src/Equilibrium/Sources/Chease.jl | 265 +++++ src/Equilibrium/Sources/GFile.jl | 128 +++ src/Equilibrium/Sources/IFile.jl | 76 ++ src/Equilibrium/Sources/IMAS.jl | 119 +++ src/Equilibrium/Sources/Ingest.jl | 222 +++++ test/runtests_equil.jl | 2 +- 32 files changed, 2471 insertions(+), 2475 deletions(-) rename src/Equilibrium/{DirectEquilibrium.jl => Axisymmetric/Direct.jl} (87%) rename src/Equilibrium/{DirectEquilibriumArcLength.jl => Axisymmetric/DirectArcLength.jl} (100%) rename src/Equilibrium/{DirectEquilibriumByInversion.jl => Axisymmetric/DirectByInversion.jl} (100%) create mode 100644 src/Equilibrium/Axisymmetric/GradShafranovResidual.jl rename src/Equilibrium/{InverseEquilibrium.jl => Axisymmetric/Inverse.jl} (100%) rename src/Equilibrium/{SeparatrixQLaw.jl => Axisymmetric/Separatrix.jl} (68%) rename src/Equilibrium/{ => Axisymmetric}/SolveTolerances.jl (100%) create mode 100644 src/Equilibrium/Config.jl rename src/Equilibrium/{ => Derived}/CoordinateInvariant.jl (100%) rename src/Equilibrium/{ => Derived}/FluxSurfaceMetrics.jl (100%) rename src/Equilibrium/{ => Derived}/GeometryProfiles.jl (100%) create mode 100644 src/Equilibrium/Derived/GlobalParameters.jl delete mode 100644 src/Equilibrium/EquilibriumTypes.jl rename src/Equilibrium/{ => Grid}/GridRefinement.jl (100%) create mode 100644 src/Equilibrium/Grid/PsiGrid.jl create mode 100644 src/Equilibrium/PlasmaEquilibrium.jl delete mode 100644 src/Equilibrium/PostProcessing.jl delete mode 100644 src/Equilibrium/ReadEquilibrium.jl create mode 100644 src/Equilibrium/Sources/Analytic/Analytic.jl create mode 100644 src/Equilibrium/Sources/Analytic/LargeAspectRatio.jl create mode 100644 src/Equilibrium/Sources/Analytic/Solovev.jl rename src/Equilibrium/{AnalyticEquilibrium.jl => Sources/Analytic/TJAnalytic.jl} (68%) create mode 100644 src/Equilibrium/Sources/Chease.jl create mode 100644 src/Equilibrium/Sources/GFile.jl create mode 100644 src/Equilibrium/Sources/IFile.jl create mode 100644 src/Equilibrium/Sources/IMAS.jl create mode 100644 src/Equilibrium/Sources/Ingest.jl diff --git a/docs/development/architecture.md b/docs/development/architecture.md index d2fce8d79..c66c4df85 100644 --- a/docs/development/architecture.md +++ b/docs/development/architecture.md @@ -32,14 +32,16 @@ Splines are provided by the external `FastInterpolations` package rather than by - `chease`, `chease2` - CHEASE equilibrium code formats - `lar` - Large Aspect Ratio analytical model - `sol` - Solovev analytical equilibrium - - Key files: - - `EquilibriumTypes.jl` - Core data structures - - `Construction.jl` - Forming a `PlasmaEquilibrium` from each source type, and re-forming it - - `PostProcessing.jl` - Separatrix, global parameters, q extrema, Grad-Shafranov residual - - `ReadEquilibrium.jl` - Parsing equilibrium files - - `DirectEquilibrium.jl` - Direct Grad-Shafranov solver - - `InverseEquilibrium.jl` - Inverse equilibrium solver - - `AnalyticEquilibrium.jl` - Analytical solutions + - Key files, by pipeline stage: + - `Config.jl` - `EquilibriumConfig`, `EquilibriumParameters` + - `PlasmaEquilibrium.jl` - The result type and its profile and geometry splines + - `Construction.jl` - Solver input for each source type, the `PlasmaEquilibrium` constructors, `reform` + - `Sources/` - One file per input format (`GFile.jl`, `Chease.jl`, `IFile.jl`, `IMAS.jl`), the ingest and + run-input types (`Ingest.jl`), and `Analytic/` (the analytic registry plus one file per model with its config and solver) + - `Grid/` - First-pass ψ grids (`PsiGrid.jl`) and the two-pass refinement (`GridRefinement.jl`) + - `Axisymmetric/` - The 2D solvers (`Direct.jl`, `DirectArcLength.jl`, `DirectByInversion.jl`, `Inverse.jl`), + the Grad-Shafranov residual, and separatrix finding + - `Derived/` - Global parameters and q extrema, flux-surface metrics, geometry profiles, coordinate-invariant weights - Status: Stable and feature-complete 3. **Vacuum** (`src/Vacuum/`) - Vacuum field calculations and Green's functions diff --git a/docs/src/conventions.md b/docs/src/conventions.md index ffaadfd29..d0e910679 100644 --- a/docs/src/conventions.md +++ b/docs/src/conventions.md @@ -40,7 +40,7 @@ kernel ``\exp(-im\theta)``; the inverse transform reconstructs with ``\exp(+im\t - Normalized from 0 (magnetic axis) to 1 (plasma boundary). - ``\psi_0 = |\psi_\mathrm{bry} - \psi_\mathrm{axis}|`` is **forced positive** on read - (`read_eq_efit` in `src/Equilibrium/ReadEquilibrium.jl`). If the EQDSK has + (`read_eq_efit` in `src/Equilibrium/Sources/GFile.jl`). If the EQDSK has ``\psi_\mathrm{bry} < \psi_\mathrm{axis}``, the 2D flux array is sign-flipped so that ``\psi`` increases from axis to edge. - The internal ``\psi`` carries **no** factor of ``2\pi`` (poloidal flux per radian). IMAS inputs, @@ -83,8 +83,8 @@ field magnitude, ``B`` is the total field magnitude, ``R`` is the major radius, | Park | 0 | 1 | 0 | 0 | The powers are set automatically from `jac_type` (`EquilibriumConfig` in -`src/Equilibrium/EquilibriumTypes.jl`) and applied in the field-line Jacobian -(`direct_fieldline_der!` in `src/Equilibrium/DirectEquilibrium.jl`). Setting `jac_type = "custom"` +`src/Equilibrium/Config.jl`) and applied in the field-line Jacobian +(`direct_fieldline_der!` in `src/Equilibrium/Axisymmetric/Direct.jl`). Setting `jac_type = "custom"` lets you specify the four exponents manually via `jac_custom_power_bp`, `jac_custom_power_b`, `jac_custom_power_r`, and `jac_custom_power_rc`. @@ -106,7 +106,7 @@ helicity = bt_sign * Int(sign(crnt)) # src/ForcingTerms/CoilFourier.jl ## ``F = R B_\phi`` The poloidal-current function ``F = R B_\phi`` from the Grad-Shafranov equation is **forced -positive** via `abs` (`read_eq_efit` in `src/Equilibrium/ReadEquilibrium.jl`): +positive** via `abs` (`read_eq_efit` in `src/Equilibrium/Sources/GFile.jl`): ```julia abs.(fpol_data) @@ -118,7 +118,7 @@ The code always works with ``|F|``; the sign of ``B_t`` is carried separately (` ## Safety Factor ``q`` The direct solver **computes** ``q`` by field-line integration for direct (EFIT) equilibria and -overwrites the file profile (`direct_run` in `src/Equilibrium/DirectEquilibrium.jl`): +overwrites the file profile (`direct_run` in `src/Equilibrium/Axisymmetric/Direct.jl`): ```math q = \frac{F}{2\pi}\oint \frac{J}{B_p}\,dl_p . @@ -129,7 +129,7 @@ q = \frac{F}{2\pi}\oint \frac{J}{B_p}\,dl_p . extrapolation yields ``q_0 \le 0``, treating it as a spline artifact rather than a physical result. - **Inverse** equilibria, for which only CHEASE and analytic options are currently supported, instead use the ``q`` profile supplied with the input - (`src/Equilibrium/InverseEquilibrium.jl`). This follows the fortran GPEC CHEASE reading convention. + (`src/Equilibrium/Axisymmetric/Inverse.jl`). This follows the fortran GPEC CHEASE reading convention. ## Mode Numbers ``m`` and ``n`` @@ -405,7 +405,7 @@ The conventions above are GPEC-native and predate the COCOS standard, as defined - ``q > 0`` and ``n > 0`` in standard operation; ``F = |F|``. The code labels this internal convention **COCOS 2** (`read_imas` in -`src/Equilibrium/ReadEquilibrium.jl`) and converts IMAS inputs accordingly: +`src/Equilibrium/Sources/IMAS.jl`) and converts IMAS inputs accordingly: ```julia imas_cocos = 11 # default: IMAS standard, ψ divided by 2π on read diff --git a/src/Equilibrium/DirectEquilibrium.jl b/src/Equilibrium/Axisymmetric/Direct.jl similarity index 87% rename from src/Equilibrium/DirectEquilibrium.jl rename to src/Equilibrium/Axisymmetric/Direct.jl index fc04bc0eb..0589aaf1e 100644 --- a/src/Equilibrium/DirectEquilibrium.jl +++ b/src/Equilibrium/Axisymmetric/Direct.jl @@ -453,77 +453,6 @@ function direct_refine(rfac::Float64, eta::Float64, psi0::Float64, params::Field atol=1e-12 * abs(psi0), rtol=1e-12, maxevals=50) end -""" - make_optimal_psi_grid(psilow, psihigh, N_core, N_mid, N_edge; psi_split_core, psi_split_edge) - -Build a three-region ψ grid with (N_core + N_mid + N_edge)+1 knots, using the counts -directly as provided — core and edge are geometric in log(ψ) and log(1−ψ) respectively; -middle is uniform in ψ. -""" -function make_optimal_psi_grid(psilow, psihigh, N_core, N_mid, N_edge; - psi_split_core=0.03, psi_split_edge=0.98) - # Core: [psilow, psi_split_core], geometric in log(ψ) - core_pts = [psilow * (psi_split_core / psilow)^(i / N_core) for i in 0:N_core] - # Middle: [psi_split_core, psi_split_edge], uniform (skip first to avoid duplicate) - mid_pts = [psi_split_core + (psi_split_edge - psi_split_core) * i / N_mid for i in 1:N_mid] - # Edge: [psi_split_edge, psihigh], geometric in log(1−ψ) (skip first to avoid duplicate) - edge_pts = [1.0 - (1.0 - psi_split_edge) * ((1.0 - psihigh) / (1.0 - psi_split_edge))^(i / N_edge) for i in 1:N_edge] - - return vcat(core_pts, mid_pts, edge_pts) -end - -""" - _build_psi_grid(equil_params, psilow, psihigh) - -Resolve `mpsi` and build `psi_nodes` for any supported `grid_type`. - -For `"auto"` (or the legacy alias `"log_asymptotic"`) with `mpsi=0`, this is the coarse pass-1 layout of -the two-pass refinement — the main driver measures the formed equilibrium and re-forms on -the `refined_psi_grid` result. Shared by `direct_fieldline_int` and `efit_by_inversion` -solvers. -""" -function _build_psi_grid(equil_params, psilow, psihigh) - mpsi = equil_params.mpsi - if equil_params.grid_type in ("auto", "log_asymptotic") && mpsi == 0 && equil_params.psi_accuracy > 0 - # Two-pass auto grid: this is the pass-1 layout; the driver measures the formed - # equilibrium's curvature and re-forms on a refined grid (GridRefinement.jl). Sized to - # resolve the gradient structure accurately so pass 2 is well-provisioned (PASS1_INTERVALS). - mpsi = PASS1_INTERVALS - @info "Auto psi grid: forming pass-1 equilibrium on $(mpsi)-interval log_asymptotic grid pending curvature-based refinement" - elseif mpsi == 0 - mpsi = 128 - end - - psi_nodes = if equil_params.grid_type in ("auto", "log_asymptotic") - psihigh > 0.98 || - @warn "grid_type = \"$(equil_params.grid_type)\" needs psihigh > 0.98 (its grid has a fixed edge region on [0.98, psihigh]); " * - "psihigh = $psihigh gives a non-monotonic ψ grid and the equilibrium will fail. Use grid_type = \"core_edge_packed\" instead." - # Distribute mpsi across the three regions by log-weights - log_core = log(0.03 / psilow) - log_mid = log(0.98 / 0.03) - log_edge = log((1.0 - 0.98) / (1.0 - psihigh)) - log_total = log_core + log_mid + log_edge - N_edge = clamp(round(Int, mpsi * log_edge / log_total), 2, mpsi ÷ 2) - N_core = round(Int, mpsi * log_core / log_total) - N_mid = mpsi - N_edge - N_core - make_optimal_psi_grid(psilow, psihigh, N_core, N_mid, N_edge) - elseif equil_params.grid_type == "core_edge_packed" - [psilow + (psihigh - psilow) * sin((ipsi / mpsi) * (π / 2))^2 for ipsi in 0:mpsi] - elseif equil_params.grid_type == "pow1" - # Fortran powspace(psilow, psihigh, 1, mpsi+1, "upper") — edge-packed grid (equil/grid.f90:92-195) - [psilow + (psihigh - psilow) * (3(ipsi / mpsi) - (ipsi / mpsi)^3) / 2 for ipsi in 0:mpsi] - elseif equil_params.grid_type == "uniform" - # Evenly spaced surfaces, e.g. for profile diagnostics at constant resolution - collect(range(psilow, psihigh; length=mpsi + 1)) - else - error("Unsupported grid_type: $(equil_params.grid_type)") - end - # Floor node spacing on the fixed grids: core_edge_packed/pow1 pack the edge as ~(π/2·mpsi)⁻² and at high - # mpsi drive it below the integration-noise scale (garbage curvature). The auto grid floors its - # refined pass-2 grid in refined_psi_grid, so it is left untouched here. - return equil_params.grid_type in ("auto", "log_asymptotic") ? psi_nodes : enforce_min_spacing(psi_nodes, MIN_KNOT_SPACING) -end - """ equilibrium_solver(raw_profile) diff --git a/src/Equilibrium/DirectEquilibriumArcLength.jl b/src/Equilibrium/Axisymmetric/DirectArcLength.jl similarity index 100% rename from src/Equilibrium/DirectEquilibriumArcLength.jl rename to src/Equilibrium/Axisymmetric/DirectArcLength.jl diff --git a/src/Equilibrium/DirectEquilibriumByInversion.jl b/src/Equilibrium/Axisymmetric/DirectByInversion.jl similarity index 100% rename from src/Equilibrium/DirectEquilibriumByInversion.jl rename to src/Equilibrium/Axisymmetric/DirectByInversion.jl diff --git a/src/Equilibrium/Axisymmetric/GradShafranovResidual.jl b/src/Equilibrium/Axisymmetric/GradShafranovResidual.jl new file mode 100644 index 000000000..8eda1f175 --- /dev/null +++ b/src/Equilibrium/Axisymmetric/GradShafranovResidual.jl @@ -0,0 +1,170 @@ +""" + equilibrium_gse!(equil::PlasmaEquilibrium) + +Diagnoses the Grad-Shafranov solution by computing the residual of the +Grad-Shafranov equation across the grid and writing diagnostic data to HDF5 files. +Performs the same function as equil_out_gse in the Fortran code. + +Returns `(; xs, ys, flux_x, flux_y, source, total, error, errori)` on the `rzphi` (ψ, θ) nodes: the two +flux-divergence terms, the source, the residual `total = flux_x - flux_y + source`, `error = |total|` over +the largest term on the grid, and `errori`, the magnitude of the θ-integrated residual on each surface. +""" +function equilibrium_gse!(equil::PlasmaEquilibrium) + + profiles = equil.profiles + mpsi = length(equil.rzphi_xs) - 1 + mtheta = length(equil.rzphi_ys) - 1 + ro, zo = equil.ro, equil.zo + psio = equil.psio + verbose = equil.params.verbose + diagnose_src = equil.params.diagnose_src + diagnose_maxima = equil.params.diagnose_maxima + + if verbose + @info "Diagnosing Grad-Shafranov solution" + end + + # Compute R, Z coordinates using nodal_derivs access + r = zeros(Float64, mpsi + 1, mtheta + 1) + z = zeros(Float64, mpsi + 1, mtheta + 1) + + for ipsi in 1:(mpsi+1) + rfac = @. sqrt(equil.rzphi_rsquared.nodal_derivs.partials[1, ipsi, :]) + angle = @. 2π * (equil.rzphi_ys + equil.rzphi_offset.nodal_derivs.partials[1, ipsi, :]) + r[ipsi, :] .= ro .+ rfac .* cos.(angle) + z[ipsi, :] .= zo .+ rfac .* sin.(angle) + end + + # Compute flux quantities using direct array access at grid points + flux_fs = zeros(Float64, mpsi + 1, mtheta + 1, 2) + flux_fsx = zeros(Float64, mpsi + 1, mtheta + 1, 2) # x-derivatives + flux_fsy = zeros(Float64, mpsi + 1, mtheta + 1, 2) # y-derivatives + for ipsi in 1:(mpsi+1) + for itheta in 1:(mtheta+1) + f1 = equil.rzphi_rsquared.nodal_derivs.partials[1, ipsi, itheta] + f2 = equil.rzphi_offset.nodal_derivs.partials[1, ipsi, itheta] + f4 = equil.rzphi_jac.nodal_derivs.partials[1, ipsi, itheta] + fy1 = equil.rzphi_rsquared.nodal_derivs.partials[3, ipsi, itheta] + fy2 = equil.rzphi_offset.nodal_derivs.partials[3, ipsi, itheta] + fx1 = equil.rzphi_rsquared.nodal_derivs.partials[2, ipsi, itheta] + fx2 = equil.rzphi_offset.nodal_derivs.partials[2, ipsi, itheta] + + flux_fs[ipsi, itheta, 1] = fy1^2 / (4π^2 * f1) + (1 + fy2)^2 * 4 * f1 + flux_fs[ipsi, itheta, 2] = fx1 * fy1 / (4π^2 * f1) + fx2 * (1 + fy2) * 4 * f1 + + flux_fs[ipsi, itheta, 1] *= 2π * psio / f4 + flux_fs[ipsi, itheta, 2] *= 2π * psio / f4 + end + end + # Create flux interpolants for Grad-Shafranov diagnostics + @views flux_fs[:, end, :] .= flux_fs[:, 1, :] + flux_opts = (bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) + flux1 = cubic_interp((equil.rzphi_xs, equil.rzphi_ys), flux_fs[:, :, 1]; flux_opts...) + flux2 = cubic_interp((equil.rzphi_xs, equil.rzphi_ys), flux_fs[:, :, 2]; flux_opts...) + + # Compute flux derivatives at all grid points for diagnostics + hint2d = (Ref(1), Ref(1)) # Shared 2D hint for hot loop optimization + for ipsi in 0:mpsi + for itheta in 0:mtheta + query_point = (equil.rzphi_xs[ipsi+1], equil.rzphi_ys[itheta+1]) + flux_fsx[ipsi+1, itheta+1, 1] = flux1(query_point; deriv=DerivOp(1, 0), hint=hint2d) + flux_fsx[ipsi+1, itheta+1, 2] = flux2(query_point; deriv=DerivOp(1, 0), hint=hint2d) + flux_fsy[ipsi+1, itheta+1, 1] = flux1(query_point; deriv=DerivOp(0, 1), hint=hint2d) + flux_fsy[ipsi+1, itheta+1, 2] = flux2(query_point; deriv=DerivOp(0, 1), hint=hint2d) + end + end + + # Compute source term using direct array access + source = zeros(Float64, mpsi + 1, mtheta + 1) + hint = Ref(1) # Linear search hint for sequential psi access + for ipsi in 1:(mpsi+1) + psi = profiles.xs[ipsi] + s1 = profiles.F_spline.y[ipsi] + s1p = profiles.F_deriv(psi; hint=hint) + s2p = profiles.P_deriv(psi; hint=hint) + for itheta in 1:(mtheta+1) + f4 = equil.rzphi_jac.nodal_derivs.partials[1, ipsi, itheta] + denom = (2π * r[ipsi, itheta])^2 + source[ipsi, itheta] = f4 / (2π * psio * π^2) * (s1 * s1p / denom + s2p) + end + end + + total = flux_fsx[:, :, 1] .- flux_fsy[:, :, 2] .+ source + error = abs.(total) ./ maximum([maximum(abs.(flux_fsx[:, :, 1])), maximum(abs.(flux_fsy[:, :, 2])), maximum(abs.(source))]) + errlog = ifelse.(error .> 0, log10.(error), 0.0) + + if diagnose_maxima + fxmax = maximum(abs.(flux_fsx[:, :, 1])) + fymax = maximum(abs.(flux_fsy[:, :, 2])) + smax = maximum(abs.(source)) + emax = maximum(abs.(error)) + lmax = maximum(errlog) + jmax = ind2sub(size(errlog), argmax(errlog)) + @info "GS residuals: fxmax = $(@sprintf("%.3e", fxmax)), fymax = $(@sprintf("%.3e", fymax)), smax = $(@sprintf("%.3e", smax)), emax = $(@sprintf("%.3e", emax)), lmax = $(@sprintf("%.3f", lmax)), maxloc = $(jmax .- 1)" + end + + # Integrated error criterion + term = zeros(Float64, mpsi + 1, 2) + for ipsi in 1:(mpsi+1) + fs_matrix = zeros(Float64, mtheta + 1, 2) + fs_matrix[:, 1] = flux_fsx[ipsi, :, 1] + fs_matrix[:, 2] = source[ipsi, :] + # Snap the repeated endpoint exactly equal to the start + fs_matrix[end, :] .= fs_matrix[1, :] + + # Compute total integral using FastInterpolations native integration + itp = cubic_interp(equil.rzphi_ys, Series(fs_matrix); bc=PeriodicBC()) + term[ipsi, :] .= FastInterpolations.integrate(itp) + end + + totali = sum(term; dims=2) + errori = abs.(totali) + errlogi = @. ifelse(errori > 0, log10(errori), 0.0) + + if diagnose_src + if verbose + @info "Writing diagnostics to HDF5 files" + end + + # Write contour data + h5open(joinpath(dirname(equil.config.eq_filename), "gsec.h5"), "w") do file + file["mpsi"] = mpsi + file["mtheta"] = mtheta + file["r"] = Float32.(r) + file["z"] = Float32.(z) + file["flux_fsx"] = Float32.(flux_fsx[:, :, 1]) + file["flux_fsy"] = Float32.(flux_fsy[:, :, 2]) + file["source"] = Float32.(source) + file["total"] = Float32.(total) + file["error"] = Float32.(error) + file["errlog"] = Float32.(errlog) + end + + # Write xy plot data + h5open(joinpath(dirname(equil.config.eq_filename), "gse.h5"), "w") do file + gse_data = Array{Float32,3}(undef, mpsi + 1, mtheta + 1, 7) + for ipsi in 0:mpsi + for itheta in 0:mtheta + gse_data[ipsi+1, itheta+1, 1] = Float32(flux.ys[itheta+1]) + gse_data[ipsi+1, itheta+1, 2] = Float32(flux.xs[ipsi+1]) + gse_data[ipsi+1, itheta+1, 3] = Float32(flux_fs[ipsi+1, itheta+1, 1]) + gse_data[ipsi+1, itheta+1, 4] = Float32(flux_fs[ipsi+1, itheta+1, 2]) + gse_data[ipsi+1, itheta+1, 5] = Float32(source[ipsi+1, itheta+1]) + gse_data[ipsi+1, itheta+1, 6] = Float32(total[ipsi+1, itheta+1]) + gse_data[ipsi+1, itheta+1, 7] = Float32(error[ipsi+1, itheta+1]) + end + end + file["gse_data"] = gse_data + end + + # Write integrated error criterion + h5open(joinpath(dirname(equil.config.eq_filename), "gsei.h5"), "w") do file + file["xs"] = Float32.(flux.xs) + file["term"] = Float32.(term) + file["totali"] = Float32.(totali) + file["errori"] = Float32.(errori) + file["errlogi"] = Float32.(errlogi) + end + end + return (; xs=equil.rzphi_xs, ys=equil.rzphi_ys, flux_x=flux_fsx[:, :, 1], flux_y=flux_fsy[:, :, 2], source, total, error, errori=vec(errori)) +end diff --git a/src/Equilibrium/InverseEquilibrium.jl b/src/Equilibrium/Axisymmetric/Inverse.jl similarity index 100% rename from src/Equilibrium/InverseEquilibrium.jl rename to src/Equilibrium/Axisymmetric/Inverse.jl diff --git a/src/Equilibrium/SeparatrixQLaw.jl b/src/Equilibrium/Axisymmetric/Separatrix.jl similarity index 68% rename from src/Equilibrium/SeparatrixQLaw.jl rename to src/Equilibrium/Axisymmetric/Separatrix.jl index 48df51aeb..25340b2ac 100644 --- a/src/Equilibrium/SeparatrixQLaw.jl +++ b/src/Equilibrium/Axisymmetric/Separatrix.jl @@ -1,7 +1,8 @@ -# SeparatrixQLaw.jl +# Separatrix.jl # # Safety factor between the outermost grid surface and the separatrix of a diverted equilibrium, -# taken from the X-point(s) of the equilibrium's own poloidal flux rather than fitted. +# taken from the X-point(s) of the equilibrium's own poloidal flux rather than fitted, and the +# separatrix's midplane and extremal points (`equilibrium_separatrix_find!`). # # Near a separatrix a field line stalls poloidally at each X-point it passes, so # @@ -170,3 +171,93 @@ function separatrix_q_psi(law, q_target::Real; psi_gap::Real=SEPARATRIX_PSI_GAP) residual(u_near) > 0 || return nothing return psi_sep - exp(find_zero(residual, (u_near, u_far), Roots.Brent())) end + +""" + equilibrium_separatrix_find!(pe::PlasmaEquilibrium) + +Finds the separatrix locations in the plasma equilibrium (rsep, zsep, rext, zext). +Performs the same function as equil_out_sep_find in the Fortran code. +""" +function equilibrium_separatrix_find!(pe::PlasmaEquilibrium) + mpsi = length(pe.rzphi_xs) - 1 + edge_idx = mpsi + 1 # Edge flux surface index + psi_edge = pe.rzphi_xs[edge_idx] + rsep = zeros(2) + + # Outboard and inboard midplane R via bracketed Brent on θ + η(θ) - η₀ = 0. + # iside=1 → outboard (η₀=0.0, θ near 0), iside=2 → inboard (η₀=0.5, θ near 0.5). + for iside in 1:2 + eta0 = (iside == 1) ? 0.0 : 0.5 + hint2d = (Ref(1), Ref(1)) + theta_lo, theta_hi = (iside == 1) ? (-0.25, 0.25) : (0.25, 0.75) + side_label = (iside == 1) ? "outboard" : "inboard" + theta = try + find_zero( + theta -> theta + pe.rzphi_offset((psi_edge, theta); hint=hint2d) - eta0, + (theta_lo, theta_hi), Roots.Brent(); + atol=1e-12, rtol=1e-12) + catch e + error("Separatrix $side_label midplane root not found in bracket " * + "[$(theta_lo), $(theta_hi)]: $(e.msg)") + end + r2 = pe.rzphi_rsquared((psi_edge, theta)) + offset = pe.rzphi_offset((psi_edge, theta)) + rsep[iside] = pe.ro + sqrt(r2) * cos(2π * (theta + offset)) + end + + # Top and bottom separatrix Z extrema via bracketed Brent on ∂z/∂θ = 0. + # iside=1 → top (θ ∈ [0.0, 0.5]), iside=2 → bottom (θ ∈ [0.5, 1.0]). + # Matches Fortran convention so (r|z)ext[1] / zsep[1] refer to the top extremum + # and (r|z)ext[2] / zsep[2] to the bottom (see equil_out.f::equil_out_sep_find). + # Splines use PeriodicBC + WrapExtrap, so θ outside [0,1] is valid. + zsep = zeros(2) + rext = zeros(2) + zext = zeros(2) + + for iside in 1:2 + hint2d = (Ref(1), Ref(1)) + + # Cache variables populated by z_deriv, read after convergence + rfac = Ref(0.0) + cos_phase = Ref(0.0) + z_val = Ref(0.0) + + # ∂z/∂θ where z(θ) = zo + √r²(θ) · sin(2π(θ + η(θ))) + function z_deriv(theta_inner) + r2 = pe.rzphi_rsquared((psi_edge, theta_inner); hint=hint2d) + r2y = pe.rzphi_rsquared((psi_edge, theta_inner); deriv=DerivOp(0, 1), hint=hint2d) + η = pe.rzphi_offset((psi_edge, theta_inner); hint=hint2d) + η1 = pe.rzphi_offset((psi_edge, theta_inner); deriv=DerivOp(0, 1), hint=hint2d) + rfac_local = sqrt(max(0.0, r2)) + rfac1 = (rfac_local > 0) ? r2y / (2 * rfac_local) : 0.0 + phase1 = 2π * (1 + η1) + sin_phase = sin(2π * (theta_inner + η)) + cos_phase_local = cos(2π * (theta_inner + η)) + + rfac[] = rfac_local + cos_phase[] = cos_phase_local + z_val[] = pe.zo + rfac_local * sin_phase + + return rfac_local * phase1 * cos_phase_local + rfac1 * sin_phase + end + + theta_lo, theta_hi = (iside == 1) ? (0.0, 0.5) : (0.5, 1.0) + side_label = (iside == 1) ? "top" : "bottom" + theta = try + find_zero(z_deriv, (theta_lo, theta_hi), Roots.Brent(); + atol=1e-12, rtol=1e-12) + catch e + error("Separatrix $side_label Z-extremum root not found in bracket " * + "[$(theta_lo), $(theta_hi)]: $(e.msg)") + end + + rext[iside] = pe.ro + rfac[] * cos_phase[] + zsep[iside] = zext[iside] = z_val[] + end + + pe.params.rsep = rsep + pe.params.zsep = zsep + pe.params.rext = rext + pe.params.zext = zext + return (rsep, zsep, rext, zext) +end diff --git a/src/Equilibrium/SolveTolerances.jl b/src/Equilibrium/Axisymmetric/SolveTolerances.jl similarity index 100% rename from src/Equilibrium/SolveTolerances.jl rename to src/Equilibrium/Axisymmetric/SolveTolerances.jl diff --git a/src/Equilibrium/Config.jl b/src/Equilibrium/Config.jl new file mode 100644 index 000000000..097f7c835 --- /dev/null +++ b/src/Equilibrium/Config.jl @@ -0,0 +1,306 @@ +function symbolize_keys(dict::Dict{String,Any}) + return Dict(Symbol(k) => v for (k, v) in dict) +end + +""" + EquilibriumConfig(...) + +An immutable struct containing configuration parameters for equilibrium reconstruction +specified in the input. + +## Fields + + - `eq_type::String` - Type of equilibrium file ("efit", "solovev", "lar", etc.) + - `eq_filename::String` - Path to equilibrium input file + - `jac_type::String` - Jacobian coordinate type ("hamada", "pest", "equal_arc", "boozer", "park", "custom") + - `power_bp::Int` - Poloidal field power exponent for Jacobian (internal, derived from `jac_type`; deprecated as a TOML input key) + - `power_b::Int` - Total field power exponent for Jacobian (internal, derived from `jac_type`; deprecated as a TOML input key) + - `power_r::Int` - Major radius power exponent for Jacobian (internal, derived from `jac_type`; deprecated as a TOML input key) + - `power_rc::Int` - Minor radius (rfac = √((R-R₀)²+(Z-Z₀)²)) power exponent for Jacobian (internal, derived from `jac_type`; deprecated as a TOML input key) + - `jac_custom_power_bp::Int` - Poloidal-field exponent for a custom Jacobian `J = Bp^bp · B^b / (R^r · rfac^rc)`; read only when `jac_type = "custom"`, ignored for named types + - `jac_custom_power_b::Int` - Total-field exponent for a custom Jacobian; read only when `jac_type = "custom"`, ignored for named types + - `jac_custom_power_r::Int` - Major-radius exponent for a custom Jacobian; read only when `jac_type = "custom"`, ignored for named types + - `jac_custom_power_rc::Int` - Minor-radius (rfac) exponent for a custom Jacobian; read only when `jac_type = "custom"`, ignored for named types + - `r0_norm::Float64` - Major radius normalization for CHEASE/EQDSK [m] + - `b0_norm::Float64` - On-axis toroidal field normalization for CHEASE/EQDSK [T] + - `grid_type::String` - Grid type for flux surface discretization ("auto" — two-pass measured-curvature + refinement when mpsi=0, three-region log layout when mpsi>0; "core_edge_packed" (sin² spacing, dense at axis and edge), "pow1", "uniform"; + inverse equilibria also "rho" (uniform in √ψ from psilow to psihigh) and "original" (the input + file's own ψ nodes off the axis; ignores psilow and mpsi); "log_asymptotic" is a legacy alias for "auto") + - `psilow::Float64` - Lower limit of normalized flux coordinate + - `psihigh::Float64` - Requested upper limit of normalized flux coordinate; the value the + equilibrium is actually formed on is `EquilibriumParameters.psihigh_resolved`. + - `mpsi::Int` - Number of radial grid intervals; 0 with grid_type="auto" selects the + two-pass auto grid: the main driver forms a coarse pass-1 equilibrium, measures its curvature, + pins knots on rational surfaces, and re-forms on the refined grid. Standalone `PlasmaEquilibrium` + callers get the coarse pass-1 grid unless they refine via `refined_psi_grid` + `override_psi_nodes`. + - `psi_accuracy::Float64` - Target relative accuracy τ of splined profile derivatives for the + two-pass auto grid (knot count scales as τ^(-1/3)) + - `mtheta::Int` - Number of poloidal grid points + - `newq0::Float64` - Target on-axis safety factor q(0); the q and F profiles are rescaled to + meet it (0 = use input value, -1 = use the axis extrapolation with its sign flipped) + - `flux_surface_rtol::Float64` - Relative tolerance of the ODE solves that trace each flux surface; their absolute tolerance + is `equil_abstol(flux_surface_rtol)`, i.e. `flux_surface_rtol` capped at `EQUIL_ABSTOL_MAX` + - `force_termination::Bool` - Terminate after equilibrium setup (skip stability calculations) +""" +@kwdef struct EquilibriumConfig + eq_type::String = "efit" + eq_filename::String = "mypath" + r0_norm::Float64 = 1.0 + b0_norm::Float64 = 1.0 + + jac_type::String = "hamada" + power_bp::Int = 0 + power_b::Int = 0 + power_r::Int = 0 + power_rc::Int = 0 + + jac_custom_power_bp::Int = 0 + jac_custom_power_b::Int = 0 + jac_custom_power_r::Int = 0 + jac_custom_power_rc::Int = 0 + + grid_type::String = "auto" + psilow::Float64 = 1e-2 + psihigh::Float64 = 0.9995 + mpsi::Int = 0 + psi_accuracy::Float64 = 0.001 + mtheta::Int = 512 + + newq0::Float64 = 0.0 + flux_surface_rtol::Float64 = 1e-10 + + force_termination::Bool = false + + # IMAS-specific: expected COCOS convention of the input dd.equilibrium (11=IMAS standard, 2=GPEC internal) + imas_cocos::Int = 11 + + """ + Modified internal constructor that enforces self consistency within the inputs + """ + # The four `_` slots are the `power_bp/power_b/power_r/power_rc` struct fields, which @kwdef + # forwards positionally. They are always derived below from `jac_type` (or `jac_custom_power_*`), + # so their incoming values are ignored (hence `_`). + function EquilibriumConfig(eq_type, eq_filename, r0_norm, b0_norm, jac_type, _, _, _, _, + jac_custom_power_bp, jac_custom_power_b, jac_custom_power_r, jac_custom_power_rc, + grid_type, psilow, psihigh, mpsi, psi_accuracy, mtheta, newq0, flux_surface_rtol, + force_termination, imas_cocos) + if jac_type == "hamada" + @info "Forcing hamada coordinate jacobian exponents: power_*" + power_b = 0 + power_bp = 0 + power_r = 0 + power_rc = 0 + elseif jac_type == "pest" + @info "Forcing pest coordinate jacobian exponents: power_*" + power_b = 0 + power_bp = 0 + power_r = 2 + power_rc = 0 + elseif jac_type == "equal_arc" + @info "Forcing equal_arc coordinate jacobian exponents: power_*" + power_b = 0 + power_bp = 1 + power_r = 0 + power_rc = 0 + elseif jac_type == "boozer" + @info "Forcing boozer coordinate jacobian exponents: power_*" + power_b = 2 + power_bp = 0 + power_r = 0 + power_rc = 0 + elseif jac_type == "park" + @info "Forcing park coordinate jacobian exponents: power_*" + power_b = 1 + power_bp = 0 + power_r = 0 + power_rc = 0 + elseif jac_type == "custom" + # Custom Jacobian: the user-facing knobs define the exponents. + power_bp = jac_custom_power_bp + power_b = jac_custom_power_b + power_r = jac_custom_power_r + power_rc = jac_custom_power_rc + # Normalize to a named type when the powers match, so fast paths are taken. + if power_b == 0 && power_bp == 0 && power_r == 0 && power_rc == 0 + jac_type = "hamada" + @info "Recognized hamada jacobian from power exponents" + elseif power_b == 0 && power_bp == 0 && power_r == 2 && power_rc == 0 + jac_type = "pest" + @info "Recognized pest jacobian from power exponents" + elseif power_b == 0 && power_bp == 1 && power_r == 0 && power_rc == 0 + jac_type = "equal_arc" + @info "Recognized equal_arc jacobian from power exponents" + elseif power_b == 2 && power_bp == 0 && power_r == 0 && power_rc == 0 + jac_type = "boozer" + @info "Recognized boozer jacobian from power exponents" + elseif power_b == 1 && power_bp == 0 && power_r == 0 && power_rc == 0 + jac_type = "park" + @info "Recognized park jacobian from power exponents" + else + @info "Using manual jacobian exponents: power b, bp, r, rc = $(power_b), $(power_bp), $(power_r), $(power_rc)" + end + else + error("Cannot recognize jac_type = $(jac_type)") + end + (isfinite(flux_surface_rtol) && flux_surface_rtol > 0) || error("flux_surface_rtol = $flux_surface_rtol must be a positive, finite relative tolerance") + if psihigh > 1.0 + @warn "psihigh = $psihigh exceeds 1.0 (separatrix); clamping to 1.0" + end + psihigh = min(psihigh, 1.0) + return new(eq_type, eq_filename, r0_norm, b0_norm, jac_type, power_bp, power_b, power_r, power_rc, + jac_custom_power_bp, jac_custom_power_b, jac_custom_power_r, jac_custom_power_rc, + grid_type, psilow, psihigh, mpsi, psi_accuracy, mtheta, newq0, flux_surface_rtol, + force_termination, imas_cocos) + end +end + +""" +Outer constructor for EquilibriumConfig from a parsed TOML dictionary +""" +function EquilibriumConfig(equil_dict::Dict{String,Any}, base_path::String="./") + # `eq_type` is always required. `eq_filename` is required for file-based + # equilibria (efit, chease, …) but optional for analytic types whose + # parameters live in an embedded `[TJ_ANALYTIC_INPUT]` / `[SOL_INPUT]` / + # `[LAR_INPUT]` section of the parent gpec.toml. + if !haskey(equil_dict, "eq_type") + error("Missing required key in [Equilibrium]: eq_type") + end + + # Filter to only known parameters + config_fields = Set(String.(fieldnames(EquilibriumConfig))) + config_data = Dict{String,Any}() + + for (k, v) in equil_dict + if k in config_fields + config_data[k] = v + else + @warn "Unknown equilibrium parameter: $k" + end + end + + # Only resolve `eq_filename` against `base_path` if the user actually supplied one + # (otherwise leave the kwdef sentinel for the embedded path). The empty string is the + # rerun path's "no input file" marker and must stay empty, not become `base_path`. + if haskey(config_data, "eq_filename") && !isempty(config_data["eq_filename"]) && !isabspath(config_data["eq_filename"]) + config_data["eq_filename"] = normpath(joinpath(base_path, config_data["eq_filename"])) + end + + return EquilibriumConfig(; symbolize_keys(config_data)...) +end + +""" + EquilibriumParameters + +A mutable struct containing computed equilibrium parameters and diagnostic flags. + +## Fields + + - `ro::Union{Nothing,Float64}` - R-coordinate of the magnetic axis [m] + - `zo::Union{Nothing,Float64}` - Z-coordinate of the magnetic axis [m] + - `psio::Union{Nothing,Float64}` - Total flux difference |ψ_axis - ψ_boundary| [Wb/rad] + - `psihigh_resolved::Union{Nothing,Float64}` - Outer flux limit the equilibrium was formed on + (the outermost ψ node); the plasma edge downstream of `PlasmaEquilibrium`. + - `rsep::Union{Nothing,Vector{Float64}}` - R-coordinates of the plasma boundary [m] + - `zsep::Union{Nothing,Vector{Float64}}` - Z-coordinates of the plasma boundary [m] + - `rext::Union{Nothing,Vector{Float64}}` - R-coordinates of the plasma edge [m] + - `zext::Union{Nothing,Vector{Float64}}` - Z-coordinates of the plasma edge [m] + - `psi0::Union{Nothing,Float64}` - Normalized poloidal flux at reference location + - `b0::Union{Nothing,Float64}` - Total magnetic field strength at the axis [T] + - `q0::Union{Nothing,Float64}` - Safety factor at the axis + - `qmin::Union{Nothing,Float64}` - Minimum safety factor in the plasma + - `qmax::Union{Nothing,Float64}` - Maximum safety factor in the plasma + - `qa::Union{Nothing,Float64}` - Safety factor at the plasma edge + - `q95::Union{Nothing,Float64}` - Safety factor at 95% flux surface + - `qextrema_psi::Union{Nothing,Vector{Float64}}` - Normalized flux values at q extrema + - `qextrema_q::Union{Nothing,Vector{Float64}}` - Safety factor values at extrema + - `mextrema::Union{Nothing,Int}` - Number of extrema in q-profile + - `psi_norm::Union{Nothing,Float64}` - Normalized poloidal flux + - `b_norm::Union{Nothing,Float64}` - Normalized magnetic field strength + - `psi_axis::Union{Nothing,Float64}` - Poloidal flux at the axis + - `psi_boundary::Union{Nothing,Float64}` - Poloidal flux at the boundary + - `psi_boundary_norm::Union{Nothing,Float64}` - Normalized boundary flux + - `psi_axis_norm::Union{Nothing,Float64}` - Normalized axis flux + - `psi_boundary_offset::Union{Nothing,Float64}` - Boundary flux offset + - `psi_axis_offset::Union{Nothing,Float64}` - Axis flux offset + - `psi_boundary_sign::Union{Nothing,Int}` - Sign of boundary flux + - `psi_axis_sign::Union{Nothing,Int}` - Sign of axis flux + - `psi_boundary_zero::Union{Nothing,Bool}` - Whether boundary flux is zero + - `rmean::Union{Nothing,Float64}` - Mean major radius [m] + - `amean::Union{Nothing,Float64}` - Mean minor radius [m] + - `aratio::Union{Nothing,Float64}` - Aspect ratio (R0/a) + - `kappa::Union{Nothing,Float64}` - Plasma elongation + - `delta1::Union{Nothing,Float64}` - Upper triangularity + - `delta2::Union{Nothing,Float64}` - Lower triangularity + - `bt0::Union{Nothing,Float64}` - Toroidal field at axis [T] + - `crnt::Union{Nothing,Float64}` - Plasma current [A] + - `bwall::Union{Nothing,Float64}` - Toroidal field at wall [T] + - `verbose::Bool` - Enable verbose output + - `diagnose_src::Bool` - Enable source data diagnostics + - `diagnose_maxima::Bool` - Enable extrema diagnostics + - `volume::Union{Nothing,Float64}` - Plasma volume [m³] + - `betat::Union{Nothing,Float64}` - Toroidal beta + - `betan::Union{Nothing,Float64}` - Normalized beta + - `betaj::Union{Nothing,Float64}` - Total beta + - `betap1::Union{Nothing,Float64}` - Poloidal beta (definition 1) + - `betap2::Union{Nothing,Float64}` - Poloidal beta (definition 2) + - `betap3::Union{Nothing,Float64}` - Poloidal beta (definition 3) + - `li1::Union{Nothing,Float64}` - Internal inductance (definition 1) + - `li2::Union{Nothing,Float64}` - Internal inductance (definition 2) + - `li3::Union{Nothing,Float64}` - Internal inductance (definition 3) +""" +@kwdef mutable struct EquilibriumParameters + ro::Union{Nothing,Float64} = nothing # R-coordinate of the magnetic axis [m] + zo::Union{Nothing,Float64} = nothing # Z-coordinate of the magnetic axis [m] + psio::Union{Nothing,Float64} = nothing # Total flux difference |ψ_axis - ψ_boundary| [Wb/rad] + psihigh_resolved::Union{Nothing,Float64} = nothing # Outer flux limit actually formed on (clamped config.psihigh) + rsep::Union{Nothing,Vector{Float64}} = nothing # R-coordinates of the plasma boundary [m] + zsep::Union{Nothing,Vector{Float64}} = nothing # Z-coordinates of the plasma boundary [m] + rext::Union{Nothing,Vector{Float64}} = nothing # R-coordinates of the plasma edge [m] + zext::Union{Nothing,Vector{Float64}} = nothing # Z-coordinates of the plasma edge [m] + psi0::Union{Nothing,Float64} = nothing # Normalized poloidal flux + b0::Union{Nothing,Float64} = nothing # Total magnetic field strength at the axis [T] + q0::Union{Nothing,Float64} = nothing # Safety factor at the axis + qmin::Union{Nothing,Float64} = nothing # Minimum safety factor in the plasma + qmax::Union{Nothing,Float64} = nothing # Maximum safety factor in the plasma + qa::Union{Nothing,Float64} = nothing # Safety factor at the plasma edge + q95::Union{Nothing,Float64} = nothing # Safety factor at 95% flux surface + qextrema_psi::Union{Nothing,Vector{Float64}} = nothing # Normalized poloidal flux values where q has extrema + qextrema_q::Union{Nothing,Vector{Float64}} = nothing # Safety factor values at the extrema points + mextrema::Union{Nothing,Int} = nothing # Number of extrema points in the q-profile + psi_norm::Union{Nothing,Float64} = nothing # Normalized poloidal flux at the axis + b_norm::Union{Nothing,Float64} = nothing # Normalized total magnetic field strength at the axis + psi_axis::Union{Nothing,Float64} = nothing # Normalized poloidal flux at the axis + psi_boundary::Union{Nothing,Float64} = nothing # Poloidal flux at the boundary + psi_boundary_norm::Union{Nothing,Float64} = nothing # Normalized poloidal flux at the boundary + psi_axis_norm::Union{Nothing,Float64} = nothing # Normalized poloidal flux at the axis + psi_boundary_offset::Union{Nothing,Float64} = nothing # Offset for the boundary poloidal flux + psi_axis_offset::Union{Nothing,Float64} = nothing # Offset for the axis poloidal flux + psi_boundary_sign::Union{Nothing,Int} = nothing # Sign of the boundary poloidal flux + psi_axis_sign::Union{Nothing,Int} = nothing # Sign of the axis poloidal flux + psi_boundary_zero::Union{Nothing,Bool} = nothing # Whether the boundary poloidal flux is zero + rmean::Union{Nothing,Float64} = nothing # Mean R-coordinate of the plasma [m] + amean::Union{Nothing,Float64} = nothing # Mean minor radius of the plasma [m] + aratio::Union{Nothing,Float64} = nothing # Aspect ratio of the plasma (R0/a) + kappa::Union{Nothing,Float64} = nothing # Elongation of the plasma cross-section + delta1::Union{Nothing,Float64} = nothing # Triangularity of the plasma cross-section (upper triangularity) + delta2::Union{Nothing,Float64} = nothing # Triangularity of the plasma cross-section (lower triangularity) + bt0::Union{Nothing,Float64} = nothing # Toroidal magnetic field at the axis [T] (always positive; sign in bt_sign) + crnt::Union{Nothing,Float64} = nothing # Plasma current at the axis [A] + bt_sign::Int = 1 # Sign of the toroidal field: +1 (positive Bt) or -1 (negative Bt, e.g. DIII-D standard) + ip_sign::Int = 1 # Sign of the plasma current as the source file states it: +1 or -1 (crnt is always positive) + bwall::Union{Nothing,Float64} = nothing # Toroidal magnetic field at the wall [T] + verbose::Bool = false # Whether to print verbose output + diagnose_src::Bool = false # Whether to diagnose source data + diagnose_maxima::Bool = false # Whether to diagnose maxima in the equilibrium + volume::Union{Nothing,Float64} = nothing # Plasma volume [m³] + betat::Union{Nothing,Float64} = nothing # Toroidal beta (normalized pressure) at the axis + betan::Union{Nothing,Float64} = nothing # Normalized beta at the axis + betaj::Union{Nothing,Float64} = nothing # Total beta at the axis + betap1::Union{Nothing,Float64} = nothing # Pressure beta at the axis + betap2::Union{Nothing,Float64} = nothing # Toroidal beta at the axis + betap3::Union{Nothing,Float64} = nothing # Poloidal beta at the axis + li1::Union{Nothing,Float64} = nothing # Internal inductance at the axis + li2::Union{Nothing,Float64} = nothing # External inductance at the axis + li3::Union{Nothing,Float64} = nothing # Total inductance at the axis +end diff --git a/src/Equilibrium/Construction.jl b/src/Equilibrium/Construction.jl index aac97140d..c715383ce 100644 --- a/src/Equilibrium/Construction.jl +++ b/src/Equilibrium/Construction.jl @@ -1,37 +1,7 @@ # Construction.jl # -# How a PlasmaEquilibrium is formed: the analytic registry, the solver input for each source type, -# the constructors, and re-forming on a new grid. - -""" - AnalyticEqSpec - -Single-source-of-truth entry describing one analytic-equilibrium kind: the `gpec.toml` -section that carries its parameters, the `*Config` type built from that section, and the -run function that turns the config into solver input. `tj_analytic` and `tj_analytic_direct` -share `TJAnalyticConfig` but bind different run functions, so `run_fn` is per-entry rather -than derived from `config_type`. - -## Fields - - - `section::String` — `gpec.toml` section key (e.g. `"SOL_INPUT"`) - - `config_type::DataType` — `*Config` type, constructed via its `(::Dict)` / `(::String)` ctor - - `run_fn::Function` — `(EquilibriumConfig, config) -> eq_input` solver entry point -""" -struct AnalyticEqSpec - section::String - config_type::DataType - run_fn::Function -end - -# Registry of analytic equilibrium kinds. Both `run_input` (fresh runs) and the rerun input -# builder dispatch off this table, so adding a new analytic kind is one new row. -const ANALYTIC_EQ = Dict( - "sol" => AnalyticEqSpec("SOL_INPUT", SolovevConfig, sol_run), - "lar" => AnalyticEqSpec("LAR_INPUT", LargeAspectRatioConfig, lar_run), - "tj_analytic" => AnalyticEqSpec("TJ_ANALYTIC_INPUT", TJAnalyticConfig, tj_analytic_run), - "tj_analytic_direct" => AnalyticEqSpec("TJ_ANALYTIC_INPUT", TJAnalyticConfig, tj_analytic_run_direct) -) +# How a PlasmaEquilibrium is formed: the solver input for each source type, the constructors, +# and re-forming on a new grid. """ run_input(eq_config::EquilibriumConfig, source) -> DirectRunInput or InverseRunInput diff --git a/src/Equilibrium/CoordinateInvariant.jl b/src/Equilibrium/Derived/CoordinateInvariant.jl similarity index 100% rename from src/Equilibrium/CoordinateInvariant.jl rename to src/Equilibrium/Derived/CoordinateInvariant.jl diff --git a/src/Equilibrium/FluxSurfaceMetrics.jl b/src/Equilibrium/Derived/FluxSurfaceMetrics.jl similarity index 100% rename from src/Equilibrium/FluxSurfaceMetrics.jl rename to src/Equilibrium/Derived/FluxSurfaceMetrics.jl diff --git a/src/Equilibrium/GeometryProfiles.jl b/src/Equilibrium/Derived/GeometryProfiles.jl similarity index 100% rename from src/Equilibrium/GeometryProfiles.jl rename to src/Equilibrium/Derived/GeometryProfiles.jl diff --git a/src/Equilibrium/Derived/GlobalParameters.jl b/src/Equilibrium/Derived/GlobalParameters.jl new file mode 100644 index 000000000..1fd98aad2 --- /dev/null +++ b/src/Equilibrium/Derived/GlobalParameters.jl @@ -0,0 +1,225 @@ +""" + equilibrium_global_parameters!(pe::PlasmaEquilibrium) + +Computes and populates global equilibrium parameters in the `PlasmaEquilibrium` +struct, such as rmean, amean, kappa, bt0, crnt, betat, betan, li1, etc. Performs +the same function as equil_out_global in the Fortran code. +""" +function equilibrium_global_parameters!(pe::PlasmaEquilibrium) + profiles = pe.profiles + mpsi = length(pe.rzphi_xs) - 1 + mtheta = length(pe.rzphi_ys) - 1 + + # Use separatrix geometry + rsep, zsep, rext, _ = equilibrium_separatrix_find!(pe) + + rmean = (rsep[1] + rsep[2]) / 2 + amean = (rsep[1] - rsep[2]) / 2 + aratio = rmean / amean + kappa = (zsep[1] - zsep[2]) / (rsep[1] - rsep[2]) + delta1 = (rmean - rext[1]) / amean + delta2 = (rmean - rext[2]) / amean + dpsi = 1.0 - pe.rzphi_xs[mpsi+1] + psi_edge = profiles.xs[end] + bt0 = (profiles.F_spline.y[end] + profiles.F_deriv(psi_edge; hint=Ref(profiles.npts_minus_1)) * dpsi) / (2π * rmean) + + pe.params.rmean = rmean + pe.params.amean = amean + pe.params.aratio = aratio + pe.params.kappa = kappa + pe.params.delta1 = delta1 + pe.params.delta2 = delta2 + pe.params.bt0 = bt0 + + psio = pe.psio + gs1 = zeros(Float64, mtheta + 1) + gs2 = zeros(Float64, mtheta + 1) + + # Direct array access at edge flux surface grid points using nodal_derivs + for itheta in 0:mtheta + # Function values (partials[1,:,:]) + r2 = pe.rzphi_rsquared.nodal_derivs.partials[1, end, itheta+1] + offset = pe.rzphi_offset.nodal_derivs.partials[1, end, itheta+1] + jac = pe.rzphi_jac.nodal_derivs.partials[1, end, itheta+1] + # Theta derivatives (partials[3,:,:] = ∂f/∂y) + r2_y = pe.rzphi_rsquared.nodal_derivs.partials[3, end, itheta+1] + offset_y = pe.rzphi_offset.nodal_derivs.partials[3, end, itheta+1] + + chi1 = 2π * psio / jac + jacfac = π / jac + rfac = sqrt(r2) + eta = 2π * (pe.rzphi_ys[itheta+1] + offset) + r = pe.ro + rfac * cos(eta) + v21 = jacfac * r2_y / (2π * rfac) + v22 = jacfac * (1 + offset_y) * (2 * rfac) + v33 = jacfac * 2π * (r / π) + dvsq = (v21^2 + v22^2) * (v33 * jac^2)^2 + gs1[itheta+1] = sqrt(dvsq) / (2π * r) + gs2[itheta+1] = chi1 * dvsq / (2π * r)^2 + end + + int1 = sum(gs1) / (mtheta + 1) + int2 = sum(gs2) / (mtheta + 1) + crnt = int2 / (1e6 * mu0) + bp0 = int2 / int1 + bwall = 1e6 * mu0 * crnt / (2π * amean) + + pe.params.crnt = crnt + pe.params.bwall = bwall + + # Flux-surface integrals over the full normalized flux [0, 1]; ExtendExtrap carries the integral past the first/last grid points (integrands are smooth/vanishing at the axis). + xs = profiles.xs + fsi(y) = FastInterpolations.integrate(cubic_interp(xs, y; extrap=ExtendExtrap()), 0.0, 1.0) + P_vals = profiles.P_spline.y + dVdpsi_vals = profiles.dVdpsi_spline.y + + fsi_pdv = fsi(P_vals .* dVdpsi_vals) # ∫ p dV/dψ + fsi_dv = fsi(dVdpsi_vals) # ∫ dV/dψ + fsi_p2dv = fsi(P_vals .^ 2 .* dVdpsi_vals) # ∫ p² dV/dψ + volume = fsi_dv # same integrand as hs col 2 in Fortran + + # Poloidal-field surface integral hs_bp2(ψ) = ψ₀² ∮dθ |∇ψ|² / (R² J). + # This is Fortran equil_out.f's hs%fs(:,3) and is the correct integrand for + # the internal inductance li1/li2/li3 (distinct from the p²·dV integrand + # that drives betaj). + hs_bp2 = zeros(Float64, mpsi + 1) + for ipsi in 0:mpsi + acc = 0.0 + for itheta in 0:mtheta + r2 = pe.rzphi_rsquared.nodal_derivs.partials[1, ipsi+1, itheta+1] + offset = pe.rzphi_offset.nodal_derivs.partials[1, ipsi+1, itheta+1] + jac = pe.rzphi_jac.nodal_derivs.partials[1, ipsi+1, itheta+1] + r2_y = pe.rzphi_rsquared.nodal_derivs.partials[3, ipsi+1, itheta+1] + offset_y = pe.rzphi_offset.nodal_derivs.partials[3, ipsi+1, itheta+1] + + jacfac = π / jac + rfac = sqrt(r2) + eta = 2π * (pe.rzphi_ys[itheta+1] + offset) + r = pe.ro + rfac * cos(eta) + v21 = jacfac * r2_y / (2π * rfac) + v22 = jacfac * (1 + offset_y) * (2 * rfac) + v33 = jacfac * 2π * (r / π) + dvsq = (v21^2 + v22^2) * (v33 * jac^2)^2 + acc += dvsq / (r^2) / jac + end + # Periodic trapezoidal rule on uniform θ grid reduces to a plain mean + # because the first and last grid points coincide — matches the int1/int2 + # pattern used for the edge-surface integrals above. + hs_bp2[ipsi+1] = (acc / (mtheta + 1)) * psio^2 + end + fsi_bp2 = fsi(hs_bp2) + + p0 = P_vals[1] - profiles.P_deriv(profiles.xs[1]; hint=Ref(1)) * profiles.xs[1] # linear extrapolation + betat = 2 * (fsi_pdv / fsi_dv) / bt0^2 + betaj = 2 * sqrt(fsi_p2dv / fsi_dv) / bwall^2 + betan = 100 * amean * bt0 * betat / crnt + betap1 = 2 * (fsi_pdv / fsi_dv) / bp0^2 + betap2 = 4 * fsi_pdv / ((1e6 * mu0 * crnt)^2 * pe.ro) + betap3 = 4 * fsi_pdv / ((1e6 * mu0 * crnt)^2 * rmean) + li1 = fsi_bp2 / fsi_dv / bp0^2 + li2 = 2 * fsi_bp2 / ((1e6 * mu0 * crnt)^2 * pe.ro) + li3 = 2 * fsi_bp2 / ((1e6 * mu0 * crnt)^2 * rmean) + + pe.params.psi0 = psio + pe.params.psi_axis = pe.psio + pe.params.psi_boundary = 1.0 + pe.params.psi_boundary_norm = 1.0 + pe.params.psi_axis_norm = 0.0 + pe.params.psi_norm = 0.0 + pe.params.psi_axis_offset = 0.0 + pe.params.psi_boundary_offset = 0.0 + pe.params.psi_axis_sign = 1 + pe.params.psi_boundary_sign = -1 + pe.params.psi_boundary_zero = false + + pe.params.q0 = profiles.q_spline.y[1] + pe.params.b0 = bt0 + + pe.params.volume = volume + pe.params.betat = betat + pe.params.betan = betan + pe.params.betaj = betaj + pe.params.betap1 = betap1 + pe.params.betap2 = betap2 + pe.params.betap3 = betap3 + pe.params.li1 = li1 + pe.params.li2 = li2 + pe.params.li3 = li3 +end + +""" + equilibrium_qfind!(equil::PlasmaEquilibrium) + +Finds the extrema of the safety factor profile q(ψ) in the plasma equilibrium +and computes derived q-values such as q0, qmin, qmax, qa, and q95. Performs +the same function as equil_out_qfind in the Fortran code. +""" +function equilibrium_qfind!(equil::PlasmaEquilibrium) + + profiles = equil.profiles + xs = profiles.xs + mpsi = length(xs) - 1 + psiexl = Float64[] + qexl = Float64[] + + # Create derivative views for q-spline + q_spline = profiles.q_spline + q_d1 = deriv1(q_spline) + q_d2 = deriv2(q_spline) + q_d3 = deriv3(q_spline) + + # Left endpoint + push!(psiexl, xs[1]) + push!(qexl, q_spline.y[1]) + + # Search for extrema in q(ψ) + for ipsi in 1:mpsi + x0 = xs[ipsi] + x1 = xs[ipsi+1] + xmax = x1 - x0 + + a = q_spline(x0) + b = q_d1(x0) + c = q_d2(x0) + d = q_d3(x0) + + if d != 0.0 + xcrit = -c / d + dx2 = xcrit^2 - 2b / d + if dx2 ≥ 0 + dx = sqrt(dx2) + for delta in (dx, -dx) + x = xcrit - delta + if 0 ≤ x < xmax + ψ = x0 + x + push!(psiexl, ψ) + push!(qexl, q_spline(ψ)) + end + end + end + end + end + + # Right endpoint + push!(psiexl, xs[end]) + push!(qexl, q_spline.y[end]) + + equil.params.qextrema_psi = psiexl + equil.params.qextrema_q = qexl + equil.params.mextrema = length(psiexl) + # Compute derived q-values + q0 = q_spline.y[1] - profiles.q_deriv(xs[1]; hint=Ref(1)) * xs[1] + qmax_edge = q_spline.y[end] + qmin = min(minimum(qexl), q0) + qmax = max(maximum(qexl), qmax_edge) + qa = q_spline.y[end] + profiles.q_deriv(xs[end]; hint=Ref(profiles.npts_minus_1)) * (1.0 - xs[end]) + + q95 = q_spline(0.95) + + # Store derived values + equil.params.q0 = q0 + equil.params.qmin = qmin + equil.params.qmax = qmax + equil.params.qa = qa + equil.params.q95 = q95 +end diff --git a/src/Equilibrium/Equilibrium.jl b/src/Equilibrium/Equilibrium.jl index c68583274..840678f34 100644 --- a/src/Equilibrium/Equilibrium.jl +++ b/src/Equilibrium/Equilibrium.jl @@ -19,22 +19,29 @@ const mu0 = 4π * 1e-7 const EFIT_KINDS = ("efit", "efit_arclength", "efit_by_inversion") # --- Internal Module Structure --- -include("EquilibriumTypes.jl") -include("SolveTolerances.jl") -include("GridRefinement.jl") -include("SeparatrixQLaw.jl") -include("FluxSurfaceMetrics.jl") -include("CoordinateInvariant.jl") -include("ReadEquilibrium.jl") -include("DirectEquilibrium.jl") -include("DirectEquilibriumArcLength.jl") -include("DirectEquilibriumByInversion.jl") -include("InverseEquilibrium.jl") -include("AnalyticEquilibrium.jl") -include("GeometryProfiles.jl") +include("Config.jl") +include("Sources/Ingest.jl") +include("Sources/Analytic/Analytic.jl") +include("PlasmaEquilibrium.jl") +include("Axisymmetric/SolveTolerances.jl") +include("Grid/GridRefinement.jl") +include("Grid/PsiGrid.jl") +include("Axisymmetric/Separatrix.jl") +include("Derived/FluxSurfaceMetrics.jl") +include("Derived/CoordinateInvariant.jl") +include("Sources/GFile.jl") +include("Sources/Chease.jl") +include("Sources/IFile.jl") +include("Sources/IMAS.jl") +include("Axisymmetric/Direct.jl") +include("Axisymmetric/DirectArcLength.jl") +include("Axisymmetric/DirectByInversion.jl") +include("Axisymmetric/Inverse.jl") +include("Derived/GeometryProfiles.jl") include("KineticProfiles.jl") include("Construction.jl") -include("PostProcessing.jl") +include("Derived/GlobalParameters.jl") +include("Axisymmetric/GradShafranovResidual.jl") # --- Expose types and functions to the user --- export reform, EquilibriumConfig, PlasmaEquilibrium, EquilibriumParameters, AnalyticEquilibriumConfig, diff --git a/src/Equilibrium/EquilibriumTypes.jl b/src/Equilibrium/EquilibriumTypes.jl deleted file mode 100644 index 2b51f7688..000000000 --- a/src/Equilibrium/EquilibriumTypes.jl +++ /dev/null @@ -1,923 +0,0 @@ -function symbolize_keys(dict::Dict{String,Any}) - return Dict(Symbol(k) => v for (k, v) in dict) -end - -""" - EquilibriumConfig(...) - -An immutable struct containing configuration parameters for equilibrium reconstruction -specified in the input. - -## Fields - - - `eq_type::String` - Type of equilibrium file ("efit", "solovev", "lar", etc.) - - `eq_filename::String` - Path to equilibrium input file - - `jac_type::String` - Jacobian coordinate type ("hamada", "pest", "equal_arc", "boozer", "park", "custom") - - `power_bp::Int` - Poloidal field power exponent for Jacobian (internal, derived from `jac_type`; deprecated as a TOML input key) - - `power_b::Int` - Total field power exponent for Jacobian (internal, derived from `jac_type`; deprecated as a TOML input key) - - `power_r::Int` - Major radius power exponent for Jacobian (internal, derived from `jac_type`; deprecated as a TOML input key) - - `power_rc::Int` - Minor radius (rfac = √((R-R₀)²+(Z-Z₀)²)) power exponent for Jacobian (internal, derived from `jac_type`; deprecated as a TOML input key) - - `jac_custom_power_bp::Int` - Poloidal-field exponent for a custom Jacobian `J = Bp^bp · B^b / (R^r · rfac^rc)`; read only when `jac_type = "custom"`, ignored for named types - - `jac_custom_power_b::Int` - Total-field exponent for a custom Jacobian; read only when `jac_type = "custom"`, ignored for named types - - `jac_custom_power_r::Int` - Major-radius exponent for a custom Jacobian; read only when `jac_type = "custom"`, ignored for named types - - `jac_custom_power_rc::Int` - Minor-radius (rfac) exponent for a custom Jacobian; read only when `jac_type = "custom"`, ignored for named types - - `r0_norm::Float64` - Major radius normalization for CHEASE/EQDSK [m] - - `b0_norm::Float64` - On-axis toroidal field normalization for CHEASE/EQDSK [T] - - `grid_type::String` - Grid type for flux surface discretization ("auto" — two-pass measured-curvature - refinement when mpsi=0, three-region log layout when mpsi>0; "core_edge_packed" (sin² spacing, dense at axis and edge), "pow1", "uniform"; - inverse equilibria also "rho" (uniform in √ψ from psilow to psihigh) and "original" (the input - file's own ψ nodes off the axis; ignores psilow and mpsi); "log_asymptotic" is a legacy alias for "auto") - - `psilow::Float64` - Lower limit of normalized flux coordinate - - `psihigh::Float64` - Requested upper limit of normalized flux coordinate; the value the - equilibrium is actually formed on is `EquilibriumParameters.psihigh_resolved`. - - `mpsi::Int` - Number of radial grid intervals; 0 with grid_type="auto" selects the - two-pass auto grid: the main driver forms a coarse pass-1 equilibrium, measures its curvature, - pins knots on rational surfaces, and re-forms on the refined grid. Standalone `PlasmaEquilibrium` - callers get the coarse pass-1 grid unless they refine via `refined_psi_grid` + `override_psi_nodes`. - - `psi_accuracy::Float64` - Target relative accuracy τ of splined profile derivatives for the - two-pass auto grid (knot count scales as τ^(-1/3)) - - `mtheta::Int` - Number of poloidal grid points - - `newq0::Float64` - Target on-axis safety factor q(0); the q and F profiles are rescaled to - meet it (0 = use input value, -1 = use the axis extrapolation with its sign flipped) - - `flux_surface_rtol::Float64` - Relative tolerance of the ODE solves that trace each flux surface; their absolute tolerance - is `equil_abstol(flux_surface_rtol)`, i.e. `flux_surface_rtol` capped at `EQUIL_ABSTOL_MAX` - - `force_termination::Bool` - Terminate after equilibrium setup (skip stability calculations) -""" -@kwdef struct EquilibriumConfig - eq_type::String = "efit" - eq_filename::String = "mypath" - r0_norm::Float64 = 1.0 - b0_norm::Float64 = 1.0 - - jac_type::String = "hamada" - power_bp::Int = 0 - power_b::Int = 0 - power_r::Int = 0 - power_rc::Int = 0 - - jac_custom_power_bp::Int = 0 - jac_custom_power_b::Int = 0 - jac_custom_power_r::Int = 0 - jac_custom_power_rc::Int = 0 - - grid_type::String = "auto" - psilow::Float64 = 1e-2 - psihigh::Float64 = 0.9995 - mpsi::Int = 0 - psi_accuracy::Float64 = 0.001 - mtheta::Int = 512 - - newq0::Float64 = 0.0 - flux_surface_rtol::Float64 = 1e-10 - - force_termination::Bool = false - - # IMAS-specific: expected COCOS convention of the input dd.equilibrium (11=IMAS standard, 2=GPEC internal) - imas_cocos::Int = 11 - - """ - Modified internal constructor that enforces self consistency within the inputs - """ - # The four `_` slots are the `power_bp/power_b/power_r/power_rc` struct fields, which @kwdef - # forwards positionally. They are always derived below from `jac_type` (or `jac_custom_power_*`), - # so their incoming values are ignored (hence `_`). - function EquilibriumConfig(eq_type, eq_filename, r0_norm, b0_norm, jac_type, _, _, _, _, - jac_custom_power_bp, jac_custom_power_b, jac_custom_power_r, jac_custom_power_rc, - grid_type, psilow, psihigh, mpsi, psi_accuracy, mtheta, newq0, flux_surface_rtol, - force_termination, imas_cocos) - if jac_type == "hamada" - @info "Forcing hamada coordinate jacobian exponents: power_*" - power_b = 0 - power_bp = 0 - power_r = 0 - power_rc = 0 - elseif jac_type == "pest" - @info "Forcing pest coordinate jacobian exponents: power_*" - power_b = 0 - power_bp = 0 - power_r = 2 - power_rc = 0 - elseif jac_type == "equal_arc" - @info "Forcing equal_arc coordinate jacobian exponents: power_*" - power_b = 0 - power_bp = 1 - power_r = 0 - power_rc = 0 - elseif jac_type == "boozer" - @info "Forcing boozer coordinate jacobian exponents: power_*" - power_b = 2 - power_bp = 0 - power_r = 0 - power_rc = 0 - elseif jac_type == "park" - @info "Forcing park coordinate jacobian exponents: power_*" - power_b = 1 - power_bp = 0 - power_r = 0 - power_rc = 0 - elseif jac_type == "custom" - # Custom Jacobian: the user-facing knobs define the exponents. - power_bp = jac_custom_power_bp - power_b = jac_custom_power_b - power_r = jac_custom_power_r - power_rc = jac_custom_power_rc - # Normalize to a named type when the powers match, so fast paths are taken. - if power_b == 0 && power_bp == 0 && power_r == 0 && power_rc == 0 - jac_type = "hamada" - @info "Recognized hamada jacobian from power exponents" - elseif power_b == 0 && power_bp == 0 && power_r == 2 && power_rc == 0 - jac_type = "pest" - @info "Recognized pest jacobian from power exponents" - elseif power_b == 0 && power_bp == 1 && power_r == 0 && power_rc == 0 - jac_type = "equal_arc" - @info "Recognized equal_arc jacobian from power exponents" - elseif power_b == 2 && power_bp == 0 && power_r == 0 && power_rc == 0 - jac_type = "boozer" - @info "Recognized boozer jacobian from power exponents" - elseif power_b == 1 && power_bp == 0 && power_r == 0 && power_rc == 0 - jac_type = "park" - @info "Recognized park jacobian from power exponents" - else - @info "Using manual jacobian exponents: power b, bp, r, rc = $(power_b), $(power_bp), $(power_r), $(power_rc)" - end - else - error("Cannot recognize jac_type = $(jac_type)") - end - (isfinite(flux_surface_rtol) && flux_surface_rtol > 0) || error("flux_surface_rtol = $flux_surface_rtol must be a positive, finite relative tolerance") - if psihigh > 1.0 - @warn "psihigh = $psihigh exceeds 1.0 (separatrix); clamping to 1.0" - end - psihigh = min(psihigh, 1.0) - return new(eq_type, eq_filename, r0_norm, b0_norm, jac_type, power_bp, power_b, power_r, power_rc, - jac_custom_power_bp, jac_custom_power_b, jac_custom_power_r, jac_custom_power_rc, - grid_type, psilow, psihigh, mpsi, psi_accuracy, mtheta, newq0, flux_surface_rtol, - force_termination, imas_cocos) - end -end - -""" -Outer constructor for EquilibriumConfig from a parsed TOML dictionary -""" -function EquilibriumConfig(equil_dict::Dict{String,Any}, base_path::String="./") - # `eq_type` is always required. `eq_filename` is required for file-based - # equilibria (efit, chease, …) but optional for analytic types whose - # parameters live in an embedded `[TJ_ANALYTIC_INPUT]` / `[SOL_INPUT]` / - # `[LAR_INPUT]` section of the parent gpec.toml. - if !haskey(equil_dict, "eq_type") - error("Missing required key in [Equilibrium]: eq_type") - end - - # Filter to only known parameters - config_fields = Set(String.(fieldnames(EquilibriumConfig))) - config_data = Dict{String,Any}() - - for (k, v) in equil_dict - if k in config_fields - config_data[k] = v - else - @warn "Unknown equilibrium parameter: $k" - end - end - - # Only resolve `eq_filename` against `base_path` if the user actually supplied one - # (otherwise leave the kwdef sentinel for the embedded path). The empty string is the - # rerun path's "no input file" marker and must stay empty, not become `base_path`. - if haskey(config_data, "eq_filename") && !isempty(config_data["eq_filename"]) && !isabspath(config_data["eq_filename"]) - config_data["eq_filename"] = normpath(joinpath(base_path, config_data["eq_filename"])) - end - - return EquilibriumConfig(; symbolize_keys(config_data)...) -end - -""" - AnalyticEquilibriumConfig - -Supertype of the parameter structs of the analytic equilibrium kinds. Each subtype is the -`config_type` of one or more `ANALYTIC_EQ` rows, which pair it with an `eq_type`, its TOML -section, and its solver. -""" -abstract type AnalyticEquilibriumConfig end - -""" - LargeAspectRatioConfig(...) - -A mutable struct holding parameters for the Large Aspect Ratio (LAR) plasma equilibrium model. - -## Fields: - - - `r0`: The major radius of the plasma [m]. - - `a`: The minor radius of the plasma [m]. - - `beta0`: The beta value on axis (normalized pressure). - - `q0`: The safety factor on axis. - - `p_pres`: The exponent for the pressure profile, defined as `p00 * (1 - (r / a)^2)^p_pres`. - - `p_sig`: The exponent that determines the shape of the current-related function profile. - - `sigma_type`: The type of sigma profile, can be "default" or "wesson". If "wesson", the sigma profile is defined as `sigma0 * (1 - (r / a)^2)^p_sig`. - - `mtheta`: The number of grid points in the poloidal direction. - - `ma`: The number of grid points in the radial direction. - - `zeroth`: If set to true, it neglects the Shafranov shift -""" -@kwdef mutable struct LargeAspectRatioConfig <: AnalyticEquilibriumConfig - r0::Float64 = 10.0 - a::Float64 = 1.0 - beta0::Float64 = 1e-3 - q0::Float64 = 1.5 - qa::Float64 = 3.6 # Edge safety factor (legacy field; not consumed by current sigma_type options) - B0::Float64 = 1.0 # On-axis toroidal field [T] (scales F and P) - p_pres::Float64 = 2.0 - p_sig::Float64 = 1.0 - sigma_type::String = "default" - mtheta::Int = 128 - ma::Int = 128 - zeroth::Bool = false -end - -""" -Build a `LargeAspectRatioConfig` from a parsed `[LAR_INPUT]` TOML table. -""" -function LargeAspectRatioConfig(input_dict::Dict{String,Any}) - return LargeAspectRatioConfig(; symbolize_keys(input_dict)...) -end - -""" - TJAnalyticConfig(...) - -Parameters for the **TJ-analytic** cylindrical large-aspect-ratio equilibrium -model — a GPEC adaptation of the analytic profile family used by -R. Fitzpatrick's TJ code (https://github.com/rfitzp/TJ). We follow the -same analytic-profile parameterization (ψ-ODE in dimensionless r/a, f₁ -for q, power-law pressure) for the inner cylindrical core and connect it -to GPEC's direct-GS pipeline; this is NOT a re-implementation of TJ. - -The model uses analytic profiles with exact control of both the on-axis -and edge safety factors. The q profile is determined by: - - f1(r) = [1 - (1-r²)^ν] / (ν·q0) - q(r) = r² / f1(r) - -where ν = qa/q0 is the current peaking parameter, q0 is the axis q, and qa -is the edge q. All lengths are normalized to R₀, fields to B₀. The pressure -profile is p₂(r) = p0·(1-r²)^μ with μ = `pressure_peaking`. - -Reference: R. Fitzpatrick, TJ code, https://github.com/rfitzp/TJ -""" -@kwdef mutable struct TJAnalyticConfig <: AnalyticEquilibriumConfig - r0::Float64 = 10.0 # Major radius R₀ [m] - a::Float64 = 1.0 # Minor radius a [m] (ε = a/R₀) - q0::Float64 = 1.5 # On-axis safety factor - qa::Float64 = 3.6 # Edge safety factor - p0::Float64 = 0.001 # Normalized on-axis pressure - pressure_peaking::Float64 = 2.0 # Pressure peaking exponent μ: p₂ = p0·(1-r²)^μ - B0::Float64 = 12.0 # On-axis toroidal field [T] - ma::Int = 128 # Radial grid points - mtheta::Int = 128 # Poloidal grid points - zeroth::Bool = false # If true, suppress Shafranov shift -end - -""" -Build a `TJAnalyticConfig` from a parsed `[TJ_ANALYTIC_INPUT]` TOML table. -""" -function TJAnalyticConfig(input_dict::Dict{String,Any}) - return TJAnalyticConfig(; symbolize_keys(input_dict)...) -end - -""" - SolovevConfig(...) - -A mutable struct holding parameters for the Solev'ev (SOL) plasma equilibrium model. - -## Fields: - - - `mr`: number of radial grid zones - - `mz`: number of axial grid zones - - `ma`: number of flux grid zones - - `elongation`: elongation - - `a`: minor radius - - `r0`: major radius - - `q0`: safety factor at the o-point - - `p0fac`: scale on-axis pressure (P-> P+P0*p0fac. beta changes. Phi,q constant) - - `b0fac`: scale toroidal field at constant beta (s*Phi,s*f,s^2*P. bt changes. Shape,beta constant) - - `f0fac`: scale toroidal field at constant pressure (s*f. beta,q changes. Phi,p,bp constant) -""" -@kwdef mutable struct SolovevConfig <: AnalyticEquilibriumConfig - mr::Int = 128 # number of radial grid zones - mz::Int = 128 # number of axial grid zones - ma::Int = 128 # number of flux grid zones - elongation::Float64 = 1.6 # elongation - a::Float64 = 0.33 # minor radius - r0::Float64 = 1.0 # major radius - q0::Float64 = 1.9 # safety factor at the o-point - p0fac::Float64 = 1 # scale on-axis pressure (P-> P+P0*p0fac. beta changes. Phi,q constant) - b0fac::Float64 = 1 # scale toroidal field at constant beta (s*Phi,s*f,s^2*P. bt changes. Shape,beta constant) - f0fac::Float64 = 1 # scale toroidal field at constant pressure (s*f. beta,q changes. Phi,p,bp constant) -end - -""" -Build a `SolovevConfig` from a parsed `[SOL_INPUT]` TOML table. -""" -function SolovevConfig(input_dict::Dict{String,Any}) - return SolovevConfig(; symbolize_keys(input_dict)...) -end - -""" - DirectIngest - -The serializable raw arrays and scalars captured by a direct-equilibrium reader -(`read_efit`, `read_imas`, `sol_run` is analytic and does not use this) — everything -needed to rebuild a `DirectRunInput`'s splines without re-reading the original g-file. -Stored on `DirectRunInput.ingest` / `PlasmaEquilibrium.ingest` and dumped to -`gpec.h5` so a run can be replayed from any directory. The interpolants themselves are -not serializable; they are reconstructed from these nodes by `build_direct_from_ingest`. - -## Fields - - - `sq_xs::Vector{Float64}` — normalized-ψ knots for the 1D profile spline - - `sq_fs::Matrix{Float64}` — 1D profile node values (F, μ₀P, q, √ψ_norm) - - `psi_xs::Vector{Float64}` — R grid for the 2D flux map [m] - - `psi_ys::Vector{Float64}` — Z grid for the 2D flux map [m] - - `psi_rz::Matrix{Float64}` — processed poloidal flux on the (R, Z) grid [Wb/rad] - - `rmin/rmax/zmin/zmax::Float64` — computational-grid bounds [m] - - `psio::Float64` — total flux difference |ψ_axis - ψ_boundary| [Wb/rad] - - `bt_sign::Int` — sign of the toroidal field (+1 or -1) - - `ip_sign::Int` — sign of the plasma current (+1 or -1) as the source file states it -""" -struct DirectIngest - sq_xs::Vector{Float64} - sq_fs::Matrix{Float64} - psi_xs::Vector{Float64} - psi_ys::Vector{Float64} - psi_rz::Matrix{Float64} - rmin::Float64 - rmax::Float64 - zmin::Float64 - zmax::Float64 - psio::Float64 - bt_sign::Int - ip_sign::Int -end - -""" - InverseIngest - -The serializable raw arrays and scalars captured by an inverse-equilibrium reader -(`read_chease_ascii`, `read_chease_binary`, `read_ifile`) — everything needed to rebuild an -`InverseRunInput`'s splines without re-reading the original CHEASE file or i-file. Stored on -`InverseRunInput.ingest` / `PlasmaEquilibrium.ingest` and reconstructed by -`build_inverse_from_ingest`. See [`DirectIngest`](@ref) for the role this plays in -the `gpec.h5` rerun snapshot. - -## Fields - - - `sq_xs::Vector{Float64}` — normalized-ψ knots for the 1D profile spline - - `sq_fs::Matrix{Float64}` — 1D profile node values - - `rz_xs::Vector{Float64}` — ψ grid for the R, Z maps - - `rz_ys::Vector{Float64}` — θ grid for the R, Z maps - - `R_nodes::Matrix{Float64}` — R node values on the (ψ, θ) grid [m] - - `Z_nodes::Matrix{Float64}` — Z node values on the (ψ, θ) grid [m] - - `ro::Float64` — R of magnetic axis [m] - - `zo::Float64` — Z of magnetic axis [m] - - `psio::Float64` — total flux difference |ψ_axis - ψ_boundary| [Wb/rad] -""" -struct InverseIngest - sq_xs::Vector{Float64} - sq_fs::Matrix{Float64} - rz_xs::Vector{Float64} - rz_ys::Vector{Float64} - R_nodes::Matrix{Float64} - Z_nodes::Matrix{Float64} - ro::Float64 - zo::Float64 - psio::Float64 -end - -# Equilibria captured for replay carry one of these; analytic equilibria carry `nothing` -# and are regenerated from their TOML section instead. -const EquilibriumIngest = Union{Nothing,DirectIngest,InverseIngest} - -""" - DirectRunInput(...) - -A container struct that bundles all necessary inputs for the `direct_run` function. -It is created by one of the equilibrium file read-in functions after processing the -raw equilibrium data and preparing the initial splines. - -## Fields - - - `config::EquilibriumConfig` - The equilibrium configuration object. - - - `sq_in` - 1D spline data versus normalized poloidal flux `psin`. - Quantities: - - 1. `F = R * B_t` — toroidal flux function [m·T] - 2. `μ₀ * Pressure` — plasma pressure (non-negative) [T²] - 3. `q` — safety factor profile - 4. `√ψ_norm` — square root of normalized flux - - `psi_in` - 2D cubic interpolant on the (R, Z) grid [m]. - The values correspond to the **poloidal flux** adjusted to be zero at the boundary [Wb/rad]. - Definitions: - - 1. `ψ(R, Z) = ψ_boundary - ψ(R, Z)` - - 2. `ψ = ψ * sign(ψ(centerR, centerZ))` - - * 1D profiles are represented by `CubicInterpolant` or `CubicSeriesInterpolant` - * 2D flux surfaces by `CubicInterpolantND` - - `psi_in_xs::Vector{Float64}` — R coordinate grid for psi_in [m] - - `psi_in_ys::Vector{Float64}` — Z coordinate grid for psi_in [m] - - `rmin::Float64` — Minimum R-coordinate of the computational grid [m] - - `rmax::Float64` — Maximum R-coordinate of the computational grid [m] - - `zmin::Float64` — Minimum Z-coordinate of the computational grid [m] - - `zmax::Float64` — Maximum Z-coordinate of the computational grid [m] - - `psio::Float64` — Total flux difference `|ψ_axis - ψ_boundary|` [Wb/rad] - - `bt_sign::Int` — Sign of the toroidal field (+1 or -1); read from fpol sign in EFIT g-files - - `ip_sign::Int` — Sign of the plasma current (+1 or -1); read from the `current` header value - in EFIT g-files and `global_quantities.ip` in IMAS. The internal flux is always made - positive at the axis, so the sign is kept here (never recovered from the computed current) - and fixes the SFL→machine toroidal-angle handedness `helicity = bt_sign × ip_sign`. - - `ingest::EquilibriumIngest` — captured raw arrays for the `gpec.h5` rerun snapshot - (a [`DirectIngest`](@ref) for file-based reads, or `nothing` for analytic equilibria) - - `psihigh_resolved::Float64` — outer flux limit the equilibrium is formed on: `config.psihigh` - clamped to the outermost closed flux surface by [`resolve_psihigh!`](@ref). Defaults to - `config.psihigh` and only differs for efit-family equilibria whose requested limit falls - outside the closed-flux region. The solvers build their ψ grid from this field. IMAS - equilibria are read into this struct but are not in `EFIT_KINDS`, so they are never - clamped and always keep the request. -""" -mutable struct DirectRunInput{S<:FastInterpolations.CubicSeriesInterpolant,I2D<:FastInterpolations.CubicInterpolantND} - config::EquilibriumConfig - sq_in::S # 1D profile spline for F, P, q - psi_in::I2D # 2D flux interpolant (CubicInterpolantND) - psi_in_xs::Vector{Float64} # R coordinates - psi_in_ys::Vector{Float64} # Z coordinates - rmin::Float64 # Minimum R-coordinate of the computational grid [m]. - rmax::Float64 # Maximum R-coordinate of the computational grid [m]. - zmin::Float64 # Minimum Z-coordinate of the computational grid [m]. - zmax::Float64 # Maximum Z-coordinate of the computational grid [m]. - psio::Float64 # The total flux difference |ψ_axis - ψ_boundary| [Weber / radian]. - bt_sign::Int # Sign of the toroidal field: +1 or -1 (from fpol sign in g-file) - ip_sign::Int # Sign of the plasma current: +1 or -1 (from the g-file current / IMAS ip) - ingest::EquilibriumIngest - psihigh_resolved::Float64 -end - -# Readers construct without a resolved psihigh; it starts at the request and `resolve_psihigh!` -# clamps it for efit-family equilibria. -DirectRunInput(config::EquilibriumConfig, sq_in, psi_in, psi_in_xs, psi_in_ys, - rmin, rmax, zmin, zmax, psio, bt_sign, ip_sign, ingest) = - DirectRunInput(config, sq_in, psi_in, psi_in_xs, psi_in_ys, - rmin, rmax, zmin, zmax, psio, bt_sign, ip_sign, ingest, config.psihigh) - -""" - InverseRunInput(...) - -A container struct for inputs to the `inverse_run` function. - -## Fields - - - `config::EquilibriumConfig` - The equilibrium configuration object - - `sq_in::CubicSeriesInterpolant` - 1D profile spline for F, P, q - - `rz_in_xs::Vector{Float64}` - ψ coordinate grid for rz_in - - `rz_in_ys::Vector{Float64}` - θ coordinate grid for rz_in - - `rz_in_R::CubicInterpolantND` - R coordinate interpolant [m] - - `rz_in_Z::CubicInterpolantND` - Z coordinate interpolant [m] - - `ro::Float64` - R-coordinate of magnetic axis [m] - - `zo::Float64` - Z-coordinate of magnetic axis [m] - - `psio::Float64` - Total flux difference |ψ_axis - ψ_boundary| [Wb/rad] - - `ingest::EquilibriumIngest` - captured raw arrays for the `gpec.h5` rerun snapshot - (an [`InverseIngest`](@ref) for file-based reads, or `nothing` for analytic equilibria) - - `psihigh_resolved::Float64` - outer flux limit the equilibrium is formed on; see - [`DirectRunInput`](@ref). `config.psihigh`, clamped to the outermost input surface for i-files and - reruns (`resolve_psihigh!`); `efit_by_inversion` forwards the clamped value from its `DirectRunInput`. -""" -mutable struct InverseRunInput{S<:FastInterpolations.CubicSeriesInterpolant,I2D<:FastInterpolations.CubicInterpolantND} - config::EquilibriumConfig - sq_in::S # 1D profile spline for F, P, q - rz_in_xs::Vector{Float64} # ψ coordinates - rz_in_ys::Vector{Float64} # θ coordinates - rz_in_R::I2D # R coordinate interpolant - rz_in_Z::I2D # Z coordinate interpolant - ro::Float64 # R axis location - zo::Float64 # Z axis location - psio::Float64 # Total flux difference |psi_axis - psi_boundary| - ingest::EquilibriumIngest - psihigh_resolved::Float64 -end - -InverseRunInput(config::EquilibriumConfig, sq_in, rz_in_xs, rz_in_ys, rz_in_R, rz_in_Z, - ro, zo, psio, ingest) = - InverseRunInput(config, sq_in, rz_in_xs, rz_in_ys, rz_in_R, rz_in_Z, - ro, zo, psio, ingest, config.psihigh) - -""" - EquilibriumParameters - -A mutable struct containing computed equilibrium parameters and diagnostic flags. - -## Fields - - - `ro::Union{Nothing,Float64}` - R-coordinate of the magnetic axis [m] - - `zo::Union{Nothing,Float64}` - Z-coordinate of the magnetic axis [m] - - `psio::Union{Nothing,Float64}` - Total flux difference |ψ_axis - ψ_boundary| [Wb/rad] - - `psihigh_resolved::Union{Nothing,Float64}` - Outer flux limit the equilibrium was formed on - (the outermost ψ node); the plasma edge downstream of `PlasmaEquilibrium`. - - `rsep::Union{Nothing,Vector{Float64}}` - R-coordinates of the plasma boundary [m] - - `zsep::Union{Nothing,Vector{Float64}}` - Z-coordinates of the plasma boundary [m] - - `rext::Union{Nothing,Vector{Float64}}` - R-coordinates of the plasma edge [m] - - `zext::Union{Nothing,Vector{Float64}}` - Z-coordinates of the plasma edge [m] - - `psi0::Union{Nothing,Float64}` - Normalized poloidal flux at reference location - - `b0::Union{Nothing,Float64}` - Total magnetic field strength at the axis [T] - - `q0::Union{Nothing,Float64}` - Safety factor at the axis - - `qmin::Union{Nothing,Float64}` - Minimum safety factor in the plasma - - `qmax::Union{Nothing,Float64}` - Maximum safety factor in the plasma - - `qa::Union{Nothing,Float64}` - Safety factor at the plasma edge - - `q95::Union{Nothing,Float64}` - Safety factor at 95% flux surface - - `qextrema_psi::Union{Nothing,Vector{Float64}}` - Normalized flux values at q extrema - - `qextrema_q::Union{Nothing,Vector{Float64}}` - Safety factor values at extrema - - `mextrema::Union{Nothing,Int}` - Number of extrema in q-profile - - `psi_norm::Union{Nothing,Float64}` - Normalized poloidal flux - - `b_norm::Union{Nothing,Float64}` - Normalized magnetic field strength - - `psi_axis::Union{Nothing,Float64}` - Poloidal flux at the axis - - `psi_boundary::Union{Nothing,Float64}` - Poloidal flux at the boundary - - `psi_boundary_norm::Union{Nothing,Float64}` - Normalized boundary flux - - `psi_axis_norm::Union{Nothing,Float64}` - Normalized axis flux - - `psi_boundary_offset::Union{Nothing,Float64}` - Boundary flux offset - - `psi_axis_offset::Union{Nothing,Float64}` - Axis flux offset - - `psi_boundary_sign::Union{Nothing,Int}` - Sign of boundary flux - - `psi_axis_sign::Union{Nothing,Int}` - Sign of axis flux - - `psi_boundary_zero::Union{Nothing,Bool}` - Whether boundary flux is zero - - `rmean::Union{Nothing,Float64}` - Mean major radius [m] - - `amean::Union{Nothing,Float64}` - Mean minor radius [m] - - `aratio::Union{Nothing,Float64}` - Aspect ratio (R0/a) - - `kappa::Union{Nothing,Float64}` - Plasma elongation - - `delta1::Union{Nothing,Float64}` - Upper triangularity - - `delta2::Union{Nothing,Float64}` - Lower triangularity - - `bt0::Union{Nothing,Float64}` - Toroidal field at axis [T] - - `crnt::Union{Nothing,Float64}` - Plasma current [A] - - `bwall::Union{Nothing,Float64}` - Toroidal field at wall [T] - - `verbose::Bool` - Enable verbose output - - `diagnose_src::Bool` - Enable source data diagnostics - - `diagnose_maxima::Bool` - Enable extrema diagnostics - - `volume::Union{Nothing,Float64}` - Plasma volume [m³] - - `betat::Union{Nothing,Float64}` - Toroidal beta - - `betan::Union{Nothing,Float64}` - Normalized beta - - `betaj::Union{Nothing,Float64}` - Total beta - - `betap1::Union{Nothing,Float64}` - Poloidal beta (definition 1) - - `betap2::Union{Nothing,Float64}` - Poloidal beta (definition 2) - - `betap3::Union{Nothing,Float64}` - Poloidal beta (definition 3) - - `li1::Union{Nothing,Float64}` - Internal inductance (definition 1) - - `li2::Union{Nothing,Float64}` - Internal inductance (definition 2) - - `li3::Union{Nothing,Float64}` - Internal inductance (definition 3) -""" -@kwdef mutable struct EquilibriumParameters - ro::Union{Nothing,Float64} = nothing # R-coordinate of the magnetic axis [m] - zo::Union{Nothing,Float64} = nothing # Z-coordinate of the magnetic axis [m] - psio::Union{Nothing,Float64} = nothing # Total flux difference |ψ_axis - ψ_boundary| [Wb/rad] - psihigh_resolved::Union{Nothing,Float64} = nothing # Outer flux limit actually formed on (clamped config.psihigh) - rsep::Union{Nothing,Vector{Float64}} = nothing # R-coordinates of the plasma boundary [m] - zsep::Union{Nothing,Vector{Float64}} = nothing # Z-coordinates of the plasma boundary [m] - rext::Union{Nothing,Vector{Float64}} = nothing # R-coordinates of the plasma edge [m] - zext::Union{Nothing,Vector{Float64}} = nothing # Z-coordinates of the plasma edge [m] - psi0::Union{Nothing,Float64} = nothing # Normalized poloidal flux - b0::Union{Nothing,Float64} = nothing # Total magnetic field strength at the axis [T] - q0::Union{Nothing,Float64} = nothing # Safety factor at the axis - qmin::Union{Nothing,Float64} = nothing # Minimum safety factor in the plasma - qmax::Union{Nothing,Float64} = nothing # Maximum safety factor in the plasma - qa::Union{Nothing,Float64} = nothing # Safety factor at the plasma edge - q95::Union{Nothing,Float64} = nothing # Safety factor at 95% flux surface - qextrema_psi::Union{Nothing,Vector{Float64}} = nothing # Normalized poloidal flux values where q has extrema - qextrema_q::Union{Nothing,Vector{Float64}} = nothing # Safety factor values at the extrema points - mextrema::Union{Nothing,Int} = nothing # Number of extrema points in the q-profile - psi_norm::Union{Nothing,Float64} = nothing # Normalized poloidal flux at the axis - b_norm::Union{Nothing,Float64} = nothing # Normalized total magnetic field strength at the axis - psi_axis::Union{Nothing,Float64} = nothing # Normalized poloidal flux at the axis - psi_boundary::Union{Nothing,Float64} = nothing # Poloidal flux at the boundary - psi_boundary_norm::Union{Nothing,Float64} = nothing # Normalized poloidal flux at the boundary - psi_axis_norm::Union{Nothing,Float64} = nothing # Normalized poloidal flux at the axis - psi_boundary_offset::Union{Nothing,Float64} = nothing # Offset for the boundary poloidal flux - psi_axis_offset::Union{Nothing,Float64} = nothing # Offset for the axis poloidal flux - psi_boundary_sign::Union{Nothing,Int} = nothing # Sign of the boundary poloidal flux - psi_axis_sign::Union{Nothing,Int} = nothing # Sign of the axis poloidal flux - psi_boundary_zero::Union{Nothing,Bool} = nothing # Whether the boundary poloidal flux is zero - rmean::Union{Nothing,Float64} = nothing # Mean R-coordinate of the plasma [m] - amean::Union{Nothing,Float64} = nothing # Mean minor radius of the plasma [m] - aratio::Union{Nothing,Float64} = nothing # Aspect ratio of the plasma (R0/a) - kappa::Union{Nothing,Float64} = nothing # Elongation of the plasma cross-section - delta1::Union{Nothing,Float64} = nothing # Triangularity of the plasma cross-section (upper triangularity) - delta2::Union{Nothing,Float64} = nothing # Triangularity of the plasma cross-section (lower triangularity) - bt0::Union{Nothing,Float64} = nothing # Toroidal magnetic field at the axis [T] (always positive; sign in bt_sign) - crnt::Union{Nothing,Float64} = nothing # Plasma current at the axis [A] - bt_sign::Int = 1 # Sign of the toroidal field: +1 (positive Bt) or -1 (negative Bt, e.g. DIII-D standard) - ip_sign::Int = 1 # Sign of the plasma current as the source file states it: +1 or -1 (crnt is always positive) - bwall::Union{Nothing,Float64} = nothing # Toroidal magnetic field at the wall [T] - verbose::Bool = false # Whether to print verbose output - diagnose_src::Bool = false # Whether to diagnose source data - diagnose_maxima::Bool = false # Whether to diagnose maxima in the equilibrium - volume::Union{Nothing,Float64} = nothing # Plasma volume [m³] - betat::Union{Nothing,Float64} = nothing # Toroidal beta (normalized pressure) at the axis - betan::Union{Nothing,Float64} = nothing # Normalized beta at the axis - betaj::Union{Nothing,Float64} = nothing # Total beta at the axis - betap1::Union{Nothing,Float64} = nothing # Pressure beta at the axis - betap2::Union{Nothing,Float64} = nothing # Toroidal beta at the axis - betap3::Union{Nothing,Float64} = nothing # Poloidal beta at the axis - li1::Union{Nothing,Float64} = nothing # Internal inductance at the axis - li2::Union{Nothing,Float64} = nothing # External inductance at the axis - li3::Union{Nothing,Float64} = nothing # Total inductance at the axis -end - -""" - ProfileSplines - -Named 1D cubic spline interpolants for equilibrium profiles. -Each profile is stored as a separate spline for code clarity. - -# Fields - - - `xs::Vector{Float64}`: Shared x-axis (normalized psi) - - `F_spline`: 2π*F (toroidal flux function, where F = R * B_toroidal) - - `P_spline`: μ₀*P (plasma pressure × μ₀) - - `dVdpsi_spline`: dV/dψ (volume derivative) - - `q_spline`: q (safety factor) - -# Derivative Interpolants (for continuous derivative evaluation) - - - `F_deriv`, `P_deriv`, `dVdpsi_deriv`, `q_deriv`: First derivative interpolants - -# Notes - - - Node values at grid points: Access via `spline.y[i]` - - Derivative at any point: Call `deriv(x)` (derivative views are callable) - - Grid: Access via `xs` field or `spline.cache.x` -""" -struct ProfileSplines{S,D} - xs::Vector{Float64} - npts::Int # length(xs), avoids redundant length() calls - npts_minus_1::Int # npts - 1, for hint at last interval - # Value interpolants - F_spline::S - P_spline::S - dVdpsi_spline::S - q_spline::S - # Derivative views (callable, share data with value interpolants) - F_deriv::D - P_deriv::D - dVdpsi_deriv::D - q_deriv::D -end - -""" - ProfileSplines(xs, F_vals, P_vals, dVdpsi_vals, q_vals; extrap=ExtendExtrap()) - -Create ProfileSplines from arrays of profile values. -Uses CubicFit boundary conditions with extension extrapolation. -""" -function ProfileSplines(xs::Vector{Float64}, - F_vals::Vector{Float64}, - P_vals::Vector{Float64}, - dVdpsi_vals::Vector{Float64}, - q_vals::Vector{Float64}; - extrap::AbstractExtrap=ExtendExtrap()) - npts = length(xs) - npts_minus_1 = npts - 1 - @assert length(F_vals) == npts - @assert length(P_vals) == npts - @assert length(dVdpsi_vals) == npts - @assert length(q_vals) == npts - - # Create value interpolants with CubicFit BC (default) for sequential psi access - F_spline = cubic_interp(xs, F_vals; extrap=extrap) - P_spline = cubic_interp(xs, P_vals; extrap=extrap) - dVdpsi_spline = cubic_interp(xs, dVdpsi_vals; extrap=extrap) - q_spline = cubic_interp(xs, q_vals; extrap=extrap) - - # Create derivative views (these share data with value interpolants, no extra storage) - F_deriv = deriv1(F_spline) - P_deriv = deriv1(P_spline) - dVdpsi_deriv = deriv1(dVdpsi_spline) - q_deriv = deriv1(q_spline) - - ProfileSplines{typeof(F_spline),typeof(F_deriv)}( - xs, npts, npts_minus_1, - F_spline, P_spline, dVdpsi_spline, q_spline, - F_deriv, P_deriv, dVdpsi_deriv, q_deriv - ) -end - -""" - GeometryProfileSplines - -Named 1D cubic spline interpolants for flux-surface-averaged geometric quantities. -Built once during equilibrium construction so they are available to every downstream -module without per-caller recomputation. - -# Fields - - - `xs::Vector{Float64}`: Shared ψ axis (matches `rzphi_xs`) - - `area_spline`: Flux-surface area ∫ dA(ψ) [m²] - - `avg_r_spline`: Surface-average minor radius ⟨r⟩(ψ) [m] - - `avg_R_spline`: Surface-average major radius ⟨R⟩(ψ) [m] -""" -struct GeometryProfileSplines{S} - xs::Vector{Float64} - npts::Int - npts_minus_1::Int - area_spline::S - avg_r_spline::S - avg_R_spline::S -end - -""" - GeometryProfileSplines(xs, area_vals, avg_r_vals, avg_R_vals; extrap=ExtendExtrap()) - -Construct `GeometryProfileSplines` from arrays of surface-averaged values defined -on the shared ψ grid `xs`. Uses CubicFit boundary conditions to match `ProfileSplines`. -""" -function GeometryProfileSplines(xs::Vector{Float64}, - area_vals::Vector{Float64}, - avg_r_vals::Vector{Float64}, - avg_R_vals::Vector{Float64}; - extrap::AbstractExtrap=ExtendExtrap()) - npts = length(xs) - @assert length(area_vals) == npts - @assert length(avg_r_vals) == npts - @assert length(avg_R_vals) == npts - - area_spline = cubic_interp(xs, area_vals; extrap=extrap) - avg_r_spline = cubic_interp(xs, avg_r_vals; extrap=extrap) - avg_R_spline = cubic_interp(xs, avg_R_vals; extrap=extrap) - - GeometryProfileSplines{typeof(area_spline)}( - xs, npts, npts - 1, - area_spline, avg_r_spline, avg_R_spline - ) -end - -""" - KineticProfileSplines - -Named 1D cubic spline interpolants for kinetic profiles (densities, temperatures, -ExB rotation, collisional diagnostics) loaded from an external `kinetic.dat` -file. Mirrors the `ProfileSplines` pattern: each quantity is its own named -spline plus a paired derivative view for the species the NTV kernel needs -(densities and temperatures, used by `wdian`/`wdiat` in `KineticForces/Torque.jl`). - -# Fields - - - `xs::Vector{Float64}`: Shared ψ axis (the regular kinetic grid) - - `ni_spline`, `ne_spline`: Ion / electron number densities [m⁻³] - - `Ti_spline`, `Te_spline`: Ion / electron temperatures [J] - - `omegaE_spline`: ExB rotation ω_E [rad/s] - - `loglam_spline`: Coulomb logarithm - - `nui_spline`, `nue_spline`: Krook collision frequencies [s⁻¹] - - `zeff_spline`: Effective charge Z_eff - - `ni_deriv`, `ne_deriv`, `Ti_deriv`, `Te_deriv`: Derivative views for ψ-derivative access -""" -struct KineticProfileSplines{S,D} - xs::Vector{Float64} - npts::Int - npts_minus_1::Int - ni_spline::S - ne_spline::S - Ti_spline::S - Te_spline::S - omegaE_spline::S - loglam_spline::S - nui_spline::S - nue_spline::S - zeff_spline::S - ni_deriv::D - ne_deriv::D - Ti_deriv::D - Te_deriv::D -end - -""" - KineticProfileSplines(xs, ni, ne, Ti, Te, omegaE, loglam, nui, nue, zeff; - extrap=ExtendExtrap()) - -Build the kinetic profile splines from arrays of values defined on the shared -ψ grid `xs`. Uses CubicFit boundary conditions to match `ProfileSplines`. -Temperatures must already be in Joules. -""" -function KineticProfileSplines(xs::Vector{Float64}, - ni::Vector{Float64}, ne::Vector{Float64}, - Ti::Vector{Float64}, Te::Vector{Float64}, - omegaE::Vector{Float64}, loglam::Vector{Float64}, - nui::Vector{Float64}, nue::Vector{Float64}, zeff::Vector{Float64}; - extrap::AbstractExtrap=ExtendExtrap()) - npts = length(xs) - @assert all(length(v) == npts for v in (ni, ne, Ti, Te, omegaE, loglam, nui, nue, zeff)) - - ni_spline = cubic_interp(xs, ni; extrap=extrap) - ne_spline = cubic_interp(xs, ne; extrap=extrap) - Ti_spline = cubic_interp(xs, Ti; extrap=extrap) - Te_spline = cubic_interp(xs, Te; extrap=extrap) - omegaE_spline = cubic_interp(xs, omegaE; extrap=extrap) - loglam_spline = cubic_interp(xs, loglam; extrap=extrap) - nui_spline = cubic_interp(xs, nui; extrap=extrap) - nue_spline = cubic_interp(xs, nue; extrap=extrap) - zeff_spline = cubic_interp(xs, zeff; extrap=extrap) - - ni_deriv = deriv1(ni_spline) - ne_deriv = deriv1(ne_spline) - Ti_deriv = deriv1(Ti_spline) - Te_deriv = deriv1(Te_spline) - - KineticProfileSplines{typeof(ni_spline),typeof(ni_deriv)}( - xs, npts, npts - 1, - ni_spline, ne_spline, Ti_spline, Te_spline, - omegaE_spline, loglam_spline, nui_spline, nue_spline, zeff_spline, - ni_deriv, ne_deriv, Ti_deriv, Te_deriv - ) -end - -""" - PlasmaEquilibrium(...) - -The final, self-contained result of the equilibrium reconstruction. -This object provides a complete representation of the processed plasma equilibrium in flux coordinates. - -# Fields - - - `config::EquilibriumConfig`: - The equilibrium configuration object used for the reconstruction. - - - `params::EquilibriumParameters`: - Computed equilibrium parameters and diagnostics. - - `profiles::ProfileSplines`: - Named 1D profile splines (F, P, dV/dψ, q) on normalized psi grid. - Access values at grid points via `profiles.F_spline.y[i]`, etc. - Access derivatives via `profiles.F_deriv.y[i]` or `profiles.F_deriv(psi)`. - - `geometry::GeometryProfileSplines`: - Named 1D splines for flux-surface-averaged geometry (area, ⟨r⟩, ⟨R⟩), - populated automatically by `compute_geometry_profiles` during construction. - - **Grid coordinates (shared by all rzphi/eqfun interpolants):** - - + `rzphi_xs::Vector{Float64}`: ψ coordinates (length mpsi+1) - + `rzphi_ys::Vector{Float64}`: θ coordinates (length mtheta+1) - - **Geometric quantities (rzphi, 4 interpolants):** - 2D cubic interpolants for flux-coordinate mapping with periodic BC in theta. - - + **x value:** normalized ψ - + **y value:** SFL poloidal angle ∈ [0, 1] - + `rzphi_rsquared::CubicInterpolantND`: r_coord² = (R - ro)² + (Z - zo)² - + `rzphi_offset::CubicInterpolantND`: η/(2π) - θₙₑw (angle offset) - + `rzphi_nu::CubicInterpolantND`: ν in ϕ = 2πζ + ν(ψ, θ) - + `rzphi_jac::CubicInterpolantND`: Jacobian - - **Physics quantities (eqfun, 3 interpolants):** - 2D cubic interpolants storing local physics and geometric quantities. - - + **x value:** normalized ψ - + **y value:** SFL poloidal angle θₙₑw - + `eqfun_B::CubicInterpolantND`: Total magnetic field strength [T] - + `eqfun_metric1::CubicInterpolantND`: (e₁⋅e₂ + q⋅e₃⋅e₁)/(J⋅B²) - + `eqfun_metric2::CubicInterpolantND`: (e₂⋅e₃ + q⋅e₃⋅e₃)/(J⋅B²) - - `ro::Float64`: R-coordinate of the magnetic axis [m] - - `zo::Float64`: Z-coordinate of the magnetic axis [m] - - `psio::Float64`: Total flux difference |Ψ_axis - Ψ_boundary| [Weber/radian] - - `ingest::EquilibriumIngest`: raw arrays forwarded from the equilibrium input for the - `gpec.h5` rerun snapshot — a [`DirectIngest`](@ref)/[`InverseIngest`](@ref) for file-based - equilibria, or `nothing` for analytic ones (regenerated from their TOML section on replay) - - `analytic::Union{Nothing,AnalyticEquilibriumConfig}`: the parameters an analytic kind - (`sol`, `lar`, `tj_analytic`, `tj_analytic_direct`) was solved from, so [`reform`](@ref) can - re-solve it on a new grid; `nothing` for file-based and IMAS equilibria. A defensive copy, so - later edits to the caller's config do not reach it. -""" -mutable struct PlasmaEquilibrium{P<:ProfileSplines,G<:GeometryProfileSplines,I2D<:FastInterpolations.CubicInterpolantND} - config::EquilibriumConfig - params::EquilibriumParameters - profiles::P - geometry::G - - # Grid coordinates (shared by all 2D interpolants) - rzphi_xs::Vector{Float64} - rzphi_ys::Vector{Float64} - - # Geometric quantities (4 interpolants) - rzphi_rsquared::I2D - rzphi_offset::I2D - rzphi_nu::I2D - rzphi_jac::I2D - - # Physics quantities (3 interpolants) - eqfun_B::I2D - eqfun_metric1::I2D - eqfun_metric2::I2D - - ro::Float64 - zo::Float64 - psio::Float64 - - ingest::EquilibriumIngest - analytic::Union{Nothing,AnalyticEquilibriumConfig} -end - -# Solvers build the equilibrium before PlasmaEquilibrium records its source, so allow -# construction without it; ingest and analytic default to nothing and are assigned post-construction. -PlasmaEquilibrium(config, params, profiles, geometry, rzphi_xs, rzphi_ys, - rzphi_rsquared, rzphi_offset, rzphi_nu, rzphi_jac, - eqfun_B, eqfun_metric1, eqfun_metric2, ro, zo, psio) = - PlasmaEquilibrium(config, params, profiles, geometry, rzphi_xs, rzphi_ys, - rzphi_rsquared, rzphi_offset, rzphi_nu, rzphi_jac, - eqfun_B, eqfun_metric1, eqfun_metric2, ro, zo, psio, nothing, nothing) diff --git a/src/Equilibrium/GridRefinement.jl b/src/Equilibrium/Grid/GridRefinement.jl similarity index 100% rename from src/Equilibrium/GridRefinement.jl rename to src/Equilibrium/Grid/GridRefinement.jl diff --git a/src/Equilibrium/Grid/PsiGrid.jl b/src/Equilibrium/Grid/PsiGrid.jl new file mode 100644 index 000000000..aafbc9063 --- /dev/null +++ b/src/Equilibrium/Grid/PsiGrid.jl @@ -0,0 +1,70 @@ +""" + make_optimal_psi_grid(psilow, psihigh, N_core, N_mid, N_edge; psi_split_core, psi_split_edge) + +Build a three-region ψ grid with (N_core + N_mid + N_edge)+1 knots, using the counts +directly as provided — core and edge are geometric in log(ψ) and log(1−ψ) respectively; +middle is uniform in ψ. +""" +function make_optimal_psi_grid(psilow, psihigh, N_core, N_mid, N_edge; + psi_split_core=0.03, psi_split_edge=0.98) + # Core: [psilow, psi_split_core], geometric in log(ψ) + core_pts = [psilow * (psi_split_core / psilow)^(i / N_core) for i in 0:N_core] + # Middle: [psi_split_core, psi_split_edge], uniform (skip first to avoid duplicate) + mid_pts = [psi_split_core + (psi_split_edge - psi_split_core) * i / N_mid for i in 1:N_mid] + # Edge: [psi_split_edge, psihigh], geometric in log(1−ψ) (skip first to avoid duplicate) + edge_pts = [1.0 - (1.0 - psi_split_edge) * ((1.0 - psihigh) / (1.0 - psi_split_edge))^(i / N_edge) for i in 1:N_edge] + + return vcat(core_pts, mid_pts, edge_pts) +end + +""" + _build_psi_grid(equil_params, psilow, psihigh) + +Resolve `mpsi` and build `psi_nodes` for any supported `grid_type`. + +For `"auto"` (or the legacy alias `"log_asymptotic"`) with `mpsi=0`, this is the coarse pass-1 layout of +the two-pass refinement — the main driver measures the formed equilibrium and re-forms on +the `refined_psi_grid` result. Shared by `direct_fieldline_int` and `efit_by_inversion` +solvers. +""" +function _build_psi_grid(equil_params, psilow, psihigh) + mpsi = equil_params.mpsi + if equil_params.grid_type in ("auto", "log_asymptotic") && mpsi == 0 && equil_params.psi_accuracy > 0 + # Two-pass auto grid: this is the pass-1 layout; the driver measures the formed + # equilibrium's curvature and re-forms on a refined grid (GridRefinement.jl). Sized to + # resolve the gradient structure accurately so pass 2 is well-provisioned (PASS1_INTERVALS). + mpsi = PASS1_INTERVALS + @info "Auto psi grid: forming pass-1 equilibrium on $(mpsi)-interval log_asymptotic grid pending curvature-based refinement" + elseif mpsi == 0 + mpsi = 128 + end + + psi_nodes = if equil_params.grid_type in ("auto", "log_asymptotic") + psihigh > 0.98 || + @warn "grid_type = \"$(equil_params.grid_type)\" needs psihigh > 0.98 (its grid has a fixed edge region on [0.98, psihigh]); " * + "psihigh = $psihigh gives a non-monotonic ψ grid and the equilibrium will fail. Use grid_type = \"core_edge_packed\" instead." + # Distribute mpsi across the three regions by log-weights + log_core = log(0.03 / psilow) + log_mid = log(0.98 / 0.03) + log_edge = log((1.0 - 0.98) / (1.0 - psihigh)) + log_total = log_core + log_mid + log_edge + N_edge = clamp(round(Int, mpsi * log_edge / log_total), 2, mpsi ÷ 2) + N_core = round(Int, mpsi * log_core / log_total) + N_mid = mpsi - N_edge - N_core + make_optimal_psi_grid(psilow, psihigh, N_core, N_mid, N_edge) + elseif equil_params.grid_type == "core_edge_packed" + [psilow + (psihigh - psilow) * sin((ipsi / mpsi) * (π / 2))^2 for ipsi in 0:mpsi] + elseif equil_params.grid_type == "pow1" + # Fortran powspace(psilow, psihigh, 1, mpsi+1, "upper") — edge-packed grid (equil/grid.f90:92-195) + [psilow + (psihigh - psilow) * (3(ipsi / mpsi) - (ipsi / mpsi)^3) / 2 for ipsi in 0:mpsi] + elseif equil_params.grid_type == "uniform" + # Evenly spaced surfaces, e.g. for profile diagnostics at constant resolution + collect(range(psilow, psihigh; length=mpsi + 1)) + else + error("Unsupported grid_type: $(equil_params.grid_type)") + end + # Floor node spacing on the fixed grids: core_edge_packed/pow1 pack the edge as ~(π/2·mpsi)⁻² and at high + # mpsi drive it below the integration-noise scale (garbage curvature). The auto grid floors its + # refined pass-2 grid in refined_psi_grid, so it is left untouched here. + return equil_params.grid_type in ("auto", "log_asymptotic") ? psi_nodes : enforce_min_spacing(psi_nodes, MIN_KNOT_SPACING) +end diff --git a/src/Equilibrium/PlasmaEquilibrium.jl b/src/Equilibrium/PlasmaEquilibrium.jl new file mode 100644 index 000000000..9c69d3218 --- /dev/null +++ b/src/Equilibrium/PlasmaEquilibrium.jl @@ -0,0 +1,295 @@ +""" + ProfileSplines + +Named 1D cubic spline interpolants for equilibrium profiles. +Each profile is stored as a separate spline for code clarity. + +# Fields + + - `xs::Vector{Float64}`: Shared x-axis (normalized psi) + - `F_spline`: 2π*F (toroidal flux function, where F = R * B_toroidal) + - `P_spline`: μ₀*P (plasma pressure × μ₀) + - `dVdpsi_spline`: dV/dψ (volume derivative) + - `q_spline`: q (safety factor) + +# Derivative Interpolants (for continuous derivative evaluation) + + - `F_deriv`, `P_deriv`, `dVdpsi_deriv`, `q_deriv`: First derivative interpolants + +# Notes + + - Node values at grid points: Access via `spline.y[i]` + - Derivative at any point: Call `deriv(x)` (derivative views are callable) + - Grid: Access via `xs` field or `spline.cache.x` +""" +struct ProfileSplines{S,D} + xs::Vector{Float64} + npts::Int # length(xs), avoids redundant length() calls + npts_minus_1::Int # npts - 1, for hint at last interval + # Value interpolants + F_spline::S + P_spline::S + dVdpsi_spline::S + q_spline::S + # Derivative views (callable, share data with value interpolants) + F_deriv::D + P_deriv::D + dVdpsi_deriv::D + q_deriv::D +end + +""" + ProfileSplines(xs, F_vals, P_vals, dVdpsi_vals, q_vals; extrap=ExtendExtrap()) + +Create ProfileSplines from arrays of profile values. +Uses CubicFit boundary conditions with extension extrapolation. +""" +function ProfileSplines(xs::Vector{Float64}, + F_vals::Vector{Float64}, + P_vals::Vector{Float64}, + dVdpsi_vals::Vector{Float64}, + q_vals::Vector{Float64}; + extrap::AbstractExtrap=ExtendExtrap()) + npts = length(xs) + npts_minus_1 = npts - 1 + @assert length(F_vals) == npts + @assert length(P_vals) == npts + @assert length(dVdpsi_vals) == npts + @assert length(q_vals) == npts + + # Create value interpolants with CubicFit BC (default) for sequential psi access + F_spline = cubic_interp(xs, F_vals; extrap=extrap) + P_spline = cubic_interp(xs, P_vals; extrap=extrap) + dVdpsi_spline = cubic_interp(xs, dVdpsi_vals; extrap=extrap) + q_spline = cubic_interp(xs, q_vals; extrap=extrap) + + # Create derivative views (these share data with value interpolants, no extra storage) + F_deriv = deriv1(F_spline) + P_deriv = deriv1(P_spline) + dVdpsi_deriv = deriv1(dVdpsi_spline) + q_deriv = deriv1(q_spline) + + ProfileSplines{typeof(F_spline),typeof(F_deriv)}( + xs, npts, npts_minus_1, + F_spline, P_spline, dVdpsi_spline, q_spline, + F_deriv, P_deriv, dVdpsi_deriv, q_deriv + ) +end + +""" + GeometryProfileSplines + +Named 1D cubic spline interpolants for flux-surface-averaged geometric quantities. +Built once during equilibrium construction so they are available to every downstream +module without per-caller recomputation. + +# Fields + + - `xs::Vector{Float64}`: Shared ψ axis (matches `rzphi_xs`) + - `area_spline`: Flux-surface area ∫ dA(ψ) [m²] + - `avg_r_spline`: Surface-average minor radius ⟨r⟩(ψ) [m] + - `avg_R_spline`: Surface-average major radius ⟨R⟩(ψ) [m] +""" +struct GeometryProfileSplines{S} + xs::Vector{Float64} + npts::Int + npts_minus_1::Int + area_spline::S + avg_r_spline::S + avg_R_spline::S +end + +""" + GeometryProfileSplines(xs, area_vals, avg_r_vals, avg_R_vals; extrap=ExtendExtrap()) + +Construct `GeometryProfileSplines` from arrays of surface-averaged values defined +on the shared ψ grid `xs`. Uses CubicFit boundary conditions to match `ProfileSplines`. +""" +function GeometryProfileSplines(xs::Vector{Float64}, + area_vals::Vector{Float64}, + avg_r_vals::Vector{Float64}, + avg_R_vals::Vector{Float64}; + extrap::AbstractExtrap=ExtendExtrap()) + npts = length(xs) + @assert length(area_vals) == npts + @assert length(avg_r_vals) == npts + @assert length(avg_R_vals) == npts + + area_spline = cubic_interp(xs, area_vals; extrap=extrap) + avg_r_spline = cubic_interp(xs, avg_r_vals; extrap=extrap) + avg_R_spline = cubic_interp(xs, avg_R_vals; extrap=extrap) + + GeometryProfileSplines{typeof(area_spline)}( + xs, npts, npts - 1, + area_spline, avg_r_spline, avg_R_spline + ) +end + +""" + KineticProfileSplines + +Named 1D cubic spline interpolants for kinetic profiles (densities, temperatures, +ExB rotation, collisional diagnostics) loaded from an external `kinetic.dat` +file. Mirrors the `ProfileSplines` pattern: each quantity is its own named +spline plus a paired derivative view for the species the NTV kernel needs +(densities and temperatures, used by `wdian`/`wdiat` in `KineticForces/Torque.jl`). + +# Fields + + - `xs::Vector{Float64}`: Shared ψ axis (the regular kinetic grid) + - `ni_spline`, `ne_spline`: Ion / electron number densities [m⁻³] + - `Ti_spline`, `Te_spline`: Ion / electron temperatures [J] + - `omegaE_spline`: ExB rotation ω_E [rad/s] + - `loglam_spline`: Coulomb logarithm + - `nui_spline`, `nue_spline`: Krook collision frequencies [s⁻¹] + - `zeff_spline`: Effective charge Z_eff + - `ni_deriv`, `ne_deriv`, `Ti_deriv`, `Te_deriv`: Derivative views for ψ-derivative access +""" +struct KineticProfileSplines{S,D} + xs::Vector{Float64} + npts::Int + npts_minus_1::Int + ni_spline::S + ne_spline::S + Ti_spline::S + Te_spline::S + omegaE_spline::S + loglam_spline::S + nui_spline::S + nue_spline::S + zeff_spline::S + ni_deriv::D + ne_deriv::D + Ti_deriv::D + Te_deriv::D +end + +""" + KineticProfileSplines(xs, ni, ne, Ti, Te, omegaE, loglam, nui, nue, zeff; + extrap=ExtendExtrap()) + +Build the kinetic profile splines from arrays of values defined on the shared +ψ grid `xs`. Uses CubicFit boundary conditions to match `ProfileSplines`. +Temperatures must already be in Joules. +""" +function KineticProfileSplines(xs::Vector{Float64}, + ni::Vector{Float64}, ne::Vector{Float64}, + Ti::Vector{Float64}, Te::Vector{Float64}, + omegaE::Vector{Float64}, loglam::Vector{Float64}, + nui::Vector{Float64}, nue::Vector{Float64}, zeff::Vector{Float64}; + extrap::AbstractExtrap=ExtendExtrap()) + npts = length(xs) + @assert all(length(v) == npts for v in (ni, ne, Ti, Te, omegaE, loglam, nui, nue, zeff)) + + ni_spline = cubic_interp(xs, ni; extrap=extrap) + ne_spline = cubic_interp(xs, ne; extrap=extrap) + Ti_spline = cubic_interp(xs, Ti; extrap=extrap) + Te_spline = cubic_interp(xs, Te; extrap=extrap) + omegaE_spline = cubic_interp(xs, omegaE; extrap=extrap) + loglam_spline = cubic_interp(xs, loglam; extrap=extrap) + nui_spline = cubic_interp(xs, nui; extrap=extrap) + nue_spline = cubic_interp(xs, nue; extrap=extrap) + zeff_spline = cubic_interp(xs, zeff; extrap=extrap) + + ni_deriv = deriv1(ni_spline) + ne_deriv = deriv1(ne_spline) + Ti_deriv = deriv1(Ti_spline) + Te_deriv = deriv1(Te_spline) + + KineticProfileSplines{typeof(ni_spline),typeof(ni_deriv)}( + xs, npts, npts - 1, + ni_spline, ne_spline, Ti_spline, Te_spline, + omegaE_spline, loglam_spline, nui_spline, nue_spline, zeff_spline, + ni_deriv, ne_deriv, Ti_deriv, Te_deriv + ) +end + +""" + PlasmaEquilibrium(...) + +The final, self-contained result of the equilibrium reconstruction. +This object provides a complete representation of the processed plasma equilibrium in flux coordinates. + +# Fields + + - `config::EquilibriumConfig`: + The equilibrium configuration object used for the reconstruction. + + - `params::EquilibriumParameters`: + Computed equilibrium parameters and diagnostics. + - `profiles::ProfileSplines`: + Named 1D profile splines (F, P, dV/dψ, q) on normalized psi grid. + Access values at grid points via `profiles.F_spline.y[i]`, etc. + Access derivatives via `profiles.F_deriv.y[i]` or `profiles.F_deriv(psi)`. + - `geometry::GeometryProfileSplines`: + Named 1D splines for flux-surface-averaged geometry (area, ⟨r⟩, ⟨R⟩), + populated automatically by `compute_geometry_profiles` during construction. + - **Grid coordinates (shared by all rzphi/eqfun interpolants):** + + + `rzphi_xs::Vector{Float64}`: ψ coordinates (length mpsi+1) + + `rzphi_ys::Vector{Float64}`: θ coordinates (length mtheta+1) + - **Geometric quantities (rzphi, 4 interpolants):** + 2D cubic interpolants for flux-coordinate mapping with periodic BC in theta. + + + **x value:** normalized ψ + + **y value:** SFL poloidal angle ∈ [0, 1] + + `rzphi_rsquared::CubicInterpolantND`: r_coord² = (R - ro)² + (Z - zo)² + + `rzphi_offset::CubicInterpolantND`: η/(2π) - θₙₑw (angle offset) + + `rzphi_nu::CubicInterpolantND`: ν in ϕ = 2πζ + ν(ψ, θ) + + `rzphi_jac::CubicInterpolantND`: Jacobian + - **Physics quantities (eqfun, 3 interpolants):** + 2D cubic interpolants storing local physics and geometric quantities. + + + **x value:** normalized ψ + + **y value:** SFL poloidal angle θₙₑw + + `eqfun_B::CubicInterpolantND`: Total magnetic field strength [T] + + `eqfun_metric1::CubicInterpolantND`: (e₁⋅e₂ + q⋅e₃⋅e₁)/(J⋅B²) + + `eqfun_metric2::CubicInterpolantND`: (e₂⋅e₃ + q⋅e₃⋅e₃)/(J⋅B²) + - `ro::Float64`: R-coordinate of the magnetic axis [m] + - `zo::Float64`: Z-coordinate of the magnetic axis [m] + - `psio::Float64`: Total flux difference |Ψ_axis - Ψ_boundary| [Weber/radian] + - `ingest::EquilibriumIngest`: raw arrays forwarded from the equilibrium input for the + `gpec.h5` rerun snapshot — a [`DirectIngest`](@ref)/[`InverseIngest`](@ref) for file-based + equilibria, or `nothing` for analytic ones (regenerated from their TOML section on replay) + - `analytic::Union{Nothing,AnalyticEquilibriumConfig}`: the parameters an analytic kind + (`sol`, `lar`, `tj_analytic`, `tj_analytic_direct`) was solved from, so [`reform`](@ref) can + re-solve it on a new grid; `nothing` for file-based and IMAS equilibria. A defensive copy, so + later edits to the caller's config do not reach it. +""" +mutable struct PlasmaEquilibrium{P<:ProfileSplines,G<:GeometryProfileSplines,I2D<:FastInterpolations.CubicInterpolantND} + config::EquilibriumConfig + params::EquilibriumParameters + profiles::P + geometry::G + + # Grid coordinates (shared by all 2D interpolants) + rzphi_xs::Vector{Float64} + rzphi_ys::Vector{Float64} + + # Geometric quantities (4 interpolants) + rzphi_rsquared::I2D + rzphi_offset::I2D + rzphi_nu::I2D + rzphi_jac::I2D + + # Physics quantities (3 interpolants) + eqfun_B::I2D + eqfun_metric1::I2D + eqfun_metric2::I2D + + ro::Float64 + zo::Float64 + psio::Float64 + + ingest::EquilibriumIngest + analytic::Union{Nothing,AnalyticEquilibriumConfig} +end + +# Solvers build the equilibrium before PlasmaEquilibrium records its source, so allow +# construction without it; ingest and analytic default to nothing and are assigned post-construction. +PlasmaEquilibrium(config, params, profiles, geometry, rzphi_xs, rzphi_ys, + rzphi_rsquared, rzphi_offset, rzphi_nu, rzphi_jac, + eqfun_B, eqfun_metric1, eqfun_metric2, ro, zo, psio) = + PlasmaEquilibrium(config, params, profiles, geometry, rzphi_xs, rzphi_ys, + rzphi_rsquared, rzphi_offset, rzphi_nu, rzphi_jac, + eqfun_B, eqfun_metric1, eqfun_metric2, ro, zo, psio, nothing, nothing) diff --git a/src/Equilibrium/PostProcessing.jl b/src/Equilibrium/PostProcessing.jl deleted file mode 100644 index 665d69e25..000000000 --- a/src/Equilibrium/PostProcessing.jl +++ /dev/null @@ -1,491 +0,0 @@ -# PostProcessing.jl -# -# Quantities derived from a formed equilibrium: separatrix, global parameters, q extrema, and the -# Grad-Shafranov residual. - -""" - equilibrium_separatrix_find!(pe::PlasmaEquilibrium) - -Finds the separatrix locations in the plasma equilibrium (rsep, zsep, rext, zext). -Performs the same function as equil_out_sep_find in the Fortran code. -""" -function equilibrium_separatrix_find!(pe::PlasmaEquilibrium) - mpsi = length(pe.rzphi_xs) - 1 - edge_idx = mpsi + 1 # Edge flux surface index - psi_edge = pe.rzphi_xs[edge_idx] - rsep = zeros(2) - - # Outboard and inboard midplane R via bracketed Brent on θ + η(θ) - η₀ = 0. - # iside=1 → outboard (η₀=0.0, θ near 0), iside=2 → inboard (η₀=0.5, θ near 0.5). - for iside in 1:2 - eta0 = (iside == 1) ? 0.0 : 0.5 - hint2d = (Ref(1), Ref(1)) - theta_lo, theta_hi = (iside == 1) ? (-0.25, 0.25) : (0.25, 0.75) - side_label = (iside == 1) ? "outboard" : "inboard" - theta = try - find_zero( - theta -> theta + pe.rzphi_offset((psi_edge, theta); hint=hint2d) - eta0, - (theta_lo, theta_hi), Roots.Brent(); - atol=1e-12, rtol=1e-12) - catch e - error("Separatrix $side_label midplane root not found in bracket " * - "[$(theta_lo), $(theta_hi)]: $(e.msg)") - end - r2 = pe.rzphi_rsquared((psi_edge, theta)) - offset = pe.rzphi_offset((psi_edge, theta)) - rsep[iside] = pe.ro + sqrt(r2) * cos(2π * (theta + offset)) - end - - # Top and bottom separatrix Z extrema via bracketed Brent on ∂z/∂θ = 0. - # iside=1 → top (θ ∈ [0.0, 0.5]), iside=2 → bottom (θ ∈ [0.5, 1.0]). - # Matches Fortran convention so (r|z)ext[1] / zsep[1] refer to the top extremum - # and (r|z)ext[2] / zsep[2] to the bottom (see equil_out.f::equil_out_sep_find). - # Splines use PeriodicBC + WrapExtrap, so θ outside [0,1] is valid. - zsep = zeros(2) - rext = zeros(2) - zext = zeros(2) - - for iside in 1:2 - hint2d = (Ref(1), Ref(1)) - - # Cache variables populated by z_deriv, read after convergence - rfac = Ref(0.0) - cos_phase = Ref(0.0) - z_val = Ref(0.0) - - # ∂z/∂θ where z(θ) = zo + √r²(θ) · sin(2π(θ + η(θ))) - function z_deriv(theta_inner) - r2 = pe.rzphi_rsquared((psi_edge, theta_inner); hint=hint2d) - r2y = pe.rzphi_rsquared((psi_edge, theta_inner); deriv=DerivOp(0, 1), hint=hint2d) - η = pe.rzphi_offset((psi_edge, theta_inner); hint=hint2d) - η1 = pe.rzphi_offset((psi_edge, theta_inner); deriv=DerivOp(0, 1), hint=hint2d) - rfac_local = sqrt(max(0.0, r2)) - rfac1 = (rfac_local > 0) ? r2y / (2 * rfac_local) : 0.0 - phase1 = 2π * (1 + η1) - sin_phase = sin(2π * (theta_inner + η)) - cos_phase_local = cos(2π * (theta_inner + η)) - - rfac[] = rfac_local - cos_phase[] = cos_phase_local - z_val[] = pe.zo + rfac_local * sin_phase - - return rfac_local * phase1 * cos_phase_local + rfac1 * sin_phase - end - - theta_lo, theta_hi = (iside == 1) ? (0.0, 0.5) : (0.5, 1.0) - side_label = (iside == 1) ? "top" : "bottom" - theta = try - find_zero(z_deriv, (theta_lo, theta_hi), Roots.Brent(); - atol=1e-12, rtol=1e-12) - catch e - error("Separatrix $side_label Z-extremum root not found in bracket " * - "[$(theta_lo), $(theta_hi)]: $(e.msg)") - end - - rext[iside] = pe.ro + rfac[] * cos_phase[] - zsep[iside] = zext[iside] = z_val[] - end - - pe.params.rsep = rsep - pe.params.zsep = zsep - pe.params.rext = rext - pe.params.zext = zext - return (rsep, zsep, rext, zext) -end - -""" - equilibrium_global_parameters!(pe::PlasmaEquilibrium) - -Computes and populates global equilibrium parameters in the `PlasmaEquilibrium` -struct, such as rmean, amean, kappa, bt0, crnt, betat, betan, li1, etc. Performs -the same function as equil_out_global in the Fortran code. -""" -function equilibrium_global_parameters!(pe::PlasmaEquilibrium) - profiles = pe.profiles - mpsi = length(pe.rzphi_xs) - 1 - mtheta = length(pe.rzphi_ys) - 1 - - # Use separatrix geometry - rsep, zsep, rext, _ = equilibrium_separatrix_find!(pe) - - rmean = (rsep[1] + rsep[2]) / 2 - amean = (rsep[1] - rsep[2]) / 2 - aratio = rmean / amean - kappa = (zsep[1] - zsep[2]) / (rsep[1] - rsep[2]) - delta1 = (rmean - rext[1]) / amean - delta2 = (rmean - rext[2]) / amean - dpsi = 1.0 - pe.rzphi_xs[mpsi+1] - psi_edge = profiles.xs[end] - bt0 = (profiles.F_spline.y[end] + profiles.F_deriv(psi_edge; hint=Ref(profiles.npts_minus_1)) * dpsi) / (2π * rmean) - - pe.params.rmean = rmean - pe.params.amean = amean - pe.params.aratio = aratio - pe.params.kappa = kappa - pe.params.delta1 = delta1 - pe.params.delta2 = delta2 - pe.params.bt0 = bt0 - - psio = pe.psio - gs1 = zeros(Float64, mtheta + 1) - gs2 = zeros(Float64, mtheta + 1) - - # Direct array access at edge flux surface grid points using nodal_derivs - for itheta in 0:mtheta - # Function values (partials[1,:,:]) - r2 = pe.rzphi_rsquared.nodal_derivs.partials[1, end, itheta+1] - offset = pe.rzphi_offset.nodal_derivs.partials[1, end, itheta+1] - jac = pe.rzphi_jac.nodal_derivs.partials[1, end, itheta+1] - # Theta derivatives (partials[3,:,:] = ∂f/∂y) - r2_y = pe.rzphi_rsquared.nodal_derivs.partials[3, end, itheta+1] - offset_y = pe.rzphi_offset.nodal_derivs.partials[3, end, itheta+1] - - chi1 = 2π * psio / jac - jacfac = π / jac - rfac = sqrt(r2) - eta = 2π * (pe.rzphi_ys[itheta+1] + offset) - r = pe.ro + rfac * cos(eta) - v21 = jacfac * r2_y / (2π * rfac) - v22 = jacfac * (1 + offset_y) * (2 * rfac) - v33 = jacfac * 2π * (r / π) - dvsq = (v21^2 + v22^2) * (v33 * jac^2)^2 - gs1[itheta+1] = sqrt(dvsq) / (2π * r) - gs2[itheta+1] = chi1 * dvsq / (2π * r)^2 - end - - int1 = sum(gs1) / (mtheta + 1) - int2 = sum(gs2) / (mtheta + 1) - crnt = int2 / (1e6 * mu0) - bp0 = int2 / int1 - bwall = 1e6 * mu0 * crnt / (2π * amean) - - pe.params.crnt = crnt - pe.params.bwall = bwall - - # Flux-surface integrals over the full normalized flux [0, 1]; ExtendExtrap carries the integral past the first/last grid points (integrands are smooth/vanishing at the axis). - xs = profiles.xs - fsi(y) = FastInterpolations.integrate(cubic_interp(xs, y; extrap=ExtendExtrap()), 0.0, 1.0) - P_vals = profiles.P_spline.y - dVdpsi_vals = profiles.dVdpsi_spline.y - - fsi_pdv = fsi(P_vals .* dVdpsi_vals) # ∫ p dV/dψ - fsi_dv = fsi(dVdpsi_vals) # ∫ dV/dψ - fsi_p2dv = fsi(P_vals .^ 2 .* dVdpsi_vals) # ∫ p² dV/dψ - volume = fsi_dv # same integrand as hs col 2 in Fortran - - # Poloidal-field surface integral hs_bp2(ψ) = ψ₀² ∮dθ |∇ψ|² / (R² J). - # This is Fortran equil_out.f's hs%fs(:,3) and is the correct integrand for - # the internal inductance li1/li2/li3 (distinct from the p²·dV integrand - # that drives betaj). - hs_bp2 = zeros(Float64, mpsi + 1) - for ipsi in 0:mpsi - acc = 0.0 - for itheta in 0:mtheta - r2 = pe.rzphi_rsquared.nodal_derivs.partials[1, ipsi+1, itheta+1] - offset = pe.rzphi_offset.nodal_derivs.partials[1, ipsi+1, itheta+1] - jac = pe.rzphi_jac.nodal_derivs.partials[1, ipsi+1, itheta+1] - r2_y = pe.rzphi_rsquared.nodal_derivs.partials[3, ipsi+1, itheta+1] - offset_y = pe.rzphi_offset.nodal_derivs.partials[3, ipsi+1, itheta+1] - - jacfac = π / jac - rfac = sqrt(r2) - eta = 2π * (pe.rzphi_ys[itheta+1] + offset) - r = pe.ro + rfac * cos(eta) - v21 = jacfac * r2_y / (2π * rfac) - v22 = jacfac * (1 + offset_y) * (2 * rfac) - v33 = jacfac * 2π * (r / π) - dvsq = (v21^2 + v22^2) * (v33 * jac^2)^2 - acc += dvsq / (r^2) / jac - end - # Periodic trapezoidal rule on uniform θ grid reduces to a plain mean - # because the first and last grid points coincide — matches the int1/int2 - # pattern used for the edge-surface integrals above. - hs_bp2[ipsi+1] = (acc / (mtheta + 1)) * psio^2 - end - fsi_bp2 = fsi(hs_bp2) - - p0 = P_vals[1] - profiles.P_deriv(profiles.xs[1]; hint=Ref(1)) * profiles.xs[1] # linear extrapolation - betat = 2 * (fsi_pdv / fsi_dv) / bt0^2 - betaj = 2 * sqrt(fsi_p2dv / fsi_dv) / bwall^2 - betan = 100 * amean * bt0 * betat / crnt - betap1 = 2 * (fsi_pdv / fsi_dv) / bp0^2 - betap2 = 4 * fsi_pdv / ((1e6 * mu0 * crnt)^2 * pe.ro) - betap3 = 4 * fsi_pdv / ((1e6 * mu0 * crnt)^2 * rmean) - li1 = fsi_bp2 / fsi_dv / bp0^2 - li2 = 2 * fsi_bp2 / ((1e6 * mu0 * crnt)^2 * pe.ro) - li3 = 2 * fsi_bp2 / ((1e6 * mu0 * crnt)^2 * rmean) - - pe.params.psi0 = psio - pe.params.psi_axis = pe.psio - pe.params.psi_boundary = 1.0 - pe.params.psi_boundary_norm = 1.0 - pe.params.psi_axis_norm = 0.0 - pe.params.psi_norm = 0.0 - pe.params.psi_axis_offset = 0.0 - pe.params.psi_boundary_offset = 0.0 - pe.params.psi_axis_sign = 1 - pe.params.psi_boundary_sign = -1 - pe.params.psi_boundary_zero = false - - pe.params.q0 = profiles.q_spline.y[1] - pe.params.b0 = bt0 - - pe.params.volume = volume - pe.params.betat = betat - pe.params.betan = betan - pe.params.betaj = betaj - pe.params.betap1 = betap1 - pe.params.betap2 = betap2 - pe.params.betap3 = betap3 - pe.params.li1 = li1 - pe.params.li2 = li2 - pe.params.li3 = li3 -end - -""" - equilibrium_qfind!(equil::PlasmaEquilibrium) - -Finds the extrema of the safety factor profile q(ψ) in the plasma equilibrium -and computes derived q-values such as q0, qmin, qmax, qa, and q95. Performs -the same function as equil_out_qfind in the Fortran code. -""" -function equilibrium_qfind!(equil::PlasmaEquilibrium) - - profiles = equil.profiles - xs = profiles.xs - mpsi = length(xs) - 1 - psiexl = Float64[] - qexl = Float64[] - - # Create derivative views for q-spline - q_spline = profiles.q_spline - q_d1 = deriv1(q_spline) - q_d2 = deriv2(q_spline) - q_d3 = deriv3(q_spline) - - # Left endpoint - push!(psiexl, xs[1]) - push!(qexl, q_spline.y[1]) - - # Search for extrema in q(ψ) - for ipsi in 1:mpsi - x0 = xs[ipsi] - x1 = xs[ipsi+1] - xmax = x1 - x0 - - a = q_spline(x0) - b = q_d1(x0) - c = q_d2(x0) - d = q_d3(x0) - - if d != 0.0 - xcrit = -c / d - dx2 = xcrit^2 - 2b / d - if dx2 ≥ 0 - dx = sqrt(dx2) - for delta in (dx, -dx) - x = xcrit - delta - if 0 ≤ x < xmax - ψ = x0 + x - push!(psiexl, ψ) - push!(qexl, q_spline(ψ)) - end - end - end - end - end - - # Right endpoint - push!(psiexl, xs[end]) - push!(qexl, q_spline.y[end]) - - equil.params.qextrema_psi = psiexl - equil.params.qextrema_q = qexl - equil.params.mextrema = length(psiexl) - # Compute derived q-values - q0 = q_spline.y[1] - profiles.q_deriv(xs[1]; hint=Ref(1)) * xs[1] - qmax_edge = q_spline.y[end] - qmin = min(minimum(qexl), q0) - qmax = max(maximum(qexl), qmax_edge) - qa = q_spline.y[end] + profiles.q_deriv(xs[end]; hint=Ref(profiles.npts_minus_1)) * (1.0 - xs[end]) - - q95 = q_spline(0.95) - - # Store derived values - equil.params.q0 = q0 - equil.params.qmin = qmin - equil.params.qmax = qmax - equil.params.qa = qa - equil.params.q95 = q95 -end - -""" - equilibrium_gse!(equil::PlasmaEquilibrium) - -Diagnoses the Grad-Shafranov solution by computing the residual of the -Grad-Shafranov equation across the grid and writing diagnostic data to HDF5 files. -Performs the same function as equil_out_gse in the Fortran code. - -Returns `(; xs, ys, flux_x, flux_y, source, total, error, errori)` on the `rzphi` (ψ, θ) nodes: the two -flux-divergence terms, the source, the residual `total = flux_x - flux_y + source`, `error = |total|` over -the largest term on the grid, and `errori`, the magnitude of the θ-integrated residual on each surface. -""" -function equilibrium_gse!(equil::PlasmaEquilibrium) - - profiles = equil.profiles - mpsi = length(equil.rzphi_xs) - 1 - mtheta = length(equil.rzphi_ys) - 1 - ro, zo = equil.ro, equil.zo - psio = equil.psio - verbose = equil.params.verbose - diagnose_src = equil.params.diagnose_src - diagnose_maxima = equil.params.diagnose_maxima - - if verbose - @info "Diagnosing Grad-Shafranov solution" - end - - # Compute R, Z coordinates using nodal_derivs access - r = zeros(Float64, mpsi + 1, mtheta + 1) - z = zeros(Float64, mpsi + 1, mtheta + 1) - - for ipsi in 1:(mpsi+1) - rfac = @. sqrt(equil.rzphi_rsquared.nodal_derivs.partials[1, ipsi, :]) - angle = @. 2π * (equil.rzphi_ys + equil.rzphi_offset.nodal_derivs.partials[1, ipsi, :]) - r[ipsi, :] .= ro .+ rfac .* cos.(angle) - z[ipsi, :] .= zo .+ rfac .* sin.(angle) - end - - # Compute flux quantities using direct array access at grid points - flux_fs = zeros(Float64, mpsi + 1, mtheta + 1, 2) - flux_fsx = zeros(Float64, mpsi + 1, mtheta + 1, 2) # x-derivatives - flux_fsy = zeros(Float64, mpsi + 1, mtheta + 1, 2) # y-derivatives - for ipsi in 1:(mpsi+1) - for itheta in 1:(mtheta+1) - f1 = equil.rzphi_rsquared.nodal_derivs.partials[1, ipsi, itheta] - f2 = equil.rzphi_offset.nodal_derivs.partials[1, ipsi, itheta] - f4 = equil.rzphi_jac.nodal_derivs.partials[1, ipsi, itheta] - fy1 = equil.rzphi_rsquared.nodal_derivs.partials[3, ipsi, itheta] - fy2 = equil.rzphi_offset.nodal_derivs.partials[3, ipsi, itheta] - fx1 = equil.rzphi_rsquared.nodal_derivs.partials[2, ipsi, itheta] - fx2 = equil.rzphi_offset.nodal_derivs.partials[2, ipsi, itheta] - - flux_fs[ipsi, itheta, 1] = fy1^2 / (4π^2 * f1) + (1 + fy2)^2 * 4 * f1 - flux_fs[ipsi, itheta, 2] = fx1 * fy1 / (4π^2 * f1) + fx2 * (1 + fy2) * 4 * f1 - - flux_fs[ipsi, itheta, 1] *= 2π * psio / f4 - flux_fs[ipsi, itheta, 2] *= 2π * psio / f4 - end - end - # Create flux interpolants for Grad-Shafranov diagnostics - @views flux_fs[:, end, :] .= flux_fs[:, 1, :] - flux_opts = (bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) - flux1 = cubic_interp((equil.rzphi_xs, equil.rzphi_ys), flux_fs[:, :, 1]; flux_opts...) - flux2 = cubic_interp((equil.rzphi_xs, equil.rzphi_ys), flux_fs[:, :, 2]; flux_opts...) - - # Compute flux derivatives at all grid points for diagnostics - hint2d = (Ref(1), Ref(1)) # Shared 2D hint for hot loop optimization - for ipsi in 0:mpsi - for itheta in 0:mtheta - query_point = (equil.rzphi_xs[ipsi+1], equil.rzphi_ys[itheta+1]) - flux_fsx[ipsi+1, itheta+1, 1] = flux1(query_point; deriv=DerivOp(1, 0), hint=hint2d) - flux_fsx[ipsi+1, itheta+1, 2] = flux2(query_point; deriv=DerivOp(1, 0), hint=hint2d) - flux_fsy[ipsi+1, itheta+1, 1] = flux1(query_point; deriv=DerivOp(0, 1), hint=hint2d) - flux_fsy[ipsi+1, itheta+1, 2] = flux2(query_point; deriv=DerivOp(0, 1), hint=hint2d) - end - end - - # Compute source term using direct array access - source = zeros(Float64, mpsi + 1, mtheta + 1) - hint = Ref(1) # Linear search hint for sequential psi access - for ipsi in 1:(mpsi+1) - psi = profiles.xs[ipsi] - s1 = profiles.F_spline.y[ipsi] - s1p = profiles.F_deriv(psi; hint=hint) - s2p = profiles.P_deriv(psi; hint=hint) - for itheta in 1:(mtheta+1) - f4 = equil.rzphi_jac.nodal_derivs.partials[1, ipsi, itheta] - denom = (2π * r[ipsi, itheta])^2 - source[ipsi, itheta] = f4 / (2π * psio * π^2) * (s1 * s1p / denom + s2p) - end - end - - total = flux_fsx[:, :, 1] .- flux_fsy[:, :, 2] .+ source - error = abs.(total) ./ maximum([maximum(abs.(flux_fsx[:, :, 1])), maximum(abs.(flux_fsy[:, :, 2])), maximum(abs.(source))]) - errlog = ifelse.(error .> 0, log10.(error), 0.0) - - if diagnose_maxima - fxmax = maximum(abs.(flux_fsx[:, :, 1])) - fymax = maximum(abs.(flux_fsy[:, :, 2])) - smax = maximum(abs.(source)) - emax = maximum(abs.(error)) - lmax = maximum(errlog) - jmax = ind2sub(size(errlog), argmax(errlog)) - @info "GS residuals: fxmax = $(@sprintf("%.3e", fxmax)), fymax = $(@sprintf("%.3e", fymax)), smax = $(@sprintf("%.3e", smax)), emax = $(@sprintf("%.3e", emax)), lmax = $(@sprintf("%.3f", lmax)), maxloc = $(jmax .- 1)" - end - - # Integrated error criterion - term = zeros(Float64, mpsi + 1, 2) - for ipsi in 1:(mpsi+1) - fs_matrix = zeros(Float64, mtheta + 1, 2) - fs_matrix[:, 1] = flux_fsx[ipsi, :, 1] - fs_matrix[:, 2] = source[ipsi, :] - # Snap the repeated endpoint exactly equal to the start - fs_matrix[end, :] .= fs_matrix[1, :] - - # Compute total integral using FastInterpolations native integration - itp = cubic_interp(equil.rzphi_ys, Series(fs_matrix); bc=PeriodicBC()) - term[ipsi, :] .= FastInterpolations.integrate(itp) - end - - totali = sum(term; dims=2) - errori = abs.(totali) - errlogi = @. ifelse(errori > 0, log10(errori), 0.0) - - if diagnose_src - if verbose - @info "Writing diagnostics to HDF5 files" - end - - # Write contour data - h5open(joinpath(dirname(equil.config.eq_filename), "gsec.h5"), "w") do file - file["mpsi"] = mpsi - file["mtheta"] = mtheta - file["r"] = Float32.(r) - file["z"] = Float32.(z) - file["flux_fsx"] = Float32.(flux_fsx[:, :, 1]) - file["flux_fsy"] = Float32.(flux_fsy[:, :, 2]) - file["source"] = Float32.(source) - file["total"] = Float32.(total) - file["error"] = Float32.(error) - file["errlog"] = Float32.(errlog) - end - - # Write xy plot data - h5open(joinpath(dirname(equil.config.eq_filename), "gse.h5"), "w") do file - gse_data = Array{Float32,3}(undef, mpsi + 1, mtheta + 1, 7) - for ipsi in 0:mpsi - for itheta in 0:mtheta - gse_data[ipsi+1, itheta+1, 1] = Float32(flux.ys[itheta+1]) - gse_data[ipsi+1, itheta+1, 2] = Float32(flux.xs[ipsi+1]) - gse_data[ipsi+1, itheta+1, 3] = Float32(flux_fs[ipsi+1, itheta+1, 1]) - gse_data[ipsi+1, itheta+1, 4] = Float32(flux_fs[ipsi+1, itheta+1, 2]) - gse_data[ipsi+1, itheta+1, 5] = Float32(source[ipsi+1, itheta+1]) - gse_data[ipsi+1, itheta+1, 6] = Float32(total[ipsi+1, itheta+1]) - gse_data[ipsi+1, itheta+1, 7] = Float32(error[ipsi+1, itheta+1]) - end - end - file["gse_data"] = gse_data - end - - # Write integrated error criterion - h5open(joinpath(dirname(equil.config.eq_filename), "gsei.h5"), "w") do file - file["xs"] = Float32.(flux.xs) - file["term"] = Float32.(term) - file["totali"] = Float32.(totali) - file["errori"] = Float32.(errori) - file["errlogi"] = Float32.(errlogi) - end - end - return (; xs=equil.rzphi_xs, ys=equil.rzphi_ys, flux_x=flux_fsx[:, :, 1], flux_y=flux_fsy[:, :, 2], source, total, error, errori=vec(errori)) -end diff --git a/src/Equilibrium/ReadEquilibrium.jl b/src/Equilibrium/ReadEquilibrium.jl deleted file mode 100644 index 31aa6e34a..000000000 --- a/src/Equilibrium/ReadEquilibrium.jl +++ /dev/null @@ -1,625 +0,0 @@ -""" -_read_1d_gfile_format(lines_block, num_values) - -Internal helper function to parse Fortran-style fixed-width numerical blocks -from a vector of strings. - -## Arguments: - - - `lines_block`: A `Vector{String}` containing the lines to parse. - - `num_values`: The total number of `Float64` values to read from the block. - -## Returns: - - - A `Vector{Float64}` containing the parsed values. -""" -function _read_1d_gfile_format(lines_block::Vector{String}, num_values::Int) - data_str = join(lines_block) - field_width = 16 - parsed_values = Float64[] - num_read = 0 - - # Ensure the string length is a multiple of the field width for safe processing - safe_len = (length(data_str) ÷ field_width) * field_width - for i in 1:field_width:safe_len - num_read >= num_values && break - val_str = strip(@view(data_str[i:(i+field_width-1)])) - if !isempty(val_str) - try - push!(parsed_values, parse(Float64, val_str)) - num_read += 1 - catch e - @warn "Parsing error for substring: '$val_str'. Error: $e. Skipping." - end - end - end - - if num_read < num_values - @warn "Expected $num_values values, but only read $num_read." - end - return parsed_values -end - -""" - _read_efit(equil_in) - -Parses an EFIT g-file, creates initial 1D and 2D splines, and bundles -them into a `DirectRunInput` object. - -## Arguments: - - - `equil_in`: The `EquilInput` object containing the filename and parameters. - -## Returns: - - - A `DirectRunInput` object ready for the direct solver. -""" -function read_efit(config::EquilibriumConfig) - @info "Processing EFIT g-file: $(config.eq_filename)" - lines = readlines(config.eq_filename) - - # --- Parse Header --- - header1_parts = split(lines[1]) - nw = parse(Int, header1_parts[end-1]) - nh = parse(Int, header1_parts[end]) - @info "Parsed from header: nw=$nw, nh=$nh" - - header_vals = _read_1d_gfile_format(lines[2:5], 20) - rdim, zdim, rcentr, rleft, zmid = header_vals[1:5] - rmaxis, zmaxis, simag, sibry = header_vals[6:9] - ip_sign = Int(sign(header_vals[11])) # the g-file's plasma current; its sign sets the helicity - ip_sign == 0 && (ip_sign = 1) - - # --- Parse Data Blocks --- - current_line_idx = 6 - function parse_block(num_pts) - num_lines = ceil(Int, num_pts / 5) - block = lines[current_line_idx:(current_line_idx+num_lines-1)] - data = _read_1d_gfile_format(block, num_pts) - current_line_idx += num_lines - return data - end - - fpol_data = parse_block(nw) - fpol_sign = Int(sign(fpol_data[end])) # sign of toroidal field (before abs is applied below) - pres_data = parse_block(nw) - ffprime_data = parse_block(nw) - pprime_data = parse_block(nw) - psi_flat_vec = parse_block(nw * nh) - qprof_data = parse_block(nw) - - psi_rz = reshape(psi_flat_vec, nw, nh) - - # --- Create 1D Profile Spline (sq_in) --- - psi_norm_grid = range(0.0, 1.0; length=nw) - sq_fs_nodes = hcat( - abs.(fpol_data), - max.(pres_data .* mu0, 0.0), - qprof_data, - sqrt.(psi_norm_grid) - ) - sq_xs = collect(psi_norm_grid) - sq_in = cubic_interp(sq_xs, Series(sq_fs_nodes); extrap=ExtendExtrap()) - - # --- Process and Normalize 2D Psi Data --- - psio_signed = sibry - simag - psi_proc = (sibry .- psi_rz) - psio = abs(psio_signed) - # Ensure psi at the magnetic axis is positive relative to the boundary - if psio_signed < 0.0 - psi_proc .*= -1.0 - end - - # --- Create 2D Psi interpolant (psi_in) --- - r_grid = range(rleft, rleft + rdim; length=nw) - z_grid = range(zmid - zdim / 2, zmid + zdim / 2; length=nh) - rmin, rmax = extrema(r_grid) - zmin, zmax = extrema(z_grid) - - psi_in_xs = collect(r_grid) - psi_in_ys = collect(z_grid) - psi_in = cubic_interp((psi_in_xs, psi_in_ys), psi_proc; extrap=ExtendExtrap()) - - # Capture the raw arrays that reconstruct sq_in and psi_in, so the rerun path - # (gpec.h5 → PlasmaEquilibrium) can skip the g-file parse entirely. - ingest = DirectIngest(sq_xs, sq_fs_nodes, psi_in_xs, psi_in_ys, psi_proc, rmin, rmax, zmin, zmax, psio, fpol_sign, ip_sign) - - return DirectRunInput(config, sq_in, psi_in, psi_in_xs, psi_in_ys, rmin, rmax, zmin, zmax, psio, fpol_sign, ip_sign, ingest) -end - - -""" - read_chease_binary(equil_config) - -Parses a binary CHEASE file, creates initial 1D and 2D splines with proper -normalization (R0, B0 scaling), and bundles them into a `InverseRunInput` object. -""" -function read_chease_binary(config::EquilibriumConfig) - @info "Reading CHEASE file (Binary): $(config.eq_filename)" - - R0EXP = config.r0_norm - B0EXP = config.b0_norm - - open(config.eq_filename, "r") do io - seekstart(io) - read(io, UInt32) - ntnova = read(io, Int32) - npsi1 = read(io, Int32) - nsym = read(io, Int32) - read(io, UInt32) - - read(io, UInt32) - axx = [read(io, Float64) for _ in 1:5] - read(io, UInt32) - - # 1D array allocation - zcpr, zcppr, zq, zdq, ztmf, ztp, zfb, zfbp, zpsi, zpsim = - [zeros(i == 1 || i == 10 ? npsi1 - 1 : npsi1) for i in 1:10] - - for arr in (zcpr, zcppr, zq, zdq, ztmf, ztp, zfb, zfbp, zpsi, zpsim) - read(io, UInt32) - read!(io, arr) - read(io, UInt32) - end - - # --- Normalization (1D) --- - ztmf .*= (R0EXP * B0EXP) - zcppr .*= (B0EXP / R0EXP^2) - psio_norm = zpsi[npsi1] - zpsi[1] - psio = psio_norm * R0EXP^2 * B0EXP - - ma = npsi1 - 1 - xs = (zpsi .- zpsi[1]) ./ psio_norm - - fs = zeros(npsi1, 4) - fs[:, 1] .= ztmf - fs[:, 2] .= zcppr - fs[:, 3] .= zq - - # Compute cumulative integral of pressure column for normalization using FastInterpolations - itp_pressure = cubic_interp(xs, fs[:, 2]) - fsi_pressure = FastInterpolations.cumulative_integrate(itp_pressure) - # Make a writable copy and normalize pressure integral - fs_copy = copy(fs) - fs_copy[:, 2] .= (fsi_pressure .- fsi_pressure[ma]) .* psio - # Create final spline with modified data - sq_in = cubic_interp(xs, Series(fs_copy); extrap=ExtendExtrap()) - - # --- 2D Geometry --- - mtau = ntnova + 1 # Same with ASCII - poloidal_start = 3 # This 3 is to skip ghost datas, which are used for derivatives - poloidal_stop = ntnova + 3 - - fs_2d = zeros(npsi1, mtau, 2) - buffer = zeros(Float64, ntnova + 3, npsi1) # CHEASE binary record size - - # Reading R data - read(io, UInt32) - read!(io, buffer) - read(io, UInt32) - buffer .*= R0EXP - - # sart reading with ASCII - ro = buffer[poloidal_start, 1] - fs_2d[:, :, 1] .= transpose(buffer[poloidal_start:poloidal_stop, :]) - - # Reading Z data - read(io, UInt32) - read!(io, buffer) - read(io, UInt32) - buffer .*= R0EXP - - zo = buffer[poloidal_start, 1] - fs_2d[:, :, 2] .= transpose(buffer[poloidal_start:poloidal_stop, :]) - - # Create separate interpolants for R and Z coordinates - rz_in_xs = xs - rz_in_ys = range(0, 1; length=mtau) |> collect - rz_in_R = cubic_interp((rz_in_xs, rz_in_ys), fs_2d[:, :, 1]; bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) - rz_in_Z = cubic_interp((rz_in_xs, rz_in_ys), fs_2d[:, :, 2]; bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) - - # Capture the raw arrays that reconstruct sq_in, rz_in_R, rz_in_Z. - ingest = InverseIngest(collect(xs), fs_copy, collect(rz_in_xs), collect(rz_in_ys), fs_2d[:, :, 1], fs_2d[:, :, 2], ro, zo, psio) - - @info "Finished reading CHEASE equilibrium (Binary)" - return InverseRunInput(config, sq_in, rz_in_xs, rz_in_ys, rz_in_R, rz_in_Z, ro, zo, psio, ingest) - end -end - - -""" - read_chease_ascii(config) - -Parses a ascii CHEASE file, creates initial 1D and 2D splines, finds magnetic axis, and bundles -them into a `InverseRunInput` object. - -## Arguments: - - - `config`: The `EquilibriumConfig` object containing the filename and parameters. - -## Returns: - - - A `InverseRunInput` object ready for the inverse solver. -""" -function read_chease_ascii(config::EquilibriumConfig) - @info "Reading CHEASE file (ASCII): $(config.eq_filename)" - lines = readlines(config.eq_filename) - R0EXP = config.r0_norm - B0EXP = config.b0_norm - - # --- Parse Header (FORMAT 10: 3I5) --- - header_parts = split(lines[1]) - ntnova = parse(Int, header_parts[1]) - npsi1 = parse(Int, header_parts[2]) - nsym = parse(Int, header_parts[3]) - - # --- Parse axx (FORMAT 20: 1E22.15) --- - axx = parse(Float64, split(lines[2])[1]) # RBOXLEN - compuational box lenth - # --- Pre-allocate Arrays --- - zcpr = zeros(npsi1 - 1) # normalized P(ψ) - zcppr = zeros(npsi1) # normalized dP/dψ - zq = zeros(npsi1) # q(ψ) - zdq = zeros(npsi1) # dq/dψ - ztmf = zeros(npsi1) # normalized F(ψ) - ztp = zeros(npsi1) # normalized dF/dψ - zfb = zeros(npsi1) # normazlied F(ψ)/q(ψ) - zfbp = zeros(npsi1) # d/dψ [F(ψ)/q(ψ) ] - zpsi = zeros(npsi1) # ψ Poloidal flux - zpsim = zeros(npsi1 - 1) # ψ mid - - zrcp = zeros(ntnova + 3, npsi1) # normalized R - zzcp = zeros(ntnova + 3, npsi1) # normalized Z - zjacm = zeros(ntnova + 3, npsi1) # 𝒥(Jacobian) - zjac = zeros(ntnova + 3, npsi1) # 𝒥(Jacobian - - # --- Helper to parse 5E22.15 data per line --- - function parse_floats(lines_range) - data = Float64[] - for line in lines[lines_range] - for i in 0:4 - s = strip(line[(22*i+1):min(end, 22 * (i + 1))]) - if !isempty(s) - push!(data, parse(Float64, s)) - end - end - end - return data - end - - # --- Compute line offsets --- - line_idx = 3 # Start after header (line 1) and axx (line 2) - - function load_vector!(vec) - count = length(vec) - lines_needed = cld(count, 5) - vec .= parse_floats(line_idx:(line_idx+lines_needed-1)) - return line_idx += lines_needed - end - - function load_matrix!(mat) - count = size(mat, 1) * size(mat, 2) - lines_needed = cld(count, 5) - data = parse_floats(line_idx:(line_idx+lines_needed-1)) - line_idx += lines_needed - # Fill column-major (Fortran-style) - for j in 1:size(mat, 2) - for i in 1:size(mat, 1) - mat[i, j] = data[(j-1)*size(mat, 1)+i] - end - end - end - - # --- Read Vectors --- - load_vector!(zcpr) - load_vector!(zcppr) - load_vector!(zq) - load_vector!(zdq) - load_vector!(ztmf) - load_vector!(ztp) - load_vector!(zfb) - load_vector!(zfbp) - load_vector!(zpsi) - load_vector!(zpsim) - - # --- Read Matrices --- - load_matrix!(zrcp) - load_matrix!(zzcp) - load_matrix!(zjacm) - load_matrix!(zjac) - @info "Parsed from header: ntnova = $ntnova, npsi1 = $npsi1, nsym = $nsym" - - # --- Apply Normalization --- - # Scale geometry - zrcp .*= R0EXP - zzcp .*= R0EXP - - # Scale flux - # Psi_phys = Psi_norm * R0^2 * B0 - psio_norm = zpsi[end] - zpsi[1] - psio = psio_norm * R0EXP^2 * B0EXP - - # Scale Toroidal Field Function F - # zfb .*= (R0EXP * B0EXP) - ztmf .*= (R0EXP * B0EXP) - - # Scale Pressure Gradient P' - zcppr .*= (B0EXP / R0EXP^2) - - # Number of spline intervals - ma = npsi1 - 1 - - # Normalize ψ to [0, 1] - xs = (zpsi .- zpsi[1]) ./ psio_norm # Use normalized range for x axis [0,1] - - fs = zeros(npsi1, 4) - # Both zq * zfb and ztmf are the same ! But I don't know why current GPEC follows zq .*zfb - JB.Cho - # fs[:, 1] .= zq .* zfb - fs[:, 1] .= ztmf - fs[:, 2] .= zcppr # normalized Pressure - fs[:, 3] .= zq # q profile - # Fit spline with extrapolation boundary condition (bctype = 3) - # Compute cumulative integral of pressure column for normalization using FastInterpolations - itp_pressure = cubic_interp(xs, fs[:, 2]) - fsi_pressure = FastInterpolations.cumulative_integrate(itp_pressure) - # Make a writable copy and normalize pressure integral - fs_copy = copy(fs) - fs_copy[:, 2] .= (fsi_pressure .- fsi_pressure[ma]) .* psio - # Create final spline with modified data - sq_in = cubic_interp(xs, Series(fs_copy); extrap=ExtendExtrap()) - - # --- Copy 2D geometry arrays --- - mtau = ntnova + 1 - poloidal_start = 3 - poloidal_stop = ntnova + 3 - ro = zrcp[poloidal_start, 1] # Already scaled - zo = zzcp[poloidal_start, 1] # Already scaled - rz_in_ys = range(0, 1; length=mtau) |> collect - # CHEASE includes 2 ghost points at the start (wrap-around); drop them. - R_data = transpose(zrcp[poloidal_start:poloidal_stop, :]) - Z_data = transpose(zzcp[poloidal_start:poloidal_stop, :]) - - # Create separate interpolants for R and Z coordinates - rz_in_xs = xs - - @views R_data[:, end] .= R_data[:, 1] - @views Z_data[:, end] .= Z_data[:, 1] - - opts2d = (bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) - rz_in_R = cubic_interp((rz_in_xs, rz_in_ys), R_data; opts2d...) - rz_in_Z = cubic_interp((rz_in_xs, rz_in_ys), Z_data; opts2d...) - # Capture the raw arrays that reconstruct sq_in, rz_in_R, rz_in_Z. - ingest = InverseIngest(collect(xs), fs_copy, collect(rz_in_xs), collect(rz_in_ys), Matrix(R_data), Matrix(Z_data), ro, zo, psio) - - @info "Finished reading CHEASE equilibrium. Magnetic axis at (ro=$(@sprintf("%.3f", ro)), zo=$(@sprintf("%.3f", zo))), psio=$(@sprintf("%.3e", psio))" - return InverseRunInput(config, sq_in, rz_in_xs, rz_in_ys, rz_in_R, rz_in_Z, ro, zo, psio, ingest) -end - - -""" - _read_fortran_reals(io, n) -> Vector{Float64} - -Read one little-endian Fortran sequential unformatted record of `n` reals, stored as real*4 or real*8 -(told apart by the record length). -""" -function _read_fortran_reals(io::IO, n::Int) - nbytes = ltoh(read(io, Int32)) - T = nbytes == 8n ? Float64 : nbytes == 4n ? Float32 : error("Fortran record holds $nbytes bytes, expected $n reals") - data = ltoh.(read!(io, Vector{T}(undef, n))) - ltoh(read(io, Int32)) == nbytes || error("Fortran record end marker does not match its start") - return Float64.(data) -end - -""" - read_ifile(config) -> InverseRunInput - -Read an inverse i-file, such as TokaMaker `save_ifile` writes: little-endian Fortran unformatted records, -reals as real*8 or real*4, axis first, θ fastest. - - - `mx, my` (int32); `psi(mx)` [Wb/rad], `f(mx)` = R·Bt [T·m], `p(mx)` [Pa], `q(mx)` - - `r(my, mx)`, `z(my, mx)` [m], with the periodic θ point duplicated - - optional `FF′(mx)`, `p′(mx)` per Wb/rad, skipped - - optional `[ψ_axis, ψ_LCFS]`: ψ_N is then normalized to this LCFS, so the last file surface (written just - inside it) sits below ψ_N = 1 - -Port of Fortran `read_eq_ldp_i`. -""" -function read_ifile(config::EquilibriumConfig) - @info "Reading inverse i-file: $(config.eq_filename)" - - psi, f, p, q, r, z, bounds = open(config.eq_filename, "r") do io - ltoh(read(io, Int32)) == 8 || error("i-file header record is not two int32") - mx, my = Int.(ltoh.(read!(io, Vector{Int32}(undef, 2)))) - read(io, Int32) - profiles = [_read_fortran_reals(io, mx) for _ in 1:4] - r, z = (reshape(_read_fortran_reals(io, my * mx), my, mx) for _ in 1:2) - # Optional records, told apart by length: FF′ and p′ (mx each, unused), then the ψ bounds (2) - bounds = nothing - while !eof(io) - nbytes = ltoh(read(io, Int32)) - if nbytes in (4mx, 8mx) - skip(io, nbytes + 4) - else - skip(io, -4) - bounds = _read_fortran_reals(io, 2) - end - end - return (profiles..., r, z, bounds) - end - mx, my = length(psi), size(r, 1) - @info "Parsed from header: mx = $mx surfaces, my = $my poloidal points" - - psi_axis, psi_lcfs = bounds === nothing ? (psi[1], psi[end]) : bounds - psio_signed = psi_lcfs - psi_axis - psio = abs(psio_signed) - xs = (psi .- psi_axis) ./ psio_signed - bounds === nothing || - @info "i-file ψ bounds: ψ_axis = $psi_axis, ψ_LCFS = $psi_lcfs; file surfaces span ψ_N ∈ [$(xs[1]), $(xs[end])]" - - sq_fs = hcat(abs.(f), p .* mu0, q, sqrt.(xs)) - sq_in = cubic_interp(xs, Series(sq_fs); extrap=ExtendExtrap()) - - R_data = Matrix(transpose(r)) - Z_data = Matrix(transpose(z)) - ro, zo = R_data[1, 1], Z_data[1, 1] - rz_in_ys = collect(range(0, 1; length=my)) - - opts2d = (bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) - rz_in_R = cubic_interp((xs, rz_in_ys), R_data; opts2d...) - rz_in_Z = cubic_interp((xs, rz_in_ys), Z_data; opts2d...) - ingest = InverseIngest(xs, sq_fs, xs, rz_in_ys, R_data, Z_data, ro, zo, psio) - - @info "Finished reading i-file. Magnetic axis at (ro=$(@sprintf("%.3f", ro)), zo=$(@sprintf("%.3f", zo))), psio=$(@sprintf("%.3e", psio))" - return InverseRunInput(config, sq_in, xs, rz_in_ys, rz_in_R, rz_in_Z, ro, zo, psio, ingest) -end - - -""" - build_direct_from_ingest(config::EquilibriumConfig, ingest::DirectIngest) -> DirectRunInput - -Rebuild a `DirectRunInput` from a [`DirectIngest`](@ref) captured by `read_efit`/`read_imas` -(or restored from `Input/RawInputs/Equilibrium/` inside `gpec.h5`). Inverse of that capture: -reconstructs the splines so the rerun path skips the g-file/IMAS parse, reusing the existing -solver dispatch. -""" -function build_direct_from_ingest(config::EquilibriumConfig, ingest::DirectIngest) - sq_in = cubic_interp(ingest.sq_xs, Series(ingest.sq_fs); extrap=ExtendExtrap()) - psi_in = cubic_interp((ingest.psi_xs, ingest.psi_ys), ingest.psi_rz; extrap=ExtendExtrap()) - return DirectRunInput(config, sq_in, psi_in, ingest.psi_xs, ingest.psi_ys, - ingest.rmin, ingest.rmax, ingest.zmin, ingest.zmax, ingest.psio, ingest.bt_sign, ingest.ip_sign, ingest) -end - -""" - build_inverse_from_ingest(config::EquilibriumConfig, ingest::InverseIngest) -> InverseRunInput - -Rebuild an `InverseRunInput` from an [`InverseIngest`](@ref) captured by -`read_chease_ascii`/`read_chease_binary`. Inverse of that capture; used by the rerun path to -skip the CHEASE parse. -""" -function build_inverse_from_ingest(config::EquilibriumConfig, ingest::InverseIngest) - sq_in = cubic_interp(ingest.sq_xs, Series(ingest.sq_fs); extrap=ExtendExtrap()) - opts2d = (bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) - rz_in_R = cubic_interp((ingest.rz_xs, ingest.rz_ys), ingest.R_nodes; opts2d...) - rz_in_Z = cubic_interp((ingest.rz_xs, ingest.rz_ys), ingest.Z_nodes; opts2d...) - return InverseRunInput(config, sq_in, ingest.rz_xs, ingest.rz_ys, rz_in_R, rz_in_Z, - ingest.ro, ingest.zo, ingest.psio, ingest) -end - -""" - read_imas(config::EquilibriumConfig, dd) - -Load an equilibrium from an IMAS data dictionary and return a `DirectRunInput`. - -The `dd.equilibrium.time_slice[]` is used (active time slice). Poloidal flux is -converted from the IMAS COCOS convention (set by `config.imas_cocos`) to the -internal COCOS 2 convention: - - - `imas_cocos = 11` (default, IMAS standard): divide ψ by 2π - - `imas_cocos = 2` (GPEC internal): no conversion - -## Arguments - - - `config`: `EquilibriumConfig` with `eq_type = "imas"` and `imas_cocos` set. - - - `dd`: populated `IMASdd.dd` with `dd.equilibrium.time_slice[]` containing: - - + `global_quantities.psi_axis`, `global_quantities.psi_boundary` - + `profiles_1d.psi`, `profiles_1d.f`, `profiles_1d.pressure`, `profiles_1d.q` - + the `profiles_2d` entry with `grid_type.index == 1` (rectangular), supplying `grid.dim1` (R), `grid.dim2` (Z) and `psi` -""" -function read_imas(config::EquilibriumConfig, dd) - @info "Processing IMAS equilibrium at global_time = $(dd.global_time) s" - - eqt = dd.equilibrium.time_slice[] - - # COCOS conversion: IMAS standard is COCOS 11 (ψ_IMAS = 2π × ψ_internal) - # config.imas_cocos controls the expected convention: - # 11 (default) → divide psi by 2π to get internal (COCOS 2) values - # 2 → data is already in internal convention, no conversion needed - if config.imas_cocos == 11 - cocos_factor = 1.0 / (2π) - @info "Converting IMAS data from COCOS 11 to internal (COCOS 2)" - elseif config.imas_cocos == 2 - cocos_factor = 1.0 - @info "IMAS data in COCOS 2 (no conversion needed)" - else - error("read_imas: unsupported imas_cocos = $(config.imas_cocos). Use 2 or 11.") - end - - psi_axis = eqt.global_quantities.psi_axis * cocos_factor - psi_boundary = eqt.global_quantities.psi_boundary * cocos_factor - - psio = abs(psi_boundary - psi_axis) - - if psio < 1e-10 - error("read_imas: |psi_axis - psi_boundary| = $psio is too small. " * - "Check that dd.equilibrium is properly populated.") - end - - # Extract 1D profiles, converting psi from COCOS 11 to internal - psi_1d = eqt.profiles_1d.psi .* cocos_factor - f_1d = eqt.profiles_1d.f # F(ψ) = R·Bt [T·m], COCOS-independent - p_1d = eqt.profiles_1d.pressure # plasma pressure P(ψ) [Pa], COCOS-independent - q_1d = eqt.profiles_1d.q # safety factor, COCOS-independent - - # Capture toroidal-field sign from the boundary F value before abs() below. - bt_sign = isempty(f_1d) ? 1 : Int(sign(f_1d[end])) - bt_sign == 0 && (bt_sign = 1) - # Plasma-current sign from the IMAS global quantity (missing or zero → +1). - # The current is optional in IMAS, and reading an absent one throws rather than returning missing. - ip_imas = getproperty(eqt.global_quantities, :ip, 0.0) - ip_sign = ip_imas == 0 ? 1 : Int(sign(ip_imas)) - - nw = length(psi_1d) - psi_norm_grid = (psi_1d .- psi_axis) ./ (psi_boundary - psi_axis) - @assert all(diff(psi_norm_grid) .> 0) "read_imas: profiles_1d.psi must increase strictly from psi_axis to psi_boundary" - @assert 0 <= psi_norm_grid[1] && psi_norm_grid[end] <= 1 "read_imas: profiles_1d.psi must lie between psi_axis and psi_boundary" - - # Build equilibrium source terms for spline interpolation - # abs(f_1d): F(ψ) can be negative depending on toroidal field direction convention; - # take absolute value since GPEC uses magnitude F = R·|Bt| - sq_fs_nodes = hcat( - abs.(f_1d), - max.(p_1d .* mu0, 0.0), - q_1d, - sqrt.(psi_norm_grid) - ) - sq_xs = collect(psi_norm_grid) - sq_in = cubic_interp(sq_xs, Series(sq_fs_nodes); extrap=ExtendExtrap()) - - # 2D ψ(R,Z) map - if isempty(eqt.profiles_2d) - error("read_imas: no profiles_2d found in equilibrium time slice. " * - "Ensure the 2D ψ(R,Z) map is stored in dd.equilibrium.") - end - prof2d = findfirst(:rectangular, eqt.profiles_2d) - prof2d === nothing && error("read_imas: no rectangular profiles_2d found.") - r_grid = prof2d.grid.dim1 - z_grid = prof2d.grid.dim2 - - psi_rz = prof2d.psi .* cocos_factor - - # Shift so psi_in = 0 at boundary, psi_in = psio at axis - psi_proc = psi_boundary .- psi_rz - if psi_boundary - psi_axis < 0.0 - psi_proc .*= -1.0 - end - - rmin, rmax = extrema(r_grid) - zmin, zmax = extrema(z_grid) - - psi_in_xs = collect(r_grid) - psi_in_ys = collect(z_grid) - psi_in = cubic_interp((psi_in_xs, psi_in_ys), psi_proc; extrap=ExtendExtrap()) - - @info "IMAS equilibrium loaded:" * - "\n psio = $(round(psio; sigdigits=5)) Wb" * - "\n 1D profile points: nw = $nw" * - "\n 2D grid: nR = $(length(r_grid)), nZ = $(length(z_grid))" * - "\n R ∈ [$(round(rmin; sigdigits=4)), $(round(rmax; sigdigits=4))] m" * - "\n Z ∈ [$(round(zmin; sigdigits=4)), $(round(zmax; sigdigits=4))] m" - - # Capture the raw arrays so an IMAS run can be replayed from gpec.h5 without the dd source. - ingest = DirectIngest(sq_xs, sq_fs_nodes, psi_in_xs, psi_in_ys, Matrix(psi_proc), rmin, rmax, zmin, zmax, psio, bt_sign, ip_sign) - - return DirectRunInput(config, sq_in, psi_in, psi_in_xs, psi_in_ys, rmin, rmax, zmin, zmax, psio, bt_sign, ip_sign, ingest) -end diff --git a/src/Equilibrium/Sources/Analytic/Analytic.jl b/src/Equilibrium/Sources/Analytic/Analytic.jl new file mode 100644 index 000000000..5621ebe9b --- /dev/null +++ b/src/Equilibrium/Sources/Analytic/Analytic.jl @@ -0,0 +1,42 @@ +""" + AnalyticEquilibriumConfig + +Supertype of the parameter structs of the analytic equilibrium kinds. Each subtype is the +`config_type` of one or more `ANALYTIC_EQ` rows, which pair it with an `eq_type`, its TOML +section, and its solver. +""" +abstract type AnalyticEquilibriumConfig end + +include("Solovev.jl") +include("LargeAspectRatio.jl") +include("TJAnalytic.jl") + +""" + AnalyticEqSpec + +Single-source-of-truth entry describing one analytic-equilibrium kind: the `gpec.toml` +section that carries its parameters, the `*Config` type built from that section, and the +run function that turns the config into solver input. `tj_analytic` and `tj_analytic_direct` +share `TJAnalyticConfig` but bind different run functions, so `run_fn` is per-entry rather +than derived from `config_type`. + +## Fields + + - `section::String` — `gpec.toml` section key (e.g. `"SOL_INPUT"`) + - `config_type::DataType` — `*Config` type, constructed via its `(::Dict)` / `(::String)` ctor + - `run_fn::Function` — `(EquilibriumConfig, config) -> eq_input` solver entry point +""" +struct AnalyticEqSpec + section::String + config_type::DataType + run_fn::Function +end + +# Registry of analytic equilibrium kinds. Both `run_input` (fresh runs) and the rerun input +# builder dispatch off this table, so adding a new analytic kind is one new row. +const ANALYTIC_EQ = Dict( + "sol" => AnalyticEqSpec("SOL_INPUT", SolovevConfig, sol_run), + "lar" => AnalyticEqSpec("LAR_INPUT", LargeAspectRatioConfig, lar_run), + "tj_analytic" => AnalyticEqSpec("TJ_ANALYTIC_INPUT", TJAnalyticConfig, tj_analytic_run), + "tj_analytic_direct" => AnalyticEqSpec("TJ_ANALYTIC_INPUT", TJAnalyticConfig, tj_analytic_run_direct) +) diff --git a/src/Equilibrium/Sources/Analytic/LargeAspectRatio.jl b/src/Equilibrium/Sources/Analytic/LargeAspectRatio.jl new file mode 100644 index 000000000..46fdc827f --- /dev/null +++ b/src/Equilibrium/Sources/Analytic/LargeAspectRatio.jl @@ -0,0 +1,271 @@ +""" + LargeAspectRatioConfig(...) + +A mutable struct holding parameters for the Large Aspect Ratio (LAR) plasma equilibrium model. + +## Fields: + + - `r0`: The major radius of the plasma [m]. + - `a`: The minor radius of the plasma [m]. + - `beta0`: The beta value on axis (normalized pressure). + - `q0`: The safety factor on axis. + - `p_pres`: The exponent for the pressure profile, defined as `p00 * (1 - (r / a)^2)^p_pres`. + - `p_sig`: The exponent that determines the shape of the current-related function profile. + - `sigma_type`: The type of sigma profile, can be "default" or "wesson". If "wesson", the sigma profile is defined as `sigma0 * (1 - (r / a)^2)^p_sig`. + - `mtheta`: The number of grid points in the poloidal direction. + - `ma`: The number of grid points in the radial direction. + - `zeroth`: If set to true, it neglects the Shafranov shift +""" +@kwdef mutable struct LargeAspectRatioConfig <: AnalyticEquilibriumConfig + r0::Float64 = 10.0 + a::Float64 = 1.0 + beta0::Float64 = 1e-3 + q0::Float64 = 1.5 + qa::Float64 = 3.6 # Edge safety factor (legacy field; not consumed by current sigma_type options) + B0::Float64 = 1.0 # On-axis toroidal field [T] (scales F and P) + p_pres::Float64 = 2.0 + p_sig::Float64 = 1.0 + sigma_type::String = "default" + mtheta::Int = 128 + ma::Int = 128 + zeroth::Bool = false +end + +""" +Build a `LargeAspectRatioConfig` from a parsed `[LAR_INPUT]` TOML table. +""" +function LargeAspectRatioConfig(input_dict::Dict{String,Any}) + return LargeAspectRatioConfig(; symbolize_keys(input_dict)...) +end + +""" + lar_init_conditions(rmin, sigma_type, params) + +Initializes the starting radius and state vector for solving the LAR ODE system. +Also evaluates the initial derivative using the analytic model. + +## Arguments: + + - `rmin`: Normalized starting radius (as a fraction of `lar_a`). + - `lar_input`: A `LargeAspectRatioConfig` object containing equilibrium parameters. + +## Returns: + + - `r`: Physical radius corresponding to `rmin * lar_a`. + - `y`: Initial state vector of length 5. +""" +function lar_init_conditions(rmin::Float64, lar_input::LargeAspectRatioConfig) + lar_a = lar_input.a + lar_r0 = lar_input.r0 + q0 = lar_input.q0 + + # Ensure rmin is not too small to avoid numerical issues in LAR equations + # Use a physically meaningful minimum (0.001% of minor radius) + if rmin < 1e-6 + @warn "Setting rmin to at least 1e-6 to avoid singularities in LAR equations." + rmin_safe = 1e-6 + else + rmin_safe = rmin + end + + r = rmin_safe * lar_a + y = zeros(5) + + y[1] = r^2 / (lar_r0 * q0) + y[2] = 1.0 + y[3] = y[1] * lar_r0 / 2 + + dy = zeros(5) + lar_der(dy, r, y, lar_input) + + y[5] = dy[5] * r / 4 + + q = r^2 * y[2] / (lar_r0 * y[1]) + y[4] = (q / r)^2 * y[5] / 2 + + return r, y +end + +""" + lar_der(dy, r, y, lar_input) + +Evaluates the spatial derivatives of the LAR (Large Aspect Ratio) equilibrium ODE system +at a given radius `r`, using the current state vector `y` and equilibrium parameters. + +## Arguments: +- `dy`: A preallocated vector where the computed derivatives will be stored (in-place). +- `r`: The radial position at which the derivative is evaluated. +- `y`: The current state vector. +- `lar_input`: A `LargeAspectRatioConfig` object containing equilibrium shaping parameters. + +## Returns: +- 0. The result is stored in-place in `dy`. +""" + +function lar_der(dy::Vector{Float64}, r::Float64, y::Vector{Float64}, lar_input::LargeAspectRatioConfig) + lar_a = lar_input.a + lar_r0 = lar_input.r0 + + q0 = lar_input.q0 + + beta0 = lar_input.beta0 + p00 = beta0 / 2.0 + p_pres = lar_input.p_pres + p_sig = lar_input.p_sig + + sigma_type = lar_input.sigma_type + x = r / lar_a + xfac = 1 - x^2 + pp = -2 * p_pres * p00 * x * xfac^(p_pres - 1) # p' + sigma0 = 2.0 / (q0 * lar_r0) + + sigma = if sigma_type == "wesson" + sigma0 * xfac^p_sig + else + sigma0 / (1 + x^(2 * p_sig))^(1 + 1 / p_sig) + end + + # LAR equations have singularities at r=0; ensure we never evaluate there + @assert r > 0.0 "LAR equations require r > 0 (called with r=$r)" + + bsq = (y[1] / r)^2 + y[2]^2 # B² + q = r^2 * y[2] / (lar_r0 * y[1]) + + dy[1] = -pp / bsq * y[1] + sigma * y[2] * r + dy[2] = -pp / bsq * y[2] - sigma * y[1] / r + dy[3] = y[1] * lar_r0 / r + dy[4] = (q / r)^2 * y[5] / r + dy[5] = y[2] * (r / q)^2 * (r / lar_r0) * (1 - 2 * (lar_r0 * q / y[2])^2 * pp) + + return 0 +end + +""" + lar_run(lar_input) + +Solves the LAR (Large Aspect Ratio) plasma equilibrium ODE system using analytic profiles +defined by `lar_input`, and returns the full solution table including derived quantities. + +## Arguments: +- `lar_input`: A `LargeAspectRatioConfig` object containing profile and geometric parameters. + +## Returns: + - working on it not yet implemented +""" + +function lar_run(equil_input::EquilibriumConfig, lar_input::LargeAspectRatioConfig) + rmin = 1e-4 + lar_a = lar_input.a + lar_r0 = lar_input.r0 + q0 = lar_input.q0 + beta0 = lar_input.beta0 + sigma_type = lar_input.sigma_type + + p00 = beta0 / 2.0 + lar_input.p_pres = max(lar_input.p_pres, 1.001) + + sigma0 = 2.0 / (q0 * lar_r0) + + ma = lar_input.ma + mtau = lar_input.mtheta + + function dydr(du, u, p, r) + lar_input = p + return lar_der(du, r, u, lar_input) + end + + r0, y0 = lar_init_conditions(rmin, lar_input) + tspan = (r0, lar_a) + p = lar_input + + prob = ODEProblem(dydr, y0, tspan, p) + + sol = solve(prob, Rosenbrock23(; autodiff=false); reltol=equil_input.flux_surface_rtol, abstol=equil_abstol(equil_input.flux_surface_rtol), dense=false) + check_equil_solve(sol, "large-aspect-ratio radial profiles") + + r_arr = sol.t + y_mat = reduce(hcat, sol.u)' + steps = length(r_arr) + + temp = zeros(steps, 9) + for i in 1:steps + r = r_arr[i] + @views y = y_mat[i, :] + x = r / lar_a + xfac = 1 - x^2 + pval = p00 * xfac^lar_input.p_pres + sigma = (sigma_type == "wesson") ? + sigma0 * xfac^lar_input.p_sig : + sigma0 / (1 + x^(2 * lar_input.p_sig))^(1 + 1 / lar_input.p_sig) + q = r^2 * y[2] / (lar_r0 * y[1]) + temp[i, 1] = r + temp[i, 2] = y[1] + temp[i, 3] = y[2] + temp[i, 4] = y[3] + temp[i, 5] = y[4] + temp[i, 6] = y[5] + temp[i, 7] = pval + temp[i, 8] = sigma + temp[i, 9] = q + end + + xs_r = temp[:, 1] + fs_r = temp[:, 2:9] + spl = cubic_interp(xs_r, Series(fs_r); extrap=ExtendExtrap()) + spl_deriv = deriv1(spl) + + dr = lar_a / (ma + 1) + r = 0.0 + psio = temp[end, 4] # ψ_edge + + sq_xs = zeros(ma + 1) + sq_fs = zeros(ma + 1, 3) + r_nodes = zeros(ma + 1) + rzphi_y_nodes = range(0.0, 1.0; length=mtau + 1) + rzphi_fs_nodes = zeros(ma + 1, mtau + 1, 2) + + hint = Ref(1) + for ia in 1:(ma+1) + r += dr + r_nodes[ia] = r + f = spl(r; hint=hint) + f1 = spl_deriv(r; hint=hint) + dψdr = f1[3] + + # Fill spline data for sq_in + sq_xs[ia] = f[3] / psio # ψ / psio + sq_fs[ia, 1] = lar_r0 * f[2] # F = R0 * Bphi + sq_fs[ia, 2] = f[6] # P + sq_fs[ia, 3] = f[8] # q + + # Compute Shafranov shift and fill rzphi grid + r2 = -(f[4] * r / f[8]) / dψdr + if lar_input.zeroth + r2 = 0.0 + end + for itau in 1:(mtau+1) + θ = 2π * (itau - 1) / mtau + cosθ, sinθ = cos(θ), sin(θ) + rfac = r + r2 * cosθ + rzphi_fs_nodes[ia, itau, 1] = lar_r0 + rfac * cosθ + rzphi_fs_nodes[ia, itau, 2] = rfac * sinθ + end + end + + sq_in = cubic_interp(sq_xs, Series(sq_fs); extrap=ExtendExtrap()) + # rz_in_xs is ψ_N (see InverseRunInput struct docs). Passing physical r + # works only by accident when lar_a ≈ 1; otherwise the inverse solver + # extrapolates the (R, Z) splines at outer surfaces. + rz_in_xs = sq_xs + rz_in_ys = collect(rzphi_y_nodes) + + @views rzphi_fs_nodes[:, end, :] .= rzphi_fs_nodes[:, 1, :] + + itp_2d_opts = (bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) + rz_in_R = cubic_interp((rz_in_xs, rz_in_ys), rzphi_fs_nodes[:, :, 1]; itp_2d_opts...) + rz_in_Z = cubic_interp((rz_in_xs, rz_in_ys), rzphi_fs_nodes[:, :, 2]; itp_2d_opts...) + + # LAR equilibrium has midplane symmetry, so magnetic axis is at Z = 0. + # Analytic: ingest=nothing — replay regenerates from the [LAR_INPUT] TOML section. + return InverseRunInput(equil_input, sq_in, rz_in_xs, rz_in_ys, rz_in_R, rz_in_Z, lar_r0, 0.0, psio, nothing) +end diff --git a/src/Equilibrium/Sources/Analytic/Solovev.jl b/src/Equilibrium/Sources/Analytic/Solovev.jl new file mode 100644 index 000000000..60b628a48 --- /dev/null +++ b/src/Equilibrium/Sources/Analytic/Solovev.jl @@ -0,0 +1,118 @@ +""" + SolovevConfig(...) + +A mutable struct holding parameters for the Solev'ev (SOL) plasma equilibrium model. + +## Fields: + + - `mr`: number of radial grid zones + - `mz`: number of axial grid zones + - `ma`: number of flux grid zones + - `elongation`: elongation + - `a`: minor radius + - `r0`: major radius + - `q0`: safety factor at the o-point + - `p0fac`: scale on-axis pressure (P-> P+P0*p0fac. beta changes. Phi,q constant) + - `b0fac`: scale toroidal field at constant beta (s*Phi,s*f,s^2*P. bt changes. Shape,beta constant) + - `f0fac`: scale toroidal field at constant pressure (s*f. beta,q changes. Phi,p,bp constant) +""" +@kwdef mutable struct SolovevConfig <: AnalyticEquilibriumConfig + mr::Int = 128 # number of radial grid zones + mz::Int = 128 # number of axial grid zones + ma::Int = 128 # number of flux grid zones + elongation::Float64 = 1.6 # elongation + a::Float64 = 0.33 # minor radius + r0::Float64 = 1.0 # major radius + q0::Float64 = 1.9 # safety factor at the o-point + p0fac::Float64 = 1 # scale on-axis pressure (P-> P+P0*p0fac. beta changes. Phi,q constant) + b0fac::Float64 = 1 # scale toroidal field at constant beta (s*Phi,s*f,s^2*P. bt changes. Shape,beta constant) + f0fac::Float64 = 1 # scale toroidal field at constant pressure (s*f. beta,q changes. Phi,p,bp constant) +end + +""" +Build a `SolovevConfig` from a parsed `[SOL_INPUT]` TOML table. +""" +function SolovevConfig(input_dict::Dict{String,Any}) + return SolovevConfig(; symbolize_keys(input_dict)...) +end + +""" +This function handles the Solovev analytical equilibrium model, transforming the input parameters +into the necessary splines and scalar values for equilibrium construction. This is a Julia version +of the Fortran code in sol.f, with no major differences except for arrays going from 0:n to 1:n+1. + +## Arguments: + + - `mr`: Number of radial grid zones + - `mz`: Number of axial grid zones + - `ma`: Number of flux grid zones + - `e`: Elongation + - `a`: Minor radius + - `r0`: Major radius + - `q0`: Safety factor at the o-point + - `p0fac`: Scales on axis pressure (s*P. beta changes. Phi,q constant) + - `b0fac`: Scales on toroidal field (s*Phi,s*f,s^2*P. bt changes. Shape,beta constant) + - `f0fac`: Scales on toroidal field (s*f. bt,q changes. Phi,p,bp constant) + +## Returns: + + - `DirectRunInput` object +""" +function sol_run(equil_inputs::EquilibriumConfig, sol_inputs::SolovevConfig) + + mr = sol_inputs.mr + mz = sol_inputs.mz + ma = sol_inputs.ma + e = sol_inputs.elongation + a = sol_inputs.a + r0 = sol_inputs.r0 + q0 = sol_inputs.q0 + p0fac = sol_inputs.p0fac + b0fac = sol_inputs.b0fac + f0fac = sol_inputs.f0fac + + # Validate inputs + if p0fac < 1.0 + @warn "Forcing p0fac ≥ 1 (no negative pressure)" + p0fac = 1.0 + end + + # Compute scalar data + f0 = r0 * b0fac + psio = e * f0 * a * a / (2 * q0 * r0) + psifac = psio / (a * r0)^2 + efac = 1 / (e * e) + pfac = 2 * psio^2 * (e * e + 1) / (a * r0 * e)^2 + rmin = r0 - 1.5 * a + rmax = r0 + 1.5 * a + zmax = 1.5 * e * a + zmin = -zmax + + # zeros, not undef: only columns 1-3 are filled, so column 4 must stay a deterministic 0. + sqfs = zeros(ma + 1, 4) + psis = [(ia / (ma + 1))^2 for ia in 1:(ma+1)] + sqfs[:, 1] .= f0 .* f0fac + sqfs[:, 2] .= pfac .* (1 .* p0fac .- psis) + sqfs[:, 3] .= 0.0 + sq_in = cubic_interp(psis, Series(sqfs); extrap=ExtendExtrap()) + + # Compute 2D data and spline + r = [rmin + i * (rmax - rmin) / mr for i in 0:mr] + z = [zmin + j * (zmax - zmin) / mz for j in 0:mz] + psifs = Array{Float64}(undef, mr + 1, mz + 1) + for iz in 1:(mz+1) + for ir in 1:(mr+1) + psifs[ir, iz] = psio - psifac * (efac * (r[ir] * z[iz])^2 + (r[ir]^2 - r0^2)^2 / 4) + end + end + psi_in_xs = r + psi_in_ys = z + psi_in = cubic_interp((psi_in_xs, psi_in_ys), psifs; extrap=ExtendExtrap()) + + # Print out equilibrium info + @info "Generating Solovev equilibrium: mr=$mr, mz=$mz, ma=$ma, e=$(@sprintf("%.3f", e)), a=$(@sprintf("%.3f", a)), r0=$(@sprintf("%.3f", r0)), q0=$(@sprintf("%.3f", q0))" + + # 1 is bt_sign=+1: Solovev has positive Bt by construction (f0 = r0 * b0fac > 0). + # Analytic equilibria carry positive field and current; ingest=nothing — replay regenerates from the [SOL_INPUT] TOML section. + return DirectRunInput(equil_inputs, sq_in, psi_in, psi_in_xs, psi_in_ys, rmin, rmax, zmin, zmax, psio, 1, 1, nothing) +end diff --git a/src/Equilibrium/AnalyticEquilibrium.jl b/src/Equilibrium/Sources/Analytic/TJAnalytic.jl similarity index 68% rename from src/Equilibrium/AnalyticEquilibrium.jl rename to src/Equilibrium/Sources/Analytic/TJAnalytic.jl index 22e246a79..94458eb24 100644 --- a/src/Equilibrium/AnalyticEquilibrium.jl +++ b/src/Equilibrium/Sources/Analytic/TJAnalytic.jl @@ -1,233 +1,43 @@ """ - lar_init_conditions(rmin, sigma_type, params) + TJAnalyticConfig(...) -Initializes the starting radius and state vector for solving the LAR ODE system. -Also evaluates the initial derivative using the analytic model. +Parameters for the **TJ-analytic** cylindrical large-aspect-ratio equilibrium +model — a GPEC adaptation of the analytic profile family used by +R. Fitzpatrick's TJ code (https://github.com/rfitzp/TJ). We follow the +same analytic-profile parameterization (ψ-ODE in dimensionless r/a, f₁ +for q, power-law pressure) for the inner cylindrical core and connect it +to GPEC's direct-GS pipeline; this is NOT a re-implementation of TJ. -## Arguments: +The model uses analytic profiles with exact control of both the on-axis +and edge safety factors. The q profile is determined by: - - `rmin`: Normalized starting radius (as a fraction of `lar_a`). - - `lar_input`: A `LargeAspectRatioConfig` object containing equilibrium parameters. + f1(r) = [1 - (1-r²)^ν] / (ν·q0) + q(r) = r² / f1(r) -## Returns: +where ν = qa/q0 is the current peaking parameter, q0 is the axis q, and qa +is the edge q. All lengths are normalized to R₀, fields to B₀. The pressure +profile is p₂(r) = p0·(1-r²)^μ with μ = `pressure_peaking`. - - `r`: Physical radius corresponding to `rmin * lar_a`. - - `y`: Initial state vector of length 5. +Reference: R. Fitzpatrick, TJ code, https://github.com/rfitzp/TJ """ -function lar_init_conditions(rmin::Float64, lar_input::LargeAspectRatioConfig) - lar_a = lar_input.a - lar_r0 = lar_input.r0 - q0 = lar_input.q0 - - # Ensure rmin is not too small to avoid numerical issues in LAR equations - # Use a physically meaningful minimum (0.001% of minor radius) - if rmin < 1e-6 - @warn "Setting rmin to at least 1e-6 to avoid singularities in LAR equations." - rmin_safe = 1e-6 - else - rmin_safe = rmin - end - - r = rmin_safe * lar_a - y = zeros(5) - - y[1] = r^2 / (lar_r0 * q0) - y[2] = 1.0 - y[3] = y[1] * lar_r0 / 2 - - dy = zeros(5) - lar_der(dy, r, y, lar_input) - - y[5] = dy[5] * r / 4 - - q = r^2 * y[2] / (lar_r0 * y[1]) - y[4] = (q / r)^2 * y[5] / 2 - - return r, y +@kwdef mutable struct TJAnalyticConfig <: AnalyticEquilibriumConfig + r0::Float64 = 10.0 # Major radius R₀ [m] + a::Float64 = 1.0 # Minor radius a [m] (ε = a/R₀) + q0::Float64 = 1.5 # On-axis safety factor + qa::Float64 = 3.6 # Edge safety factor + p0::Float64 = 0.001 # Normalized on-axis pressure + pressure_peaking::Float64 = 2.0 # Pressure peaking exponent μ: p₂ = p0·(1-r²)^μ + B0::Float64 = 12.0 # On-axis toroidal field [T] + ma::Int = 128 # Radial grid points + mtheta::Int = 128 # Poloidal grid points + zeroth::Bool = false # If true, suppress Shafranov shift end """ - lar_der(dy, r, y, lar_input) - -Evaluates the spatial derivatives of the LAR (Large Aspect Ratio) equilibrium ODE system -at a given radius `r`, using the current state vector `y` and equilibrium parameters. - -## Arguments: -- `dy`: A preallocated vector where the computed derivatives will be stored (in-place). -- `r`: The radial position at which the derivative is evaluated. -- `y`: The current state vector. -- `lar_input`: A `LargeAspectRatioConfig` object containing equilibrium shaping parameters. - -## Returns: -- 0. The result is stored in-place in `dy`. +Build a `TJAnalyticConfig` from a parsed `[TJ_ANALYTIC_INPUT]` TOML table. """ - -function lar_der(dy::Vector{Float64}, r::Float64, y::Vector{Float64}, lar_input::LargeAspectRatioConfig) - lar_a = lar_input.a - lar_r0 = lar_input.r0 - - q0 = lar_input.q0 - - beta0 = lar_input.beta0 - p00 = beta0 / 2.0 - p_pres = lar_input.p_pres - p_sig = lar_input.p_sig - - sigma_type = lar_input.sigma_type - x = r / lar_a - xfac = 1 - x^2 - pp = -2 * p_pres * p00 * x * xfac^(p_pres - 1) # p' - sigma0 = 2.0 / (q0 * lar_r0) - - sigma = if sigma_type == "wesson" - sigma0 * xfac^p_sig - else - sigma0 / (1 + x^(2 * p_sig))^(1 + 1 / p_sig) - end - - # LAR equations have singularities at r=0; ensure we never evaluate there - @assert r > 0.0 "LAR equations require r > 0 (called with r=$r)" - - bsq = (y[1] / r)^2 + y[2]^2 # B² - q = r^2 * y[2] / (lar_r0 * y[1]) - - dy[1] = -pp / bsq * y[1] + sigma * y[2] * r - dy[2] = -pp / bsq * y[2] - sigma * y[1] / r - dy[3] = y[1] * lar_r0 / r - dy[4] = (q / r)^2 * y[5] / r - dy[5] = y[2] * (r / q)^2 * (r / lar_r0) * (1 - 2 * (lar_r0 * q / y[2])^2 * pp) - - return 0 -end - -""" - lar_run(lar_input) - -Solves the LAR (Large Aspect Ratio) plasma equilibrium ODE system using analytic profiles -defined by `lar_input`, and returns the full solution table including derived quantities. - -## Arguments: -- `lar_input`: A `LargeAspectRatioConfig` object containing profile and geometric parameters. - -## Returns: - - working on it not yet implemented -""" - -function lar_run(equil_input::EquilibriumConfig, lar_input::LargeAspectRatioConfig) - rmin = 1e-4 - lar_a = lar_input.a - lar_r0 = lar_input.r0 - q0 = lar_input.q0 - beta0 = lar_input.beta0 - sigma_type = lar_input.sigma_type - - p00 = beta0 / 2.0 - lar_input.p_pres = max(lar_input.p_pres, 1.001) - - sigma0 = 2.0 / (q0 * lar_r0) - - ma = lar_input.ma - mtau = lar_input.mtheta - - function dydr(du, u, p, r) - lar_input = p - return lar_der(du, r, u, lar_input) - end - - r0, y0 = lar_init_conditions(rmin, lar_input) - tspan = (r0, lar_a) - p = lar_input - - prob = ODEProblem(dydr, y0, tspan, p) - - sol = solve(prob, Rosenbrock23(; autodiff=false); reltol=equil_input.flux_surface_rtol, abstol=equil_abstol(equil_input.flux_surface_rtol), dense=false) - check_equil_solve(sol, "large-aspect-ratio radial profiles") - - r_arr = sol.t - y_mat = reduce(hcat, sol.u)' - steps = length(r_arr) - - temp = zeros(steps, 9) - for i in 1:steps - r = r_arr[i] - @views y = y_mat[i, :] - x = r / lar_a - xfac = 1 - x^2 - pval = p00 * xfac^lar_input.p_pres - sigma = (sigma_type == "wesson") ? - sigma0 * xfac^lar_input.p_sig : - sigma0 / (1 + x^(2 * lar_input.p_sig))^(1 + 1 / lar_input.p_sig) - q = r^2 * y[2] / (lar_r0 * y[1]) - temp[i, 1] = r - temp[i, 2] = y[1] - temp[i, 3] = y[2] - temp[i, 4] = y[3] - temp[i, 5] = y[4] - temp[i, 6] = y[5] - temp[i, 7] = pval - temp[i, 8] = sigma - temp[i, 9] = q - end - - xs_r = temp[:, 1] - fs_r = temp[:, 2:9] - spl = cubic_interp(xs_r, Series(fs_r); extrap=ExtendExtrap()) - spl_deriv = deriv1(spl) - - dr = lar_a / (ma + 1) - r = 0.0 - psio = temp[end, 4] # ψ_edge - - sq_xs = zeros(ma + 1) - sq_fs = zeros(ma + 1, 3) - r_nodes = zeros(ma + 1) - rzphi_y_nodes = range(0.0, 1.0; length=mtau + 1) - rzphi_fs_nodes = zeros(ma + 1, mtau + 1, 2) - - hint = Ref(1) - for ia in 1:(ma+1) - r += dr - r_nodes[ia] = r - f = spl(r; hint=hint) - f1 = spl_deriv(r; hint=hint) - dψdr = f1[3] - - # Fill spline data for sq_in - sq_xs[ia] = f[3] / psio # ψ / psio - sq_fs[ia, 1] = lar_r0 * f[2] # F = R0 * Bphi - sq_fs[ia, 2] = f[6] # P - sq_fs[ia, 3] = f[8] # q - - # Compute Shafranov shift and fill rzphi grid - r2 = -(f[4] * r / f[8]) / dψdr - if lar_input.zeroth - r2 = 0.0 - end - for itau in 1:(mtau+1) - θ = 2π * (itau - 1) / mtau - cosθ, sinθ = cos(θ), sin(θ) - rfac = r + r2 * cosθ - rzphi_fs_nodes[ia, itau, 1] = lar_r0 + rfac * cosθ - rzphi_fs_nodes[ia, itau, 2] = rfac * sinθ - end - end - - sq_in = cubic_interp(sq_xs, Series(sq_fs); extrap=ExtendExtrap()) - # rz_in_xs is ψ_N (see InverseRunInput struct docs). Passing physical r - # works only by accident when lar_a ≈ 1; otherwise the inverse solver - # extrapolates the (R, Z) splines at outer surfaces. - rz_in_xs = sq_xs - rz_in_ys = collect(rzphi_y_nodes) - - @views rzphi_fs_nodes[:, end, :] .= rzphi_fs_nodes[:, 1, :] - - itp_2d_opts = (bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) - rz_in_R = cubic_interp((rz_in_xs, rz_in_ys), rzphi_fs_nodes[:, :, 1]; itp_2d_opts...) - rz_in_Z = cubic_interp((rz_in_xs, rz_in_ys), rzphi_fs_nodes[:, :, 2]; itp_2d_opts...) - - # LAR equilibrium has midplane symmetry, so magnetic axis is at Z = 0. - # Analytic: ingest=nothing — replay regenerates from the [LAR_INPUT] TOML section. - return InverseRunInput(equil_input, sq_in, rz_in_xs, rz_in_ys, rz_in_R, rz_in_Z, lar_r0, 0.0, psio, nothing) +function TJAnalyticConfig(input_dict::Dict{String,Any}) + return TJAnalyticConfig(; symbolize_keys(input_dict)...) end """ @@ -541,7 +351,7 @@ function tj_analytic_run(equil_input::EquilibriumConfig, tj::TJAnalyticConfig) end sq_in = cubic_interp(sq_xs, Series(sq_fs); extrap=ExtendExtrap()) - # InverseRunInput's rz_in_xs is specified as ψ_N (see EquilibriumTypes.jl docs); + # InverseRunInput's rz_in_xs is specified as ψ_N (see its docstring in Sources/Ingest.jl); # the inverse solver queries (R, Z) splines at ψ_N values from sq_xs. Passing # physical r here happens to work when a ≈ 1 (r and ψ_N cover the same range) # but extrapolates the (R, Z) splines for any a < 1, corrupting outer surfaces. @@ -771,84 +581,3 @@ function tj_analytic_run_direct(equil_input::EquilibriumConfig, tj::TJAnalyticCo return DirectRunInput(equil_input, sq_in, psi_in, psi_in_xs, psi_in_ys, rmin_grid, rmax_grid, zmin_grid, zmax_grid, psio, 1, 1, nothing) end - -""" -This function handles the Solovev analytical equilibrium model, transforming the input parameters -into the necessary splines and scalar values for equilibrium construction. This is a Julia version -of the Fortran code in sol.f, with no major differences except for arrays going from 0:n to 1:n+1. - -## Arguments: - - - `mr`: Number of radial grid zones - - `mz`: Number of axial grid zones - - `ma`: Number of flux grid zones - - `e`: Elongation - - `a`: Minor radius - - `r0`: Major radius - - `q0`: Safety factor at the o-point - - `p0fac`: Scales on axis pressure (s*P. beta changes. Phi,q constant) - - `b0fac`: Scales on toroidal field (s*Phi,s*f,s^2*P. bt changes. Shape,beta constant) - - `f0fac`: Scales on toroidal field (s*f. bt,q changes. Phi,p,bp constant) - -## Returns: - - - `DirectRunInput` object -""" -function sol_run(equil_inputs::EquilibriumConfig, sol_inputs::SolovevConfig) - - mr = sol_inputs.mr - mz = sol_inputs.mz - ma = sol_inputs.ma - e = sol_inputs.elongation - a = sol_inputs.a - r0 = sol_inputs.r0 - q0 = sol_inputs.q0 - p0fac = sol_inputs.p0fac - b0fac = sol_inputs.b0fac - f0fac = sol_inputs.f0fac - - # Validate inputs - if p0fac < 1.0 - @warn "Forcing p0fac ≥ 1 (no negative pressure)" - p0fac = 1.0 - end - - # Compute scalar data - f0 = r0 * b0fac - psio = e * f0 * a * a / (2 * q0 * r0) - psifac = psio / (a * r0)^2 - efac = 1 / (e * e) - pfac = 2 * psio^2 * (e * e + 1) / (a * r0 * e)^2 - rmin = r0 - 1.5 * a - rmax = r0 + 1.5 * a - zmax = 1.5 * e * a - zmin = -zmax - - # zeros, not undef: only columns 1-3 are filled, so column 4 must stay a deterministic 0. - sqfs = zeros(ma + 1, 4) - psis = [(ia / (ma + 1))^2 for ia in 1:(ma+1)] - sqfs[:, 1] .= f0 .* f0fac - sqfs[:, 2] .= pfac .* (1 .* p0fac .- psis) - sqfs[:, 3] .= 0.0 - sq_in = cubic_interp(psis, Series(sqfs); extrap=ExtendExtrap()) - - # Compute 2D data and spline - r = [rmin + i * (rmax - rmin) / mr for i in 0:mr] - z = [zmin + j * (zmax - zmin) / mz for j in 0:mz] - psifs = Array{Float64}(undef, mr + 1, mz + 1) - for iz in 1:(mz+1) - for ir in 1:(mr+1) - psifs[ir, iz] = psio - psifac * (efac * (r[ir] * z[iz])^2 + (r[ir]^2 - r0^2)^2 / 4) - end - end - psi_in_xs = r - psi_in_ys = z - psi_in = cubic_interp((psi_in_xs, psi_in_ys), psifs; extrap=ExtendExtrap()) - - # Print out equilibrium info - @info "Generating Solovev equilibrium: mr=$mr, mz=$mz, ma=$ma, e=$(@sprintf("%.3f", e)), a=$(@sprintf("%.3f", a)), r0=$(@sprintf("%.3f", r0)), q0=$(@sprintf("%.3f", q0))" - - # 1 is bt_sign=+1: Solovev has positive Bt by construction (f0 = r0 * b0fac > 0). - # Analytic equilibria carry positive field and current; ingest=nothing — replay regenerates from the [SOL_INPUT] TOML section. - return DirectRunInput(equil_inputs, sq_in, psi_in, psi_in_xs, psi_in_ys, rmin, rmax, zmin, zmax, psio, 1, 1, nothing) -end diff --git a/src/Equilibrium/Sources/Chease.jl b/src/Equilibrium/Sources/Chease.jl new file mode 100644 index 000000000..22960e5f2 --- /dev/null +++ b/src/Equilibrium/Sources/Chease.jl @@ -0,0 +1,265 @@ +""" + read_chease_binary(equil_config) + +Parses a binary CHEASE file, creates initial 1D and 2D splines with proper +normalization (R0, B0 scaling), and bundles them into a `InverseRunInput` object. +""" +function read_chease_binary(config::EquilibriumConfig) + @info "Reading CHEASE file (Binary): $(config.eq_filename)" + + R0EXP = config.r0_norm + B0EXP = config.b0_norm + + open(config.eq_filename, "r") do io + seekstart(io) + read(io, UInt32) + ntnova = read(io, Int32) + npsi1 = read(io, Int32) + nsym = read(io, Int32) + read(io, UInt32) + + read(io, UInt32) + axx = [read(io, Float64) for _ in 1:5] + read(io, UInt32) + + # 1D array allocation + zcpr, zcppr, zq, zdq, ztmf, ztp, zfb, zfbp, zpsi, zpsim = + [zeros(i == 1 || i == 10 ? npsi1 - 1 : npsi1) for i in 1:10] + + for arr in (zcpr, zcppr, zq, zdq, ztmf, ztp, zfb, zfbp, zpsi, zpsim) + read(io, UInt32) + read!(io, arr) + read(io, UInt32) + end + + # --- Normalization (1D) --- + ztmf .*= (R0EXP * B0EXP) + zcppr .*= (B0EXP / R0EXP^2) + psio_norm = zpsi[npsi1] - zpsi[1] + psio = psio_norm * R0EXP^2 * B0EXP + + ma = npsi1 - 1 + xs = (zpsi .- zpsi[1]) ./ psio_norm + + fs = zeros(npsi1, 4) + fs[:, 1] .= ztmf + fs[:, 2] .= zcppr + fs[:, 3] .= zq + + # Compute cumulative integral of pressure column for normalization using FastInterpolations + itp_pressure = cubic_interp(xs, fs[:, 2]) + fsi_pressure = FastInterpolations.cumulative_integrate(itp_pressure) + # Make a writable copy and normalize pressure integral + fs_copy = copy(fs) + fs_copy[:, 2] .= (fsi_pressure .- fsi_pressure[ma]) .* psio + # Create final spline with modified data + sq_in = cubic_interp(xs, Series(fs_copy); extrap=ExtendExtrap()) + + # --- 2D Geometry --- + mtau = ntnova + 1 # Same with ASCII + poloidal_start = 3 # This 3 is to skip ghost datas, which are used for derivatives + poloidal_stop = ntnova + 3 + + fs_2d = zeros(npsi1, mtau, 2) + buffer = zeros(Float64, ntnova + 3, npsi1) # CHEASE binary record size + + # Reading R data + read(io, UInt32) + read!(io, buffer) + read(io, UInt32) + buffer .*= R0EXP + + # sart reading with ASCII + ro = buffer[poloidal_start, 1] + fs_2d[:, :, 1] .= transpose(buffer[poloidal_start:poloidal_stop, :]) + + # Reading Z data + read(io, UInt32) + read!(io, buffer) + read(io, UInt32) + buffer .*= R0EXP + + zo = buffer[poloidal_start, 1] + fs_2d[:, :, 2] .= transpose(buffer[poloidal_start:poloidal_stop, :]) + + # Create separate interpolants for R and Z coordinates + rz_in_xs = xs + rz_in_ys = range(0, 1; length=mtau) |> collect + rz_in_R = cubic_interp((rz_in_xs, rz_in_ys), fs_2d[:, :, 1]; bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) + rz_in_Z = cubic_interp((rz_in_xs, rz_in_ys), fs_2d[:, :, 2]; bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) + + # Capture the raw arrays that reconstruct sq_in, rz_in_R, rz_in_Z. + ingest = InverseIngest(collect(xs), fs_copy, collect(rz_in_xs), collect(rz_in_ys), fs_2d[:, :, 1], fs_2d[:, :, 2], ro, zo, psio) + + @info "Finished reading CHEASE equilibrium (Binary)" + return InverseRunInput(config, sq_in, rz_in_xs, rz_in_ys, rz_in_R, rz_in_Z, ro, zo, psio, ingest) + end +end + + +""" + read_chease_ascii(config) + +Parses a ascii CHEASE file, creates initial 1D and 2D splines, finds magnetic axis, and bundles +them into a `InverseRunInput` object. + +## Arguments: + + - `config`: The `EquilibriumConfig` object containing the filename and parameters. + +## Returns: + + - A `InverseRunInput` object ready for the inverse solver. +""" +function read_chease_ascii(config::EquilibriumConfig) + @info "Reading CHEASE file (ASCII): $(config.eq_filename)" + lines = readlines(config.eq_filename) + R0EXP = config.r0_norm + B0EXP = config.b0_norm + + # --- Parse Header (FORMAT 10: 3I5) --- + header_parts = split(lines[1]) + ntnova = parse(Int, header_parts[1]) + npsi1 = parse(Int, header_parts[2]) + nsym = parse(Int, header_parts[3]) + + # --- Parse axx (FORMAT 20: 1E22.15) --- + axx = parse(Float64, split(lines[2])[1]) # RBOXLEN - compuational box lenth + # --- Pre-allocate Arrays --- + zcpr = zeros(npsi1 - 1) # normalized P(ψ) + zcppr = zeros(npsi1) # normalized dP/dψ + zq = zeros(npsi1) # q(ψ) + zdq = zeros(npsi1) # dq/dψ + ztmf = zeros(npsi1) # normalized F(ψ) + ztp = zeros(npsi1) # normalized dF/dψ + zfb = zeros(npsi1) # normazlied F(ψ)/q(ψ) + zfbp = zeros(npsi1) # d/dψ [F(ψ)/q(ψ) ] + zpsi = zeros(npsi1) # ψ Poloidal flux + zpsim = zeros(npsi1 - 1) # ψ mid + + zrcp = zeros(ntnova + 3, npsi1) # normalized R + zzcp = zeros(ntnova + 3, npsi1) # normalized Z + zjacm = zeros(ntnova + 3, npsi1) # 𝒥(Jacobian) + zjac = zeros(ntnova + 3, npsi1) # 𝒥(Jacobian + + # --- Helper to parse 5E22.15 data per line --- + function parse_floats(lines_range) + data = Float64[] + for line in lines[lines_range] + for i in 0:4 + s = strip(line[(22*i+1):min(end, 22 * (i + 1))]) + if !isempty(s) + push!(data, parse(Float64, s)) + end + end + end + return data + end + + # --- Compute line offsets --- + line_idx = 3 # Start after header (line 1) and axx (line 2) + + function load_vector!(vec) + count = length(vec) + lines_needed = cld(count, 5) + vec .= parse_floats(line_idx:(line_idx+lines_needed-1)) + return line_idx += lines_needed + end + + function load_matrix!(mat) + count = size(mat, 1) * size(mat, 2) + lines_needed = cld(count, 5) + data = parse_floats(line_idx:(line_idx+lines_needed-1)) + line_idx += lines_needed + # Fill column-major (Fortran-style) + for j in 1:size(mat, 2) + for i in 1:size(mat, 1) + mat[i, j] = data[(j-1)*size(mat, 1)+i] + end + end + end + + # --- Read Vectors --- + load_vector!(zcpr) + load_vector!(zcppr) + load_vector!(zq) + load_vector!(zdq) + load_vector!(ztmf) + load_vector!(ztp) + load_vector!(zfb) + load_vector!(zfbp) + load_vector!(zpsi) + load_vector!(zpsim) + + # --- Read Matrices --- + load_matrix!(zrcp) + load_matrix!(zzcp) + load_matrix!(zjacm) + load_matrix!(zjac) + @info "Parsed from header: ntnova = $ntnova, npsi1 = $npsi1, nsym = $nsym" + + # --- Apply Normalization --- + # Scale geometry + zrcp .*= R0EXP + zzcp .*= R0EXP + + # Scale flux + # Psi_phys = Psi_norm * R0^2 * B0 + psio_norm = zpsi[end] - zpsi[1] + psio = psio_norm * R0EXP^2 * B0EXP + + # Scale Toroidal Field Function F + # zfb .*= (R0EXP * B0EXP) + ztmf .*= (R0EXP * B0EXP) + + # Scale Pressure Gradient P' + zcppr .*= (B0EXP / R0EXP^2) + + # Number of spline intervals + ma = npsi1 - 1 + + # Normalize ψ to [0, 1] + xs = (zpsi .- zpsi[1]) ./ psio_norm # Use normalized range for x axis [0,1] + + fs = zeros(npsi1, 4) + # Both zq * zfb and ztmf are the same ! But I don't know why current GPEC follows zq .*zfb - JB.Cho + # fs[:, 1] .= zq .* zfb + fs[:, 1] .= ztmf + fs[:, 2] .= zcppr # normalized Pressure + fs[:, 3] .= zq # q profile + # Fit spline with extrapolation boundary condition (bctype = 3) + # Compute cumulative integral of pressure column for normalization using FastInterpolations + itp_pressure = cubic_interp(xs, fs[:, 2]) + fsi_pressure = FastInterpolations.cumulative_integrate(itp_pressure) + # Make a writable copy and normalize pressure integral + fs_copy = copy(fs) + fs_copy[:, 2] .= (fsi_pressure .- fsi_pressure[ma]) .* psio + # Create final spline with modified data + sq_in = cubic_interp(xs, Series(fs_copy); extrap=ExtendExtrap()) + + # --- Copy 2D geometry arrays --- + mtau = ntnova + 1 + poloidal_start = 3 + poloidal_stop = ntnova + 3 + ro = zrcp[poloidal_start, 1] # Already scaled + zo = zzcp[poloidal_start, 1] # Already scaled + rz_in_ys = range(0, 1; length=mtau) |> collect + # CHEASE includes 2 ghost points at the start (wrap-around); drop them. + R_data = transpose(zrcp[poloidal_start:poloidal_stop, :]) + Z_data = transpose(zzcp[poloidal_start:poloidal_stop, :]) + + # Create separate interpolants for R and Z coordinates + rz_in_xs = xs + + @views R_data[:, end] .= R_data[:, 1] + @views Z_data[:, end] .= Z_data[:, 1] + + opts2d = (bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) + rz_in_R = cubic_interp((rz_in_xs, rz_in_ys), R_data; opts2d...) + rz_in_Z = cubic_interp((rz_in_xs, rz_in_ys), Z_data; opts2d...) + # Capture the raw arrays that reconstruct sq_in, rz_in_R, rz_in_Z. + ingest = InverseIngest(collect(xs), fs_copy, collect(rz_in_xs), collect(rz_in_ys), Matrix(R_data), Matrix(Z_data), ro, zo, psio) + + @info "Finished reading CHEASE equilibrium. Magnetic axis at (ro=$(@sprintf("%.3f", ro)), zo=$(@sprintf("%.3f", zo))), psio=$(@sprintf("%.3e", psio))" + return InverseRunInput(config, sq_in, rz_in_xs, rz_in_ys, rz_in_R, rz_in_Z, ro, zo, psio, ingest) +end diff --git a/src/Equilibrium/Sources/GFile.jl b/src/Equilibrium/Sources/GFile.jl new file mode 100644 index 000000000..84027a667 --- /dev/null +++ b/src/Equilibrium/Sources/GFile.jl @@ -0,0 +1,128 @@ +""" +_read_1d_gfile_format(lines_block, num_values) + +Internal helper function to parse Fortran-style fixed-width numerical blocks +from a vector of strings. + +## Arguments: + + - `lines_block`: A `Vector{String}` containing the lines to parse. + - `num_values`: The total number of `Float64` values to read from the block. + +## Returns: + + - A `Vector{Float64}` containing the parsed values. +""" +function _read_1d_gfile_format(lines_block::Vector{String}, num_values::Int) + data_str = join(lines_block) + field_width = 16 + parsed_values = Float64[] + num_read = 0 + + # Ensure the string length is a multiple of the field width for safe processing + safe_len = (length(data_str) ÷ field_width) * field_width + for i in 1:field_width:safe_len + num_read >= num_values && break + val_str = strip(@view(data_str[i:(i+field_width-1)])) + if !isempty(val_str) + try + push!(parsed_values, parse(Float64, val_str)) + num_read += 1 + catch e + @warn "Parsing error for substring: '$val_str'. Error: $e. Skipping." + end + end + end + + if num_read < num_values + @warn "Expected $num_values values, but only read $num_read." + end + return parsed_values +end + +""" + _read_efit(equil_in) + +Parses an EFIT g-file, creates initial 1D and 2D splines, and bundles +them into a `DirectRunInput` object. + +## Arguments: + + - `equil_in`: The `EquilInput` object containing the filename and parameters. + +## Returns: + + - A `DirectRunInput` object ready for the direct solver. +""" +function read_efit(config::EquilibriumConfig) + @info "Processing EFIT g-file: $(config.eq_filename)" + lines = readlines(config.eq_filename) + + # --- Parse Header --- + header1_parts = split(lines[1]) + nw = parse(Int, header1_parts[end-1]) + nh = parse(Int, header1_parts[end]) + @info "Parsed from header: nw=$nw, nh=$nh" + + header_vals = _read_1d_gfile_format(lines[2:5], 20) + rdim, zdim, rcentr, rleft, zmid = header_vals[1:5] + rmaxis, zmaxis, simag, sibry = header_vals[6:9] + ip_sign = Int(sign(header_vals[11])) # the g-file's plasma current; its sign sets the helicity + ip_sign == 0 && (ip_sign = 1) + + # --- Parse Data Blocks --- + current_line_idx = 6 + function parse_block(num_pts) + num_lines = ceil(Int, num_pts / 5) + block = lines[current_line_idx:(current_line_idx+num_lines-1)] + data = _read_1d_gfile_format(block, num_pts) + current_line_idx += num_lines + return data + end + + fpol_data = parse_block(nw) + fpol_sign = Int(sign(fpol_data[end])) # sign of toroidal field (before abs is applied below) + pres_data = parse_block(nw) + ffprime_data = parse_block(nw) + pprime_data = parse_block(nw) + psi_flat_vec = parse_block(nw * nh) + qprof_data = parse_block(nw) + + psi_rz = reshape(psi_flat_vec, nw, nh) + + # --- Create 1D Profile Spline (sq_in) --- + psi_norm_grid = range(0.0, 1.0; length=nw) + sq_fs_nodes = hcat( + abs.(fpol_data), + max.(pres_data .* mu0, 0.0), + qprof_data, + sqrt.(psi_norm_grid) + ) + sq_xs = collect(psi_norm_grid) + sq_in = cubic_interp(sq_xs, Series(sq_fs_nodes); extrap=ExtendExtrap()) + + # --- Process and Normalize 2D Psi Data --- + psio_signed = sibry - simag + psi_proc = (sibry .- psi_rz) + psio = abs(psio_signed) + # Ensure psi at the magnetic axis is positive relative to the boundary + if psio_signed < 0.0 + psi_proc .*= -1.0 + end + + # --- Create 2D Psi interpolant (psi_in) --- + r_grid = range(rleft, rleft + rdim; length=nw) + z_grid = range(zmid - zdim / 2, zmid + zdim / 2; length=nh) + rmin, rmax = extrema(r_grid) + zmin, zmax = extrema(z_grid) + + psi_in_xs = collect(r_grid) + psi_in_ys = collect(z_grid) + psi_in = cubic_interp((psi_in_xs, psi_in_ys), psi_proc; extrap=ExtendExtrap()) + + # Capture the raw arrays that reconstruct sq_in and psi_in, so the rerun path + # (gpec.h5 → PlasmaEquilibrium) can skip the g-file parse entirely. + ingest = DirectIngest(sq_xs, sq_fs_nodes, psi_in_xs, psi_in_ys, psi_proc, rmin, rmax, zmin, zmax, psio, fpol_sign, ip_sign) + + return DirectRunInput(config, sq_in, psi_in, psi_in_xs, psi_in_ys, rmin, rmax, zmin, zmax, psio, fpol_sign, ip_sign, ingest) +end diff --git a/src/Equilibrium/Sources/IFile.jl b/src/Equilibrium/Sources/IFile.jl new file mode 100644 index 000000000..4e036cfef --- /dev/null +++ b/src/Equilibrium/Sources/IFile.jl @@ -0,0 +1,76 @@ +""" + _read_fortran_reals(io, n) -> Vector{Float64} + +Read one little-endian Fortran sequential unformatted record of `n` reals, stored as real*4 or real*8 +(told apart by the record length). +""" +function _read_fortran_reals(io::IO, n::Int) + nbytes = ltoh(read(io, Int32)) + T = nbytes == 8n ? Float64 : nbytes == 4n ? Float32 : error("Fortran record holds $nbytes bytes, expected $n reals") + data = ltoh.(read!(io, Vector{T}(undef, n))) + ltoh(read(io, Int32)) == nbytes || error("Fortran record end marker does not match its start") + return Float64.(data) +end + +""" + read_ifile(config) -> InverseRunInput + +Read an inverse i-file, such as TokaMaker `save_ifile` writes: little-endian Fortran unformatted records, +reals as real*8 or real*4, axis first, θ fastest. + + - `mx, my` (int32); `psi(mx)` [Wb/rad], `f(mx)` = R·Bt [T·m], `p(mx)` [Pa], `q(mx)` + - `r(my, mx)`, `z(my, mx)` [m], with the periodic θ point duplicated + - optional `FF′(mx)`, `p′(mx)` per Wb/rad, skipped + - optional `[ψ_axis, ψ_LCFS]`: ψ_N is then normalized to this LCFS, so the last file surface (written just + inside it) sits below ψ_N = 1 + +Port of Fortran `read_eq_ldp_i`. +""" +function read_ifile(config::EquilibriumConfig) + @info "Reading inverse i-file: $(config.eq_filename)" + + psi, f, p, q, r, z, bounds = open(config.eq_filename, "r") do io + ltoh(read(io, Int32)) == 8 || error("i-file header record is not two int32") + mx, my = Int.(ltoh.(read!(io, Vector{Int32}(undef, 2)))) + read(io, Int32) + profiles = [_read_fortran_reals(io, mx) for _ in 1:4] + r, z = (reshape(_read_fortran_reals(io, my * mx), my, mx) for _ in 1:2) + # Optional records, told apart by length: FF′ and p′ (mx each, unused), then the ψ bounds (2) + bounds = nothing + while !eof(io) + nbytes = ltoh(read(io, Int32)) + if nbytes in (4mx, 8mx) + skip(io, nbytes + 4) + else + skip(io, -4) + bounds = _read_fortran_reals(io, 2) + end + end + return (profiles..., r, z, bounds) + end + mx, my = length(psi), size(r, 1) + @info "Parsed from header: mx = $mx surfaces, my = $my poloidal points" + + psi_axis, psi_lcfs = bounds === nothing ? (psi[1], psi[end]) : bounds + psio_signed = psi_lcfs - psi_axis + psio = abs(psio_signed) + xs = (psi .- psi_axis) ./ psio_signed + bounds === nothing || + @info "i-file ψ bounds: ψ_axis = $psi_axis, ψ_LCFS = $psi_lcfs; file surfaces span ψ_N ∈ [$(xs[1]), $(xs[end])]" + + sq_fs = hcat(abs.(f), p .* mu0, q, sqrt.(xs)) + sq_in = cubic_interp(xs, Series(sq_fs); extrap=ExtendExtrap()) + + R_data = Matrix(transpose(r)) + Z_data = Matrix(transpose(z)) + ro, zo = R_data[1, 1], Z_data[1, 1] + rz_in_ys = collect(range(0, 1; length=my)) + + opts2d = (bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) + rz_in_R = cubic_interp((xs, rz_in_ys), R_data; opts2d...) + rz_in_Z = cubic_interp((xs, rz_in_ys), Z_data; opts2d...) + ingest = InverseIngest(xs, sq_fs, xs, rz_in_ys, R_data, Z_data, ro, zo, psio) + + @info "Finished reading i-file. Magnetic axis at (ro=$(@sprintf("%.3f", ro)), zo=$(@sprintf("%.3f", zo))), psio=$(@sprintf("%.3e", psio))" + return InverseRunInput(config, sq_in, xs, rz_in_ys, rz_in_R, rz_in_Z, ro, zo, psio, ingest) +end diff --git a/src/Equilibrium/Sources/IMAS.jl b/src/Equilibrium/Sources/IMAS.jl new file mode 100644 index 000000000..b6b898495 --- /dev/null +++ b/src/Equilibrium/Sources/IMAS.jl @@ -0,0 +1,119 @@ +""" + read_imas(config::EquilibriumConfig, dd) + +Load an equilibrium from an IMAS data dictionary and return a `DirectRunInput`. + +The `dd.equilibrium.time_slice[]` is used (active time slice). Poloidal flux is +converted from the IMAS COCOS convention (set by `config.imas_cocos`) to the +internal COCOS 2 convention: + + - `imas_cocos = 11` (default, IMAS standard): divide ψ by 2π + - `imas_cocos = 2` (GPEC internal): no conversion + +## Arguments + + - `config`: `EquilibriumConfig` with `eq_type = "imas"` and `imas_cocos` set. + + - `dd`: populated `IMASdd.dd` with `dd.equilibrium.time_slice[]` containing: + + + `global_quantities.psi_axis`, `global_quantities.psi_boundary` + + `profiles_1d.psi`, `profiles_1d.f`, `profiles_1d.pressure`, `profiles_1d.q` + + the `profiles_2d` entry with `grid_type.index == 1` (rectangular), supplying `grid.dim1` (R), `grid.dim2` (Z) and `psi` +""" +function read_imas(config::EquilibriumConfig, dd) + @info "Processing IMAS equilibrium at global_time = $(dd.global_time) s" + + eqt = dd.equilibrium.time_slice[] + + # COCOS conversion: IMAS standard is COCOS 11 (ψ_IMAS = 2π × ψ_internal) + # config.imas_cocos controls the expected convention: + # 11 (default) → divide psi by 2π to get internal (COCOS 2) values + # 2 → data is already in internal convention, no conversion needed + if config.imas_cocos == 11 + cocos_factor = 1.0 / (2π) + @info "Converting IMAS data from COCOS 11 to internal (COCOS 2)" + elseif config.imas_cocos == 2 + cocos_factor = 1.0 + @info "IMAS data in COCOS 2 (no conversion needed)" + else + error("read_imas: unsupported imas_cocos = $(config.imas_cocos). Use 2 or 11.") + end + + psi_axis = eqt.global_quantities.psi_axis * cocos_factor + psi_boundary = eqt.global_quantities.psi_boundary * cocos_factor + + psio = abs(psi_boundary - psi_axis) + + if psio < 1e-10 + error("read_imas: |psi_axis - psi_boundary| = $psio is too small. " * + "Check that dd.equilibrium is properly populated.") + end + + # Extract 1D profiles, converting psi from COCOS 11 to internal + psi_1d = eqt.profiles_1d.psi .* cocos_factor + f_1d = eqt.profiles_1d.f # F(ψ) = R·Bt [T·m], COCOS-independent + p_1d = eqt.profiles_1d.pressure # plasma pressure P(ψ) [Pa], COCOS-independent + q_1d = eqt.profiles_1d.q # safety factor, COCOS-independent + + # Capture toroidal-field sign from the boundary F value before abs() below. + bt_sign = isempty(f_1d) ? 1 : Int(sign(f_1d[end])) + bt_sign == 0 && (bt_sign = 1) + # Plasma-current sign from the IMAS global quantity (missing or zero → +1). + # The current is optional in IMAS, and reading an absent one throws rather than returning missing. + ip_imas = getproperty(eqt.global_quantities, :ip, 0.0) + ip_sign = ip_imas == 0 ? 1 : Int(sign(ip_imas)) + + nw = length(psi_1d) + psi_norm_grid = (psi_1d .- psi_axis) ./ (psi_boundary - psi_axis) + @assert all(diff(psi_norm_grid) .> 0) "read_imas: profiles_1d.psi must increase strictly from psi_axis to psi_boundary" + @assert 0 <= psi_norm_grid[1] && psi_norm_grid[end] <= 1 "read_imas: profiles_1d.psi must lie between psi_axis and psi_boundary" + + # Build equilibrium source terms for spline interpolation + # abs(f_1d): F(ψ) can be negative depending on toroidal field direction convention; + # take absolute value since GPEC uses magnitude F = R·|Bt| + sq_fs_nodes = hcat( + abs.(f_1d), + max.(p_1d .* mu0, 0.0), + q_1d, + sqrt.(psi_norm_grid) + ) + sq_xs = collect(psi_norm_grid) + sq_in = cubic_interp(sq_xs, Series(sq_fs_nodes); extrap=ExtendExtrap()) + + # 2D ψ(R,Z) map + if isempty(eqt.profiles_2d) + error("read_imas: no profiles_2d found in equilibrium time slice. " * + "Ensure the 2D ψ(R,Z) map is stored in dd.equilibrium.") + end + prof2d = findfirst(:rectangular, eqt.profiles_2d) + prof2d === nothing && error("read_imas: no rectangular profiles_2d found.") + r_grid = prof2d.grid.dim1 + z_grid = prof2d.grid.dim2 + + psi_rz = prof2d.psi .* cocos_factor + + # Shift so psi_in = 0 at boundary, psi_in = psio at axis + psi_proc = psi_boundary .- psi_rz + if psi_boundary - psi_axis < 0.0 + psi_proc .*= -1.0 + end + + rmin, rmax = extrema(r_grid) + zmin, zmax = extrema(z_grid) + + psi_in_xs = collect(r_grid) + psi_in_ys = collect(z_grid) + psi_in = cubic_interp((psi_in_xs, psi_in_ys), psi_proc; extrap=ExtendExtrap()) + + @info "IMAS equilibrium loaded:" * + "\n psio = $(round(psio; sigdigits=5)) Wb" * + "\n 1D profile points: nw = $nw" * + "\n 2D grid: nR = $(length(r_grid)), nZ = $(length(z_grid))" * + "\n R ∈ [$(round(rmin; sigdigits=4)), $(round(rmax; sigdigits=4))] m" * + "\n Z ∈ [$(round(zmin; sigdigits=4)), $(round(zmax; sigdigits=4))] m" + + # Capture the raw arrays so an IMAS run can be replayed from gpec.h5 without the dd source. + ingest = DirectIngest(sq_xs, sq_fs_nodes, psi_in_xs, psi_in_ys, Matrix(psi_proc), rmin, rmax, zmin, zmax, psio, bt_sign, ip_sign) + + return DirectRunInput(config, sq_in, psi_in, psi_in_xs, psi_in_ys, rmin, rmax, zmin, zmax, psio, bt_sign, ip_sign, ingest) +end diff --git a/src/Equilibrium/Sources/Ingest.jl b/src/Equilibrium/Sources/Ingest.jl new file mode 100644 index 000000000..a8bb7d665 --- /dev/null +++ b/src/Equilibrium/Sources/Ingest.jl @@ -0,0 +1,222 @@ +""" + DirectIngest + +The serializable raw arrays and scalars captured by a direct-equilibrium reader +(`read_efit`, `read_imas`, `sol_run` is analytic and does not use this) — everything +needed to rebuild a `DirectRunInput`'s splines without re-reading the original g-file. +Stored on `DirectRunInput.ingest` / `PlasmaEquilibrium.ingest` and dumped to +`gpec.h5` so a run can be replayed from any directory. The interpolants themselves are +not serializable; they are reconstructed from these nodes by `build_direct_from_ingest`. + +## Fields + + - `sq_xs::Vector{Float64}` — normalized-ψ knots for the 1D profile spline + - `sq_fs::Matrix{Float64}` — 1D profile node values (F, μ₀P, q, √ψ_norm) + - `psi_xs::Vector{Float64}` — R grid for the 2D flux map [m] + - `psi_ys::Vector{Float64}` — Z grid for the 2D flux map [m] + - `psi_rz::Matrix{Float64}` — processed poloidal flux on the (R, Z) grid [Wb/rad] + - `rmin/rmax/zmin/zmax::Float64` — computational-grid bounds [m] + - `psio::Float64` — total flux difference |ψ_axis - ψ_boundary| [Wb/rad] + - `bt_sign::Int` — sign of the toroidal field (+1 or -1) + - `ip_sign::Int` — sign of the plasma current (+1 or -1) as the source file states it +""" +struct DirectIngest + sq_xs::Vector{Float64} + sq_fs::Matrix{Float64} + psi_xs::Vector{Float64} + psi_ys::Vector{Float64} + psi_rz::Matrix{Float64} + rmin::Float64 + rmax::Float64 + zmin::Float64 + zmax::Float64 + psio::Float64 + bt_sign::Int + ip_sign::Int +end + +""" + InverseIngest + +The serializable raw arrays and scalars captured by an inverse-equilibrium reader +(`read_chease_ascii`, `read_chease_binary`, `read_ifile`) — everything needed to rebuild an +`InverseRunInput`'s splines without re-reading the original CHEASE file or i-file. Stored on +`InverseRunInput.ingest` / `PlasmaEquilibrium.ingest` and reconstructed by +`build_inverse_from_ingest`. See [`DirectIngest`](@ref) for the role this plays in +the `gpec.h5` rerun snapshot. + +## Fields + + - `sq_xs::Vector{Float64}` — normalized-ψ knots for the 1D profile spline + - `sq_fs::Matrix{Float64}` — 1D profile node values + - `rz_xs::Vector{Float64}` — ψ grid for the R, Z maps + - `rz_ys::Vector{Float64}` — θ grid for the R, Z maps + - `R_nodes::Matrix{Float64}` — R node values on the (ψ, θ) grid [m] + - `Z_nodes::Matrix{Float64}` — Z node values on the (ψ, θ) grid [m] + - `ro::Float64` — R of magnetic axis [m] + - `zo::Float64` — Z of magnetic axis [m] + - `psio::Float64` — total flux difference |ψ_axis - ψ_boundary| [Wb/rad] +""" +struct InverseIngest + sq_xs::Vector{Float64} + sq_fs::Matrix{Float64} + rz_xs::Vector{Float64} + rz_ys::Vector{Float64} + R_nodes::Matrix{Float64} + Z_nodes::Matrix{Float64} + ro::Float64 + zo::Float64 + psio::Float64 +end + +# Equilibria captured for replay carry one of these; analytic equilibria carry `nothing` +# and are regenerated from their TOML section instead. +const EquilibriumIngest = Union{Nothing,DirectIngest,InverseIngest} + +""" + DirectRunInput(...) + +A container struct that bundles all necessary inputs for the `direct_run` function. +It is created by one of the equilibrium file read-in functions after processing the +raw equilibrium data and preparing the initial splines. + +## Fields + + - `config::EquilibriumConfig` + The equilibrium configuration object. + + - `sq_in` + 1D spline data versus normalized poloidal flux `psin`. + Quantities: + + 1. `F = R * B_t` — toroidal flux function [m·T] + 2. `μ₀ * Pressure` — plasma pressure (non-negative) [T²] + 3. `q` — safety factor profile + 4. `√ψ_norm` — square root of normalized flux + - `psi_in` + 2D cubic interpolant on the (R, Z) grid [m]. + The values correspond to the **poloidal flux** adjusted to be zero at the boundary [Wb/rad]. + Definitions: + + 1. `ψ(R, Z) = ψ_boundary - ψ(R, Z)` + + 2. `ψ = ψ * sign(ψ(centerR, centerZ))` + + * 1D profiles are represented by `CubicInterpolant` or `CubicSeriesInterpolant` + * 2D flux surfaces by `CubicInterpolantND` + - `psi_in_xs::Vector{Float64}` — R coordinate grid for psi_in [m] + - `psi_in_ys::Vector{Float64}` — Z coordinate grid for psi_in [m] + - `rmin::Float64` — Minimum R-coordinate of the computational grid [m] + - `rmax::Float64` — Maximum R-coordinate of the computational grid [m] + - `zmin::Float64` — Minimum Z-coordinate of the computational grid [m] + - `zmax::Float64` — Maximum Z-coordinate of the computational grid [m] + - `psio::Float64` — Total flux difference `|ψ_axis - ψ_boundary|` [Wb/rad] + - `bt_sign::Int` — Sign of the toroidal field (+1 or -1); read from fpol sign in EFIT g-files + - `ip_sign::Int` — Sign of the plasma current (+1 or -1); read from the `current` header value + in EFIT g-files and `global_quantities.ip` in IMAS. The internal flux is always made + positive at the axis, so the sign is kept here (never recovered from the computed current) + and fixes the SFL→machine toroidal-angle handedness `helicity = bt_sign × ip_sign`. + - `ingest::EquilibriumIngest` — captured raw arrays for the `gpec.h5` rerun snapshot + (a [`DirectIngest`](@ref) for file-based reads, or `nothing` for analytic equilibria) + - `psihigh_resolved::Float64` — outer flux limit the equilibrium is formed on: `config.psihigh` + clamped to the outermost closed flux surface by [`resolve_psihigh!`](@ref). Defaults to + `config.psihigh` and only differs for efit-family equilibria whose requested limit falls + outside the closed-flux region. The solvers build their ψ grid from this field. IMAS + equilibria are read into this struct but are not in `EFIT_KINDS`, so they are never + clamped and always keep the request. +""" +mutable struct DirectRunInput{S<:FastInterpolations.CubicSeriesInterpolant,I2D<:FastInterpolations.CubicInterpolantND} + config::EquilibriumConfig + sq_in::S # 1D profile spline for F, P, q + psi_in::I2D # 2D flux interpolant (CubicInterpolantND) + psi_in_xs::Vector{Float64} # R coordinates + psi_in_ys::Vector{Float64} # Z coordinates + rmin::Float64 # Minimum R-coordinate of the computational grid [m]. + rmax::Float64 # Maximum R-coordinate of the computational grid [m]. + zmin::Float64 # Minimum Z-coordinate of the computational grid [m]. + zmax::Float64 # Maximum Z-coordinate of the computational grid [m]. + psio::Float64 # The total flux difference |ψ_axis - ψ_boundary| [Weber / radian]. + bt_sign::Int # Sign of the toroidal field: +1 or -1 (from fpol sign in g-file) + ip_sign::Int # Sign of the plasma current: +1 or -1 (from the g-file current / IMAS ip) + ingest::EquilibriumIngest + psihigh_resolved::Float64 +end + +# Readers construct without a resolved psihigh; it starts at the request and `resolve_psihigh!` +# clamps it for efit-family equilibria. +DirectRunInput(config::EquilibriumConfig, sq_in, psi_in, psi_in_xs, psi_in_ys, + rmin, rmax, zmin, zmax, psio, bt_sign, ip_sign, ingest) = + DirectRunInput(config, sq_in, psi_in, psi_in_xs, psi_in_ys, + rmin, rmax, zmin, zmax, psio, bt_sign, ip_sign, ingest, config.psihigh) + +""" + InverseRunInput(...) + +A container struct for inputs to the `inverse_run` function. + +## Fields + + - `config::EquilibriumConfig` - The equilibrium configuration object + - `sq_in::CubicSeriesInterpolant` - 1D profile spline for F, P, q + - `rz_in_xs::Vector{Float64}` - ψ coordinate grid for rz_in + - `rz_in_ys::Vector{Float64}` - θ coordinate grid for rz_in + - `rz_in_R::CubicInterpolantND` - R coordinate interpolant [m] + - `rz_in_Z::CubicInterpolantND` - Z coordinate interpolant [m] + - `ro::Float64` - R-coordinate of magnetic axis [m] + - `zo::Float64` - Z-coordinate of magnetic axis [m] + - `psio::Float64` - Total flux difference |ψ_axis - ψ_boundary| [Wb/rad] + - `ingest::EquilibriumIngest` - captured raw arrays for the `gpec.h5` rerun snapshot + (an [`InverseIngest`](@ref) for file-based reads, or `nothing` for analytic equilibria) + - `psihigh_resolved::Float64` - outer flux limit the equilibrium is formed on; see + [`DirectRunInput`](@ref). `config.psihigh`, clamped to the outermost input surface for i-files and + reruns (`resolve_psihigh!`); `efit_by_inversion` forwards the clamped value from its `DirectRunInput`. +""" +mutable struct InverseRunInput{S<:FastInterpolations.CubicSeriesInterpolant,I2D<:FastInterpolations.CubicInterpolantND} + config::EquilibriumConfig + sq_in::S # 1D profile spline for F, P, q + rz_in_xs::Vector{Float64} # ψ coordinates + rz_in_ys::Vector{Float64} # θ coordinates + rz_in_R::I2D # R coordinate interpolant + rz_in_Z::I2D # Z coordinate interpolant + ro::Float64 # R axis location + zo::Float64 # Z axis location + psio::Float64 # Total flux difference |psi_axis - psi_boundary| + ingest::EquilibriumIngest + psihigh_resolved::Float64 +end + +InverseRunInput(config::EquilibriumConfig, sq_in, rz_in_xs, rz_in_ys, rz_in_R, rz_in_Z, + ro, zo, psio, ingest) = + InverseRunInput(config, sq_in, rz_in_xs, rz_in_ys, rz_in_R, rz_in_Z, + ro, zo, psio, ingest, config.psihigh) + +""" + build_direct_from_ingest(config::EquilibriumConfig, ingest::DirectIngest) -> DirectRunInput + +Rebuild a `DirectRunInput` from a [`DirectIngest`](@ref) captured by `read_efit`/`read_imas` +(or restored from `Input/RawInputs/Equilibrium/` inside `gpec.h5`). Inverse of that capture: +reconstructs the splines so the rerun path skips the g-file/IMAS parse, reusing the existing +solver dispatch. +""" +function build_direct_from_ingest(config::EquilibriumConfig, ingest::DirectIngest) + sq_in = cubic_interp(ingest.sq_xs, Series(ingest.sq_fs); extrap=ExtendExtrap()) + psi_in = cubic_interp((ingest.psi_xs, ingest.psi_ys), ingest.psi_rz; extrap=ExtendExtrap()) + return DirectRunInput(config, sq_in, psi_in, ingest.psi_xs, ingest.psi_ys, + ingest.rmin, ingest.rmax, ingest.zmin, ingest.zmax, ingest.psio, ingest.bt_sign, ingest.ip_sign, ingest) +end + +""" + build_inverse_from_ingest(config::EquilibriumConfig, ingest::InverseIngest) -> InverseRunInput + +Rebuild an `InverseRunInput` from an [`InverseIngest`](@ref) captured by +`read_chease_ascii`/`read_chease_binary`. Inverse of that capture; used by the rerun path to +skip the CHEASE parse. +""" +function build_inverse_from_ingest(config::EquilibriumConfig, ingest::InverseIngest) + sq_in = cubic_interp(ingest.sq_xs, Series(ingest.sq_fs); extrap=ExtendExtrap()) + opts2d = (bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap())) + rz_in_R = cubic_interp((ingest.rz_xs, ingest.rz_ys), ingest.R_nodes; opts2d...) + rz_in_Z = cubic_interp((ingest.rz_xs, ingest.rz_ys), ingest.Z_nodes; opts2d...) + return InverseRunInput(config, sq_in, ingest.rz_xs, ingest.rz_ys, rz_in_R, rz_in_Z, + ingest.ro, ingest.zo, ingest.psio, ingest) +end diff --git a/test/runtests_equil.jl b/test/runtests_equil.jl index 46a064867..886666513 100644 --- a/test/runtests_equil.jl +++ b/test/runtests_equil.jl @@ -701,7 +701,7 @@ using Statistics @testset "field-line ODE failure is an error, not silent truncation" begin # Checked in the source, since a non-closing field line is costly to synthesize; the EFIT # testsets above are the control that neither guard false-fires. - src = read(joinpath(dirname(@__DIR__), "src", "Equilibrium", "DirectEquilibrium.jl"), String) + src = read(joinpath(dirname(@__DIR__), "src", "Equilibrium", "Axisymmetric", "Direct.jl"), String) @test occursin("sol.retcode != ReturnCode.Success", src) @test occursin("isapprox(sol.t[end], 2π", src) end From f5e1f9459c7dba11b1d832525a6b194afea3636e Mon Sep 17 00:00:00 2001 From: Matthew Pharr Date: Tue, 27 Oct 2026 11:10:23 +0100 Subject: [PATCH 3/3] Equilibrium - DOCS - State where code goes in the Equilibrium module docstring Co-Authored-By: Claude Opus 5.5 --- src/Equilibrium/Equilibrium.jl | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/src/Equilibrium/Equilibrium.jl b/src/Equilibrium/Equilibrium.jl index 69633b685..6c8443787 100644 --- a/src/Equilibrium/Equilibrium.jl +++ b/src/Equilibrium/Equilibrium.jl @@ -1,3 +1,22 @@ +""" + Equilibrium + +Reads or solves an MHD equilibrium and maps it to straight-field-line flux coordinates as a +[`PlasmaEquilibrium`](@ref). + +## Where code goes + + - `Config.jl`: the run configuration (`EquilibriumConfig`) and global parameters (`EquilibriumParameters`). + - `PlasmaEquilibrium.jl`: the result type and its spline containers. + - `Construction.jl`: `run_input` for each source type, the `PlasmaEquilibrium` constructors, and `reform`. + - `Sources/`: one file per input format, each building a `DirectRunInput` or `InverseRunInput` + (`Ingest.jl` holds those and the ingest types). A new analytic model is one file in + `Sources/Analytic/` holding its config and solver, plus one `ANALYTIC_EQ` row. + - `Grid/`: radial ψ grids, first-pass and refined. + - `Axisymmetric/`: code that assumes axisymmetry (the 2D solvers, the Grad-Shafranov residual, + separatrix finding). A non-axisymmetric equilibrium gets a sibling folder. + - `Derived/`: quantities computed from a formed `PlasmaEquilibrium`. +""" module Equilibrium # --- Module-level Dependencies ---