Equilibrium - BUGFIX - Make the magnetic-axis search robust to an unreliable psi-spline Hessian

src/Equilibrium/DirectEquilibrium.jl

@@ -161,28 +161,88 @@ 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
-        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
+    r_march, z_march = r, z
+
+    # Now, use Newton iteration to find the O-point (magnetic axis) where Br=0 and Bz=0.
+    # The 2-D ψ spline's second derivatives are not reliable at every point near the
+    # axis (∂B_z/∂R can pass through zero at isolated R), and a single near-singular
+    # Hessian sends an undamped Newton step across the whole box. Cap each step at one
+    # march step: inactive for well-behaved iterations (their steps are ≲ dr/2), it
+    # only keeps a bad iterate inside the axis neighbourhood until the Hessian recovers.
+    step_cap = dr
+    function _newton(r0, z0)
+        local dr, dz   # not the enclosing march step
+        r, z = r0, z0
+        for _ in 1:max_iterations
+            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
+            abs(det) < 1e-20 && return r, z, false   # singular Hessian: let the caller fall back
+            # Δx = -J⁻¹ F
+            dr = (bfield.brz * bfield.bz - bfield.bzz * bfield.br) / det
+            dz = (bfield.bzr * bfield.br - bfield.brr * bfield.bz) / det
+            step = hypot(dr, dz)
+            if step > step_cap
+                dr *= step_cap / step
+                dz *= step_cap / step
+            end
+            r += dr
+            z += dz
+            if abs(dr) <= 1e-12 * abs(r) && abs(dz) <= 1e-12 * abs(r)
+                return r, z, true
+            end
         end
+        return r, z, false
     end
-
-    if !(abs(dr) <= 1e-12 * abs(r) && abs(dz) <= 1e-12 * abs(r))
-        error("Failed to find magnetic axis after $max_iterations iterations.")
+    r, z, converged = _newton(r, z)
+
+    if !converged
+        # Fallback, reached only when Newton from the midplane guess cycles or hits a
+        # singular Hessian. Locate the axis with first derivatives only: B_z changes
+        # sign across the axis along the midplane and B_r changes sign across it along
+        # a column, so alternating 1-D bisections converge to ∇ψ = 0 without touching
+        # the spline's second derivatives. Newton is then retried from that point as a
+        # polish; if even that fails, the bisection point itself is a converged zero of
+        # the first derivatives and is used as is.
+        _bz(rr, zz) = (direct_get_bfield!(bfield, rr, zz, raw_profile.psi_in, raw_profile.sq_in, sq_in_deriv, raw_profile.psio; derivs=1); bfield.bz)
+        _br(rr, zz) = (direct_get_bfield!(bfield, rr, zz, raw_profile.psi_in, raw_profile.sq_in, sq_in_deriv, raw_profile.psio; derivs=1); bfield.br)
+        function _bisect(f, lo, hi)
+            flo = f(lo)
+            for _ in 1:80
+                mid = 0.5 * (lo + hi)
+                fmid = f(mid)
+                (fmid == 0.0 || hi - lo < 4eps(mid)) && return mid
+                if sign(fmid) == sign(flo)
+                    lo, flo = mid, fmid
+                else
+                    hi = mid
+                end
+            end
+            return 0.5 * (lo + hi)
+        end
+        r, z = r_march, z_march
+        for _ in 1:4
+            # B_z: negative at r - dr (the march crossed there), non-negative at r
+            r = _bisect(rr -> _bz(rr, z), r - step_cap, r)
+            # B_r along the column through r; widen the bracket until it straddles zero
+            w = step_cap
+            while sign(_br(r, z - w)) == sign(_br(r, z + w)) && w < 20 * step_cap
+                w *= 2
+            end
+            sign(_br(r, z - w)) != sign(_br(r, z + w)) && (z = _bisect(zz -> _br(r, zz), z - w, z + w))
+            # r bracket for the next pass: re-establish the sign change around the new r
+            (_bz(r - step_cap, z) < 0 <= _bz(r + step_cap, z)) && (r += step_cap)
+        end
+        r_bis, z_bis = r, z
+        b_bis = hypot(_br(r, z), _bz(r, z))
+        r, z, converged = _newton(r_bis, z_bis)
+        if converged
+            @info "Magnetic axis found at R = $(@sprintf("%.3f", r)), Z = $(@sprintf("%.3f", z)) (Newton restarted from a first-derivative bisection)"
+        else
+            hypot(_br(r, z), _bz(r, z)) > b_bis && ((r, z) = (r_bis, z_bis))
+            @info "Magnetic axis from first-derivative bisection at R = $(@sprintf("%.3f", r)), Z = $(@sprintf("%.3f", z)) (Newton did not converge; |B| = $(@sprintf("%.2e", b_bis)))"
+        end
+    else
+        @info "Magnetic axis found at R = $(@sprintf("%.3f", r)), Z = $(@sprintf("%.3f", z))"
     end
 
     ro = r

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_grid_refinement.jl")
     include("./runtests_coordinate_invariant.jl")
     include("./runtests_eulerlagrange.jl")

test/runtests_equil_axis_newton.jl

@@ -0,0 +1,45 @@
+@testset "Direct equilibrium: magnetic-axis search robustness" begin
+    using GeneralizedPerturbedEquilibrium.Equilibrium
+    using GeneralizedPerturbedEquilibrium.Equilibrium: EquilibriumConfig, read_efit, direct_position!
+
+    data_dir = joinpath(@__DIR__, "test_data")
+
+    # A 257x257 TJ circular geqdsk (R0 = 2 m) on which the 2-D psi spline's
+    # d(Bz)/dR passes through zero at Newton's second iterate. Undamped Newton
+    # previously stepped -4.3 m and died with "Jacobian matrix is singular".
+    # The step cap keeps the iterate near the axis until the Hessian recovers, and
+    # capped Newton then converges on its own (no fallback in the log).
+    @testset "previously divergent geqdsk now converges to the axis" begin
+        cfg = EquilibriumConfig(;
+            eq_filename=joinpath(data_dir, "TJ_circular_axis_newton_regression.geqdsk"),
+            eq_type="efit")
+        rp = read_efit(cfg)
+        ro, zo, _, _ = @test_logs (:info, r"^Magnetic axis found at R = [0-9.]+, Z = -?[0-9.]+$") match_mode=:any direct_position!(rp)
+        @test isapprox(ro, 2.0; atol=1e-3)
+        @test abs(zo) < 1e-3
+    end
+
+    # A second failure mode: the Hessian is indefinite over several cells around the
+    # axis (dB_z/dR swings between ~5 and ~0), so even a capped Newton cycles without
+    # converging. The first-derivative bisection stage resolves it.
+    @testset "cycling Newton is rescued by the bisection stage" begin
+        cfg = EquilibriumConfig(;
+            eq_filename=joinpath(data_dir, "TJ_circular_axis_newton_cycling.geqdsk"),
+            eq_type="efit")
+        rp = read_efit(cfg)
+        ro, zo, _, _ = @test_logs (:info, r"first-derivative bisection") match_mode=:any direct_position!(rp)
+        @test isapprox(ro, 2.0; atol=1e-3)
+        @test abs(zo) < 1e-3
+    end
+
+    # A healthy file must take the plain Newton path (the log carries no fallback
+    # suffix), so its axis -- and everything downstream -- is unchanged.
+    @testset "well-behaved geqdsk takes the plain Newton path" begin
+        cfg = EquilibriumConfig(;
+            eq_filename=joinpath(data_dir, "CHEASE_test_data", "EQDSK_COCOS_02"), eq_type="efit")
+        rp = read_efit(cfg)
+        ro, zo, _, _ = @test_logs (:info, r"^Magnetic axis found at R = [0-9.]+, Z = -?[0-9.]+$") match_mode=:any direct_position!(rp)
+        @test 6.5 < ro < 7.5
+        @test isfinite(zo)
+    end
+end