def test_beam_focusing(show=False):
"""
Runs the simulation of a focusing charged beam, in a boosted-frame,
with and without the injection through a plane.
The value of the RMS radius at focus is automatically checked.
"""
# Simulate beam focusing with injection through plane or not
simulate_beam_focusing(None, 'direct')
simulate_beam_focusing(z_focus, 'through_plane')
# Analyze the results and show that the beam reaches
# the right RMS radius at focus
ts1 = OpenPMDTimeSeries('./direct/hdf5/')
r1 = get_rms_radius(ts1)
ts2 = OpenPMDTimeSeries('./through_plane/hdf5/')
r2 = get_rms_radius(ts2)
if show:
import matplotlib.pyplot as plt
plt.plot(1.e3 * c * ts1.t, 1.e6 * r1)
plt.plot(1.e3 * c * ts2.t, 1.e6 * r2)
plt.xlabel('z (mm)')
plt.ylabel('RMS radius (microns)')
plt.show()
# Find the index of the output at z_focus
i = np.argmin(abs(c * ts2.t - z_focus))
# With injection through plane, we get the right RMS value at focus
assert abs(r2[i] - sigma_r) < 0.05e-6
# Without injection through plane, the RMS value is significantly different
assert abs(r1[i] - sigma_r) > 0.5e-6
# Clean up the data folders
shutil.rmtree('direct')
shutil.rmtree('through_plane')
def __init__(self, path_to_dir, check_all_files=True, backend=None):
"""
Initialize an OpenPMD time series with various methods to diagnose the
data
Parameter
---------
path_to_dir : string
The path to the directory where the openPMD files are.
For the moment, only HDF5 files are supported. There should be
one file per iteration, and the name of the files should end
with the iteration number, followed by '.h5' (e.g. data0005000.h5)
check_all_files: bool, optional
Check that all the files in the timeseries are consistent
(i.e. that they contain the same fields and particles,
with the same metadata)
For fast access to the files, this can be changed to False.
backend: string
Backend to be used for data reading. Can be `openpmd-api`
or `h5py`. If not provided will use `openpmd-api` if available
and `h5py` otherwise.
"""
OpenPMDTimeSeries.__init__(self,
path_to_dir,
check_all_files=check_all_files,
backend=backend)
class Backend:
''' Use openPMD-viewer as the backend reader to read openPMD files
'''
def __init__(self, filename):
''' Constructor: store the dataset object
'''
self.dataset = OpenPMDTimeSeries(filename)
def fields_list(self):
''' Return the list of fields defined on the grid
'''
return self.dataset.avail_fields
def species_list(self):
''' Return the list of species in the dataset
'''
return self.dataset.avail_species
def n_levels(self):
''' Return the number of MR levels in the dataset
'''
return 1
def get_field_checksum(self, lev, field, test_name):
''' Calculate the checksum for a given field at a given level in the dataset
'''
Q = self.dataset.get_field(field=field,
iteration=self.dataset.iterations[-1])[0]
return np.sum(np.abs(Q))
def get_species_attributes(self, species):
''' Return the list of attributes for a given species in the dataset
'''
return self.dataset.avail_record_components[species]
def get_species_checksum(self, species, attribute):
''' Calculate the checksum for a given attribute of a given species in the dataset
'''
Q = self.dataset.get_particle(var_list=[attribute],
species=species,
iteration=self.dataset.iterations[-1])
# JSON complains with numpy integers, so if the quantity is a np.int64, convert to int
checksum = np.sum(np.abs(Q))
if type(checksum) in [np.int64, np.uint64]:
return int(checksum)
return checksum
def plot_snapshots(job: Job) -> None:
"""
Plot a snapshot of
a. the electric field Ex
b. the 1D weighted particle energy histogram
corresponding to ``iteration``.
:param job: the job instance is a handle to the data of a unique statepoint
"""
e0 = electric_field_amplitude_norm(lambda0=job.sp.lambda0)
h5_path: Union[bytes, str] = os.path.join(job.ws, "diags", "hdf5")
time_series: OpenPMDTimeSeries = OpenPMDTimeSeries(
h5_path, check_all_files=False)
# plot electric field and save to disk
field_snapshot(
tseries=time_series,
it=iteration,
field_name="E",
coord="x",
normalization_factor=1.0 / e0,
path=job.ws,
zlabel=r"$E_x/E_0$",
vmin=-4, # CHANGEME
vmax=4, # CHANGEME
hslice_val=
0.0, # do a 1D slice through the middle of the simulation box
)
def fbpic_ran(job: Job) -> bool:
"""
Check if ``fbpic`` produced all the output .h5 files.
:param job: the job instance is a handle to the data of a unique statepoint
:return: True if all output files are in {job_dir}/diags/hdf5, False otherwise
"""
h5_path: Union[bytes, str] = os.path.join(job.ws, "diags", "hdf5")
if not os.path.isdir(h5_path):
# {job_dir}/diags/hdf5 not present, ``fbpic`` didn't run
did_it_run = False
return did_it_run
time_series: OpenPMDTimeSeries = OpenPMDTimeSeries(h5_path,
check_all_files=True)
iterations: np.ndarray = time_series.iterations
# estimate iteration array based on input parameters
estimated_iterations = np.arange(0,
job.sp.N_step,
job.sp.diag_period,
dtype=np.int)
# check if iterations array corresponds to input params
did_it_run = np.array_equal(estimated_iterations, iterations)
return did_it_run
def test_boosted_frame_sim_twoproc():
"Test the example input script with two procs in `docs/source/example_input`"
temporary_dir = './tests/tmp_test_dir'
# Create a temporary directory for the simulation
# and copy the example script into this directory
if os.path.exists(temporary_dir):
shutil.rmtree(temporary_dir)
os.mkdir(temporary_dir)
shutil.copy('./docs/source/example_input/boosted_frame_script.py',
temporary_dir)
# Shortcut for the script file, which is repeatedly changed
script_filename = os.path.join(temporary_dir, 'boosted_frame_script.py')
# Read the script
with open(script_filename) as f:
script = f.read()
# Change default N_step
script = replace_string(script, 'N_step = .*', 'N_step = 101')
# Modify the script so as to enable finite order
script = replace_string(script, 'n_order = -1', 'n_order = 16')
script = replace_string(script, 'track_bunch = False',
'track_bunch = True')
with open(script_filename, 'w') as f:
f.write(script)
# Launch the script from the OS
command_line = 'cd %s; NUMBA_NUM_THREADS=1 MKL_NUM_THREADS=1 ' % temporary_dir
command_line += 'mpirun -np 2 python boosted_frame_script.py'
response = os.system(command_line)
assert response == 0
# Check that the particle ids are unique at each iterations
ts = OpenPMDTimeSeries(os.path.join(temporary_dir, 'lab_diags/hdf5'))
print('Checking particle ids...')
start_time = time.time()
for iteration in ts.iterations:
pid, = ts.get_particle(["id"], iteration=iteration)
assert len(np.unique(pid)) == len(pid)
end_time = time.time()
print("%.2f seconds" % (end_time - start_time))
# Suppress the temporary directory
shutil.rmtree(temporary_dir)
def check_theory_gaussian():
"""
Check that the transverse E and B field are close to the high-gamma
theory for a gaussian bunch
"""
ts = OpenPMDTimeSeries( os.path.join(temporary_dir, 'diags_serial/hdf5/') )
Ex, info = ts.get_field( 'E', 'x', iteration=0 )
By, info = ts.get_field( 'B', 'y', iteration=0 )
r, z = np.meshgrid( info.r, info.z, indexing='ij' )
# High-gamma theory for Gaussian bunch
Eth = -Q/(2*np.pi)**1.5/sig_z/epsilon_0/r * \
(1 - np.exp(-0.5*r**2/sig_r**2)) * \
np.exp( -0.5*(z-zf)**2/sig_z**2)
Bth = Eth/c
# Check that the fields agree
assert np.allclose( Ex, Eth, atol=0.1*Eth.max() )
assert np.allclose( By, Bth, atol=0.1*Bth.max() )
def get_a0(
tseries: OpenPMDTimeSeries,
t: Optional[float] = None,
it: Optional[int] = None,
coord="x",
m="all",
slicing_dir="y",
theta=0.0,
lambda0=0.8e-6,
) -> Tuple[float, float, float, float]:
"""
Compute z₀, a₀, w₀, cτ.
:param tseries: whole simulation time series
:param t: time (in seconds) at which to obtain the data
:param it: time step at which to obtain the data
:param coord: which component of the field to extract
:param m: 'all' for extracting the sum of all the modes
:param slicing_dir: the direction along which to slice the data eg., 'x', 'y' or 'z'
:param theta: the angle of the plane of observation, with respect to the 'x' axis
:param lambda0: laser wavelength (meters)
:return: z₀, a₀, w₀, cτ
"""
# get E_x field in V/m
electric_field_x, info_electric_field_x = tseries.get_field(
field="E",
coord=coord,
t=t,
iteration=it,
m=m,
theta=theta,
slicing_dir=slicing_dir,
)
# normalized vector potential
e0 = electric_field_amplitude_norm(lambda0=lambda0)
a0 = electric_field_x / e0
# get pulse envelope
envelope = np.abs(hilbert(a0, axis=1))
envelope_z = envelope[envelope.shape[0] // 2, :]
a0_max = np.amax(envelope_z)
# index of peak
z_idx = np.argmax(envelope_z)
# pulse peak position
z0 = info_electric_field_x.z[z_idx]
# FWHM perpendicular size of beam, proportional to w0
fwhm_a0_w0 = (np.sum(np.greater_equal(envelope[:, z_idx], a0_max / 2)) *
info_electric_field_x.dr)
# FWHM longitudinal size of the beam, proportional to ctau
fwhm_a0_ctau = (np.sum(np.greater_equal(envelope_z, a0_max / 2)) *
info_electric_field_x.dz)
return z0, a0_max, fwhm_a0_w0, fwhm_a0_ctau
def _opmd_time_series(data_file):
prev = None
try:
prev = main.list_h5_files
main.list_h5_files = lambda x: ([data_file.filename], [data_file.iteration])
return OpenPMDTimeSeries(py.path.local(data_file.filename).dirname)
finally:
if prev:
main.list_h5_files = prev
def field_snapshot(
tseries: OpenPMDTimeSeries,
it: int,
field_name: str,
normalization_factor=1,
coord: Optional[str] = None,
m="all",
theta=0.0,
chop: Optional[List[float]] = None,
path="./",
**kwargs,
) -> None:
"""
Plot the ``field_name`` field from ``tseries`` at step ``iter``.
:param path: path to output file
:param tseries: whole simulation time series
:param it: time step in the simulation
:param field_name: which field to extract, eg. 'rho', 'E', 'B' or 'J'
:param normalization_factor: normalization factor for the extracted field
:param coord: which component of the field to extract, eg. 'r', 't' or 'z'
:param m: 'all' for extracting the sum of all the azimuthal modes
:param theta: the angle of the plane of observation, with respect to the 'x' axis
:param chop: adjusting extent of simulation box plot
:param kwargs: extra plotting arguments, eg. labels, data limits etc.
:return: saves field plot image to disk
"""
if chop is None: # how much to cut out from simulation domain
chop = [40, -20, 15, -15] # CHANGEME
field, info = tseries.get_field(field=field_name,
coord=coord,
iteration=it,
m=m,
theta=theta)
field *= normalization_factor
plot = sliceplots.Plot2D(
arr2d=field,
h_axis=info.z * 1e6,
v_axis=info.r * 1e6,
xlabel=r"${} \;(\mu m)$".format(info.axes[1]),
ylabel=r"${} \;(\mu m)$".format(info.axes[0]),
extent=(
info.zmin * 1e6 + chop[0],
info.zmax * 1e6 + chop[1],
info.rmin * 1e6 + chop[2],
info.rmax * 1e6 + chop[3],
),
cbar=True,
text=f"iteration {it}",
**kwargs,
)
filename = os.path.join(path, f"{field_name}{it:06d}.png")
plot.canvas.print_figure(filename)
def __init__( self, path_to_dir, check_all_files=True ):
"""
Initialize an OpenPMD time series with various methods to diagnose the
data
Parameter
---------
path_to_dir : string
The path to the directory where the openPMD files are.
For the moment, only HDF5 files are supported. There should be
one file per iteration, and the name of the files should end
with the iteration number, followed by '.h5' (e.g. data0005000.h5)
check_all_files: bool, optional
Check that all the files in the timeseries are consistent
(i.e. that they contain the same fields and particles,
with the same metadata)
For fast access to the files, this can be changed to False.
"""
OpenPMDTimeSeries.__init__( self, path_to_dir,
check_all_files=check_all_files )
def check_identical_fields( folder1, folder2 ):
ts1 = OpenPMDTimeSeries( folder1 )
ts2 = OpenPMDTimeSeries( folder2 )
# Check the vector fields
for field, coord in [("J", "z"), ("E","r"), ("E","z"), ("B","t")]:
print("Checking %s%s" %(field, coord))
field1, info = ts1.get_field(field, coord, iteration=0)
field2, info = ts2.get_field(field, coord, iteration=0)
# For 0 fields, do not use allclose
if abs(field1).max() == 0:
assert abs(field2).max() == 0
else:
assert np.allclose(
field1/abs(field1).max(), field2/abs(field2).max() )
# Check the rho field
print("Checking rho")
field1, info = ts1.get_field("rho", iteration=0)
field2, info = ts2.get_field("rho", iteration=0)
assert np.allclose( field1/abs(field1).max(), field2/abs(field2).max() )
def check_theory_pml(show, boundaries):
"""
Check that the transverse E and B field are close to the high-gamma
theory for a gaussian bunch
"""
ts = OpenPMDTimeSeries(os.path.join(temporary_dir, 'diags/hdf5/'))
for iteration in ts.iterations:
compare_E(ts,
'x',
m=0,
iteration=iteration,
rtol=rtol0,
show=show,
boundaries=boundaries)
compare_E(ts,
'x',
m=1,
iteration=iteration,
rtol=rtol1,
show=show,
boundaries=boundaries)
def particle_energy_histogram(
tseries: OpenPMDTimeSeries,
it: int,
energy_min=1,
energy_max=800,
delta_energy=1,
cutoff=1, # CHANGEME
):
"""
Compute the weighted particle energy histogram from ``tseries`` at step ``iteration``.
:param tseries: whole simulation time series
:param it: time step in the simulation
:param energy_min: lower energy threshold (MeV)
:param energy_max: upper energy threshold (MeV)
:param delta_energy: size of each energy bin (MeV)
:param cutoff: upper threshold for the histogram, in pC / MeV
:return: histogram values and bin edges
"""
nbins = (energy_max - energy_min) // delta_energy
energy_bins = np.linspace(start=energy_min, stop=energy_max, num=nbins + 1)
ux, uy, uz, w = tseries.get_particle(["ux", "uy", "uz", "w"], iteration=it)
energy = mc2 * np.sqrt(1 + ux**2 + uy**2 + uz**2)
# Explanation of weights:
# 1. convert electron charge from C to pC (factor 1e12)
# 2. multiply by weight w to get real number of electrons
# 3. divide by energy bin size delta_energy to get charge / MeV
hist, _ = np.histogram(energy,
bins=energy_bins,
weights=q_e * 1e12 / delta_energy * w)
# cut off histogram
np.clip(hist, a_min=None, a_max=cutoff, out=hist)
return hist, energy_bins, nbins
# This script compares a field from a simulation with
# full IO and from a simulation with only slice IO
import matplotlib.pyplot as plt
import scipy.constants as scc
import matplotlib
import sys
import numpy as np
import math
import argparse
from openpmd_viewer import OpenPMDTimeSeries
do_plot = False
field = 'Ez'
ts1 = OpenPMDTimeSeries('full_io')
F_full = ts1.get_field(field=field, iteration=ts1.iterations[-1])[0]
F_full = np.swapaxes(F_full, 0, 2)
F_full_xz = (F_full[:, F_full.shape[1] // 2, :].squeeze() +
F_full[:, F_full.shape[1] // 2 - 1, :].squeeze()) / 2.
F_full_yz = (F_full[F_full.shape[0] // 2, :, :].squeeze() +
F_full[F_full.shape[0] // 2 - 1, :, :].squeeze()) / 2.
ts2 = OpenPMDTimeSeries('slice_io_xz')
F_slice_xz = ts2.get_field(field=field,
iteration=ts2.iterations[-1])[0].transpose()
ts3 = OpenPMDTimeSeries('slice_io_yz')
F_slice_yz = ts3.get_field(field=field,
iteration=ts3.iterations[-1])[0].transpose()
def restart_from_checkpoint(sim,
iteration=None,
checkpoint_dir='./checkpoints'):
"""
Fills the Simulation object `sim` with data saved in a checkpoint.
More precisely, the following data from `sim` is overwritten:
- Current time and iteration number of the simulation
- Position of the boundaries of the simulation box
- Values of the field arrays
- Size and values of the particle arrays
Any other information (e.g. diagnostics of the simulation, presence of a
moving window, presence of a laser antenna, etc.) need to be set by hand.
For this reason, a successful restart will often require to modify the
original input script that produced the checkpoint, rather than to start
a new input script from scratch.
NB: This function should always be called *before* the initialization
of the moving window, since the moving window infers the position of
particle injection from the existing particle data.
Parameters
----------
sim: a Simulation object
The Simulation object into which the checkpoint should be loaded
iteration: integer, optional
The iteration number of the checkpoint from which to restart
If None, the latest checkpoint available will be used.
checkpoint_dir: string, optional
The path to the directory that contains the checkpoints to be loaded.
(When running a simulation with several MPI ranks, use the
same path for all ranks.)
"""
# Try to import openPMD-viewer, version 1
try:
from openpmd_viewer import OpenPMDTimeSeries
openpmd_viewer_version = 1
except ImportError:
# If not available, try to import openPMD-viewer, version 0
try:
from opmd_viewer import OpenPMDTimeSeries
openpmd_viewer_version = 0
except ImportError:
openpmd_viewer_version = None
# Otherwise, raise an error
if openpmd_viewer_version is None:
raise ImportError(
'The package openPMD-viewer is required to restart from checkpoints.'
'\nPlease install it from https://github.com/openPMD/openPMD-viewer'
)
# Verify that the restart is valid (only for the first processor)
# (Use the global MPI communicator instead of the `BoundaryCommunicator`,
# so that this also works for `use_all_ranks=False`)
if comm.rank == 0:
check_restart(sim, iteration, checkpoint_dir)
comm.barrier()
# Choose the name of the directory from which to restart:
# one directory per processor
data_dir = os.path.join(checkpoint_dir, 'proc%d/hdf5' % comm.rank)
ts = OpenPMDTimeSeries(data_dir)
# Select the iteration, and its index
if iteration is None:
iteration = ts.iterations[-1]
# Find the index of the closest iteration
i_iteration = np.argmin(abs(np.array(ts.iterations) - iteration))
# Modify parameters of the simulation
sim.iteration = iteration
sim.time = ts.t[i_iteration]
# Export available species as a list
avail_species = ts.avail_species
if avail_species is None:
avail_species = []
# Load the particles
# Loop through the different species
if len(avail_species) == len(sim.ptcl):
for i in range(len(sim.ptcl)):
name = 'species %d' % i
load_species(sim.ptcl[i], name, ts, iteration, sim.comm,
openpmd_viewer_version)
else:
raise RuntimeError( \
"""Species numbers in checkpoint and simulation should be same, but
got {:d} and {:d}. Use add_new_species method to add species to
simulation or sim.ptcl = [] to remove them""".format(len(avail_species),
len(sim.ptcl)) )
# Record position of grid before restart
zmin_old = sim.fld.interp[0].zmin
# Load the fields
# Loop through the different modes
for m in range(sim.fld.Nm):
# Load the fields E and B
for fieldtype in ['E', 'B']:
for coord in ['r', 't', 'z']:
load_fields(sim.fld.interp[m], fieldtype, coord, ts, iteration)
# Load the PML components if needed
if sim.use_pml:
for fieldtype in ['Er_pml', 'Et_pml', 'Br_pml', 'Bt_pml']:
load_fields(sim.fld.interp[m], fieldtype, None, ts, iteration)
# Record position after restart (`zmin` is modified by `load_fields`)
# and shift the global domain position in the BoundaryCommunicator
zmin_new = sim.fld.interp[0].zmin
sim.comm.shift_global_domain_positions(zmin_new - zmin_old)
def run_cpu_gpu_deposition(show=False, particle_shape='cubic'):
# Skip this test if cuda is not installed
if not cuda_installed:
return
# Perform deposition for a few timesteps, with both the CPU and GPU
for hardware in ['cpu', 'gpu']:
if hardware == 'cpu':
use_cuda = False
elif hardware == 'gpu':
use_cuda = True
# Initialize the simulation object
sim = Simulation(Nz,
zmax,
Nr,
rmax,
Nm,
dt,
zmin=zmin,
use_cuda=use_cuda,
particle_shape=particle_shape)
sim.ptcl = []
# Add an electron bunch (set the random seed first)
np.random.seed(0)
add_elec_bunch_gaussian(sim, sig_r, sig_z, n_emit, gamma0, sig_gamma,
Q, N)
# Add a field diagnostic
sim.diags = [
FieldDiagnostic(diag_period,
sim.fld,
fieldtypes=['rho', 'J'],
comm=sim.comm,
write_dir=os.path.join('tests', hardware))
]
### Run the simulation
sim.step(N_step)
# Check that the results are identical
ts_cpu = OpenPMDTimeSeries('tests/cpu/hdf5')
ts_gpu = OpenPMDTimeSeries('tests/gpu/hdf5')
for iteration in ts_cpu.iterations:
for field, coord in [('rho', ''), ('J', 'x'), ('J', 'z')]:
# Jy is not tested because it is zero
print('Testing %s at iteration %d' % (field + coord, iteration))
F_cpu, info = ts_cpu.get_field(field, coord, iteration=iteration)
F_gpu, info = ts_gpu.get_field(field, coord, iteration=iteration)
tolerance = 1.e-13 * (abs(F_cpu).max() + abs(F_gpu).max())
if not show:
assert np.allclose(F_cpu, F_gpu, atol=tolerance)
else:
if not np.allclose(F_cpu, F_gpu, atol=tolerance):
plot_difference(field, coord, iteration, F_cpu, F_gpu,
info)
# Remove the files used
shutil.rmtree('tests/cpu')
shutil.rmtree('tests/gpu')
dest='norm_units',
action='store_true',
default=False,
help='Run the analysis in normalized units')
parser.add_argument('--do-plot',
dest='do_plot',
action='store_true',
default=False,
help='Plot figures and save them to file')
parser.add_argument('--output-dir',
dest='output_dir',
default='diags/hdf5',
help='Path to the directory containing output files')
args = parser.parse_args()
ts = OpenPMDTimeSeries(args.output_dir)
if args.norm_units:
c = 1.
jz0 = -1.
rho0 = -1.
mu_0 = 1.
eps_0 = 1.
R = 1.
else:
# Density of the can beam
dens = 2.8239587008591567e23 # at this density, 1/kp = 10um, allowing for an easy comparison with normalized units
# Define array for transverse coordinate and theory for By and Bx
jz0 = - scc.e * scc.c * dens
rho0 = - scc.e * dens
c = scc.c
dest='norm_units',
action='store_true',
default=False,
help='Run the analysis in normalized units')
parser.add_argument('--do-plot',
dest='do_plot',
action='store_true',
default=False,
help='Plot figures and save them to file')
parser.add_argument('--output-dir',
dest='output_dir',
default='diags/hdf5',
help='Path to the directory containing output files')
args = parser.parse_args()
ts = OpenPMDTimeSeries(args.output_dir)
if args.norm_units:
c = 1.
jz0 = -1.
rho0 = -1.
mu_0 = 1.
eps_0 = 1.
R = 1.
else:
# Density of the can beam
dens = 2.8239587008591567e23 # at this density, 1/kp = 10um, allowing for an easy comparison with normalized units
# Define array for transverse coordinate and theory for By and Bx
jz0 = -scc.e * scc.c * dens
rho0 = -scc.e * dens
c = scc.c
#
# This file is part of HiPACE++.
#
# Authors: MaxThevenet, Severin Diederichs
# License: BSD-3-Clause-LBNL
# This script calculates the sum of jz.
# The beam current and the grid current should cancel each other.
import argparse
import numpy as np
from openpmd_viewer import OpenPMDTimeSeries
parser = argparse.ArgumentParser(description='Script to analyze the correctness of the beam in vacuum')
parser.add_argument('--output-dir',
dest='output_dir',
default='diags/hdf5',
help='Path to the directory containing output files')
args = parser.parse_args()
ts = OpenPMDTimeSeries(args.output_dir)
# Load Hipace data for jz
jz_sim, jz_info = ts.get_field(field='jz', iteration=1)
# Assert that the grid current and the beam current cancel each other
error_jz = np.sum( (jz_sim)**2)
print("sum of jz**2: " + str(error_jz) + " (tolerance = 3e-3)")
assert(error_jz < 3e-3)
dest='si_data',
required=True,
help='Path to the data of the SI units run')
parser.add_argument(
'--si-fixed-weight-data',
dest='si_fixed_weight_data',
required=True,
help='Path to the data of the SI units run with a fixed weight beam')
parser.add_argument('--do-plot',
dest='do_plot',
action='store_true',
default=False,
help='Plot figures and save them to file')
args = parser.parse_args()
ts_norm = OpenPMDTimeSeries(args.norm_data)
ts_si = OpenPMDTimeSeries(args.si_data)
ts_si_fixed_weight = OpenPMDTimeSeries(args.si_fixed_weight_data)
elec_density = 2.8239587008591567e23 # [1/m^3]
# calculation of the plasma frequency
omega_p = np.sqrt(elec_density * (scc.e**2) / (scc.epsilon_0 * scc.m_e))
E_0 = omega_p * scc.m_e * scc.c / scc.e
kp = omega_p / scc.c # 1./10.e-6
# Load HiPACE++ data for Ez in both normalized and SI units
Ez_along_z_norm, meta_norm = ts_norm.get_field(field='Ez',
iteration=1,
slice_across=['x', 'y'],
slice_relative_position=[0, 0])
fn = sys.argv[1]
ds = yt.load(fn)
ad = ds.all_data()
x = ad['electron', 'particle_position_x'].v
error = len(x) - 512
print('error = ', error)
print('tolerance = ', tolerance)
assert (error == tolerance)
# Check that all the removed particles are properly recorded
# by making sure that, at each iteration, the sum of the number of
# remaining particles and scraped particles is equal to the
# original number of particles
ts_full = OpenPMDTimeSeries('./diags/diag2/')
ts_scraping = OpenPMDTimeSeries('./diags/diag3/')
def n_remaining_particles(iteration):
w, = ts_full.get_particle(['w'], iteration=iteration)
return len(w)
def n_scraped_particles(iteration):
timestamp = ts_scraping.get_particle(['timestamp'])
return (timestamp <= iteration).sum()
n_remaining = np.array(
[n_remaining_particles(iteration) for iteration in ts_full.iterations])
# License: BSD-3-Clause-LBNL
# This script compares a field from a simulation with full IO and from a simulation with coarse IO
import matplotlib.pyplot as plt
import scipy.constants as scc
import matplotlib
import sys
import numpy as np
import math
import argparse
from openpmd_viewer import OpenPMDTimeSeries
fields = ['Ez', 'ExmBy', 'EypBx']
ts1 = OpenPMDTimeSeries('fine_io')
ts2 = OpenPMDTimeSeries('coarse_io')
for field in fields:
F_full = ts1.get_field(field=field, iteration=ts1.iterations[-1])[0]
F_full_coarse = (F_full[2::5,1::4,1::3] + F_full[2::5,2::4,1::3])/2
F_coarse = ts2.get_field(field=field, iteration=ts2.iterations[-1])[0]
print("F_full.shape =", F_full.shape)
print("F_full_coarse.shape =", F_full_coarse.shape)
print("F_coarse.shape =", F_coarse.shape)
error = np.max(np.abs(F_coarse-F_full_coarse)) / np.max(np.abs(F_full_coarse))
print("error =", error)
parser = argparse.ArgumentParser(description='Script to analyze the correctness of the beam in vacuum')
parser.add_argument('--output-dir',
dest='output_dir',
default='diags/hdf5',
help='Path to the directory containing output files')
args = parser.parse_args()
# Numerical parameters of the simulation
field_strength = 0.5
gamma = 1000.
x_std_initial = 1./2.
omega_beta = np.sqrt(field_strength/gamma)
# Load beam particle data
ts = OpenPMDTimeSeries(args.output_dir)
xp, yp, uzp, wp = ts.get_particle(species='beam', iteration=ts.iterations[-1],
var_list=['x', 'y', 'uz', 'w'])
std_theory = x_std_initial * np.abs(np.cos(omega_beta * ts.current_t))
std_sim_x = np.sqrt(np.sum(xp**2*wp)/np.sum(wp))
std_sim_y = np.sqrt(np.sum(yp**2*wp)/np.sum(wp))
if do_plot:
plt.figure()
plt.plot(xp, yp, '.')
plt.xlim(-2, 2)
plt.ylim(-2, 2)
plt.xlabel('x')
plt.ylabel('y')
plt.savefig('image.pdf', bbox_inches='tight')
def __init__(self, filename):
''' Constructor: store the dataset object
'''
self.dataset = OpenPMDTimeSeries(filename)
import matplotlib
import sys
import numpy as np
import math
import argparse
from openpmd_viewer import OpenPMDTimeSeries
parser = argparse.ArgumentParser(
description='Script to analyze the equality of two simulations')
parser.add_argument('--first',
dest='first',
required=True)
parser.add_argument('--second',
dest='second',
required=True)
args = parser.parse_args()
# Replace the string below, to point to your data
tss = OpenPMDTimeSeries(args.first)
tsp = OpenPMDTimeSeries(args.second)
iteration = 0
for field in ['Bx', 'By', 'Ez', 'ExmBy', 'EypBx']:
Fs, ms = tss.get_field(iteration=iteration, field=field)
Fp, mp = tsp.get_field(iteration=iteration, field=field)
error = np.sum((Fp-Fs)**2) / np.sum(Fs**2)
print(field, 'error = np.sum((Fp-Fs)**2) / np.sum(Fs**2) = ' +str(error))
assert(error<0.006)
dest='do_plot',
action='store_true',
default=False,
help='Plot figures and save them to file')
parser.add_argument('--gaussian-beam',
dest='gaussian_beam',
action='store_true',
default=False,
help='Run the analysis on the Gaussian beam')
parser.add_argument('--output-dir',
dest='output_dir',
default='diags/hdf5',
help='Path to the directory containing output files')
args = parser.parse_args()
ts = OpenPMDTimeSeries(args.output_dir)
if args.norm_units:
kp = 1.
ne = 1.
q_e = 1.
else:
kp = 1./10.e-6
ne = scc.c**2 / scc.e**2 * scc.m_e * scc.epsilon_0 * kp**2
q_e = scc.e
rho_along_z, rho_meta = ts.get_field(field='rho', iteration=ts.iterations[-1],
slice_across=['x','y'], slice_relative_position=[0,0])
zeta_array = rho_meta.z
dzeta = rho_meta.dz
nz = len(rho_meta.z)
def add_particle_bunch_openPMD(sim,
q,
m,
ts_path,
z_off=0.,
species=None,
select=None,
iteration=None,
boost=None,
z_injection_plane=None,
initialize_self_field=True):
"""
Introduce a relativistic particle bunch in the simulation,
along with its space charge field, loading particles from an openPMD
timeseries.
Parameters
----------
sim : a Simulation object
The structure that contains the simulation.
q : float (in Coulomb)
Charge of the particle species
m : float (in kg)
Mass of the particle species
ts_path : string
The path to the directory where the openPMD files are.
For the moment, only HDF5 files are supported. There should be
one file per iteration, and the name of the files should end
with the iteration number, followed by '.h5' (e.g. data0005000.h5)
z_off: float (in meters)
Shift the particle positions in z by z_off. By default the initialized
phasespace is centered at z=0.
species: string
A string indicating the name of the species
This is optional if there is only one species
select: dict, optional
Either None or a dictionary of rules
to select the particles, of the form
'x' : [-4., 10.] (Particles having x between -4 and 10 microns)
'ux' : [-0.1, 0.1] (Particles having ux between -0.1 and 0.1 mc)
'uz' : [5., None] (Particles with uz above 5 mc)
iteration: integer (optional)
The iteration number of the openPMD file from which to extract the
particles.
boost : a BoostConverter object, optional
A BoostConverter object defining the Lorentz boost of
the simulation.
z_injection_plane: float (in meters) or None
When `z_injection_plane` is not None, then particles have a ballistic
motion for z<z_injection_plane. This is sometimes useful in
boosted-frame simulations.
`z_injection_plane` is always given in the lab frame.
initialize_self_field: bool, optional
Whether to calculate the initial space charge fields of the bunch
and add these fields to the fields on the grid (Default: True)
"""
# Try to import openPMD-viewer, version 1
try:
from openpmd_viewer import OpenPMDTimeSeries
openpmd_viewer_version = 1
except ImportError:
# If not available, try to import openPMD-viewer, version 0
try:
from opmd_viewer import OpenPMDTimeSeries
openpmd_viewer_version = 0
except ImportError:
openpmd_viewer_version = None
# Otherwise, raise an error
if openpmd_viewer_version is None:
raise ImportError(
'The package openPMD-viewer is required to load a particle bunch from on openPMD file.'
'\nPlease install it from https://github.com/openPMD/openPMD-viewer'
)
ts = OpenPMDTimeSeries(ts_path)
# Extract phasespace and particle weights
x, y, z, ux, uy, uz, w = ts.get_particle(
['x', 'y', 'z', 'ux', 'uy', 'uz', 'w'],
iteration=iteration,
species=species,
select=select)
if openpmd_viewer_version == 0:
# Convert the positions from microns to meters
x *= 1.e-6
y *= 1.e-6
z *= 1.e-6
# Shift the center of the phasespace to z_off
z = z - np.average(z, weights=w) + z_off
# Add the electrons to the simulation, and calculate the space charge
ptcl_bunch = add_particle_bunch_from_arrays(
sim,
q,
m,
x,
y,
z,
ux,
uy,
uz,
w,
boost=boost,
z_injection_plane=z_injection_plane,
initialize_self_field=initialize_self_field)
return ptcl_bunch
def run_sim(script_name, n_MPI, checked_fields, test_checkpoint_dir=False):
"""
Runs the script `script_name` from the folder docs/source/example_input,
with `n_MPI` MPI processes. The simulation is then restarted with
the same number of processes ; the code checks that the restarted results
are identical.
More precisely:
- The first simulation is run for N_step, then the random seed is reset
(for reproducibility) and the code runs for N_step more steps.
- Then a second simulation is launched, which reruns the last N_step.
"""
temporary_dir = './tests/tmp_test_dir'
# Create a temporary directory for the simulation
# and copy the example script into this directory
if os.path.exists(temporary_dir):
shutil.rmtree(temporary_dir)
os.mkdir(temporary_dir)
shutil.copy('./docs/source/example_input/%s' % script_name, temporary_dir)
# Shortcut for the script file, which is repeatedly changed
script_filename = os.path.join(temporary_dir, script_name)
# Read the script and check
with open(script_filename) as f:
script = f.read()
# Change default N_step, diag_period and checkpoint_period
script = replace_string(script, 'N_step = .*', 'N_step = 200')
script = replace_string(script, 'diag_period = 50', 'diag_period = 10')
script = replace_string(script, 'checkpoint_period = 100',
'checkpoint_period = 50')
# For MPI simulations: modify the script to use finite-order
if n_MPI > 1:
script = replace_string(script, 'n_order = -1', 'n_order = 16')
# Modify the script so as to enable checkpoints
script = replace_string(script, 'save_checkpoints = False',
'save_checkpoints = True')
if test_checkpoint_dir:
# Try to change the name of the checkpoint directory
checkpoint_dir = './test_chkpt'
script = replace_string(
script, 'set_periodic_checkpoint\( sim, checkpoint_period \)',
'set_periodic_checkpoint( sim, checkpoint_period, checkpoint_dir="%s" )'
% checkpoint_dir)
script = replace_string(
script, 'restart_from_checkpoint\( sim \)',
'restart_from_checkpoint( sim, checkpoint_dir="%s" )' %
checkpoint_dir)
else:
checkpoint_dir = './checkpoints'
script = replace_string(script, 'track_electrons = False',
'track_electrons = True')
# Modify the script to perform N_step, enforce the random seed
# (should be the same when restarting, for exact comparison),
# and perform again N_step.
script = replace_string(
script, 'sim.step\( N_step \)',
'sim.step( N_step ); np.random.seed(0); sim.step( N_step )')
with open(script_filename, 'w') as f:
f.write(script)
# Launch the script from the OS
command_line = 'cd %s' % temporary_dir
if n_MPI == 1:
command_line += '; python %s' % script_name
else:
# Use only one thread for multiple MPI
command_line += '; NUMBA_NUM_THREADS=1 MKL_NUM_THREADS=1 '
command_line += 'mpirun -np %d python %s' % (n_MPI, script_name)
response = os.system(command_line)
assert response == 0
# Move diagnostics (for later comparison with the restarted simulation)
shutil.move(os.path.join(temporary_dir, 'diags'),
os.path.join(temporary_dir, 'original_diags'))
# Keep only the checkpoints from the first N_step
N_step = int(get_string('N_step = (\d+)', script))
period = int(get_string('checkpoint_period = (\d+)', script))
for i_MPI in range(n_MPI):
for step in range(N_step + period, 2 * N_step + period, period):
os.remove(
os.path.join(
temporary_dir, '%s/proc%d/hdf5/data%08d.h5' %
(checkpoint_dir, i_MPI, step)))
# Modify the script so as to enable restarts
script = replace_string(script, 'use_restart = False',
'use_restart = True')
# Redo only the last N_step
script = replace_string(
script,
'sim.step\( N_step \); np.random.seed\(0\); sim.step\( N_step \)',
'np.random.seed(0); sim.step( N_step )',
)
with open(script_filename, 'w') as f:
f.write(script)
# Launch the modified script from the OS, with 2 proc
response = os.system(command_line)
assert response == 0
# Check that restarted simulation gives the same results
# as the original simulation
print('Checking restarted simulation...')
start_time = time.time()
ts1 = OpenPMDTimeSeries(os.path.join(temporary_dir, 'diags/hdf5'))
ts2 = OpenPMDTimeSeries(os.path.join(temporary_dir, 'original_diags/hdf5'))
compare_simulations(ts1, ts2, checked_fields)
end_time = time.time()
print("%.2f seconds" % (end_time - start_time))
# Check that the particle IDs are unique
print('Checking particle ids...')
start_time = time.time()
for iteration in ts1.iterations:
pid, = ts1.get_particle(["id"],
iteration=iteration,
species="electrons")
assert len(np.unique(pid)) == len(pid)
end_time = time.time()
print("%.2f seconds" % (end_time - start_time))
# Suppress the temporary directory
shutil.rmtree(temporary_dir)
def test_boosted_output(gamma_boost=10.):
"""
# TODO
Parameters
----------
gamma_boost: float
The Lorentz factor of the frame in which the simulation is carried out.
"""
# The simulation box
Nz = 500 # Number of gridpoints along z
zmax_lab = 0.e-6 # Length of the box along z (meters)
zmin_lab = -20.e-6
Nr = 10 # Number of gridpoints along r
rmax = 10.e-6 # Length of the box along r (meters)
Nm = 2 # Number of modes used
# Number of timesteps
N_steps = 500
diag_period = 20 # Period of the diagnostics in number of timesteps
dt_lab = (zmax_lab - zmin_lab) / Nz * 1. / c
T_sim_lab = N_steps * dt_lab
# Move into directory `tests`
os.chdir('./tests')
# Initialize the simulation object
sim = Simulation(
Nz,
zmax_lab,
Nr,
rmax,
Nm,
dt_lab,
0,
0, # No electrons get created because we pass p_zmin=p_zmax=0
0,
rmax,
1,
1,
4,
n_e=0,
zmin=zmin_lab,
initialize_ions=False,
gamma_boost=gamma_boost,
v_comoving=-0.9999 * c,
boundaries='open',
use_cuda=use_cuda)
sim.set_moving_window(v=c)
# Remove the electron species
sim.ptcl = []
# Add a Gaussian electron bunch
# Note: the total charge is 0 so all fields should remain 0
# throughout the simulation. As a consequence, the motion of the beam
# is a mere translation.
N_particles = 3000
add_elec_bunch_gaussian(sim,
sig_r=1.e-6,
sig_z=1.e-6,
n_emit=0.,
gamma0=100,
sig_gamma=0.,
Q=0.,
N=N_particles,
zf=0.5 * (zmax_lab + zmin_lab),
boost=BoostConverter(gamma_boost))
sim.ptcl[0].track(sim.comm)
# openPMD diagnostics
sim.diags = [
BackTransformedParticleDiagnostic(zmin_lab,
zmax_lab,
v_lab=c,
dt_snapshots_lab=T_sim_lab / 3.,
Ntot_snapshots_lab=3,
gamma_boost=gamma_boost,
period=diag_period,
fldobject=sim.fld,
species={"bunch": sim.ptcl[0]},
comm=sim.comm)
]
# Run the simulation
sim.step(N_steps)
# Check consistency of the back-transformed openPMD diagnostics:
# Make sure that all the particles were retrived by checking particle IDs
ts = OpenPMDTimeSeries('./lab_diags/hdf5/')
ref_pid = np.sort(sim.ptcl[0].tracker.id)
for iteration in ts.iterations:
pid, = ts.get_particle(['id'], iteration=iteration)
pid = np.sort(pid)
assert len(pid) == N_particles
assert np.all(ref_pid == pid)
# Remove openPMD files
shutil.rmtree('./lab_diags/')
os.chdir('../')
