Equilibrium - REFACTOR - Derive every rzphi |∇ψ| from one shared metric formula

src/Equilibrium/Derived/FluxSurfaceMetrics.jl

@@ -1,52 +1,61 @@
 """
-    flux_surface_metric(equil, psi, theta; hint=(Ref(1), Ref(1))) -> (; r, jac, delpsi)
+    rzphi_primitives(equil, psi, theta; hint=(Ref(1), Ref(1))) -> (; r2, deta, jac, r2_y, deta_y)
 
-Evaluate the flux-surface geometry at `(psi, theta)`: major radius `r`, Jacobian `jac`, and the
-flux-gradient magnitude `delpsi = |∇ψ|`. The metric components `w11`/`w12` follow GPEC's gpout.f
-convention for the rzphi straight-field-line coordinates. This is the single source of truth for the
-`|∇ψ|` computation reused across the FFS and PE modules.
+Evaluate the smooth `rzphi` primitives at `(psi, theta)`: `r²`, the `η`-offset, the Jacobian and the
+θ-derivatives of `r²` and the offset. Every consumer that derives flux-surface geometry from `rzphi`
+(the metric and area helpers here, the perturbed-field reconstruction transform) samples this one
+function, and `metric_from_primitives` turns them into the gpout.f metric. ψ-derivatives are
+deliberately not provided: the bicubic is only C¹ in ψ, so callers differentiate a ψ-resample instead.
 """
-function flux_surface_metric(equil::PlasmaEquilibrium, psi::Float64, theta::Float64; hint=(Ref(1), Ref(1)))
-    r2 = equil.rzphi_rsquared((psi, theta); hint=hint)
-    jac = equil.rzphi_jac((psi, theta); hint=hint)
-    deta = equil.rzphi_offset((psi, theta); hint=hint)
-    r2_y = equil.rzphi_rsquared((psi, theta); deriv=DerivOp(0, 1), hint=hint)
-    deta_y = equil.rzphi_offset((psi, theta); deriv=DerivOp(0, 1), hint=hint)
+function rzphi_primitives(equil::PlasmaEquilibrium, psi::Float64, theta::Float64; hint=(Ref(1), Ref(1)))
+    pt = (psi, theta)
+    r2 = equil.rzphi_rsquared(pt; hint=hint)
+    deta = equil.rzphi_offset(pt; hint=hint)
+    jac = equil.rzphi_jac(pt; hint=hint)
+    r2_y = equil.rzphi_rsquared(pt; deriv=DerivOp(0, 1), hint=hint)
+    deta_y = equil.rzphi_offset(pt; deriv=DerivOp(0, 1), hint=hint)
+    return (; r2, deta, jac, r2_y, deta_y)
+end
 
-    rfac = sqrt(abs(r2))
-    eta = 2π * (theta + deta)
-    r = equil.ro + rfac * cos(eta)
+"""
+    metric_from_primitives(ro, theta, p) -> (; rfac, r, jac, delpsi)
+
+The gpout.f flux-surface metric at one point from the `rzphi` primitives `p` (the fields of
+`rzphi_primitives`): minor radius `rfac = √r²`, major radius `r`, Jacobian `jac` and
+`delpsi = |∇ψ|`. Every `|∇ψ|` derived from `rzphi` goes through this one formula, whether the
+primitives come from evaluating the bicubic or from its nodal values.
+"""
+function metric_from_primitives(ro::Float64, theta::Float64, p)
+    rfac = sqrt(abs(p.r2))
+    r = ro + rfac * cos(2π * (theta + p.deta))
+    w11 = (1.0 + p.deta_y) * (2π)^2 * rfac * r / p.jac
+    w12 = rfac > 0 ? -p.r2_y * π * r / (rfac * p.jac) : 0.0
+    return (; rfac, r, jac=p.jac, delpsi=sqrt(w11^2 + w12^2))
+end
 
-    # Metric components for |∇ψ| in flux coordinates (GPEC gpout.f)
-    w11 = (1.0 + deta_y) * (2π)^2 * rfac * r / jac
-    w12 = -r2_y * π * r / (rfac * jac)
-    delpsi = sqrt(w11^2 + w12^2)
+"""
+    flux_surface_metric(equil, psi, theta; hint=(Ref(1), Ref(1))) -> (; r, jac, delpsi)
 
+Evaluate the flux-surface geometry at `(psi, theta)`: major radius `r`, Jacobian `jac`, and the
+flux-gradient magnitude `delpsi = |∇ψ|`, by `metric_from_primitives` on the bicubic at that point.
+"""
+function flux_surface_metric(equil::PlasmaEquilibrium, psi::Float64, theta::Float64; hint=(Ref(1), Ref(1)))
+    (; r, jac, delpsi) = metric_from_primitives(equil.ro, theta, rzphi_primitives(equil, psi, theta; hint=hint))
     return (; r, jac, delpsi)
 end
 
 """
     flux_surface_area(equil, psi, mtheta) -> Float64
 
-Compute ∫ J·|∇ψ| dθ over the flux surface at `psi` using `mtheta` equally-spaced θ points and the
-trapezoidal rule with periodic end correction.
+Compute ⟨J·|∇ψ|⟩_θ on the flux surface at `psi`: the periodic trapezoidal rule on `mtheta` equally
+spaced θ points (θ = 1 duplicates θ = 0).
 """
 function flux_surface_area(equil::PlasmaEquilibrium, psi::Float64, mtheta::Int)
+    hint = (Ref(1), Ref(1))
     area = 0.0
-    last_contrib = 0.0
-
-    hint = (Ref(1), Ref(1))  # Shared 2D hint for hot loop optimization
-    for itheta in 0:mtheta
-        theta = itheta / mtheta
-        m = flux_surface_metric(equil, psi, theta; hint=hint)
-        contrib = m.jac * m.delpsi / mtheta
-        area += contrib
-
-        if itheta == mtheta
-            last_contrib = contrib
-        end
+    for itheta in 0:(mtheta-1)
+        m = flux_surface_metric(equil, psi, itheta / mtheta; hint=hint)
+        area += m.jac * m.delpsi
     end
-
-    # Trapezoidal rule end correction
-    return area - last_contrib
+    return area / mtheta
 end

src/Equilibrium/Derived/GeometryProfiles.jl

@@ -17,9 +17,8 @@ equilibrium's 2D rzphi interpolants. Uses trapezoidal sums in θ across the
 non-periodic θ nodes (the trailing duplicate at θ = 1 is skipped, mthsurf =
 length(rzphi_ys) - 1).
 
-The integrand `delpsi = |∇ψ|` (in geometry units) is built from the same
-nodal derivatives the GS-residual diagnostic uses; see `Equilibrium.jl:407-419`
-and the Fortran reference at `pentrc/dcon_interface.f:1199-1254`.
+The integrand `delpsi = |∇ψ|` (in geometry units) applies `metric_from_primitives`
+to the bicubic's nodal values.
 """
 function compute_geometry_profiles(
     rzphi_xs::Vector{Float64},
@@ -41,24 +40,15 @@ function compute_geometry_profiles(
         r_int = 0.0
         R_int = 0.0
         for itheta in 1:mthsurf
-            f1 = rzphi_rsquared.nodal_derivs.partials[1, ipsi, itheta]
-            offset = rzphi_offset.nodal_derivs.partials[1, ipsi, itheta]
-            jac = rzphi_jac.nodal_derivs.partials[1, ipsi, itheta]
-            fy1 = rzphi_rsquared.nodal_derivs.partials[3, ipsi, itheta]
-            fy2 = rzphi_offset.nodal_derivs.partials[3, ipsi, itheta]
-
-            rfac = sqrt(f1)
-            eta = 2π * (rzphi_ys[itheta] + offset)
-            R = ro + rfac * cos(eta)
-
-            w1 = (1 + fy2) * (2π)^2 * rfac * R / jac
-            w2 = (rfac > 0) ? (-fy1 * π * R / (rfac * jac)) : 0.0
-            delpsi = sqrt(w1^2 + w2^2)
+            p = (r2=rzphi_rsquared.nodal_derivs.partials[1, ipsi, itheta], deta=rzphi_offset.nodal_derivs.partials[1, ipsi, itheta],
+                jac=rzphi_jac.nodal_derivs.partials[1, ipsi, itheta], r2_y=rzphi_rsquared.nodal_derivs.partials[3, ipsi, itheta],
+                deta_y=rzphi_offset.nodal_derivs.partials[3, ipsi, itheta])
+            (; rfac, r, jac, delpsi) = metric_from_primitives(ro, rzphi_ys[itheta], p)
             weight = jac * delpsi / mthsurf
 
             area_psi += weight                  # wegt=0, ave=0
             r_int += weight * rfac              # wegt=3, ave=1
-            R_int += weight * R                 # wegt=1, ave=1
+            R_int += weight * r                 # wegt=1, ave=1
         end
         area_vals[ipsi] = area_psi
         avg_r_vals[ipsi] = (area_psi > 0) ? r_int / area_psi : 0.0

src/PerturbedEquilibrium/FieldReconstruction.jl

@@ -140,32 +140,11 @@ function reconstruct_physical_fields(
 
         # Compute b^ψ / area for HDF5 output — matches Fortran gpout_xbnormal fast path
         # area = <J·|∇ψ|>_θ (Hamada: J is const over θ)
-        ro = equil.ro
         mthsurf = length(equil.rzphi_ys) - 1
-        thetas_for_avg = [(k - 1) / mthsurf for k in 1:mthsurf]
         Jb_psi_modes = similar(b_psi_modes)
         # Surfaces are independent; :static keeps threadid() stable for any per-thread state.
         Threads.@threads :static for ipsi in 1:npsi
-            psi = psi_grid[ipsi]
-            hint2d = (Ref(1), Ref(1))
-            area = 0.0
-            for k in 1:mthsurf
-                theta = thetas_for_avg[k]
-                r2 = equil.rzphi_rsquared((psi, theta); hint=hint2d)
-                deta = equil.rzphi_offset((psi, theta); hint=hint2d)
-                jac = equil.rzphi_jac((psi, theta); hint=hint2d)
-                r2_y = equil.rzphi_rsquared((psi, theta); deriv=DerivOp(0, 1), hint=hint2d)
-                deta_y = equil.rzphi_offset((psi, theta); deriv=DerivOp(0, 1), hint=hint2d)
-                rfac = sqrt(abs(r2))
-                eta = 2π * (theta + deta)
-                r = ro + rfac * cos(eta)
-                w11 = (1.0 + deta_y) * 4π^2 * rfac * r / jac
-                w12 = -r2_y * π * r / (rfac * jac)
-                delpsi = sqrt(w11^2 + w12^2)
-                area += jac * delpsi
-            end
-            area /= mthsurf
-            Jb_psi_modes[ipsi, :] = b_psi_modes[ipsi, :] ./ area
+            Jb_psi_modes[ipsi, :] = b_psi_modes[ipsi, :] ./ Equilibrium.flux_surface_area(equil, psi_grid[ipsi], mthsurf)
         end
 
         # R,Z,φ cylindrical components (Fortran gpeq_rzphi)
@@ -794,7 +773,6 @@ function compute_b_n_xi_n_modes(
     npsi, numpert_total = size(b_psi_modes)
     (; mpert, mlow, npert) = ffs
     mthsurf = length(equil.rzphi_ys) - 1
-    ro = equil.ro
     twopi = 2π
 
     b_n_modes = zeros(ComplexF64, npsi, numpert_total)
@@ -814,20 +792,9 @@ function compute_b_n_xi_n_modes(
         delpsis = zeros(Float64, mthsurf)
         jacs = zeros(Float64, mthsurf)
         for k in 1:mthsurf
-            theta = thetas[k]
-            r2 = equil.rzphi_rsquared((psi, theta); hint=hint2d_psi)
-            deta = equil.rzphi_offset((psi, theta); hint=hint2d_psi)
-            jac = equil.rzphi_jac((psi, theta); hint=hint2d_psi)
-            r2_y = equil.rzphi_rsquared((psi, theta); deriv=DerivOp(0, 1), hint=hint2d_psi)
-            deta_y = equil.rzphi_offset((psi, theta); deriv=DerivOp(0, 1), hint=hint2d_psi)
-            rfac = sqrt(abs(r2))
-            eta = twopi * (theta + deta)
-            r = ro + rfac * cos(eta)
-            w11 = (1.0 + deta_y) * twopi^2 * rfac * r / jac
-            w12 = -r2_y * π * r / (rfac * jac)
-            delpsi = sqrt(w11^2 + w12^2)
-            delpsis[k] = delpsi
-            jacs[k] = jac
+            m = Equilibrium.flux_surface_metric(equil, psi, thetas[k]; hint=hint2d_psi)
+            delpsis[k] = m.delpsi
+            jacs[k] = m.jac
         end
 
         jd = jacs .* delpsis
@@ -901,13 +868,13 @@ function _build_rzphi_geometry(
 
         for itheta in 1:mtheta
             theta = (itheta - 1) / mtheta  # SFL theta ∈ [0, 1)
-            pt = (psi, theta)
+            g = Equilibrium.rzphi_primitives(equil, psi, theta; hint=hint)
 
-            knot_vals[ik, 0*mtheta+itheta] = equil.rzphi_rsquared(pt; hint=hint)
-            knot_vals[ik, 1*mtheta+itheta] = equil.rzphi_offset(pt; hint=hint)
-            knot_vals[ik, 2*mtheta+itheta] = equil.rzphi_jac(pt; hint=hint)
-            knot_vals[ik, 3*mtheta+itheta] = equil.rzphi_rsquared(pt; deriv=DerivOp(0, 1), hint=hint)
-            knot_vals[ik, 4*mtheta+itheta] = equil.rzphi_offset(pt; deriv=DerivOp(0, 1), hint=hint)
+            knot_vals[ik, 0*mtheta+itheta] = g.r2
+            knot_vals[ik, 1*mtheta+itheta] = g.deta
+            knot_vals[ik, 2*mtheta+itheta] = g.jac
+            knot_vals[ik, 3*mtheta+itheta] = g.r2_y
+            knot_vals[ik, 4*mtheta+itheta] = g.deta_y
         end
     end