From 8a8a3cf17abaef18e10c10cbd4d5b99f3b07e00f Mon Sep 17 00:00:00 2001
From: StuartBenjamin <stubenj9@mit.edu>
Date: Wed, 7 Oct 2026 09:34:36 -0700
Subject: [PATCH] ForceFreeStates.Galerkin - BUGFIX! - Jacobi-scale the banded
 Galerkin LU solve

The assembled Galerkin matrix diagonal spans ~2e-15..2e10, and unscaled banded LU
(gbtrf/gbtrs) lost the small resonant coefficients that carry Delta': on DIII-D
g147131 n=1 the 4/1 row varied erratically with gal_nx and disagreed with Riccati.
Symmetrically scale A <- D A D, d_j = |A_jj|^(-1/2) (1 for a zero diagonal), before
factoring, and un-scale every solution column (incl. rpec coil columns).

Add test/runtests_galerkin.jl: a badly scaled banded system that unscaled LU solves
to ~1e-4 relative error, now ~1e-15.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/ForceFreeStates/Galerkin/GalerkinSolve.jl | 25 ++++++++++-
 test/runtests.jl                              |  1 +
 test/runtests_galerkin.jl                     | 43 +++++++++++++++++++
 3 files changed, 67 insertions(+), 2 deletions(-)
 create mode 100644 test/runtests_galerkin.jl

diff --git a/src/ForceFreeStates/Galerkin/GalerkinSolve.jl b/src/ForceFreeStates/Galerkin/GalerkinSolve.jl
index 963c107f0..a06e37e82 100644
--- a/src/ForceFreeStates/Galerkin/GalerkinSolve.jl
+++ b/src/ForceFreeStates/Galerkin/GalerkinSolve.jl
@@ -146,8 +146,7 @@ function galerkin_solve(ctrl::ForceFreeStatesControl, equil, mats::MatrixSplines
     ws.sol .= ws.rhs
     if ws.solver == "LU"
         ctrl.verbose && @info "Galerkin LU banded factorization + solve"
-        ab, ipiv = LinearAlgebra.LAPACK.gbtrf!(ws.kl, ws.ku, ws.ndim, ws.mat)
-        LinearAlgebra.LAPACK.gbtrs!('N', ws.kl, ws.ku, ws.ndim, ab, ipiv, ws.sol)
+        gal_scaled_lu_solve!(ws.mat, ws.sol, ws.kl, ws.ku)
     else
         ctrl.verbose && @info "Galerkin Cholesky banded factorization + solve"
         LinearAlgebra.LAPACK.pbtrf!('L', ws.kl, ws.mat)
@@ -204,6 +203,28 @@ function galerkin_solve(ctrl::ForceFreeStatesControl, equil, mats::MatrixSplines
     return GalerkinResult(msing, sing_psi, sing_q, sing_m, sing_n, di, alpha, solution, match), dp
 end
 
+"""
+    gal_scaled_lu_solve!(mat, sol, kl, ku) -> sol
+
+Banded LU solve in place: `mat` (LAPACK `gbtrf!` band storage, `ldab = 2kl + ku + 1`) is overwritten by
+its factor and `sol` (the right-hand sides on entry) by the solution. The system is first symmetrically
+Jacobi-scaled, `D·A·D` with `d_j = |A_jj|^(-1/2)` (1 for a zero diagonal): the assembled diagonal spans
+~25 decades, and unscaled banded LU loses the small resonant coefficients that carry Δ′.
+"""
+function gal_scaled_lu_solve!(mat::Matrix{ComplexF64}, sol::AbstractVecOrMat{ComplexF64}, kl::Int, ku::Int)
+    n = size(mat, 2)
+    off = kl + ku + 1  # band row holding the diagonal
+    d = [iszero(mat[off, j]) ? 1.0 : 1 / sqrt(abs(mat[off, j])) for j in 1:n]
+    for j in 1:n, i in max(1, j - ku):min(n, j + kl)
+        mat[off+i-j, j] *= d[i] * d[j]
+    end
+    sol .*= d
+    _, ipiv = LAPACK.gbtrf!(kl, ku, n, mat)
+    LAPACK.gbtrs!('N', kl, ku, n, mat, ipiv, sol)
+    sol .*= d
+    return sol
+end
+
 """
     gal_pest3_blocks(delta, msing) -> (Ap, Bp, Gammap, Deltap)
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 389370209..e1256e56d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -29,6 +29,7 @@ else
     include("./runtests_coordinate_invariant.jl")
     include("./runtests_eulerlagrange.jl")
     include("./runtests_riccati.jl")
+    include("./runtests_galerkin.jl")
     include("./runtests_parallel_integration.jl")
     include("./runtests_result_struct.jl")
     include("./runtests_solve_api.jl")
diff --git a/test/runtests_galerkin.jl b/test/runtests_galerkin.jl
new file mode 100644
index 000000000..6165b34ab
--- /dev/null
+++ b/test/runtests_galerkin.jl
@@ -0,0 +1,43 @@
+using Test
+using LinearAlgebra
+
+# The Galerkin banded LU solve must survive the ~25-decade diagonal spread of the assembled system:
+# A = D·A0·D with A0 well conditioned, weakly coupled, and D alternating 4e-8 / 1e5 drives
+# unscaled partial-pivoting LU to ~1e-4 relative error; Jacobi scaling recovers ~1e-15.
+@testset "Galerkin scaled banded LU" begin
+    FFS = GeneralizedPerturbedEquilibrium.ForceFreeStates
+
+    n, kl = 24, 2
+    A0 = zeros(ComplexF64, n, n)
+    for j in 1:n, i in max(1, j - kl):min(n, j + kl)
+        A0[i, j] = i == j ? 2 + 0.1 * sin(j) : 1e-5 * cis(0.3i + 0.7j)
+    end
+    A0 = (A0 + A0') / 2
+    d = [isodd(j) ? 4e-8 : 1e5 for j in 1:n]
+    A = Diagonal(d) * A0 * Diagonal(d)
+    x = ComplexF64[cis(0.1j) for j in 1:n] ./ d
+    b = A * hcat(x, 2x)
+
+    # LAPACK gbtrf! band storage, as GalWorkspace lays it out for "LU"
+    function band(A)
+        m = size(A, 2)
+        ab = zeros(ComplexF64, 3kl + 1, m)
+        for j in 1:m, i in max(1, j - kl):min(m, j + kl)
+            ab[2kl+1+i-j, j] = A[i, j]
+        end
+        return ab
+    end
+
+    sol = copy(b)
+    FFS.gal_scaled_lu_solve!(band(A), sol, kl, kl)
+    @test maximum(abs.(sol[:, 1] .- x) ./ abs.(x)) < 1e-12
+    @test sol[:, 2] ≈ 2 * sol[:, 1]  # every right-hand side is un-scaled
+
+    @testset "zero diagonal entry" begin
+        Z = ComplexF64[0 1 0; 1 1 1; 0 1 3]
+        z = ComplexF64[1, -2, 0.5]
+        sol = Z * z
+        FFS.gal_scaled_lu_solve!(band(Z), sol, kl, kl)
+        @test sol ≈ z
+    end
+end