1#!/usr/bin/env python3
2"""
3Production-Quality 1D Vlasov-Poisson Solver
4
5This script implements a high-quality 1D electrostatic Vlasov-Poisson solver
6using spectral methods and operator splitting.
7
8Algorithm:
9- Strang splitting: half x-advect → full v-advect → half x-advect
10- Cubic spline interpolation for advection
11- FFT Poisson solver with periodic boundaries
12
13Test cases:
141. Linear plasma oscillation (verify ω = ωpe)
152. Landau damping (verify γ matches theory)
163. Two-stream instability (growth + saturation)
174. Bump-on-tail (plateau formation)
18
19Author: Plasma Physics Examples
20License: MIT
21"""
22
23import numpy as np
24import matplotlib.pyplot as plt
25from matplotlib.gridspec import GridSpec
26from scipy.interpolate import CubicSpline
27from scipy.fft import fft, ifft, fftfreq
28
29# Physical constants
30QE = 1.602176634e-19 # C
31ME = 9.10938356e-31 # kg
32EPS0 = 8.854187817e-12 # F/m
33
34class VlasovPoisson1D:
35 """1D Vlasov-Poisson solver."""
36
37 def __init__(self, nx, nv, Lx, vmax, dt):
38 """
39 Initialize solver.
40
41 Parameters:
42 -----------
43 nx : int
44 Number of spatial grid points
45 nv : int
46 Number of velocity grid points
47 Lx : float
48 Spatial domain length [m]
49 vmax : float
50 Maximum velocity [m/s]
51 dt : float
52 Time step [s]
53 """
54 self.nx = nx
55 self.nv = nv
56 self.Lx = Lx
57 self.vmax = vmax
58 self.dt = dt
59
60 # Spatial grid (periodic)
61 self.x = np.linspace(0, Lx, nx, endpoint=False)
62 self.dx = Lx / nx
63
64 # Velocity grid (symmetric)
65 self.v = np.linspace(-vmax, vmax, nv)
66 self.dv = 2 * vmax / (nv - 1)
67
68 # Phase space distribution
69 self.f = np.zeros((nx, nv))
70
71 # Electric field and potential
72 self.E = np.zeros(nx)
73 self.phi = np.zeros(nx)
74
75 # Diagnostics storage
76 self.diagnostics = {
77 'time': [],
78 'field_energy': [],
79 'kinetic_energy': [],
80 'total_energy': [],
81 'momentum': [],
82 'entropy': [],
83 'E_max': []
84 }
85
86 def initialize_maxwellian(self, n0, v0, vth):
87 """
88 Initialize with a shifted Maxwellian.
89
90 f(x,v) = (n0 / sqrt(2π vth²)) exp(-(v-v0)²/(2vth²))
91
92 Parameters:
93 -----------
94 n0 : float
95 Density [m^-3]
96 v0 : float
97 Drift velocity [m/s]
98 vth : float
99 Thermal velocity [m/s]
100 """
101 for i in range(self.nx):
102 self.f[i, :] = (n0 / np.sqrt(2 * np.pi * vth**2)) * \
103 np.exp(-(self.v - v0)**2 / (2 * vth**2))
104
105 def initialize_two_stream(self, n0, v0, vth):
106 """
107 Initialize two counter-streaming beams.
108
109 Parameters:
110 -----------
111 n0 : float
112 Density per beam [m^-3]
113 v0 : float
114 Beam velocity [m/s]
115 vth : float
116 Thermal spread [m/s]
117 """
118 for i in range(self.nx):
119 self.f[i, :] = (n0 / (2 * np.sqrt(2 * np.pi * vth**2))) * \
120 (np.exp(-(self.v - v0)**2 / (2 * vth**2)) +
121 np.exp(-(self.v + v0)**2 / (2 * vth**2)))
122
123 def initialize_bump_on_tail(self, n_bulk, n_beam, v_beam, vth_bulk, vth_beam):
124 """
125 Initialize bump-on-tail distribution.
126
127 Parameters:
128 -----------
129 n_bulk : float
130 Bulk density [m^-3]
131 n_beam : float
132 Beam density [m^-3]
133 v_beam : float
134 Beam velocity [m/s]
135 vth_bulk : float
136 Bulk thermal velocity [m/s]
137 vth_beam : float
138 Beam thermal spread [m/s]
139 """
140 for i in range(self.nx):
141 # Bulk
142 f_bulk = (n_bulk / np.sqrt(2 * np.pi * vth_bulk**2)) * \
143 np.exp(-self.v**2 / (2 * vth_bulk**2))
144
145 # Beam
146 f_beam = (n_beam / np.sqrt(2 * np.pi * vth_beam**2)) * \
147 np.exp(-(self.v - v_beam)**2 / (2 * vth_beam**2))
148
149 self.f[i, :] = f_bulk + f_beam
150
151 def add_perturbation(self, k, amplitude):
152 """
153 Add sinusoidal perturbation in x.
154
155 f → f * (1 + amplitude * cos(k*x))
156
157 Parameters:
158 -----------
159 k : float
160 Wavenumber [rad/m]
161 amplitude : float
162 Perturbation amplitude
163 """
164 for j in range(self.nv):
165 self.f[:, j] *= (1 + amplitude * np.cos(k * self.x))
166
167 def compute_moments(self):
168 """
169 Compute density and current from distribution function.
170
171 Returns:
172 --------
173 n, j : density [m^-3] and current density [A/m^2]
174 """
175 # Density: n = ∫ f dv
176 n = np.trapz(self.f, x=self.v, axis=1)
177
178 # Current: j = q ∫ v f dv
179 j = QE * np.trapz(self.v[None, :] * self.f, x=self.v, axis=1)
180
181 return n, j
182
183 def solve_poisson_fft(self, ne):
184 """
185 Solve Poisson equation using FFT: ∇²φ = e(ne - ni)/ε0.
186
187 Assumes ni = n0 (background ions).
188
189 Parameters:
190 -----------
191 ne : array
192 Electron density [m^-3]
193
194 Returns:
195 --------
196 phi, E : potential [V] and electric field [V/m]
197 """
198 # Background ion density (uniform)
199 n0 = np.mean(ne)
200
201 # Charge density
202 rho = -QE * (ne - n0)
203
204 # FFT
205 rho_k = fft(rho)
206 k = 2 * np.pi * fftfreq(self.nx, d=self.dx)
207
208 # Solve in Fourier space: -k²φ_k = rho_k/ε0
209 # Avoid division by zero at k=0
210 phi_k = np.zeros_like(rho_k, dtype=complex)
211 phi_k[k != 0] = -rho_k[k != 0] / (k[k != 0]**2 * EPS0)
212
213 # Inverse FFT
214 phi = np.real(ifft(phi_k))
215
216 # Electric field: E = -dφ/dx
217 E_k = 1j * k * phi_k
218 E = np.real(ifft(E_k))
219
220 return phi, E
221
222 def advect_x(self, f, dt_frac):
223 """
224 Advect in x direction: ∂f/∂t + v·∂f/∂x = 0.
225
226 Uses cubic spline interpolation.
227
228 Parameters:
229 -----------
230 f : array (nx, nv)
231 Distribution function
232 dt_frac : float
233 Time step (can be fractional for splitting)
234
235 Returns:
236 --------
237 f_new : array (nx, nv)
238 """
239 f_new = np.zeros_like(f)
240
241 for j in range(self.nv):
242 # Displacement
243 dx = self.v[j] * dt_frac
244
245 # Interpolate (periodic boundary)
246 cs = CubicSpline(self.x, f[:, j], bc_type='periodic')
247
248 x_new = (self.x - dx) % self.Lx # Periodic
249 f_new[:, j] = cs(x_new)
250
251 return f_new
252
253 def advect_v(self, f, E, dt_frac):
254 """
255 Advect in v direction: ∂f/∂t + a·∂f/∂v = 0.
256
257 where a = qE/m.
258
259 Parameters:
260 -----------
261 f : array (nx, nv)
262 Distribution function
263 E : array (nx,)
264 Electric field [V/m]
265 dt_frac : float
266 Time step
267
268 Returns:
269 --------
270 f_new : array (nx, nv)
271 """
272 f_new = np.zeros_like(f)
273
274 # Acceleration
275 a = -QE * E / ME # Electron
276
277 for i in range(self.nx):
278 # Velocity displacement
279 dv = a[i] * dt_frac
280
281 # Interpolate
282 cs = CubicSpline(self.v, f[i, :], bc_type='natural')
283
284 v_new = self.v - dv
285
286 # Clip to velocity domain
287 v_new = np.clip(v_new, self.v[0], self.v[-1])
288
289 f_new[i, :] = cs(v_new)
290
291 return f_new
292
293 def step(self):
294 """
295 Advance one time step using Strang splitting.
296
297 Strang splitting:
298 1. Half x-advect
299 2. Full v-advect
300 3. Half x-advect
301 """
302 # 1. Half x-advect
303 self.f = self.advect_x(self.f, 0.5 * self.dt)
304
305 # 2. Compute field
306 ne, _ = self.compute_moments()
307 self.phi, self.E = self.solve_poisson_fft(ne)
308
309 # 3. Full v-advect
310 self.f = self.advect_v(self.f, self.E, self.dt)
311
312 # 4. Half x-advect
313 self.f = self.advect_x(self.f, 0.5 * self.dt)
314
315 # Update field for diagnostics
316 ne, _ = self.compute_moments()
317 self.phi, self.E = self.solve_poisson_fft(ne)
318
319 def compute_diagnostics(self, t):
320 """
321 Compute and store diagnostic quantities.
322
323 Parameters:
324 -----------
325 t : float
326 Current time [s]
327 """
328 ne, _ = self.compute_moments()
329
330 # Field energy: (ε0/2) ∫ E² dx
331 E_field = 0.5 * EPS0 * np.sum(self.E**2) * self.dx
332
333 # Kinetic energy: (m/2) ∫∫ v² f dx dv
334 v2_grid = self.v[None, :]**2
335 E_kinetic = 0.5 * ME * np.sum(v2_grid * self.f) * self.dx * self.dv
336
337 # Total energy
338 E_total = E_field + E_kinetic
339
340 # Momentum: ∫∫ m v f dx dv
341 momentum = ME * np.sum(self.v[None, :] * self.f) * self.dx * self.dv
342
343 # Entropy: -∫∫ f ln(f) dx dv
344 f_safe = self.f + 1e-30 # Avoid log(0)
345 entropy = -np.sum(self.f * np.log(f_safe)) * self.dx * self.dv
346
347 # Maximum E field
348 E_max = np.max(np.abs(self.E))
349
350 # Store
351 self.diagnostics['time'].append(t)
352 self.diagnostics['field_energy'].append(E_field)
353 self.diagnostics['kinetic_energy'].append(E_kinetic)
354 self.diagnostics['total_energy'].append(E_total)
355 self.diagnostics['momentum'].append(momentum)
356 self.diagnostics['entropy'].append(entropy)
357 self.diagnostics['E_max'].append(E_max)
358
359 def run(self, t_max, save_every=10):
360 """
361 Run simulation.
362
363 Parameters:
364 -----------
365 t_max : float
366 Maximum simulation time [s]
367 save_every : int
368 Save diagnostics every N steps
369
370 Returns:
371 --------
372 snapshots : list of (t, f, E) tuples
373 """
374 n_steps = int(t_max / self.dt)
375 snapshots = []
376
377 print(f"Running simulation: {n_steps} steps...")
378
379 for step in range(n_steps):
380 t = step * self.dt
381
382 # Save snapshot
383 if step % save_every == 0:
384 snapshots.append((t, self.f.copy(), self.E.copy()))
385 self.compute_diagnostics(t)
386
387 if step % (save_every * 10) == 0:
388 print(f" Step {step}/{n_steps} (t = {t:.2e} s)")
389
390 # Advance
391 self.step()
392
393 print("Simulation complete!")
394 return snapshots
395
396def test_plasma_oscillation():
397 """Test case 1: Linear plasma oscillation."""
398 print("\n" + "=" * 70)
399 print("Test 1: Linear Plasma Oscillation")
400 print("=" * 70)
401
402 # Parameters
403 n0 = 1e16 # m^-3
404 Te = 1.0 # eV
405 vth = np.sqrt(2 * Te * QE / ME)
406
407 # Plasma frequency
408 omega_pe = np.sqrt(n0 * QE**2 / (ME * EPS0))
409 T_pe = 2 * np.pi / omega_pe
410
411 print(f"Density: {n0:.2e} m^-3")
412 print(f"Plasma frequency: {omega_pe/(2*np.pi)/1e9:.3f} GHz")
413 print(f"Plasma period: {T_pe*1e9:.3f} ns")
414
415 # Setup grid
416 Lx = 1e-2 # 1 cm
417 k = 2 * np.pi / Lx
418 vmax = 6 * vth
419
420 solver = VlasovPoisson1D(nx=128, nv=256, Lx=Lx, vmax=vmax, dt=0.1 * T_pe / 100)
421
422 # Initialize
423 solver.initialize_maxwellian(n0, v0=0, vth=vth)
424 solver.add_perturbation(k, amplitude=0.01)
425
426 # Run
427 snapshots = solver.run(t_max=5 * T_pe, save_every=5)
428
429 return solver, snapshots, omega_pe
430
431def test_landau_damping():
432 """Test case 2: Landau damping."""
433 print("\n" + "=" * 70)
434 print("Test 2: Landau Damping")
435 print("=" * 70)
436
437 # Parameters
438 n0 = 1e16 # m^-3
439 Te = 1.0 # eV
440 vth = np.sqrt(2 * Te * QE / ME)
441
442 omega_pe = np.sqrt(n0 * QE**2 / (ME * EPS0))
443
444 # Choose k*λD = 0.5 for moderate damping
445 lambda_D = vth / omega_pe
446 k = 0.5 / lambda_D
447 Lx = 2 * np.pi / k
448
449 print(f"Debye length: {lambda_D*1e6:.2f} μm")
450 print(f"k·λD = {k*lambda_D:.2f}")
451
452 # Setup
453 vmax = 6 * vth
454 T_pe = 2 * np.pi / omega_pe
455
456 solver = VlasovPoisson1D(nx=128, nv=256, Lx=Lx, vmax=vmax, dt=0.1 * T_pe / 100)
457
458 # Initialize
459 solver.initialize_maxwellian(n0, v0=0, vth=vth)
460 solver.add_perturbation(k, amplitude=0.01)
461
462 # Run
463 snapshots = solver.run(t_max=10 * T_pe, save_every=5)
464
465 return solver, snapshots, omega_pe
466
467def plot_results(solver, snapshots, omega_pe, test_name):
468 """Plot simulation results."""
469 fig = plt.figure(figsize=(16, 10))
470 gs = GridSpec(3, 3, figure=fig, hspace=0.35, wspace=0.35)
471
472 # Extract diagnostics
473 t = np.array(solver.diagnostics['time'])
474 E_max = np.array(solver.diagnostics['E_max'])
475 E_field = np.array(solver.diagnostics['field_energy'])
476 E_total = np.array(solver.diagnostics['total_energy'])
477
478 # Plot 1: Electric field vs time
479 ax1 = fig.add_subplot(gs[0, :])
480 ax1.semilogy(t * omega_pe / (2 * np.pi), E_max, 'b-', linewidth=2)
481 ax1.set_xlabel(r'Time ($\omega_{pe} t / 2\pi$)', fontsize=11)
482 ax1.set_ylabel('Max |E| (V/m)', fontsize=11)
483 ax1.set_title(f'{test_name}: Electric Field Evolution',
484 fontsize=12, fontweight='bold')
485 ax1.grid(True, alpha=0.3)
486
487 # Plot 2: Energy conservation
488 ax2 = fig.add_subplot(gs[1, 0])
489 E_total_norm = (E_total - E_total[0]) / E_total[0] * 100
490 ax2.plot(t * omega_pe / (2 * np.pi), E_total_norm, 'r-', linewidth=2)
491 ax2.set_xlabel(r'Time ($\omega_{pe} t / 2\pi$)', fontsize=11)
492 ax2.set_ylabel('Energy Error (%)', fontsize=11)
493 ax2.set_title('Energy Conservation', fontsize=11, fontweight='bold')
494 ax2.grid(True, alpha=0.3)
495
496 # Plot 3-5: Phase space snapshots
497 snapshot_indices = [0, len(snapshots)//2, -1]
498 axes = [fig.add_subplot(gs[1, 1]), fig.add_subplot(gs[1, 2]),
499 fig.add_subplot(gs[2, 0])]
500
501 for idx, ax in zip(snapshot_indices, axes):
502 t_snap, f_snap, E_snap = snapshots[idx]
503
504 im = ax.pcolormesh(solver.x * 1e3, solver.v / 1e6, f_snap.T,
505 shading='auto', cmap='viridis')
506 ax.set_xlabel('x (mm)', fontsize=10)
507 ax.set_ylabel('v (Mm/s)', fontsize=10)
508 ax.set_title(f't = {t_snap*omega_pe/(2*np.pi):.1f} / ωpe',
509 fontsize=10, fontweight='bold')
510 plt.colorbar(im, ax=ax, label='f(x,v)')
511
512 # Plot 6: Final density and E field
513 ax6 = fig.add_subplot(gs[2, 1:])
514 t_final, f_final, E_final = snapshots[-1]
515 ne_final, _ = solver.compute_moments()
516
517 ax6_twin = ax6.twinx()
518 ax6.plot(solver.x * 1e3, ne_final / 1e16, 'b-', linewidth=2, label='Density')
519 ax6_twin.plot(solver.x * 1e3, E_final, 'r-', linewidth=2, label='E field')
520
521 ax6.set_xlabel('x (mm)', fontsize=11)
522 ax6.set_ylabel(r'Density ($10^{16}$ m$^{-3}$)', fontsize=11, color='b')
523 ax6_twin.set_ylabel('E (V/m)', fontsize=11, color='r')
524 ax6.tick_params(axis='y', labelcolor='b')
525 ax6_twin.tick_params(axis='y', labelcolor='r')
526 ax6.set_title('Final State', fontsize=11, fontweight='bold')
527 ax6.grid(True, alpha=0.3)
528
529 plt.suptitle(f'1D Vlasov-Poisson: {test_name}',
530 fontsize=14, fontweight='bold')
531
532 filename = test_name.lower().replace(' ', '_') + '.png'
533 plt.savefig(filename, dpi=150, bbox_inches='tight')
534 print(f"Plot saved as '{filename}'")
535
536 plt.show()
537
538if __name__ == "__main__":
539 # Run test case 1: Plasma oscillation
540 solver1, snapshots1, omega_pe1 = test_plasma_oscillation()
541 plot_results(solver1, snapshots1, omega_pe1, "Plasma Oscillation")
542
543 # Run test case 2: Landau damping
544 solver2, snapshots2, omega_pe2 = test_landau_damping()
545 plot_results(solver2, snapshots2, omega_pe2, "Landau Damping")
546
547 print("\n" + "=" * 70)
548 print("All tests complete!")
549 print("=" * 70)