Equilibrium - BUGFIX! - ⚠️ Recover the magnetic axis when psi is not smooth near it, without relaxing the convergence test

src/Equilibrium/Axisymmetric/Direct.jl

@@ -116,15 +116,99 @@ the Julia spline implementation.
     bf_out.bzz = -bf_out.psirz / r
 end
 
+# Newton iteration for the O-point (magnetic axis), where B_r = B_z = 0, with each step capped at `cap`.
+# Returns `(r, z, status)`, status `:converged` (step test), `:singular_jacobian` (|det J| < `singular_det`)
+# or `:iteration_cap`.
+function _axis_newton(rp::DirectRunInput, bfield::DirectBField, sq_in_deriv, r::Float64, z::Float64, cap::Float64;
+    max_iterations::Int=200, singular_det::Float64=1e-20)
+    for _ in 1:max_iterations
+        direct_get_bfield!(bfield, r, z, rp.psi_in, rp.sq_in, sq_in_deriv, rp.psio; derivs=2)
+        det = bfield.brr * bfield.bzz - bfield.brz * bfield.bzr
+        abs(det) < singular_det && return r, z, :singular_jacobian
+        # Δx = -J⁻¹ F
+        δr = (bfield.brz * bfield.bz - bfield.bzz * bfield.br) / det
+        δz = (bfield.bzr * bfield.br - bfield.brr * bfield.bz) / det
+        step = hypot(δr, δz)
+        step > cap && ((δr, δz) = (δr * cap / step, δz * cap / step))
+        r += δr
+        z += δz
+        abs(δr) <= 1e-12 * abs(r) && abs(δz) <= 1e-12 * abs(r) && return r, z, :converged
+    end
+    return r, z, :iteration_cap
+end
+
+# Seed for the Newton restart: alternate 1-D bisections on B_z along R and B_r along Z (a nonlinear
+# Gauss–Seidel sweep, contracting by ψ_RZ²/(ψ_RR ψ_ZZ) per pass) until a pass moves the point by less than a
+# tenth of a grid cell, which leaves it within about a cell of the axis for a contraction up to 0.9. Returns `(r, z)`,
+# or a message if a sign change is missing.
+function _axis_bisection_seed(rp::DirectRunInput, bfield::DirectBField, sq_in_deriv, r::Float64, z::Float64, w::Float64;
+    max_passes::Int=50)
+    _bz(rr, zz) = (direct_get_bfield!(bfield, rr, zz, rp.psi_in, rp.sq_in, sq_in_deriv, rp.psio; derivs=1); bfield.bz)
+    _br(rr, zz) = (direct_get_bfield!(bfield, rr, zz, rp.psi_in, rp.sq_in, sq_in_deriv, rp.psio; derivs=1); bfield.br)
+    # Nearest sign change of f around x, doubling the half-width up to the grid edges; `nothing` if none.
+    function _bracketed_root(f, x, xmin, xmax)
+        hw = w
+        while true
+            lo, hi = max(x - hw, xmin), min(x + hw, xmax)
+            sign(f(lo)) * sign(f(hi)) <= 0 && return find_zero(f, (lo, hi), Roots.Bisection())
+            (lo == xmin && hi == xmax) && return nothing
+            hw *= 2
+        end
+    end
+    hr = (rp.rmax - rp.rmin) / (length(rp.psi_in_xs) - 1)
+    hz = (rp.zmax - rp.zmin) / (length(rp.psi_in_ys) - 1)
+    for _ in 1:max_passes
+        r_new = _bracketed_root(rr -> _bz(rr, z), r, rp.rmin, rp.rmax)
+        r_new === nothing && return "B_z has no sign change along Z = $z"
+        z_new = _bracketed_root(zz -> _br(r_new, zz), z, rp.zmin, rp.zmax)
+        z_new === nothing && return "B_r has no sign change along R = $r_new"
+        moved = max(abs(r_new - r) / hr, abs(z_new - z) / hz)
+        r, z = r_new, z_new
+        moved < 0.1 && break
+    end
+    return r, z
+end
+
+# Why (r, z) is not an acceptable magnetic axis, or `nothing`. The axis must lie on the ψ grid (a root off it is an
+# artifact of spline extrapolation) and be the maximum of ψ that the readers place there, ψ = ψ₀ > 0: an O-point has
+# det ∂(B_R,B_Z)/∂(R,Z) = det(Hess ψ)/R² > 0 (an X-point has det < 0), and a maximum also has ψ_RR < 0. With
+# `bound_residual`, |B_p| must also be below `1e-10` × B_scale = |ψ₀|/(R (R_max − R_min)).
+function _axis_failure(rp::DirectRunInput, bfield::DirectBField, sq_in_deriv, r::Float64, z::Float64; bound_residual::Bool)
+    (rp.rmin <= r <= rp.rmax && rp.zmin <= z <= rp.zmax) ||
+        return "converged at (R, Z) = ($r, $z), outside the ψ grid R ∈ [$(rp.rmin), $(rp.rmax)], Z ∈ [$(rp.zmin), $(rp.zmax)]"
+    direct_get_bfield!(bfield, r, z, rp.psi_in, rp.sq_in, sq_in_deriv, rp.psio; derivs=2)
+    det = bfield.brr * bfield.bzz - bfield.brz * bfield.bzr
+    residual = hypot(bfield.br, bfield.bz)
+    det > 0 || return "converged at (R, Z) = ($r, $z), which is not an O-point (det ∂(B_R,B_Z)/∂(R,Z) = $det ≤ 0, |B_p| = $residual T)"
+    (bfield.psirr < 0 && bfield.psi > 0) ||
+        return "converged at (R, Z) = ($r, $z), an O-point that is not the maximum of ψ (ψ = $(bfield.psi), ψ_RR = $(bfield.psirr))"
+    if bound_residual
+        b_scale = abs(rp.psio) / (r * (rp.rmax - rp.rmin))
+        residual <= 1e-10 * b_scale ||
+            return "converged at (R, Z) = ($r, $z) with |B_p| = $residual T, above 1e-10 × B_scale = $(1e-10 * b_scale) T"
+    end
+    return nothing
+end
+
 """
     direct_position!(raw_profile)
 
 Finds the key geometric locations of the equilibrium: the magnetic axis (O-point)
 and the inboard/outboard separatrix crossings on the midplane. It also updates the
 spline representing the poloidal flux `ψ(R,Z)` based on the new magnetic axis location.
-This function performs the same overall function as the Fortran `direct_position`
-subroutine with better iteration control and error handling. We have also added a
-helper function for separatrix finding.
+The returned axis is a root of `∇ψ = 0` that meets the Newton step-size test
+`|δR|, |δZ| ≤ 1e-12 R`. Newton starts from a midplane guess, found by marching outward along
+the grid's midplane `Z` until `B_z ≥ 0`. The axis must
+
+  - lie inside the ψ grid;
+  - be an O-point, `det ∂(B_R,B_Z)/∂(R,Z) > 0`;
+  - be the maximum of ψ, `ψ_RR < 0` and `ψ > 0`.
+
+If Newton stops on a singular Jacobian or the iteration cap, or converges to a point that fails
+these checks, it is restarted with steps capped at one march step from a first-derivative
+bisection seed. An axis found by the restart must also satisfy `|B_p| ≤ 1e-10 B_scale` with
+`B_scale = |ψ₀| / (R_axis (R_max - R_min))`, and a warning reports `|B_p|`. If the restart also
+fails, the search errors.
 
 ## Arguments:
 
@@ -161,29 +245,35 @@ function direct_position!(raw_profile::DirectRunInput)
 
     # If we never exited early, the loop failed to find bz = 0
     !(bfield.bz >= 0) && error("Took too many iterations to get bz=0.")
-
-    # Now, use Newton iteration to find the O-point (magnetic axis) where Br=0 and Bz=0
-    dr, dz = 0.0, 0.0
-    for _ in 1:max_iterations
+    r_march, z_march = r, z
+
+    max_newton_iterations = 200
+    singular_det = 1e-20
+    _stop_reason(status) = status === :singular_jacobian ? "stopped on a singular Jacobian, |det| < $singular_det" :
+                           "stopped without convergence in $max_newton_iterations iterations"
+    r, z, status = _axis_newton(raw_profile, bfield, sq_in_deriv, r_march, z_march, Inf;
+        max_iterations=max_newton_iterations, singular_det=singular_det)
+    failure = status === :converged ? _axis_failure(raw_profile, bfield, sq_in_deriv, r, z; bound_residual=false) : _stop_reason(status)
+    if failure !== nothing
+        seed = _axis_bisection_seed(raw_profile, bfield, sq_in_deriv, r_march, z_march, dr)
+        seed isa String && error("Failed to find magnetic axis: Newton from the midplane guess $failure, and the restart's seed failed: $seed.")
+        rb, zb = seed
+        r, z, restart_status = _axis_newton(raw_profile, bfield, sq_in_deriv, rb, zb, dr;
+            max_iterations=max_newton_iterations, singular_det=singular_det)
+        restart_failure = restart_status === :converged ? _axis_failure(raw_profile, bfield, sq_in_deriv, r, z; bound_residual=true) :
+                          _stop_reason(restart_status)
+        restart_failure === nothing || error(
+            "Failed to find magnetic axis: Newton from the midplane guess $failure, and the restart from the bisected point " *
+            "($rb, $zb) $restart_failure."
+        )
         direct_get_bfield!(bfield, r, z, raw_profile.psi_in, raw_profile.sq_in, sq_in_deriv, raw_profile.psio; derivs=2)
-        det = bfield.brr * bfield.bzz - bfield.brz * bfield.bzr
-        if abs(det) < 1e-20
-            error("Jacobian matrix is singular near ($r, $z).")
-        end
-        # Δx = -J⁻¹ F
-        dr = (bfield.brz * bfield.bz - bfield.bzz * bfield.br) / det
-        dz = (bfield.bzr * bfield.br - bfield.brr * bfield.bz) / det
-        r += dr
-        z += dz
-        if abs(dr) <= 1e-12 * abs(r) && abs(dz) <= 1e-12 * abs(r)
-            @info "Magnetic axis found at R = $(@sprintf("%.3f", r)), Z = $(@sprintf("%.3f", z))"
-            break
-        end
-    end
-
-    if !(abs(dr) <= 1e-12 * abs(r) && abs(dz) <= 1e-12 * abs(r))
-        error("Failed to find magnetic axis after $max_iterations iterations.")
+        residual = hypot(bfield.br, bfield.bz)
+        b_scale = abs(raw_profile.psio) / (r * (raw_profile.rmax - raw_profile.rmin))
+        @warn "Newton from the midplane guess $failure; magnetic axis found by restarting Newton from a first-derivative " *
+              "bisection, at (R, Z) = ($r, $z) with |B_p| = $residual T ($(residual / b_scale) × B_scale) and " *
+              "det ∂(B_R,B_Z)/∂(R,Z) = $(bfield.brr * bfield.bzz - bfield.brz * bfield.bzr) T²/m². Check ψ near the axis in the equilibrium file."
     end
+    @info "Magnetic axis found at R = $(@sprintf("%.3f", r)), Z = $(@sprintf("%.3f", z))"
 
     ro = r
     zo = z
@@ -205,6 +295,7 @@ function direct_position!(raw_profile::DirectRunInput)
     function find_separatrix_crossing(r_edge, label)
         # On each failed attempt, shift the initial guess (a weighted blend of r_edge and the axis
         # ro) from near the grid edge toward the axis (Fortran direct.f restart counter ird = 0..5).
+        local r, dr
         for restart in 0:5
             r = ((3 - 0.5 * restart) * r_edge + ro) / (4 - 0.5 * restart)
             for _ in 1:max_separatrix_steps

test/runtests.jl

@@ -25,6 +25,7 @@ else
     include("./runtests_fouriertransforms.jl")
     include("./runtests_vacuum.jl")
     include("./runtests_equil.jl")
+    include("./runtests_equil_axis_newton.jl")
     include("./runtests_sfl_sampling.jl")
     include("./runtests_grid_refinement.jl")
     include("./runtests_core_capped_knots.jl")

test/runtests_equil_axis_newton.jl

@@ -0,0 +1,152 @@
+@testset "Direct equilibrium: magnetic-axis search robustness" begin
+    using GeneralizedPerturbedEquilibrium.Equilibrium
+    using GeneralizedPerturbedEquilibrium.Equilibrium: EquilibriumConfig, SolovevConfig, DirectRunInput, sol_run, read_efit, direct_position!
+    using FastInterpolations
+
+    fallback = r"restarting Newton from a first-derivative bisection"
+    # Relative Newton step tolerance of the axis search, reused as the bound on the residual |∇ψ|.
+    newton_tol = 1e-12
+
+    # Up-down symmetric Solovev ψ whose axis sits off the grid's R-midpoint and off every node, plus a
+    # conical cusp -A·ρ·exp(-(ρ/w)⁴) at the axis: the cells there carry a curvature spike, and the
+    # cusp strength A is set so that a Newton step overshoots the axis by `ncell` grid cells.
+    # `r_in`, `r_out` set the grid's extent about r0 in units of a; with r_in > r_out the grid's R-midpoint lies inboard
+    # of the axis, so the midplane march takes a step and brackets the axis in R. `z_shift` moves the grid, and with it
+    # the march's midplane, in Z in units of a.
+    function cusp_solovev(; ncell, nr=96, r_in=1.4, r_out=1.6, z_shift=0.0)
+        sol = SolovevConfig(; mr=nr, mz=nr)
+        cfg = EquilibriumConfig(; eq_type="sol")
+        (; elongation, a, r0, q0, b0fac) = sol
+        e = elongation
+        psio = e * r0 * b0fac * a^2 / (2 * q0 * r0)
+        psifac = psio / (a * r0)^2
+        rs = collect(range(r0 - r_in * a, r0 + r_out * a; length=nr + 1))
+        zh = collect(range(0.0, 1.5e * a; length=nr ÷ 2 + 1))
+        zs = z_shift * a .+ vcat(-reverse(zh[2:end]), zh)
+        h = rs[2] - rs[1]
+        A = ncell * 2 * psifac * r0^2 * h
+        cusp(ρ) = A * ρ * exp(-(ρ / (0.5a))^4)
+        psi = [psio - psifac * ((R * Z)^2 / e^2 + (R^2 - r0^2)^2 / 4) - cusp(hypot(R - r0, Z)) for R in rs, Z in zs]
+        psi_in = cubic_interp((rs, zs), psi; extrap=ExtendExtrap())
+        rp = DirectRunInput(cfg, sol_run(cfg, sol).sq_in, psi_in, rs, zs, rs[1], rs[end], zs[1], zs[end], psio, 1, 1, nothing)
+        return rp, r0, h
+    end
+
+    # The returned axis must be an O-point of the ψ spline with a vanishing gradient.
+    function test_axis_is_o_point(rp, ro, zo)
+        d(i, j) = rp.psi_in((ro, zo); deriv=DerivOp(i, j))
+        @test d(2, 0) * d(0, 2) - d(1, 1)^2 > 0
+        @test hypot(d(1, 0), d(0, 1)) < newton_tol * rp.psio / (rp.rmax - rp.rmin)
+    end
+
+    # Newton on the cusp is chaotic: whether it converges, cycles or leaves the grid at a given cusp strength depends
+    # on the machine's rounding. Whatever it does, the search must return the O-point at the cusp or fail loudly.
+    @testset "curvature spike at the axis: a verified axis or a loud failure" begin
+        n_fallback = 0
+        grids = ((;), (; r_in=1.6, r_out=1.4), (; r_in=1.7, r_out=1.3, z_shift=0.7))
+        for grid in grids, ncell in 1.0:0.25:3.0
+            rp, r0, h = cusp_solovev(; ncell, grid...)
+            logger = Test.TestLogger(; min_level=Base.CoreLogging.Warn)
+            outcome = try
+                Base.CoreLogging.with_logger(() -> direct_position!(rp), logger)
+            catch err
+                err
+            end
+            if outcome isa Exception
+                @test occursin("Failed to find magnetic axis", sprint(showerror, outcome))
+            else
+                ro, zo = outcome[1], outcome[2]
+                test_axis_is_o_point(rp, ro, zo)
+                @test hypot(ro - r0, zo) < h
+                n_fallback += any(occursin(fallback, log.message) for log in logger.logs)
+            end
+        end
+        @test n_fallback > 0
+    end
+
+    @testset "smooth ψ takes the plain Newton path" begin
+        rp, r0, h = cusp_solovev(; ncell=0.0)
+        ro, zo, _, _ = @test_logs min_level = Base.CoreLogging.Warn direct_position!(rp)
+        test_axis_is_o_point(rp, ro, zo)
+        @test abs(zo) <= eps(ro)
+        @test abs(ro - r0) < h
+    end
+
+    @testset "well-behaved geqdsk takes the plain Newton path" begin
+        cfg = EquilibriumConfig(; eq_filename=joinpath(@__DIR__, "test_data", "CHEASE_test_data", "EQDSK_COCOS_02"), eq_type="efit")
+        rp = read_efit(cfg)
+        ro, zo, _, _ = @test_logs min_level = Base.CoreLogging.Warn direct_position!(rp)
+        test_axis_is_o_point(rp, ro, zo)
+        @test 6.5 < ro < 7.5
+    end
+
+    @testset "midplane march that brackets the axis takes the plain Newton path" begin
+        rp, r0, h = cusp_solovev(; ncell=0.0, r_in=1.6, r_out=1.4)
+        ro, zo, _, _ = @test_logs min_level = Base.CoreLogging.Warn direct_position!(rp)
+        test_axis_is_o_point(rp, ro, zo)
+        @test abs(zo) <= eps(ro)
+        @test abs(ro - r0) < h
+    end
+
+    @testset "axis outside the R-bracket of the midplane march is found" begin
+        # Off the axis's midplane the Solovev B_z changes sign at R² = r0² - 2Z²/e², inboard of the axis, and the
+        # march stops between that point and the axis, so the axis lies outside the march's last step.
+        (; elongation, a) = SolovevConfig()
+        e = elongation
+        rp, r0, h = cusp_solovev(; ncell=0.0, r_in=1.7, r_out=1.3, z_shift=0.7)
+        r_mid, z_mid, dr = (rp.rmin + rp.rmax) / 2, (rp.zmin + rp.zmax) / 2, (rp.rmax - rp.rmin) / 20
+        @test r_mid < sqrt(r0^2 - 2 * z_mid^2 / e^2) < r_mid + dr < r0
+        ro, zo, _, _ = @test_logs min_level = Base.CoreLogging.Warn direct_position!(rp)
+        test_axis_is_o_point(rp, ro, zo)
+        @test hypot(ro - r0, zo) < h
+    end
+
+    @testset "Newton converging off the ψ grid is an error" begin
+        # A smooth ψ on a grid that lies wholly outboard of the axis: Newton follows the spline's extrapolation inboard.
+        rp, r0, _ = cusp_solovev(; ncell=0.0, r_in=-0.2, r_out=1.6)
+        @test rp.rmin > r0
+        @test_throws r"outside the ψ grid" direct_position!(rp)
+    end
+
+    # Quadratic ψ = ψ₀ − c (x² + 2b·x·y + y²/e²) about (r0, 0), exact on the cubic spline; `sgn = -1` turns the
+    # axis into a minimum of ψ. With b·e near 1 the Hessian is strongly tilted (ψ_RZ ≠ 0), where the Gauss–Seidel
+    # bisection sweep contracts by (b e)² per pass.
+    function quadratic_psi(; b, sgn=1.0, nr=64)
+        sol = SolovevConfig(; mr=nr, mz=nr)
+        cfg = EquilibriumConfig(; eq_type="sol")
+        (; elongation, a, r0) = sol
+        e = elongation
+        psio = 1.0
+        c = psio / a^2
+        rs = collect(range(r0 - 1.5a, r0 + 1.5a; length=nr + 1))
+        zs = collect(range(-1.5e * a, 1.5e * a; length=nr + 1))
+        psi = [psio - sgn * c * ((R - r0)^2 + 2b * (R - r0) * Z + Z^2 / e^2) for R in rs, Z in zs]
+        psi_in = cubic_interp((rs, zs), psi; extrap=ExtendExtrap())
+        rp = DirectRunInput(cfg, sol_run(cfg, sol).sq_in, psi_in, rs, zs, rs[1], rs[end], zs[1], zs[end], psio, 1, 1, nothing)
+        return rp, r0, rs[2] - rs[1], e
+    end
+
+    @testset "axis checks reject a minimum of ψ and accept the maximum" begin
+        EQ = GeneralizedPerturbedEquilibrium.Equilibrium
+        rp, r0, _, e = quadratic_psi(; b=0.0)
+        @test EQ._axis_failure(rp, EQ.DirectBField(), deriv1(rp.sq_in), r0, 0.0; bound_residual=true) === nothing
+        rp_min, _, _, _ = quadratic_psi(; b=0.0, sgn=-1.0)
+        msg = EQ._axis_failure(rp_min, EQ.DirectBField(), deriv1(rp_min.sq_in), r0, 0.0; bound_residual=false)
+        @test msg isa String && occursin("not the maximum of ψ", msg)
+    end
+
+    @testset "restart seed and capped Newton recover a tilted axis" begin
+        EQ = GeneralizedPerturbedEquilibrium.Equilibrium
+        rp, r0, h, e = quadratic_psi(; b=0.9 / 1.6)
+        bf, sqd = EQ.DirectBField(), deriv1(rp.sq_in)
+        w = (rp.rmax - rp.rmin) / 20
+        # Start ten cells away in both directions; four fixed passes at (b e)² = 0.81 would leave it several cells off.
+        seed = EQ._axis_bisection_seed(rp, bf, sqd, r0 + 10h, 10h, w)
+        @test seed isa Tuple
+        rb, zb = seed
+        @test hypot(rb - r0, zb) < h
+        ro, zo, status = EQ._axis_newton(rp, bf, sqd, rb, zb, w)
+        @test status === :converged
+        @test abs(ro - r0) < 1e-10 && abs(zo) < 1e-10
+    end
+end