-
Notifications
You must be signed in to change notification settings - Fork 200
/
Copy pathvisualize_densities_and_fields_2d.py
173 lines (135 loc) · 5.62 KB
/
visualize_densities_and_fields_2d.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
#!/usr/bin/env python3
# Copyright 2023 The WarpX Community
#
# This file is part of WarpX.
#
# Authors: Marco Garten
# License: BSD-3-Clause-LBNL
#
# This script plots the densities and fields of a 2D laser-ion acceleration simulation.
import argparse
import os
from matplotlib.colors import TwoSlopeNorm
import matplotlib.pyplot as plt
import numpy as np
from openpmd_viewer import OpenPMDTimeSeries
import scipy.constants as sc
def create_analysis_dir(directory):
if not os.path.exists(directory):
os.makedirs(directory)
def visualize_density_iteration(ts, iteration, out_dir):
"""
Visualize densities and fields of a single iteration.
:param ts: OpenPMDTimeSeries
:param iteration: Output iteration (simulation timestep)
:param out_dir: Directory for PNG output
:return:
"""
# Physics parameters
lambda_L = 800e-9 # Laser wavelength in meters
omega_L = 2 * np.pi * sc.c / lambda_L # Laser frequency in seconds
n_c = sc.m_e * sc.epsilon_0 * omega_L**2 / sc.elementary_charge**2 # Critical plasma density in meters^(-3)
# Simulation parameters
n_e0 = 30
n_max = 2 * n_e0
nr = 1 # Number to decrease resolution
# Data fetching
it = iteration
ii = np.where(ts.iterations == it)[0][0]
time = ts.t[ii]
rho_e, rho_e_info = ts.get_field(field="rho_electrons", iteration=it)
rho_d, rho_d_info = ts.get_field(field="rho_hydrogen", iteration=it)
# Rescale to critical density
rho_e = rho_e / (sc.elementary_charge * n_c)
rho_d = rho_d / (sc.elementary_charge * n_c)
# Axes setup
fig, axs = plt.subplots(3, 1, figsize=(5, 8))
xax, zax = rho_e_info.x, rho_e_info.z
# Plotting
# Electron density
im0 = axs[0].pcolormesh(zax[::nr]/1e-6, xax[::nr]/1e-6, -rho_e.T[::nr, ::nr],
vmin=0, vmax=n_max, cmap="Reds", rasterized=True)
plt.colorbar(im0, ax=axs[0], label=r"$n_\mathrm{\,e}\ (n_\mathrm{c})$")
# Hydrogen density
im1 = axs[1].pcolormesh(zax[::nr]/1e-6, xax[::nr]/1e-6, rho_d.T[::nr, ::nr],
vmin=0, vmax=n_max, cmap="Blues", rasterized=True)
plt.colorbar(im1, ax=axs[1], label=r"$n_\mathrm{\,H}\ (n_\mathrm{c})$")
# Masked electron density
divnorm = TwoSlopeNorm(vmin=-7., vcenter=0., vmax=2)
masked_data = np.ma.masked_where(rho_e.T == 0, rho_e.T)
my_cmap = plt.cm.PiYG_r.copy()
my_cmap.set_bad(color='black')
im2 = axs[2].pcolormesh(zax[::nr]/1e-6, xax[::nr]/1e-6, np.log(-masked_data[::nr, ::nr]),
norm=divnorm, cmap=my_cmap, rasterized=True)
plt.colorbar(im2, ax=axs[2], ticks=[-6, -3, 0, 1, 2], extend='both',
label=r"$\log n_\mathrm{\,e}\ (n_\mathrm{c})$")
# Axis labels and title
for ax in axs:
ax.set_aspect(1.0)
ax.set_ylabel(r"$x$ ($\mu$m)")
for ax in axs[:-1]:
ax.set_xticklabels([])
axs[2].set_xlabel(r"$z$ ($\mu$m)")
fig.suptitle(f"Iteration: {it}, Time: {time/1e-12:.3f} ps")
plt.tight_layout()
plt.savefig(f"{out_dir}/densities_{it:06d}.png")
def visualize_field_iteration(ts, iteration, out_dir):
# Additional parameters
nr = 1 # Number to decrease resolution
# Data fetching
it = iteration
ii = np.where(ts.iterations == it)[0][0]
time = ts.t[ii]
Ex, Ex_info = ts.get_field(field="E", coord="x", iteration=it)
Exmax = np.max(np.abs([np.min(Ex),np.max(Ex)]))
By, By_info = ts.get_field(field="B", coord="y", iteration=it)
Bymax = np.max(np.abs([np.min(By),np.max(By)]))
Ez, Ez_info = ts.get_field(field="E", coord="z", iteration=it)
Ezmax = np.max(np.abs([np.min(Ez),np.max(Ez)]))
# Axes setup
fig,axs = plt.subplots(3, 1, figsize=(5, 8))
xax, zax = Ex_info.x, Ex_info.z
# Plotting
im0 = axs[0].pcolormesh(
zax[::nr]/1e-6,xax[::nr]/1e-6,Ex.T[::nr,::nr],
vmin=-Exmax, vmax=Exmax,
cmap="RdBu", rasterized=True)
plt.colorbar(im0,ax=axs[00], label=r"$E_x$ (V/m)")
im1 = axs[1].pcolormesh(
zax[::nr]/1e-6,xax[::nr]/1e-6,By.T[::nr,::nr],
vmin=-Bymax, vmax=Bymax,
cmap="RdBu", rasterized=True)
plt.colorbar(im1,ax=axs[1], label=r"$B_y$ (T)")
im2 = axs[2].pcolormesh(
zax[::nr]/1e-6,xax[::nr]/1e-6,Ez.T[::nr,::nr],
vmin=-Ezmax, vmax=Ezmax,
cmap="RdBu", rasterized=True)
plt.colorbar(im2,ax=axs[2],label=r"$E_z$ (V/m)")
# Axis labels and title
for ax in axs:
ax.set_aspect(1.0)
ax.set_ylabel(r"$x$ ($\mu$m)")
for ax in axs[:-1]:
ax.set_xticklabels([])
axs[2].set_xlabel(r"$z$ ($\mu$m)")
fig.suptitle(f"Iteration: {it}, Time: {time/1e-12:.3f} ps")
plt.tight_layout()
plt.savefig(f"{out_dir}/fields_{it:06d}.png")
if __name__ == "__main__":
# Argument parsing
parser = argparse.ArgumentParser(description='Visualize Laser-Ion Accelerator Densities and Fields')
parser.add_argument('-d', '--diag_dir', type=str, default='./diags/diag1', help='Directory containing density and field diagnostics')
parser.add_argument('-i', '--iteration', type=int, default=0, help='Specific iteration to visualize')
args = parser.parse_args()
# Create analysis directory
analysis_dir = 'analysis'
create_analysis_dir(analysis_dir)
# Loading the time series
ts = OpenPMDTimeSeries(args.diag_dir)
if args.iteration is not None:
visualize_density_iteration(ts, args.iteration, analysis_dir)
visualize_field_iteration(ts, args.iteration, analysis_dir)
else:
for it in ts.iterations:
visualize_density_iteration(ts, it, analysis_dir)
visualize_field_iteration(ts, args.iteration, analysis_dir)