1#!/usr/bin/env python3
2"""
3Tokamak Equilibrium Calculation
4
5This module performs simplified tokamak equilibrium calculations, computing
6key parameters for magnetic confinement fusion.
7
8Key concepts:
9- Safety factor q(r): measures field line winding
10- Shafranov shift: outward displacement of flux surfaces due to pressure
11- Beta values: ratio of plasma to magnetic pressure
12- Troyon beta limit: operational constraint
13
14The safety factor is crucial for stability:
15 q = r B_toroidal / (R B_poloidal)
16
17Author: MHD Course Examples
18License: MIT
19"""
20
21import numpy as np
22import matplotlib.pyplot as plt
23from scipy.integrate import cumtrapz
24
25
26class TokamakEquilibrium:
27 """
28 Simplified tokamak equilibrium solver.
29
30 Attributes:
31 R0 (float): Major radius (m)
32 a (float): Minor radius (m)
33 B_tor (float): Toroidal field at R0 (T)
34 I_p (float): Plasma current (MA)
35 beta_0 (float): Central beta value
36 """
37
38 def __init__(self, R0=3.0, a=1.0, B_tor=5.0, I_p=15.0, beta_0=0.02):
39 """
40 Initialize tokamak parameters.
41
42 Parameters:
43 R0 (float): Major radius (m)
44 a (float): Minor radius (m)
45 B_tor (float): Toroidal field (T)
46 I_p (float): Plasma current (MA)
47 beta_0 (float): Central beta
48 """
49 self.R0 = R0
50 self.a = a
51 self.B_tor = B_tor
52 self.I_p = I_p * 1e6 # Convert to Amperes
53 self.beta_0 = beta_0
54
55 # Compute derived quantities
56 self.epsilon = a / R0 # Inverse aspect ratio
57 self.mu_0 = 4 * np.pi * 1e-7
58
59 # Grid
60 self.nr = 100
61 self.r = np.linspace(0, a, self.nr)
62
63 def current_density_profile(self, r, alpha=2.0):
64 """
65 Parameterized current density profile.
66
67 j(r) = j_0 (1 - (r/a)²)^α
68
69 Parameters:
70 r (ndarray): Minor radius
71 alpha (float): Profile parameter
72
73 Returns:
74 ndarray: Current density (A/m²)
75 """
76 rho = r / self.a
77 j = (1 - rho**2)**alpha
78
79 # Normalize to give total current I_p
80 # I_p = ∫ j(r) 2πr dr
81 integral = 2 * np.pi * np.trapz(j * r, r)
82 j_0 = self.I_p / integral
83
84 return j_0 * j
85
86 def pressure_profile(self, r, alpha_p=2.0):
87 """
88 Pressure profile.
89
90 p(r) = p_0 (1 - (r/a)²)^α_p
91
92 Parameters:
93 r (ndarray): Minor radius
94 alpha_p (float): Pressure profile parameter
95
96 Returns:
97 ndarray: Pressure (Pa)
98 """
99 rho = r / self.a
100 p_norm = (1 - rho**2)**alpha_p
101
102 # p_0 from beta_0
103 # β = 2μ₀p/B²
104 p_0 = self.beta_0 * self.B_tor**2 / (2 * self.mu_0)
105
106 return p_0 * p_norm
107
108 def compute_q_profile(self):
109 """
110 Compute safety factor q(r).
111
112 q(r) = (r B_tor) / (R₀ B_pol)
113 where B_pol comes from enclosed current.
114
115 Returns:
116 ndarray: Safety factor profile
117 """
118 r = self.r
119 j = self.current_density_profile(r)
120
121 # Poloidal field from Ampere's law
122 # B_pol(r) = μ₀ I_enclosed / (2π r)
123
124 # Compute enclosed current
125 I_enclosed = np.zeros_like(r)
126 for i in range(1, len(r)):
127 I_enclosed[i] = 2 * np.pi * np.trapz(j[:i+1] * r[:i+1], r[:i+1])
128
129 # Poloidal field
130 B_pol = np.zeros_like(r)
131 B_pol[1:] = self.mu_0 * I_enclosed[1:] / (2 * np.pi * r[1:])
132 B_pol[0] = B_pol[1] # Avoid singularity
133
134 # Safety factor
135 q = r * self.B_tor / (self.R0 * B_pol)
136 q[0] = q[1] # Fix q(0)
137
138 return q, B_pol
139
140 def compute_shafranov_shift(self):
141 """
142 Compute Shafranov shift Δ(r).
143
144 Simplified formula:
145 Δ(r) ≈ (r²/R₀) × (β_pol + l_i/2)
146
147 where β_pol is poloidal beta and l_i is internal inductance.
148
149 Returns:
150 ndarray: Shafranov shift
151 """
152 r = self.r
153 p = self.pressure_profile(r)
154 q, B_pol = self.compute_q_profile()
155
156 # Poloidal beta
157 beta_pol_avg = np.trapz(p * r, r) / np.trapz(0.5 * B_pol**2 / self.mu_0 * r, r)
158
159 # Internal inductance (approximate)
160 l_i = 0.5 # Typical value
161
162 # Shafranov shift
163 Delta = (r**2 / self.R0) * (beta_pol_avg + l_i / 2)
164
165 return Delta, beta_pol_avg
166
167 def compute_beta_values(self):
168 """
169 Compute various beta values.
170
171 Returns:
172 dict: Beta values
173 """
174 r = self.r
175 p = self.pressure_profile(r)
176 q, B_pol = self.compute_q_profile()
177
178 # Volume-averaged pressure
179 p_avg = np.trapz(p * r, r) / np.trapz(r, r)
180
181 # Toroidal beta
182 beta_tor = 2 * self.mu_0 * p_avg / self.B_tor**2
183
184 # Poloidal beta
185 B_pol_avg = np.sqrt(np.trapz(B_pol**2 * r, r) / np.trapz(r, r))
186 beta_pol = 2 * self.mu_0 * p_avg / B_pol_avg**2
187
188 # Total beta
189 B_total_avg = np.sqrt(self.B_tor**2 + B_pol_avg**2)
190 beta_total = 2 * self.mu_0 * p_avg / B_total_avg**2
191
192 # Normalized beta (Troyon)
193 beta_N = beta_tor * 100 * self.a * self.B_tor / (self.I_p / 1e6)
194
195 return {
196 'beta_tor': beta_tor,
197 'beta_pol': beta_pol,
198 'beta_total': beta_total,
199 'beta_N': beta_N
200 }
201
202 def check_troyon_limit(self):
203 """
204 Check Troyon beta limit: β_N < β_{N,max} ≈ 3
205
206 Returns:
207 bool: True if within limit
208 """
209 betas = self.compute_beta_values()
210 beta_N_max = 3.0
211
212 within_limit = betas['beta_N'] < beta_N_max
213
214 return within_limit, betas['beta_N'], beta_N_max
215
216 def plot_equilibrium(self):
217 """
218 Plot equilibrium profiles and flux surfaces.
219 """
220 fig = plt.figure(figsize=(14, 10))
221 gs = fig.add_gridspec(2, 2)
222
223 ax1 = fig.add_subplot(gs[0, 0])
224 ax2 = fig.add_subplot(gs[0, 1])
225 ax3 = fig.add_subplot(gs[1, 0])
226 ax4 = fig.add_subplot(gs[1, 1])
227
228 r = self.r
229 q, B_pol = self.compute_q_profile()
230 p = self.pressure_profile(r)
231 j = self.current_density_profile(r)
232 Delta, beta_pol = self.compute_shafranov_shift()
233
234 # Safety factor
235 ax1.plot(r / self.a, q, 'b-', linewidth=2)
236 ax1.axhline(y=1, color='r', linestyle='--', alpha=0.5,
237 label='q=1 (sawtooth)')
238 ax1.axhline(y=2, color='orange', linestyle='--', alpha=0.5,
239 label='q=2 (m=2 kink)')
240 ax1.set_xlabel(r'$r/a$', fontsize=11)
241 ax1.set_ylabel('Safety Factor q', fontsize=11)
242 ax1.set_title('Safety Factor Profile', fontsize=13)
243 ax1.grid(True, alpha=0.3)
244 ax1.legend(fontsize=10)
245 ax1.text(0.05, 0.95, f'q(0) = {q[0]:.2f}\nq(a) = {q[-1]:.2f}',
246 transform=ax1.transAxes, verticalalignment='top',
247 bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
248
249 # Pressure and current
250 ax2_twin = ax2.twinx()
251 ax2.plot(r / self.a, p / 1e3, 'b-', linewidth=2, label='Pressure')
252 ax2_twin.plot(r / self.a, j / 1e6, 'r--', linewidth=2, label='Current density')
253 ax2.set_xlabel(r'$r/a$', fontsize=11)
254 ax2.set_ylabel('Pressure (kPa)', fontsize=11, color='b')
255 ax2_twin.set_ylabel(r'Current Density (MA/m$^2$)', fontsize=11, color='r')
256 ax2.set_title('Pressure and Current Profiles', fontsize=13)
257 ax2.grid(True, alpha=0.3)
258 ax2.tick_params(axis='y', labelcolor='b')
259 ax2_twin.tick_params(axis='y', labelcolor='r')
260
261 # Flux surfaces with Shafranov shift
262 theta = np.linspace(0, 2*np.pi, 100)
263 colors = plt.cm.viridis(np.linspace(0, 1, 10))
264
265 for i, r_flux in enumerate(np.linspace(0.1, 0.9, 10) * self.a):
266 idx = np.argmin(np.abs(r - r_flux))
267 shift = Delta[idx]
268
269 R = self.R0 + shift + r_flux * np.cos(theta)
270 Z = r_flux * np.sin(theta)
271
272 ax3.plot(R, Z, color=colors[i], linewidth=1.5)
273
274 # First wall
275 R_wall = self.R0 + self.a * np.cos(theta)
276 Z_wall = self.a * np.sin(theta)
277 ax3.plot(R_wall, Z_wall, 'k-', linewidth=2.5, label='First wall')
278
279 ax3.set_xlabel('R (m)', fontsize=11)
280 ax3.set_ylabel('Z (m)', fontsize=11)
281 ax3.set_title(f'Flux Surfaces with Shafranov Shift\nΔ(a) = {Delta[-1]*100:.1f} cm',
282 fontsize=13)
283 ax3.set_aspect('equal')
284 ax3.grid(True, alpha=0.3)
285 ax3.legend(fontsize=10)
286
287 # Beta values
288 betas = self.compute_beta_values()
289 within_limit, beta_N, beta_N_max = self.check_troyon_limit()
290
291 beta_names = ['β_tor', 'β_pol', 'β_total', 'β_N']
292 beta_vals = [betas['beta_tor']*100, betas['beta_pol']*100,
293 betas['beta_total']*100, betas['beta_N']]
294 colors_beta = ['blue', 'red', 'green', 'purple']
295
296 bars = ax4.bar(beta_names, beta_vals, color=colors_beta, alpha=0.7)
297 ax4.axhline(y=beta_N_max, color='orange', linestyle='--',
298 linewidth=2, label=f'Troyon limit (β_N={beta_N_max})')
299 ax4.set_ylabel('Value (%)', fontsize=11)
300 ax4.set_title('Beta Values', fontsize=13)
301 ax4.grid(True, alpha=0.3, axis='y')
302 ax4.legend(fontsize=10)
303
304 # Add values on bars
305 for bar, val in zip(bars, beta_vals):
306 height = bar.get_height()
307 ax4.text(bar.get_x() + bar.get_width()/2., height,
308 f'{val:.2f}',
309 ha='center', va='bottom', fontsize=10)
310
311 # Status text
312 status = "WITHIN LIMIT" if within_limit else "EXCEEDS LIMIT"
313 color = 'green' if within_limit else 'red'
314 ax4.text(0.5, 0.95, f'Troyon: {status}',
315 transform=ax4.transAxes, ha='center', va='top',
316 fontsize=11, weight='bold', color=color,
317 bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
318
319 plt.tight_layout()
320 return fig
321
322
323def main():
324 """
325 Main function demonstrating tokamak equilibrium calculation.
326 """
327 print("=" * 60)
328 print("Tokamak Equilibrium Calculation")
329 print("=" * 60)
330
331 # ITER-like parameters
332 tokamak = TokamakEquilibrium(
333 R0=6.2, # Major radius (m)
334 a=2.0, # Minor radius (m)
335 B_tor=5.3, # Toroidal field (T)
336 I_p=15.0, # Plasma current (MA)
337 beta_0=0.025 # Central beta
338 )
339
340 print(f"\nParameters:")
341 print(f" Major radius R₀ = {tokamak.R0:.1f} m")
342 print(f" Minor radius a = {tokamak.a:.1f} m")
343 print(f" Aspect ratio A = {tokamak.R0/tokamak.a:.1f}")
344 print(f" Toroidal field = {tokamak.B_tor:.1f} T")
345 print(f" Plasma current = {tokamak.I_p/1e6:.1f} MA")
346
347 # Compute equilibrium
348 q, B_pol = tokamak.compute_q_profile()
349 betas = tokamak.compute_beta_values()
350 within_limit, beta_N, beta_N_max = tokamak.check_troyon_limit()
351
352 print(f"\nEquilibrium properties:")
353 print(f" q(0) = {q[0]:.2f}")
354 print(f" q(a) = {q[-1]:.2f}")
355 print(f" β_toroidal = {betas['beta_tor']*100:.2f}%")
356 print(f" β_poloidal = {betas['beta_pol']*100:.2f}%")
357 print(f" β_N = {betas['beta_N']:.2f}")
358 print(f"\nTroyon limit check:")
359 print(f" β_N = {beta_N:.2f} < {beta_N_max:.1f}: {within_limit}")
360
361 # Plot
362 tokamak.plot_equilibrium()
363
364 plt.savefig('/tmp/tokamak_equilibrium.png', dpi=150, bbox_inches='tight')
365 print("\nPlot saved to /tmp/tokamak_equilibrium.png")
366
367 plt.show()
368
369
370if __name__ == "__main__":
371 main()