The simulation runs in a boosted frame, and the analysis is done in the lab
frame, i.e., on the back-transformed diagnostics.
'''
import numpy as np
import read_raw_data
import yt
yt.funcs.mylog.setLevel(0)
# Read data from back-transformed diagnostics
snapshot = './lab_frame_data/snapshots/snapshot00001'
header = './lab_frame_data/snapshots/Header'
allrd, info = read_raw_data.read_lab_snapshot(snapshot, header)
z = np.mean(read_raw_data.get_particle_field(snapshot, 'beam', 'z'))
w = np.std(read_raw_data.get_particle_field(snapshot, 'beam', 'x'))
# initial parameters
z0 = 20.e-6
w0 = 1.e-6
theta0 = np.arcsin(0.1)
# Theoretical beam width after propagation if rigid ON
wth = np.sqrt(w0**2 + (z - z0)**2 * theta0**2)
error_rel = np.abs((w - wth) / wth)
tolerance_rel = 0.03
# Print error and assert small error
print("Beam position: " + str(z))
print("Beam width : " + str(w))
number_list = range(nfiles)
comm_world = MPI.COMM_WORLD
me = comm_world.Get_rank()
nrank = comm_world.Get_size()
# List of files to process for current proc
my_list = file_list[(me * nfiles) / nrank:((me + 1) * nfiles) / nrank]
# List if file numbers for current proc
my_number_list = number_list[(me * nfiles) / nrank:((me + 1) * nfiles) / nrank]
# --- Run parallel analysis --- #
# ----------------------------- #
# Each MPI rank reads roughly (nb snapshots)/(nb ranks) snapshots.
# Works with any number of snapshots.
for count, filename in enumerate(my_list):
zwin, a0 = get_a0(res_dir, filename)
uzlab = read_raw_data.get_particle_field(filename, species, 'uz') / scc.c
select_particles = (uzlab > 5.)
uzlab = uzlab[select_particles]
uzmean = np.mean(uzlab)
# --- gather and rank 0 plots --- #
# ------------------------------- #
# Gather particle quantities to rank 0 to plot history of quantities.
UZMEAN = comm_world.gather(uzmean, root=0)
ZWIN = comm_world.gather(zwin, root=0)
A0 = comm_world.gather(a0, root=0)
# Rank 0 does the plot.
if me == 0:
# Convert to numpy arrays
UZMEAN, ZWIN, A0 = convert_to_np_array([UZMEAN, ZWIN, A0])
# Plot and save