Equilibrium - REFACTOR - Arrange the Equilibrium module by pipeline stage

docs/development/architecture.md

@@ -32,12 +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
-     - `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

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

src/Equilibrium/DirectEquilibrium.jl → 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)
 

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