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

src/Equilibrium/EquilibriumTypes.jl

@@ -25,7 +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";
-    "log_asymptotic" is a legacy alias for "auto")
+    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`.

src/Equilibrium/InverseEquilibrium.jl

@@ -124,14 +124,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 +157,16 @@ 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 off the axis; psilow and mpsi are ignored, psihigh still sets psilim
+        sq_xs = filter(>(0.0), rz_in_xs)
+        mpsi = length(sq_xs) - 1
     elseif grid_type == "uniform"
         if mpsi == 0
             mpsi = 128
@@ -174,9 +176,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

test/runtests_equil.jl

@@ -174,6 +174,24 @@
             (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 off the axis (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[2:end]
+        @test finite(eq_o)
+
+        # 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.