Equilibrium - BUGFIX - Support original and rho grids for inverse equilibria

src/Equilibrium/Equilibrium.jl

@@ -536,6 +536,7 @@ function equilibrium_gse!(equil::PlasmaEquilibrium)
             flux_fs[ipsi, itheta, 2] *= 2π * psio / f4
         end
     end
+    equil.rzphi_xs[1] == 0.0 && extrap_axis_row!(flux_fs, equil.rzphi_xs, 1:2)
     # Create flux interpolants for Grad-Shafranov diagnostics
     @views flux_fs[:, end, :] .= flux_fs[:, 1, :]
     flux_opts = (bc=(CubicFit(), PeriodicBC()), extrap=(ExtendExtrap(), WrapExtrap()))

src/Equilibrium/EquilibriumTypes.jl

@@ -25,6 +25,8 @@ specified in the input.
   - `b0exp::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; "ldp", "pow1", "uniform";
+    inverse equilibria also "rho" (uniform in √ψ from psilow to psihigh) and "original" (the input
+    file's own ψ nodes; 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

src/Equilibrium/InverseEquilibrium.jl

@@ -39,6 +39,20 @@ function inverse_extrap(xx::Matrix{Float64}, ff::Matrix{Float64}, x::Float64)::V
     return f
 end
 
+"""
+    extrap_axis_row!(fs, xs, ks; me=3) -> fs
+
+Fill the axis row (`xs[1] == 0`, where the θ-dependent SFL and GSE terms are 0/0) of `fs[:, :, k]`
+for each `k` in `ks`, extrapolating from the next `me` surfaces in √ψ (∝ r).
+"""
+function extrap_axis_row!(fs::Array{Float64,3}, xs::AbstractVector{Float64}, ks; me::Int=3)
+    rho = repeat(sqrt.(xs[2:(me+1)]), 1, size(fs, 2))
+    for k in ks
+        fs[1, :, k] .= inverse_extrap(rho, fs[2:(me+1), :, k], 0.0)
+    end
+    return fs
+end
+
 
 
 function equilibrium_solver(input::InverseRunInput; override_psi_nodes::Union{Nothing,Vector{Float64}}=nothing)
@@ -124,14 +138,6 @@ function equilibrium_solver(input::InverseRunInput; override_psi_nodes::Union{No
     # Create 2D interpolants for r² and dη
     rz_rsq = cubic_interp((rz_in_xs, rz_in_ys), r2; itp_opts2d...)
     rz_deta = cubic_interp((rz_in_xs, rz_in_ys), deta; itp_opts2d...)
-    # c-----------------------------------------------------------------------
-    # c     prepare new spline type for surface quantities.
-    # c-----------------------------------------------------------------------
-
-    if grid_type in ["original", "orig"]
-        mpsi = size(sq_in._fs, 1) - 1
-    end
-
     # c-----------------------------------------------------------------------
     # c     set up radial grid
     # c-----------------------------------------------------------------------
@@ -165,6 +171,17 @@ function equilibrium_solver(input::InverseRunInput; override_psi_nodes::Union{No
             mpsi = 128
         end
         sq_xs = [psilow + (psihigh - psilow) * (3(i / mpsi) - (i / mpsi)^3) / 2 for i in 0:mpsi]
+    elseif grid_type == "rho"
+        if mpsi == 0
+            mpsi = 128
+        end
+        # Uniform in √ψ from psilow to psihigh (Fortran's rho ignores psilow)
+        sq_xs = [(sqrt(psilow) + (sqrt(psihigh) - sqrt(psilow)) * i / mpsi)^2 for i in 0:mpsi]
+    elseif grid_type in ("original", "orig")
+        # The input's own ψ nodes; psilow and mpsi are ignored, psihigh still sets psilim (as Fortran,
+        # which takes the 1D input's nodes; they differ from these only for efit_by_inversion)
+        sq_xs = copy(rz_in_xs)
+        mpsi = length(sq_xs) - 1
     elseif grid_type == "uniform"
         if mpsi == 0
             mpsi = 128
@@ -174,9 +191,9 @@ function equilibrium_solver(input::InverseRunInput; override_psi_nodes::Union{No
         error("Unsupported grid_type: $grid_type")
     end
     # Floor node spacing on the fixed grids (ldp/pow1 pack the edge below the integration-noise
-    # scale at high mpsi); the auto grid floors its refined grid in refined_psi_grid, and an
-    # override grid arrives already floored.
-    if override_psi_nodes === nothing && !(grid_type in ("auto", "log_asymptotic"))
+    # scale at high mpsi); the auto grid floors its refined grid in refined_psi_grid, an override
+    # grid arrives already floored, and the original grid keeps the input's nodes.
+    if override_psi_nodes === nothing && !(grid_type in ("auto", "log_asymptotic", "original", "orig"))
         sq_xs = enforce_min_spacing(sq_xs, MIN_KNOT_SPACING)
         mpsi = length(sq_xs) - 1
     end
@@ -215,6 +232,14 @@ function equilibrium_solver(input::InverseRunInput; override_psi_nodes::Union{No
     for ipsi in 0:mpsi
         psifac = rzphi_xs[ipsi+1]
         sq_in(f_sq_in_buf, psifac; hint=sq_in_hint)
+        sq_fs[ipsi+1, 1] = f_sq_in_buf[1] * twopi
+        sq_fs[ipsi+1, 2] = f_sq_in_buf[2]
+        # Use the input q profile directly (from LAR ODE or CHEASE), matching the
+        # Fortran `inverse_chease4_run` convention (sq%fs(ipsi,4) = sq_in%f(3)).
+        # The field-line-integration-based q formula (spl_fsi * F / (2*twopi*psio))
+        # is inaccurate for cylindrical LAR geometry.
+        sq_fs[ipsi+1, 4] = f_sq_in_buf[3]  # q from input profile
+        psifac == 0.0 && continue  # axis node: filled by extrapolation below
         spl_xs .= rzphi_ys
         for itheta in 0:mtheta
             theta = rzphi_ys[itheta+1]
@@ -276,15 +301,13 @@ function equilibrium_solver(input::InverseRunInput; override_psi_nodes::Union{No
             spl_post(spl_post_buf, rzphi_ys[itheta+1]; hint=hint_post)
             @views rzphi_fs[ipsi+1, itheta+1, :] .= spl_post_buf[1:4]
         end
-
-        sq_fs[ipsi+1, 1] = f_sq_in_buf[1] * twopi
-        sq_fs[ipsi+1, 2] = f_sq_in_buf[2]
         sq_fs[ipsi+1, 3] = spl_fsi[mtheta+1, 3] * twopi * pi # dV/d(psi)
-        # Use the input q profile directly (from LAR ODE or CHEASE), matching the
-        # Fortran `inverse_chease4_run` convention (sq%fs(ipsi,4) = sq_in%f(3)).
-        # The field-line-integration-based q formula (spl_fsi * F / (2*twopi*psio))
-        # is inaccurate for cylindrical LAR geometry.
-        sq_fs[ipsi+1, 4] = f_sq_in_buf[3]  # q from input profile
+    end
+    # An axis node (ψ = 0) has r² = 0, where the SFL weights are 0/0: extrapolate it from the innermost
+    # surfaces as deta is above, θ-dependent fields in √ψ (∝ r) and dV/dψ in ψ.
+    if sq_xs[1] == 0.0
+        extrap_axis_row!(rzphi_fs, sq_xs, 2:4; me)
+        sq_fs[1, 3] = inverse_extrap(reshape(sq_xs[2:(me+1)], :, 1), sq_fs[2:(me+1), 3:3], 0.0)[1]
     end
 
     sq = cubic_interp(sq_xs, Series(sq_fs); extrap=ExtendExtrap())
@@ -327,6 +350,7 @@ function equilibrium_solver(input::InverseRunInput; override_psi_nodes::Union{No
 
     v = zeros(Float64, 3, 3)
     for ipsi in 0:mpsi
+        sq_xs[ipsi+1] == 0.0 && continue  # axis node: extrapolated below
         f_sq_vec = @view sq.y[ipsi+1, :]
         for itheta in 0:mtheta
             # Evaluate rzphi interpolants at grid points using nodal_derivs
@@ -371,6 +395,7 @@ function equilibrium_solver(input::InverseRunInput; override_psi_nodes::Union{No
             eqfun_fs[ipsi+1, itheta+1, 3] = (v[2, 3] * v[3, 3] + f_sq_vec[4] * v[3, 3]^2) / (jacfac * eqfun_fs[ipsi+1, itheta+1, 1]^2)
         end
     end
+    sq_xs[1] == 0.0 && extrap_axis_row!(eqfun_fs, sq_xs, 1:3; me)
     # Create 2D interpolants for physics quantities (eqfun)
     eqfun_grid2d = (eqfun_xs, theta_range)
     @views eqfun_fs[:, end, :] .= eqfun_fs[:, 1, :]

test/runtests_equil.jl

@@ -174,6 +174,28 @@
             (Symbol(k) => v for (k, v) in ffs_table)...) isa GeneralizedPerturbedEquilibrium.ForceFreeStates.ForceFreeStatesControl
     end
 
+    @testset "Inverse grid types (original, rho)" begin
+        Eq = GeneralizedPerturbedEquilibrium.Equilibrium
+        cfg(grid) = Eq.EquilibriumConfig(; eq_type="chease_binary", eq_filename=joinpath(data_dir, "INP1_binary"), jac_type="boozer",
+            grid_type=grid, mpsi=16, psilow=0.01, psihigh=0.99, r0exp=6.8, b0exp=7.4)
+        splines(eq) = (eq.rzphi_rsquared, eq.rzphi_offset, eq.rzphi_nu, eq.rzphi_jac, eq.eqfun_B, eq.eqfun_metric1, eq.eqfun_metric2)
+        finite(eq) = all(isfinite, eq.profiles.dVdpsi_spline.y) && all(all(isfinite, s.nodal_derivs.partials) for s in splines(eq))
+
+        # original: the file's own ψ nodes (this file starts at the axis), whatever psilow/mpsi
+        eq_o = Eq.setup_equilibrium(cfg("original"))
+        @test eq_o.rzphi_xs == Eq.read_chease_binary(cfg("original")).rz_in_xs
+        @test eq_o.rzphi_xs[1] == 0.0
+        @test finite(eq_o)
+        # The extrapolated axis node: B close to F/R₀
+        b_axis = abs(eq_o.profiles.F_spline.y[1]) / (2π * eq_o.ro)
+        @test maximum(abs.(eq_o.eqfun_B.nodal_derivs.partials[1, 1, :] ./ b_axis .- 1)) < 1e-2
+
+        # rho: starts at psilow (Fortran's rho starts at psihigh/(mpsi+1)²)
+        eq_r = Eq.setup_equilibrium(cfg("rho"))
+        @test eq_r.rzphi_xs[[1, end]] ≈ [0.01, 0.99]
+        @test finite(eq_r)
+    end
+
     @testset "EFIT Method Consistency" begin
         # All three methods solve the same equilibrium — q-profiles should broadly agree.
         # Tolerance is 10% to allow for method-specific discretisation differences.