Notebooks · posts/grad-shafranov-solovev

Grad-Shafranov by finite differences, against Solov'ev's closed form

Adam Zettel · September 25, 2026 · written with an agent

A tokamak equilibrium solved on a grid and graded against an exact one. Solov'ev's $ψ = (R^{2} - 1)^{2} /8 + R^{2} Z^{2} / (2 κ^{2})$ solves the Grad-Shafranov equation with constant pressure gradient; the five-point solver, relaxed by SOR until the residual is below $1 0^{- 11}$ , misses it by $1.68 \times 1 0^{- 4}$ , $4.21 \times 1 0^{- 5}$ and $1.06 \times 1 0^{- 5}$ on grids of spacing 0.1, 0.05 and 0.025, an observed order of 1.995 and then 1.994. Those errors are scipy's direct solve of the same equations, to $1 0^{- 10}$ .

Δ^{*} ψ \equiv R \frac{\partial}{\partial R} (\frac{1}{R} \frac{\partial ψ}{\partial R}) + \frac{\partial ^{2} ψ}{\partial Z ^{2}} = - μ_{0} R^{2} p^{'} (ψ) - F F^{'} (ψ)

A tokamak holds its plasma in a magnetic field shaped like nested doughnuts. In equilibrium the magnetic force balances the pressure gradient, $J \times B = \nabla p$ . With the machine's rotational symmetry about its vertical axis, that balance reduces to one equation for one function, the poloidal flux $ψ (R, Z)$ , where $R$ is the distance from the axis of symmetry and $Z$ the height. The surfaces $ψ = const$ are the flux surfaces: the field lines lie in them, and the pressure is constant on each. The equation is the Grad-Shafranov equation above. Its right side holds two free functions of $ψ$ , the pressure $p$ and $F = R B_{ϕ}$ .

When $p^{'}$ and $F F^{'}$ are constants the equation is linear, and Solov'ev (1968) wrote down exact solutions. This post takes $F F^{'} = 0$ and $- μ_{0} p^{'} = 1 + 1/ κ^{2}$ (lengths in units of the major radius), for which

ψ (R, Z) = \frac{( R ^{2} - 1 ) ^{2}}{8} + \frac{R ^{2} Z ^{2}}{2 κ ^{2}}

is exact: its flux surfaces are nested about the magnetic axis at $(R, Z) = (1, 0)$ , and near the axis they are ellipses of elongation $κ = 1.7$ . posts/refs/grad-shafranov-solovev.py checks the closed form with sympy, then solves the same finite-difference equations directly with scipy's sparse LU, which shares no code or method with the SOR iteration here.

Look inside this program

Atlas
The program’s map: every function and which calls which.
There’s no recorded run of this program yet.

Shift+Enter in a cell runs it and moves to the next. Running a cell runs every cell above it too, as one program.

[1]lines 36–54

PI = 3.141592653589793
KAPPA = 1.7
R_MIN = 0.6
R_MAX = 1.4
Z_HALF = 0.7
# SOR stops when the largest residual of the difference equations is below this
RES_TOL = 1.0e-11
MAX_SWEEPS = 20000

#| Solov'ev's flux function, the exact solution.
def psi_exact(r, z) {
    t = r * r - 1.0
    return 0.125 * t * t + r * r * z * z / (2.0 * KAPPA * KAPPA)
}

#| The right side (1 + 1/kappa^2) R^2.
def source(r) {
    return (1.0 + 1.0 / (KAPPA * KAPPA)) * r * r
}

The discrete equation and SOR

On a grid with spacing $h_{R}$ in $R$ and $h_{Z}$ in $Z$ , the centered differences for $ψ_{R R}$ , $ψ_{R}$ and $ψ_{Z Z}$ give at each interior node

(\frac{1}{h _{R}^{2}} - \frac{1}{2 R h _{R}}) ψ_{i + 1, j} + (\frac{1}{h _{R}^{2}} + \frac{1}{2 R h _{R}}) ψ_{i - 1, j} + \frac{ψ _{i, j + 1} + ψ _{i, j - 1}}{h _{Z}^{2}} - (\frac{2}{h _{R}^{2}} + \frac{2}{h _{Z}^{2}}) ψ_{i, j} = (1 + κ^{- 2}) R_{i}^{2},

with $ψ$ fixed to the closed form on the edge of the box. The $- ψ_{R} / R$ term makes the equation unsymmetric in $R$ ; that is the only difference from Poisson's equation.

Successive over-relaxation (SOR) sweeps the grid, solves each node's equation for $ψ_{i, j}$ with its neighbors' newest values, and moves $ω$ times that far. The program takes the $ω$ that is optimal for Poisson's equation on the same box, $ω = 2/ (1 + 1 - ρ^{2})$ , with $ρ$ the Jacobi iteration's spectral radius $(cos (π / n_{R}) + cos (π / n_{Z})) /2$ for $n_{R}$ and $n_{Z}$ intervals of equal spacing. It sweeps until the largest residual is below $1 0^{- 11}$ , so what is left is the discretization error.

[2]lines 79–120

#| Solve on an nr x nz grid by SOR into psi[nr, nz]. Returns
#| [largest error against the closed form, sweeps taken, final residual].
def solve_gs(psi: Arr, nr, nz) {
    hr = (R_MAX - R_MIN) / (nr - 1)
    hz = 2.0 * Z_HALF / (nz - 1)
    for i to nr {
        for j to nz {
            psi[i, j] = 0.0
            if i == 0 or i == nr - 1 or j == 0 or j == nz - 1 {
                psi[i, j] = psi_exact(R_MIN + i * hr, -Z_HALF + j * hz)
            }
        }
    }
    rho = 0.5 * (cos(PI / (nr - 1)) + cos(PI / (nz - 1)))
    omega = 2.0 / (1.0 + sqrt(1.0 - rho * rho))
    aN = 1.0 / (hz * hz)
    aC = 2.0 / (hr * hr) + 2.0 / (hz * hz)
    sweeps = 0
    resid = 1.0
    while resid > RES_TOL and sweeps < MAX_SWEEPS {
        resid = 0.0
        for i from 1 to nr - 1 {
            r = R_MIN + i * hr
            aE = 1.0 / (hr * hr) - 1.0 / (2.0 * r * hr)
            aW = 1.0 / (hr * hr) + 1.0 / (2.0 * r * hr)
            f = source(r)
            for j from 1 to nz - 1 {
                res = aE * psi[i + 1, j] + aW * psi[i - 1, j] + aN * (psi[i, j + 1] + psi[i, j - 1]) - aC * psi[i, j] - f
                resid = max(resid, abs(res))
                psi[i, j] = psi[i, j] + omega * res / aC
            }
        }
        sweeps = sweeps + 1
    }
    err = 0.0
    for i to nr {
        for j to nz {
            err = max(err, abs(psi[i, j] - psi_exact(R_MIN + i * hr, -Z_HALF + j * hz)))
        }
    }
    return [err, sweeps, resid]
}

#| Solve on an nr x nz grid by SOR into psi[nr, nz]. Returns
#| [largest error against the closed form, sweeps taken, final residual].
def solve_gs(psi: Arr, nr, nz) {
    hr = (R_MAX - R_MIN) / (nr - 1)
    hz = 2.0 * Z_HALF / (nz - 1)
    for i to nr {
        for j to nz {
            psi[i, j] = 0.0
            if i == 0 or i == nr - 1 or j == 0 or j == nz - 1 {
                psi[i, j] = psi_exact(R_MIN + i * hr, -Z_HALF + j * hz)
            }
        }
    }
    rho = 0.5 * (cos(PI / (nr - 1)) + cos(PI / (nz - 1)))
    omega = 2.0 / (1.0 + sqrt(1.0 - rho * rho))
    aN = 1.0 / (hz * hz)
    aC = 2.0 / (hr * hr) + 2.0 / (hz * hz)
    sweeps = 0
    resid = 1.0
    while resid > RES_TOL and sweeps < MAX_SWEEPS {
        resid = 0.0
        for i from 1 to nr - 1 {
            r = R_MIN + i * hr
            aE = 1.0 / (hr * hr) - 1.0 / (2.0 * r * hr)
            aW = 1.0 / (hr * hr) + 1.0 / (2.0 * r * hr)
            f = source(r)
            for j from 1 to nz - 1 {
                res = aE * psi[i + 1, j] + aW * psi[i - 1, j] + aN * (psi[i, j + 1] + psi[i, j - 1]) - aC * psi[i, j] - f
                resid = max(resid, abs(res))
                psi[i, j] = psi[i, j] + omega * res / aC
            }
        }
        sweeps = sweeps + 1
    }
    err = 0.0
    for i to nr {
        for j to nz {
            err = max(err, abs(psi[i, j] - psi_exact(R_MIN + i * hr, -Z_HALF + j * hz)))
        }
    }
    return [err, sweeps, resid]
}

Three grids

The three grids have 8, 16 and 32 intervals across $R$ and 14, 28 and 56 across $Z$ , so $h_{R} = h_{Z} = h$ = 0.1, 0.05 and 0.025. (The reference script also solves a fourth grid, $h = 0.0125$ , where the order is 1.9998; it takes too long in a browser tab.)

The error of the difference equations is second order, and here it can be written down. The closed form's $R^{2} Z^{2}$ term is a quadratic in each variable, so the centered differences are exact on it. On $(R^{2} - 1)^{2} /8$ the second difference is off by $h^{2} ψ_{R R R R} /12 = h^{2} /4$ and the first difference over $R$ by $h^{2} ψ_{R R R} / (6 R) = h^{2} /2$ , so the difference operator applied to the exact $ψ$ misses the source by $- h^{2} /4$ , the same constant at every node. The error $e = ψ_{h} - ψ$ therefore solves the difference equation with source $h^{2} /4$ and zero boundary values: it is $h^{2}$ times a fixed function, up to that function's own discretization, and each halving of $h$ divides it by 4.

[3]lines 142–220

NR_GRID = [9, 17, 33]
NZ_GRID = [15, 29, 57]
errs = zeros(len(NR_GRID))
hs = zeros(len(NR_GRID))
sweeps = zeros(len(NR_GRID))
resids = zeros(len(NR_GRID))
# the finest solution is drawn and read on the midplane
psi_fine = zeros(NR_GRID[2], NZ_GRID[2])
for g to len(NR_GRID) {
    if g == 2 {
        out = solve_gs(psi_fine, NR_GRID[g], NZ_GRID[g])
    } else {
        out = solve_gs(zeros(NR_GRID[g], NZ_GRID[g]), NR_GRID[g], NZ_GRID[g])
    }
    errs[g] = out[0]
    sweeps[g] = out[1]
    resids[g] = out[2]
    hs[g] = (R_MAX - R_MIN) / (NR_GRID[g] - 1)
}
order_1 = log(errs[0] / errs[1]) / log(2.0)
order_2 = log(errs[1] / errs[2]) / log(2.0)
# intervals across R and Z (a grid holds one more point each way), and what
# each halving of h divides the error by
ints_r = zeros(len(NR_GRID))
ints_z = zeros(len(NR_GRID))
for g to len(NR_GRID) {
    ints_r[g] = NR_GRID[g] - 1
    ints_z[g] = NZ_GRID[g] - 1
}
halving_1 = errs[0] / errs[1]
halving_2 = errs[1] / errs[2]
show ints_r
show ints_z
show hs
show halving_1
show halving_2
show errs
show sweeps
show resids
show order_1
show order_2

test "SOR brought every grid's residual below 1e-11" {
    for g to len(NR_GRID) {
        assert_lt(resids[g], RES_TOL)
    }
}
test "the errors are scipy's direct solve of the same equations: 1.6771831e-4, 4.2061527e-5, 1.0558059e-5" {
    assert_near(errs[0], 1.6771831249e-04, 1.0e-10)
    assert_near(errs[1], 4.2061527279e-05, 1.0e-10)
    assert_near(errs[2], 1.0558058854e-05, 1.0e-10)
}
test "the three grids have 8, 16 and 32 intervals across R and 14, 28 and 56 across Z, so h = 0.1, 0.05 and 0.025" {
    for g to len(NR_GRID) {
        assert_near(ints_r[g], 8 * pow(2.0, g), 0.0)
        assert_near(ints_z[g], 14 * pow(2.0, g), 0.0)
        assert_near(hs[g], 0.1 / pow(2.0, g), 1.0e-15)
        assert_near((2.0 * Z_HALF) / ints_z[g], hs[g], 1.0e-15)
    }
}
test "each halving of h divides the error by 4, to within 0.02" {
    assert_near(halving_1, 4.0, 0.02)
    assert_near(halving_2, 4.0, 0.02)
}
test "the observed order is 1.995 and then 1.994" {
    assert_near(order_1, 1.995467, 1.0e-4)
    assert_near(order_2, 1.994157, 1.0e-4)
}

h2 = zeros(len(NR_GRID))
for g to len(NR_GRID) {
    h2[g] = errs[2] * (hs[g] / hs[2]) * (hs[g] / hs[2])
}
figure("Grad-Shafranov error against the Solov'ev solution", "grid spacing h", "$\\max |\\psi - \\psi_{\\mathrm{exact}}|$")
xscale("log")
yscale("log")
plot(hs, errs, "five-point difference equations, solved by SOR")
plot(hs, h2, "slope 2")
style("slope 2", "dash")

NR_GRID = [9, 17, 33]
NZ_GRID = [15, 29, 57]
errs = zeros(len(NR_GRID))
hs = zeros(len(NR_GRID))
sweeps = zeros(len(NR_GRID))
resids = zeros(len(NR_GRID))
# the finest solution is drawn and read on the midplane
psi_fine = zeros(NR_GRID[2], NZ_GRID[2])
for g to len(NR_GRID) {
    if g == 2 {
        out = solve_gs(psi_fine, NR_GRID[g], NZ_GRID[g])
    } else {
        out = solve_gs(zeros(NR_GRID[g], NZ_GRID[g]), NR_GRID[g], NZ_GRID[g])
    }
    errs[g] = out[0]
    sweeps[g] = out[1]
    resids[g] = out[2]
    hs[g] = (R_MAX - R_MIN) / (NR_GRID[g] - 1)
}
order_1 = log(errs[0] / errs[1]) / log(2.0)
order_2 = log(errs[1] / errs[2]) / log(2.0)
# intervals across R and Z (a grid holds one more point each way), and what
# each halving of h divides the error by
ints_r = zeros(len(NR_GRID))
ints_z = zeros(len(NR_GRID))
for g to len(NR_GRID) {
    ints_r[g] = NR_GRID[g] - 1
    ints_z[g] = NZ_GRID[g] - 1
}
halving_1 = errs[0] / errs[1]
halving_2 = errs[1] / errs[2]
show ints_r
show ints_z
show hs
show halving_1
show halving_2
show errs
show sweeps
show resids
show order_1
show order_2

test "SOR brought every grid's residual below 1e-11" {
    for g to len(NR_GRID) {
        assert_lt(resids[g], RES_TOL)
    }
}
test "the errors are scipy's direct solve of the same equations: 1.6771831e-4, 4.2061527e-5, 1.0558059e-5" {
    assert_near(errs[0], 1.6771831249e-04, 1.0e-10)
    assert_near(errs[1], 4.2061527279e-05, 1.0e-10)
    assert_near(errs[2], 1.0558058854e-05, 1.0e-10)
}
test "the three grids have 8, 16 and 32 intervals across R and 14, 28 and 56 across Z, so h = 0.1, 0.05 and 0.025" {
    for g to len(NR_GRID) {
        assert_near(ints_r[g], 8 * pow(2.0, g), 0.0)
        assert_near(ints_z[g], 14 * pow(2.0, g), 0.0)
        assert_near(hs[g], 0.1 / pow(2.0, g), 1.0e-15)
        assert_near((2.0 * Z_HALF) / ints_z[g], hs[g], 1.0e-15)
    }
}
test "each halving of h divides the error by 4, to within 0.02" {
    assert_near(halving_1, 4.0, 0.02)
    assert_near(halving_2, 4.0, 0.02)
}
test "the observed order is 1.995 and then 1.994" {
    assert_near(order_1, 1.995467, 1.0e-4)
    assert_near(order_2, 1.994157, 1.0e-4)
}

h2 = zeros(len(NR_GRID))
for g to len(NR_GRID) {
    h2[g] = errs[2] * (hs[g] / hs[2]) * (hs[g] / hs[2])
}
figure("Grad-Shafranov error against the Solov'ev solution", "grid spacing h", "$\\max |\\psi - \\psi_{\\mathrm{exact}}|$")
xscale("log")
yscale("log")
plot(hs, errs, "five-point difference equations, solved by SOR")
plot(hs, h2, "slope 2")
style("slope 2", "dash")

The flux surfaces, and the Shafranov shift

The heatmap is the solution on the 33 by 57 grid, with the flux surfaces drawn at five values of $ψ$ . They are nested about the magnetic axis at $(1, 0)$ , taller than wide near the axis ( $κ = 1.7$ ), and pushed toward small $R$ farther out.

On the midplane $Z = 0$ the closed form gives the crossings of a surface directly: $ψ = (R^{2} - 1)^{2} /8$ puts them at $R = 1 \pm 8 ψ$ . The surface through $R = 1.25$ crosses again at $R = 0.66144$ , so its center is at $0.95572$ : the magnetic axis sits $0.0443$ outboard of the center of this surface. That offset is the Shafranov shift, and it appears whenever there is pressure: the plasma pushes itself outward, where the field is weaker. The program reads the same crossing off the 33 by 57 grid by linear interpolation along the midplane and gets 0.661489, which is the reference script's interpolation on its own solution of the same grid, and $5 \times 1 0^{- 5}$ from the closed form.

[4]lines 242–276

nr_f = NR_GRID[2]
nz_f = NZ_GRID[2]
h_f = hs[2]
jz0 = (nz_f - 1) // 2
psi_b = psi_exact(1.25, 0.0)
r_in = 0.0
for i from 1 to nr_f {
    a = psi_fine[i - 1, jz0]
    b = psi_fine[i, jz0]
    if a >= psi_b and b < psi_b {
        r_in = R_MIN + (i - 1 + (a - psi_b) / (a - b)) * h_f
    }
}
center = 0.5 * (1.25 + r_in)
show r_in
show center

test "the surface through R = 1.25 crosses the midplane again at 0.661489, the scipy solution's crossing, and within 1e-4 of the exact 0.661438" {
    assert_near(r_in, 0.6614893199, 1.0e-9)
    assert_near(r_in, 0.6614378278, 1.0e-4)
}
test "its center is 0.0443 inboard of the magnetic axis at R = 1" {
    assert_near(1.0 - center, 1.0 - 0.9557189139, 1.0e-4)
}

LEVELS = [0.002, 0.008, 0.02, 0.04, 0.07]
figure("Solov'ev equilibrium on a 33 x 57 grid: poloidal flux $\\psi(R, Z)$", "R", "Z", R_MIN, R_MAX, -Z_HALF, Z_HALF)
plot2d(transpose(psi_fine), "$\\psi$")
origin("lower")
colormap("viridis")
levels(LEVELS)
label("axis", "magnetic axis", 1.0, -0.05)

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors

nr_f = NR_GRID[2]
nz_f = NZ_GRID[2]
h_f = hs[2]
jz0 = (nz_f - 1) // 2
psi_b = psi_exact(1.25, 0.0)
r_in = 0.0
for i from 1 to nr_f {
    a = psi_fine[i - 1, jz0]
    b = psi_fine[i, jz0]
    if a >= psi_b and b < psi_b {
        r_in = R_MIN + (i - 1 + (a - psi_b) / (a - b)) * h_f
    }
}
center = 0.5 * (1.25 + r_in)
show r_in
show center

test "the surface through R = 1.25 crosses the midplane again at 0.661489, the scipy solution's crossing, and within 1e-4 of the exact 0.661438" {
    assert_near(r_in, 0.6614893199, 1.0e-9)
    assert_near(r_in, 0.6614378278, 1.0e-4)
}
test "its center is 0.0443 inboard of the magnetic axis at R = 1" {
    assert_near(1.0 - center, 1.0 - 0.9557189139, 1.0e-4)
}

LEVELS = [0.002, 0.008, 0.02, 0.04, 0.07]
figure("Solov'ev equilibrium on a 33 x 57 grid: poloidal flux $\\psi(R, Z)$", "R", "Z", R_MIN, R_MAX, -Z_HALF, Z_HALF)
plot2d(transpose(psi_fine), "$\\psi$")
origin("lower")
colormap("viridis")
levels(LEVELS)
label("axis", "magnetic axis", 1.0, -0.05)

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors