1#!/usr/bin/env python3
2"""
3Grad-Shafranov Equation Solver
4
5Solves the Grad-Shafranov equation for axisymmetric toroidal MHD equilibria:
6 Δ*ψ = -μ₀ R² dp/dψ - F dF/dψ
7
8where:
9 Δ* = R ∂/∂R (1/R ∂/∂R) + ∂²/∂Z²
10 ψ is the poloidal flux function
11 p(ψ) is the pressure profile
12 F(ψ) = R*B_φ is the poloidal current function
13
14Uses iterative Picard iteration and verifies against Solovev analytical solution.
15
16Author: Claude
17Date: 2026-02-14
18"""
19
20import numpy as np
21import matplotlib.pyplot as plt
22from scipy.sparse import diags, linalg as splinalg
23from scipy.interpolate import interp1d
24
25# Physical constants
26MU_0 = 4 * np.pi * 1e-7 # Permeability of free space
27
28
29class GradShafranovSolver:
30 """
31 Solver for the Grad-Shafranov equation.
32
33 Attributes
34 ----------
35 R : ndarray
36 Major radius grid
37 Z : ndarray
38 Vertical position grid
39 psi : ndarray
40 Poloidal flux function
41 """
42
43 def __init__(self, R_range, Z_range, nR, nZ):
44 """
45 Initialize solver grid.
46
47 Parameters
48 ----------
49 R_range : tuple
50 (R_min, R_max) in meters
51 Z_range : tuple
52 (Z_min, Z_max) in meters
53 nR, nZ : int
54 Number of grid points
55 """
56 self.R_arr = np.linspace(R_range[0], R_range[1], nR)
57 self.Z_arr = np.linspace(Z_range[0], Z_range[1], nZ)
58 self.R, self.Z = np.meshgrid(self.R_arr, self.Z_arr, indexing='ij')
59
60 self.dR = self.R_arr[1] - self.R_arr[0]
61 self.dZ = self.Z_arr[1] - self.Z_arr[0]
62
63 self.nR = nR
64 self.nZ = nZ
65
66 self.psi = np.zeros((nR, nZ))
67
68 def source_term(self, psi):
69 """
70 Compute source term: S(ψ) = -μ₀ R² dp/dψ - F dF/dψ
71
72 Parameters
73 ----------
74 psi : ndarray
75 Current flux function
76
77 Returns
78 -------
79 source : ndarray
80 Source term
81 """
82 # Normalize psi to [0, 1] for profile functions
83 psi_min, psi_max = psi.min(), psi.max()
84 if psi_max > psi_min:
85 psi_norm = (psi - psi_min) / (psi_max - psi_min)
86 else:
87 psi_norm = np.zeros_like(psi)
88
89 # Pressure profile: p(ψ) = p0 * (1 - ψ_norm)²
90 p0 = 1e5 # 100 kPa
91 dpdpsi = -2 * p0 / (psi_max - psi_min + 1e-10)
92
93 # F profile: F(ψ) = F0 (constant for simplicity)
94 F0 = 1.0 # R*B_phi at magnetic axis
95 dFdpsi = 0.0
96
97 source = -MU_0 * self.R**2 * dpdpsi - F0 * dFdpsi
98
99 return source
100
101 def delta_star_operator(self, psi):
102 """
103 Apply Δ* operator: Δ*ψ = R ∂/∂R (1/R ∂ψ/∂R) + ∂²ψ/∂Z²
104
105 Parameters
106 ----------
107 psi : ndarray
108 Input flux function
109
110 Returns
111 -------
112 result : ndarray
113 Result of Δ* operation
114 """
115 result = np.zeros_like(psi)
116
117 for i in range(1, self.nR - 1):
118 for j in range(1, self.nZ - 1):
119 R = self.R[i, j]
120
121 # ∂ψ/∂R at i±1/2
122 dpsi_dR_plus = (psi[i+1, j] - psi[i, j]) / self.dR
123 dpsi_dR_minus = (psi[i, j] - psi[i-1, j]) / self.dR
124
125 # R values at i±1/2
126 R_plus = (self.R[i+1, j] + R) / 2
127 R_minus = (self.R[i-1, j] + R) / 2
128
129 # R ∂/∂R (1/R ∂ψ/∂R)
130 term1 = R * (dpsi_dR_plus / R_plus - dpsi_dR_minus / R_minus) / self.dR
131
132 # ∂²ψ/∂Z²
133 term2 = (psi[i, j+1] - 2*psi[i, j] + psi[i, j-1]) / self.dZ**2
134
135 result[i, j] = term1 + term2
136
137 return result
138
139 def solve_picard(self, max_iter=100, tol=1e-6, omega=0.5):
140 """
141 Solve Grad-Shafranov using Picard iteration.
142
143 Parameters
144 ----------
145 max_iter : int
146 Maximum iterations
147 tol : float
148 Convergence tolerance
149 omega : float
150 Relaxation parameter
151
152 Returns
153 -------
154 converged : bool
155 Whether solver converged
156 """
157 print("Starting Picard iteration...")
158
159 for iteration in range(max_iter):
160 psi_old = self.psi.copy()
161
162 # Compute source term with current psi
163 source = self.source_term(self.psi)
164
165 # Solve Δ*ψ = source using finite differences
166 # This is a simplified direct update; a real solver would use
167 # a linear system solver
168 for i in range(1, self.nR - 1):
169 for j in range(1, self.nZ - 1):
170 R = self.R[i, j]
171
172 # Coefficients for finite difference stencil
173 R_plus = (self.R[i+1, j] + R) / 2
174 R_minus = (self.R[i-1, j] + R) / 2
175
176 a_R = R / (R_plus * self.dR**2)
177 b_R = -R / (R_minus * self.dR**2)
178 c_Z = 1.0 / self.dZ**2
179
180 coeff = a_R + b_R + 2*c_Z
181
182 rhs = (a_R * psi_old[i+1, j] +
183 b_R * psi_old[i-1, j] +
184 c_Z * (psi_old[i, j+1] + psi_old[i, j-1]) -
185 source[i, j])
186
187 self.psi[i, j] = rhs / coeff
188
189 # Apply relaxation
190 self.psi = omega * self.psi + (1 - omega) * psi_old
191
192 # Boundary conditions: psi = 0 at boundaries
193 self.psi[0, :] = 0
194 self.psi[-1, :] = 0
195 self.psi[:, 0] = 0
196 self.psi[:, -1] = 0
197
198 # Check convergence
199 error = np.max(np.abs(self.psi - psi_old))
200
201 if iteration % 10 == 0:
202 print(f" Iteration {iteration}: error = {error:.2e}")
203
204 if error < tol:
205 print(f"Converged after {iteration} iterations")
206 return True
207
208 print(f"Did not converge after {max_iter} iterations")
209 return False
210
211 def solovev_solution(self, epsilon=0.32, kappa=1.0, delta=0.0):
212 """
213 Analytical Solovev solution for verification.
214
215 Parameters
216 ----------
217 epsilon : float
218 Inverse aspect ratio a/R0
219 kappa : float
220 Elongation
221 delta : float
222 Triangularity
223
224 Returns
225 -------
226 psi_solovev : ndarray
227 Analytical flux function
228 """
229 R0 = (self.R_arr.max() + self.R_arr.min()) / 2
230 a = epsilon * R0
231
232 # Solovev parameters
233 c = [1.0, 0.0, -1.0/(2*a**2), 0.0] # Simplified
234
235 psi_sol = np.zeros_like(self.psi)
236
237 for i in range(self.nR):
238 for j in range(self.nZ):
239 R = self.R[i, j]
240 Z = self.Z[i, j]
241
242 x = (R - R0) / a
243 y = Z / a
244
245 # Solovev form: ψ = c[0]*x² + c[2]*y² + higher order terms
246 psi_sol[i, j] = (c[0] * x**2 + c[2] * y**2 +
247 c[1] * x + c[3] * (x**2 * y - y**3/3))
248
249 # Normalize
250 psi_sol = psi_sol - psi_sol.min()
251 psi_sol = psi_sol / psi_sol.max()
252
253 return psi_sol
254
255
256def compute_current_density(solver):
257 """
258 Compute toroidal current density: J_φ = R Δ*ψ / μ₀
259
260 Parameters
261 ----------
262 solver : GradShafranovSolver
263 Solver instance
264
265 Returns
266 -------
267 J_phi : ndarray
268 Toroidal current density
269 """
270 delta_star_psi = solver.delta_star_operator(solver.psi)
271 J_phi = solver.R * delta_star_psi / MU_0
272
273 return J_phi
274
275
276def plot_equilibrium(solver, J_phi):
277 """
278 Plot flux surfaces and current density.
279
280 Parameters
281 ----------
282 solver : GradShafranovSolver
283 Solver instance
284 J_phi : ndarray
285 Current density
286 """
287 fig, axes = plt.subplots(1, 3, figsize=(15, 5))
288
289 # Flux surfaces
290 levels = np.linspace(solver.psi.min(), solver.psi.max(), 20)
291 cs1 = axes[0].contour(solver.R, solver.Z, solver.psi, levels=levels,
292 colors='blue', linewidths=1.5)
293 axes[0].set_xlabel('R (m)', fontsize=12)
294 axes[0].set_ylabel('Z (m)', fontsize=12)
295 axes[0].set_title('Flux Surfaces (ψ contours)', fontsize=13, fontweight='bold')
296 axes[0].set_aspect('equal')
297 axes[0].grid(True, alpha=0.3)
298
299 # Current density
300 im2 = axes[1].contourf(solver.R, solver.Z, J_phi / 1e6,
301 levels=20, cmap='RdBu_r')
302 axes[1].contour(solver.R, solver.Z, solver.psi, levels=10,
303 colors='black', linewidths=0.5, alpha=0.3)
304 axes[1].set_xlabel('R (m)', fontsize=12)
305 axes[1].set_ylabel('Z (m)', fontsize=12)
306 axes[1].set_title('Toroidal Current Density J_φ', fontsize=13, fontweight='bold')
307 axes[1].set_aspect('equal')
308 plt.colorbar(im2, ax=axes[1], label='J_φ (MA/m²)')
309
310 # Pressure (computed from psi)
311 psi_norm = (solver.psi - solver.psi.min()) / (solver.psi.max() - solver.psi.min() + 1e-10)
312 p = 1e5 * (1 - psi_norm)**2
313 im3 = axes[2].contourf(solver.R, solver.Z, p / 1e3, levels=20, cmap='hot')
314 axes[2].contour(solver.R, solver.Z, solver.psi, levels=10,
315 colors='blue', linewidths=0.5, alpha=0.3)
316 axes[2].set_xlabel('R (m)', fontsize=12)
317 axes[2].set_ylabel('Z (m)', fontsize=12)
318 axes[2].set_title('Pressure', fontsize=13, fontweight='bold')
319 axes[2].set_aspect('equal')
320 plt.colorbar(im3, ax=axes[2], label='Pressure (kPa)')
321
322 plt.suptitle('Grad-Shafranov Equilibrium Solution',
323 fontsize=14, fontweight='bold')
324 plt.tight_layout()
325 plt.savefig('grad_shafranov_solution.png', dpi=150, bbox_inches='tight')
326 plt.show()
327
328
329def main():
330 """Main execution function."""
331 # Create solver
332 R_range = (0.5, 1.5) # 0.5 to 1.5 meters
333 Z_range = (-0.5, 0.5) # -0.5 to 0.5 meters
334 nR, nZ = 80, 80
335
336 solver = GradShafranovSolver(R_range, Z_range, nR, nZ)
337
338 print("Grad-Shafranov Equation Solver")
339 print("=" * 50)
340 print(f"Grid: {nR} x {nZ}")
341 print(f"R range: {R_range[0]:.2f} to {R_range[1]:.2f} m")
342 print(f"Z range: {Z_range[0]:.2f} to {Z_range[1]:.2f} m")
343 print()
344
345 # Solve using Picard iteration
346 converged = solver.solve_picard(max_iter=200, tol=1e-6, omega=0.5)
347
348 if not converged:
349 print("Warning: Solver did not fully converge")
350
351 # Compute current density
352 J_phi = compute_current_density(solver)
353
354 print(f"\nResults:")
355 print(f" ψ range: {solver.psi.min():.3e} to {solver.psi.max():.3e}")
356 print(f" Peak current density: {np.max(np.abs(J_phi))/1e6:.2f} MA/m²")
357
358 # Compare with Solovev solution (for verification)
359 psi_solovev = solver.solovev_solution(epsilon=0.3)
360 difference = np.max(np.abs(solver.psi/solver.psi.max() - psi_solovev))
361 print(f" Difference from Solovev solution: {difference:.3f}")
362
363 # Plot results
364 plot_equilibrium(solver, J_phi)
365 print("\nPlot saved as 'grad_shafranov_solution.png'")
366
367
368if __name__ == '__main__':
369 main()