03 — Grad-Shafranov Equilibrium Solver

The Grad-Shafranov equation is the fundamental equation of tokamak plasma equilibrium. This tutorial demonstrates:

1. Setting up an ITER-like tokamak configuration
2. Computing vacuum magnetic flux from external coils
3. Solving the free-boundary equilibrium
4. Visualising flux surfaces and key plasma features

License: © 1998–2026 Miroslav Šotek. GNU AGPL v3.

---

import numpy as np
import matplotlib.pyplot as plt
import json, tempfile, os

Step 1: Define the Tokamak Configuration

ITER-like parameters: R₀ = 6.2 m, a = 2.0 m, I_p = 15 MA

config = {
    "reactor_name": "ITER-like",
    "grid_resolution": [65, 65],
    "dimensions": {
        "R_min": 1.0, "R_max": 9.0,
        "Z_min": -5.0, "Z_max": 5.0
    },
    "physics": {
        "plasma_current_target": 15.0,
        "vacuum_permeability": 1.2566370614e-6
    },
    "coils": [
        {"R": 3.5, "Z": 4.0, "current": 5.0},
        {"R": 3.5, "Z": -4.0, "current": 5.0},
        {"R": 9.0, "Z": 4.0, "current": -3.0},
        {"R": 9.0, "Z": -4.0, "current": -3.0},
        {"R": 6.2, "Z": 5.5, "current": -1.5},
        {"R": 6.2, "Z": -5.5, "current": -1.5}
    ],
    "solver": {
        "max_iterations": 100,
        "convergence_threshold": 1e-6,
        "relaxation_factor": 0.1
    }
}

# Write config to temp file for the kernel
fd, config_path = tempfile.mkstemp(suffix=".json")
os.close(fd)
with open(config_path, 'w') as f:
    json.dump(config, f)

print("Configuration:")
print(f"  Grid: {config['grid_resolution'][0]}×{config['grid_resolution'][1]}")
print(f"  R range: [{config['dimensions']['R_min']}, {config['dimensions']['R_max']}] m")
print(f"  Z range: [{config['dimensions']['Z_min']}, {config['dimensions']['Z_max']}] m")
print(f"  Coils: {len(config['coils'])}")

Step 2: Initialize and Compute Vacuum Field

The vacuum field is computed from the external PF coils using complete elliptic integrals (Green's function method).

from scpn_fusion.core import FusionKernel

kernel = FusionKernel(config_path)
kernel.initialize_grid()
kernel.calculate_vacuum_field()

# Extract grid and vacuum psi for plotting
R = kernel.R
Z = kernel.Z
RR, ZZ = np.meshgrid(R, Z)
psi_vac = kernel.Psi.copy()

print(f"Grid: {len(R)}×{len(Z)} points")
print(f"Psi_vac range: [{psi_vac.min():.4f}, {psi_vac.max():.4f}]")

Step 3: Visualise the Vacuum Flux

The vacuum field shows the magnetic geometry before any plasma current.

fig, ax = plt.subplots(figsize=(6, 8))
levels = np.linspace(psi_vac.min(), psi_vac.max(), 30)
cs = ax.contour(RR, ZZ, psi_vac, levels=levels, cmap='RdBu_r')
ax.set_xlabel('R [m]')
ax.set_ylabel('Z [m]')
ax.set_title('Vacuum Poloidal Flux (Ψ_vac)')
ax.set_aspect('equal')

# Mark coil positions
for coil in config['coils']:
    marker = 'r^' if coil['current'] > 0 else 'bv'
    ax.plot(coil['R'], coil['Z'], marker, markersize=10)

plt.colorbar(cs, label='Ψ [Wb/rad]')
plt.tight_layout()
plt.show()

Step 4: Solve Full Equilibrium

The Picard iteration alternates between:

1. Solving the Grad-Shafranov PDE (SOR solver) for Ψ
2. Updating the plasma current source J_φ based on Ψ

until convergence.

kernel.solve_equilibrium()

psi_eq = kernel.Psi.copy()

# Find magnetic axis
axis_r, axis_z, psi_axis = kernel.find_x_point(psi_eq)
print(f"Magnetic axis: R={axis_r:.3f} m, Z={axis_z:.3f} m")
print(f"Ψ at axis: {psi_axis:.6f} Wb/rad")
print(f"Ψ range: [{psi_eq.min():.4f}, {psi_eq.max():.4f}]")

Benchmark

%%timeit -n1 -r3
_k = FusionKernel(config_path)
_k.initialize_grid()
_k.calculate_vacuum_field()
_k.solve_equilibrium()

Step 5: Compare Vacuum vs Equilibrium

fig, axes = plt.subplots(1, 2, figsize=(12, 8))

# Vacuum
levels_v = np.linspace(psi_vac.min(), psi_vac.max(), 25)
axes[0].contour(RR, ZZ, psi_vac, levels=levels_v, cmap='RdBu_r')
axes[0].set_title('Vacuum Flux')
axes[0].set_xlabel('R [m]')
axes[0].set_ylabel('Z [m]')
axes[0].set_aspect('equal')

# Equilibrium
levels_e = np.linspace(psi_eq.min(), psi_eq.max(), 25)
axes[1].contour(RR, ZZ, psi_eq, levels=levels_e, cmap='RdBu_r')
axes[1].set_title('Equilibrium Flux (with plasma)')
axes[1].set_xlabel('R [m]')
axes[1].set_ylabel('Z [m]')
axes[1].set_aspect('equal')

plt.suptitle('Grad-Shafranov Equilibrium', fontsize=14)
plt.tight_layout()
plt.show()

# Cleanup
os.unlink(config_path)

Performance Benchmarks

Timing the key computations in this notebook:

1. Vacuum field calculation (calculate_vacuum_field)
2. Full Grad-Shafranov equilibrium solve (solve_equilibrium)
3. X-point / magnetic axis search (find_x_point)

import timeit

# Re-create a fresh kernel for benchmarking (config was cleaned up above)
_bench_config = {
    "reactor_name": "ITER-like",
    "grid_resolution": [65, 65],
    "dimensions": {"R_min": 1.0, "R_max": 9.0, "Z_min": -5.0, "Z_max": 5.0},
    "physics": {"plasma_current_target": 15.0, "vacuum_permeability": 1.2566370614e-6},
    "coils": [
        {"R": 3.5, "Z": 4.0, "current": 5.0},
        {"R": 3.5, "Z": -4.0, "current": 5.0},
        {"R": 9.0, "Z": 4.0, "current": -3.0},
        {"R": 9.0, "Z": -4.0, "current": -3.0},
        {"R": 6.2, "Z": 5.5, "current": -1.5},
        {"R": 6.2, "Z": -5.5, "current": -1.5},
    ],
    "solver": {"max_iterations": 100, "convergence_threshold": 1e-6, "relaxation_factor": 0.1},
}

import tempfile
_fd, _bench_path = tempfile.mkstemp(suffix=".json")
os.close(_fd)
with open(_bench_path, 'w') as _f:
    json.dump(_bench_config, _f)

# 1. Vacuum field calculation
def bench_vacuum():
    k = FusionKernel(_bench_path)
    k.initialize_grid()
    k.calculate_vacuum_field()

t_vac = timeit.repeat(bench_vacuum, number=3, repeat=3)
print(f"calculate_vacuum_field (65x65 grid, 3 calls):")
print(f"  Mean: {np.mean(t_vac)*1000:.1f} ms +/- {np.std(t_vac)*1000:.1f} ms")
print(f"  Per call: {np.mean(t_vac)/3*1000:.1f} ms")

# 2. Full equilibrium solve (Picard + SOR)
def bench_equilibrium():
    k = FusionKernel(_bench_path)
    k.initialize_grid()
    k.calculate_vacuum_field()
    k.solve_equilibrium()

t_eq = timeit.repeat(bench_equilibrium, number=1, repeat=3)
print(f"\nsolve_equilibrium (full Picard iteration, 65x65 grid):")
print(f"  Mean: {np.mean(t_eq)*1000:.1f} ms +/- {np.std(t_eq)*1000:.1f} ms")

# 3. X-point / magnetic axis search
_k_bench = FusionKernel(_bench_path)
_k_bench.initialize_grid()
_k_bench.calculate_vacuum_field()
_k_bench.solve_equilibrium()
_psi_bench = _k_bench.Psi.copy()

def bench_xpoint():
    _k_bench.find_x_point(_psi_bench)

t_xp = timeit.repeat(bench_xpoint, number=100, repeat=5)
print(f"\nfind_x_point (100 calls):")
print(f"  Mean: {np.mean(t_xp)*1000:.1f} ms +/- {np.std(t_xp)*1000:.1f} ms")
print(f"  Per call: {np.mean(t_xp)/100*1000:.2f} ms")

# Cleanup
os.unlink(_bench_path)

Summary

The Grad-Shafranov solver:

• Uses a Red-Black SOR scheme with 5-point stencil and toroidal 1/R corrections
• Picard iteration couples PDE solve with nonlinear source update
• When the Rust acceleration is available (scpn_fusion_rs), the SOR step runs at native speed (~10-100× faster than NumPy for large grids)

Next: See 04_disruption_predictor.ipynb for ML-based disruption early warning.