1#!/usr/bin/env python3
2"""
3Flux Surface Visualization
4
5Visualizes nested flux surfaces from Grad-Shafranov solution, including:
6- Nested flux surfaces (contours of ψ)
7- Last Closed Flux Surface (LCFS)
8- Magnetic axis (ψ maximum)
9- X-point location (for divertor configurations)
10- Safety factor q on each surface
11
12Author: Claude
13Date: 2026-02-14
14"""
15
16import numpy as np
17import matplotlib.pyplot as plt
18from scipy.interpolate import RectBivariateSpline
19from scipy.optimize import minimize_scalar
20
21
22def solovev_flux_function(R, Z, R0=1.0, a=0.3, kappa=1.0, delta=0.0):
23 """
24 Compute Solovev analytical flux function.
25
26 Parameters
27 ----------
28 R, Z : array_like
29 Major radius and vertical position
30 R0 : float
31 Major radius of magnetic axis
32 a : float
33 Minor radius
34 kappa : float
35 Elongation
36 delta : float
37 Triangularity
38
39 Returns
40 -------
41 psi : ndarray
42 Poloidal flux
43 """
44 x = (R - R0) / a
45 y = Z / (a * kappa)
46
47 # Solovev solution parameters
48 c1 = 1.0
49 c2 = -1.0 / (2 * a**2)
50
51 # Include triangularity effect
52 c3 = delta / a**2
53
54 psi = c1 * x**2 + c2 * y**2 + c3 * x * y**2
55
56 return psi
57
58
59def find_magnetic_axis(R, Z, psi):
60 """
61 Find magnetic axis (maximum of ψ).
62
63 Parameters
64 ----------
65 R, Z : ndarray
66 Grid coordinates
67 psi : ndarray
68 Flux function
69
70 Returns
71 -------
72 R_axis, Z_axis : float
73 Magnetic axis location
74 psi_axis : float
75 Flux at axis
76 """
77 idx = np.unravel_index(np.argmax(psi), psi.shape)
78 R_axis = R[idx]
79 Z_axis = Z[idx]
80 psi_axis = psi[idx]
81
82 return R_axis, Z_axis, psi_axis
83
84
85def find_x_point(R, Z, psi):
86 """
87 Find X-point (saddle point of ψ).
88
89 For a divertor configuration, the X-point is a saddle point
90 where ∂ψ/∂R = ∂ψ/∂Z = 0 and the Hessian has opposite-sign eigenvalues.
91
92 Parameters
93 ----------
94 R, Z : ndarray
95 Grid coordinates
96 psi : ndarray
97 Flux function
98
99 Returns
100 -------
101 R_x, Z_x : float or None
102 X-point location (None if not found)
103 psi_x : float or None
104 Flux at X-point
105 """
106 # Compute gradients
107 dR = R[1, 0] - R[0, 0]
108 dZ = Z[0, 1] - Z[0, 0]
109
110 dpsi_dR = np.gradient(psi, dR, axis=0)
111 dpsi_dZ = np.gradient(psi, dZ, axis=1)
112
113 # Find points where gradient is small
114 grad_mag = np.sqrt(dpsi_dR**2 + dpsi_dZ**2)
115 threshold = 0.1 * np.max(grad_mag)
116
117 candidates = grad_mag < threshold
118
119 # Look for saddle points (below magnetic axis)
120 Z_mid = (Z.min() + Z.max()) / 2
121 candidates = candidates & (Z < Z_mid)
122
123 if np.any(candidates):
124 # Take the point with lowest Z
125 idx_candidates = np.where(candidates)
126 min_z_idx = np.argmin(Z[idx_candidates])
127 i, j = idx_candidates[0][min_z_idx], idx_candidates[1][min_z_idx]
128
129 R_x = R[i, j]
130 Z_x = Z[i, j]
131 psi_x = psi[i, j]
132
133 return R_x, Z_x, psi_x
134 else:
135 return None, None, None
136
137
138def compute_safety_factor(R, Z, psi, R0, B0):
139 """
140 Compute safety factor q(ψ) for each flux surface.
141
142 q = (1/2π) ∮ (B·∇φ)/(B·∇θ) dθ
143
144 For circular surfaces: q ≈ r*B_φ/(R*B_θ)
145
146 Parameters
147 ----------
148 R, Z : ndarray
149 Grid coordinates
150 psi : ndarray
151 Flux function
152 R0 : float
153 Major radius
154 B0 : float
155 Toroidal field at R0
156
157 Returns
158 -------
159 psi_values : ndarray
160 Flux values
161 q_values : ndarray
162 Safety factor values
163 """
164 # Select flux surfaces
165 psi_min, psi_max = psi.min(), psi.max()
166 psi_levels = np.linspace(psi_min, psi_max * 0.95, 10)
167
168 q_values = []
169
170 for psi_val in psi_levels:
171 # Find contour at this psi value
172 # Approximate q using circular approximation
173 # Find average radius of this flux surface
174 mask = (psi >= psi_val * 0.99) & (psi <= psi_val * 1.01)
175 if np.any(mask):
176 R_surf = R[mask]
177 r_minor = np.mean(np.abs(R_surf - R0))
178
179 # Toroidal field: B_phi ~ B0 * R0 / R
180 B_phi = B0 * R0 / R0
181
182 # Poloidal field estimate from psi
183 dpsi = psi_max - psi_min
184 B_poloidal = dpsi / (2 * np.pi * r_minor * R0) # Rough estimate
185
186 q = r_minor * B_phi / (R0 * B_poloidal + 1e-10)
187 q_values.append(q)
188 else:
189 q_values.append(np.nan)
190
191 return psi_levels, np.array(q_values)
192
193
194def plot_flux_surfaces(R, Z, psi, n_surfaces=15):
195 """
196 Plot nested flux surfaces with annotations.
197
198 Parameters
199 ----------
200 R, Z : ndarray
201 Grid coordinates
202 psi : ndarray
203 Flux function
204 n_surfaces : int
205 Number of flux surfaces to plot
206 """
207 fig, axes = plt.subplots(1, 2, figsize=(14, 6))
208
209 # Find magnetic axis and X-point
210 R_axis, Z_axis, psi_axis = find_magnetic_axis(R, Z, psi)
211 R_x, Z_x, psi_x = find_x_point(R, Z, psi)
212
213 # Plot 1: Flux surfaces with annotations
214 ax = axes[0]
215
216 # Flux surface levels
217 psi_min, psi_max = psi.min(), psi.max()
218 levels = np.linspace(psi_min, psi_max, n_surfaces)
219
220 # Plot flux surfaces
221 cs = ax.contour(R, Z, psi, levels=levels, colors='blue', linewidths=1.5)
222 ax.clabel(cs, inline=True, fontsize=8, fmt='%.3f')
223
224 # Mark LCFS (outermost closed surface)
225 lcfs_level = psi_max * 0.95
226 ax.contour(R, Z, psi, levels=[lcfs_level], colors='red',
227 linewidths=3, linestyles='--')
228
229 # Mark magnetic axis
230 ax.plot(R_axis, Z_axis, 'r*', markersize=20, label='Magnetic Axis')
231 ax.text(R_axis + 0.05, Z_axis + 0.05, f'Axis\n({R_axis:.2f}, {Z_axis:.2f})',
232 fontsize=10, color='red', fontweight='bold')
233
234 # Mark X-point if found
235 if R_x is not None:
236 ax.plot(R_x, Z_x, 'kx', markersize=15, markeredgewidth=3, label='X-point')
237 ax.text(R_x + 0.05, Z_x - 0.05, f'X-point\n({R_x:.2f}, {Z_x:.2f})',
238 fontsize=10, color='black', fontweight='bold')
239
240 ax.set_xlabel('R (m)', fontsize=12)
241 ax.set_ylabel('Z (m)', fontsize=12)
242 ax.set_title('Flux Surfaces', fontsize=13, fontweight='bold')
243 ax.set_aspect('equal')
244 ax.legend(loc='upper right', fontsize=10)
245 ax.grid(True, alpha=0.3)
246
247 # Plot 2: Safety factor q(ψ)
248 ax2 = axes[1]
249
250 R0 = R_axis
251 B0 = 2.0 # Tesla
252 psi_vals, q_vals = compute_safety_factor(R, Z, psi, R0, B0)
253
254 # Normalize psi for plotting
255 psi_norm = (psi_vals - psi_min) / (psi_max - psi_min)
256
257 ax2.plot(psi_norm, q_vals, 'b-o', linewidth=2, markersize=6)
258 ax2.axhline(y=1, color='r', linestyle='--', linewidth=1, alpha=0.5, label='q=1')
259 ax2.axhline(y=2, color='g', linestyle='--', linewidth=1, alpha=0.5, label='q=2')
260 ax2.axhline(y=3, color='m', linestyle='--', linewidth=1, alpha=0.5, label='q=3')
261
262 ax2.set_xlabel('Normalized flux ψ_N', fontsize=12)
263 ax2.set_ylabel('Safety factor q', fontsize=12)
264 ax2.set_title('Safety Factor Profile', fontsize=13, fontweight='bold')
265 ax2.grid(True, alpha=0.3)
266 ax2.legend(loc='best', fontsize=10)
267 ax2.set_xlim([0, 1])
268
269 plt.suptitle('Tokamak Flux Surface Analysis',
270 fontsize=14, fontweight='bold')
271 plt.tight_layout()
272 plt.savefig('flux_surfaces.png', dpi=150, bbox_inches='tight')
273 plt.show()
274
275
276def plot_surface_quantities(R, Z, psi):
277 """
278 Plot additional surface quantities.
279
280 Parameters
281 ----------
282 R, Z : ndarray
283 Grid coordinates
284 psi : ndarray
285 Flux function
286 """
287 fig, axes = plt.subplots(2, 2, figsize=(12, 10))
288
289 # Compute gradients for field components
290 dR = R[1, 0] - R[0, 0]
291 dZ = Z[0, 1] - Z[0, 0]
292
293 dpsi_dR = np.gradient(psi, dR, axis=0)
294 dpsi_dZ = np.gradient(psi, dZ, axis=1)
295
296 # Poloidal field components: B_R = -1/(R) * ∂ψ/∂Z, B_Z = 1/(R) * ∂ψ/∂R
297 B_R = -dpsi_dZ / (R + 1e-10)
298 B_Z = dpsi_dR / (R + 1e-10)
299 B_poloidal = np.sqrt(B_R**2 + B_Z**2)
300
301 # Plot B_R
302 im1 = axes[0, 0].contourf(R, Z, B_R, levels=20, cmap='RdBu_r')
303 axes[0, 0].contour(R, Z, psi, levels=10, colors='black',
304 linewidths=0.5, alpha=0.3)
305 axes[0, 0].set_xlabel('R (m)', fontsize=11)
306 axes[0, 0].set_ylabel('Z (m)', fontsize=11)
307 axes[0, 0].set_title('B_R (radial field)', fontsize=12, fontweight='bold')
308 axes[0, 0].set_aspect('equal')
309 plt.colorbar(im1, ax=axes[0, 0], label='B_R (T)')
310
311 # Plot B_Z
312 im2 = axes[0, 1].contourf(R, Z, B_Z, levels=20, cmap='RdBu_r')
313 axes[0, 1].contour(R, Z, psi, levels=10, colors='black',
314 linewidths=0.5, alpha=0.3)
315 axes[0, 1].set_xlabel('R (m)', fontsize=11)
316 axes[0, 1].set_ylabel('Z (m)', fontsize=11)
317 axes[0, 1].set_title('B_Z (vertical field)', fontsize=12, fontweight='bold')
318 axes[0, 1].set_aspect('equal')
319 plt.colorbar(im2, ax=axes[0, 1], label='B_Z (T)')
320
321 # Plot |B_poloidal|
322 im3 = axes[1, 0].contourf(R, Z, B_poloidal, levels=20, cmap='viridis')
323 axes[1, 0].contour(R, Z, psi, levels=10, colors='white',
324 linewidths=0.5, alpha=0.5)
325 axes[1, 0].set_xlabel('R (m)', fontsize=11)
326 axes[1, 0].set_ylabel('Z (m)', fontsize=11)
327 axes[1, 0].set_title('|B_poloidal|', fontsize=12, fontweight='bold')
328 axes[1, 0].set_aspect('equal')
329 plt.colorbar(im3, ax=axes[1, 0], label='|B_p| (T)')
330
331 # Plot flux surface spacing (proxy for confinement quality)
332 dpsi_norm = np.gradient(psi, axis=0)**2 + np.gradient(psi, axis=1)**2
333 dpsi_norm = np.sqrt(dpsi_norm)
334
335 im4 = axes[1, 1].contourf(R, Z, dpsi_norm, levels=20, cmap='hot')
336 axes[1, 1].contour(R, Z, psi, levels=10, colors='blue',
337 linewidths=0.5, alpha=0.3)
338 axes[1, 1].set_xlabel('R (m)', fontsize=11)
339 axes[1, 1].set_ylabel('Z (m)', fontsize=11)
340 axes[1, 1].set_title('|∇ψ| (flux surface spacing)', fontsize=12, fontweight='bold')
341 axes[1, 1].set_aspect('equal')
342 plt.colorbar(im4, ax=axes[1, 1], label='|∇ψ|')
343
344 plt.suptitle('Poloidal Field Components',
345 fontsize=14, fontweight='bold')
346 plt.tight_layout()
347 plt.savefig('flux_surface_fields.png', dpi=150, bbox_inches='tight')
348 plt.show()
349
350
351def main():
352 """Main execution function."""
353 # Create grid
354 R = np.linspace(0.5, 1.5, 150)
355 Z = np.linspace(-0.6, 0.6, 150)
356 R_grid, Z_grid = np.meshgrid(R, Z, indexing='ij')
357
358 # Generate Solovev solution
359 R0 = 1.0 # Major radius
360 a = 0.3 # Minor radius
361 kappa = 1.5 # Elongation
362 delta = 0.3 # Triangularity
363
364 print("Flux Surface Visualization")
365 print("=" * 50)
366 print(f"Configuration parameters:")
367 print(f" Major radius R0: {R0:.2f} m")
368 print(f" Minor radius a: {a:.2f} m")
369 print(f" Elongation κ: {kappa:.2f}")
370 print(f" Triangularity δ: {delta:.2f}")
371 print()
372
373 psi = solovev_flux_function(R_grid, Z_grid, R0, a, kappa, delta)
374
375 # Find key locations
376 R_axis, Z_axis, psi_axis = find_magnetic_axis(R_grid, Z_grid, psi)
377 R_x, Z_x, psi_x = find_x_point(R_grid, Z_grid, psi)
378
379 print(f"Magnetic axis: R = {R_axis:.3f} m, Z = {Z_axis:.3f} m")
380 print(f" ψ_axis = {psi_axis:.3f}")
381
382 if R_x is not None:
383 print(f"X-point: R = {R_x:.3f} m, Z = {Z_x:.3f} m")
384 print(f" ψ_x = {psi_x:.3f}")
385 else:
386 print("No X-point found (limiter configuration)")
387
388 print()
389
390 # Plot flux surfaces
391 plot_flux_surfaces(R_grid, Z_grid, psi, n_surfaces=15)
392 print("Flux surface plot saved as 'flux_surfaces.png'")
393
394 # Plot field components
395 plot_surface_quantities(R_grid, Z_grid, psi)
396 print("Field components plot saved as 'flux_surface_fields.png'")
397
398
399if __name__ == '__main__':
400 main()