def test_lpa_sim_twoproc_restart():
"Test the checkpoint/restart mechanism with two proc"
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/lwfa_script.py', temporary_dir )
# Enter the temporary directory
os.chdir( temporary_dir )
# Read the script and check that the targeted lines are present
with open('lwfa_script.py') as f:
script = f.read()
# Check that the targeted lines are present
if script.find('save_checkpoints = False') == -1 \
or script.find('use_restart = False') == -1 \
or script.find('track_electrons = False') == -1 \
or script.find('N_step = 200') == -1:
raise RuntimeError('Did not find expected lines in lwfa_script.py')
# Modify the script so as to enable checkpoints and particle tracking
script = script.replace('save_checkpoints = False',
'save_checkpoints = True')
script = script.replace('track_electrons = False',
'track_electrons = True')
script = script.replace('N_step = 200', 'N_step = 101')
with open('lwfa_script.py', 'w') as f:
f.write(script)
# Launch the modified script from the OS, with 2 proc
response = os.system( 'mpirun -np 2 python lwfa_script.py' )
assert response==0
# Modify the script so as to enable restarts
script = script.replace('use_restart = False',
'use_restart = True')
script = script.replace('save_checkpoints = True',
'save_checkpoints = False')
with open('lwfa_script.py', 'w') as f:
f.write(script)
# Launch the modified script from the OS, with 2 proc
response = os.system( 'mpirun -np 2 python lwfa_script.py' )
assert response==0
# Check that the particle ids are unique at each iterations
ts = OpenPMDTimeSeries('./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))
# Exit the temporary directory and suppress it
os.chdir('../../')
shutil.rmtree( temporary_dir )
class OpenPMDRawDataReader(RawDataReaderBase):
def __init__(self, location, speciesName, dataName, internalName,
firstTimeStep):
# Store an openPMD timeseries object
# (Its API is used in order to conveniently extract data from the file)
self.openpmd_ts = OpenPMDTimeSeries(location, check_all_files=False)
# Initialize the instance
RawDataReaderBase.__init__(self, location, speciesName, dataName,
internalName, firstTimeStep)
def _ReadData(self, timeStep):
data, = self.openpmd_ts.get_particle([self.internalName],
species=self.speciesName,
iteration=timeStep)
self.currentTime = self._ReadTime(timeStep)
return data
def _ReadTime(self, timeStep):
# The line below sets the attribute `_current_i` of openpmd_ts
self.openpmd_ts._find_output(None, timeStep)
# This sets the corresponding time
self.currentTime = self.openpmd_ts.t[self.openpmd_ts._current_i]
def _ReadUnits(self):
# OpenPMD data always provide conversion to SI units
self.dataUnits = "" # TODO:
self.timeUnits = "s"
def _OpenFile(self, timeStep):
# The line below sets the attribute `_current_i` of openpmd_ts
self.openpmd_ts._find_output(None, timeStep)
# This finds the full path to the corresponding file
fileName = self.openpmd_ts.h5_files[self.openpmd_ts._current_i]
file_content = H5File(fileName, 'r')
return file_content
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 = int(T_interact/sim.dt)',
'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)
class OpenPMDRawDataReader(RawDataReaderBase):
def __init__(self, location, speciesName, dataName, internalName, firstTimeStep):
# First check whether openPMD is installed
if not openpmd_installed:
raise RunTimeError("You need to install openPMD-viewer, e.g. with:\n"
"pip install openPMD-viewer")
# Store an openPMD timeseries object
# (Its API is used in order to conveniently extract data from the file)
self.openpmd_ts = OpenPMDTimeSeries( location, check_all_files=False )
# Initialize the instance
RawDataReaderBase.__init__(self, location, speciesName, dataName, internalName, firstTimeStep)
def _ReadData(self, timeStep):
data, = self.openpmd_ts.get_particle( [self.internalName],
species=self.speciesName, iteration=timeStep )
self.currentTime = self._ReadTime(timeStep)
return data
def _ReadTime(self, timeStep):
# The line below sets the attribute `_current_i` of openpmd_ts
self.openpmd_ts._find_output( None, timeStep )
# This sets the corresponding time
self.currentTime = self.openpmd_ts.t[ self.openpmd_ts._current_i ]
def _ReadUnits(self):
# OpenPMD data always provide conversion to SI units
# TODO: Get the units from file
self.dataUnits = "arb.u."
self.timeUnits = "s"
def _OpenFile(self, timeStep):
# The line below sets the attribute `_current_i` of openpmd_ts
self.openpmd_ts._find_output( None, timeStep )
# This finds the full path to the corresponding file
fileName = self.openpmd_ts.h5_files[ self.openpmd_ts._current_i ]
file_content = H5File(fileName, 'r')
return file_content
def _ReadSimulationProperties(self, file_content):
# TODO: Add the proper resolution
self.grid_resolution = None
self.grid_size = None
self.grid_units = 'm'
def _ReadBasicData(self):
file_content = self._OpenFile(self.firstTimeStep)
self._ReadSimulationProperties(file_content)
file_content.close()
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 and check that the targeted lines are present
with open(script_filename) as f:
script = f.read()
# Check that the targeted lines are present
if script.find('n_order = -1') == -1 \
or script.find('track_bunch = False') == -1:
raise RuntimeError('Did not find expected lines in \
boosted_frame_script.py')
# Modify the script so as to enable finite order
script = script.replace('n_order = -1', 'n_order = 16')
script = script.replace('track_bunch = False', 'track_bunch = True')
with open(script_filename, 'w') as f:
f.write(script)
# Launch the script from the OS
response = os.system(
'cd %s; mpirun -np 2 python boosted_frame_script.py' %temporary_dir )
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 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 )
# Enter the temporary directory
os.chdir( temporary_dir )
# Read the script and check that the targeted lines are present
with open('boosted_frame_script.py') as f:
script = f.read()
# Modify the script so as to enable particle tracking
script = script.replace('track_bunch = False', 'track_bunch = True')
with open('boosted_frame_script.py', 'w') as f:
f.write(script)
# Launch the script from the OS
response = os.system( 'mpirun -np 2 python boosted_frame_script.py' )
assert response==0
# Check that the particle ids are unique at each iterations
ts = OpenPMDTimeSeries('./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))
# Exit the temporary directory and suppress it
os.chdir('../../')
shutil.rmtree( temporary_dir )
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 = int(T_interact/sim.dt)',
'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 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)
"""
# Import openPMD viewer
try:
from opmd_viewer import OpenPMDTimeSeries
except ImportError:
raise ImportError(
'The package `opmd_viewer` is required to restart from checkpoints.'
'\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)
# 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_simulation(gamma_boost, use_separate_electron_species):
"""
Run a simulation with a laser pulse going through a gas jet of ionizable
N5+ atoms, and check the fraction of atoms that are in the N5+ state.
Parameters
----------
gamma_boost: float
The Lorentz factor of the frame in which the simulation is carried out.
use_separate_electron_species: bool
Whether to use separate electron species for each level, or
a single electron species for all levels.
"""
# The simulation box
zmax_lab = 20.e-6 # Length of the box along z (meters)
zmin_lab = 0.e-6
Nr = 3 # Number of gridpoints along r
rmax = 10.e-6 # Length of the box along r (meters)
Nm = 2 # Number of modes used
# The particles of the plasma
p_zmin = 5.e-6 # Position of the beginning of the plasma (meters)
p_zmax = 15.e-6
p_rmin = 0. # Minimal radial position of the plasma (meters)
p_rmax = 100.e-6 # Maximal radial position of the plasma (meters)
n_atoms = 0.2 # The atomic density is chosen very low,
# to avoid collective effects
p_nz = 2 # Number of particles per cell along z
p_nr = 1 # Number of particles per cell along r
p_nt = 4 # Number of particles per cell along theta
# Boosted frame
boost = BoostConverter(gamma_boost)
# Boost the different quantities
beta_boost = np.sqrt(1. - 1. / gamma_boost**2)
zmin, zmax = boost.static_length([zmin_lab, zmax_lab])
p_zmin, p_zmax = boost.static_length([p_zmin, p_zmax])
n_atoms, = boost.static_density([n_atoms])
# Increase the number of particles per cell in order to keep sufficient
# statistics for the evaluation of the ionization fraction
if gamma_boost > 1:
p_nz = int(2 * gamma_boost * (1 + beta_boost) * p_nz)
# The laser
a0 = 1.8 # Laser amplitude
lambda0_lab = 0.8e-6 # Laser wavelength
# Boost the laser wavelength before calculating the laser amplitude
lambda0, = boost.copropag_length([lambda0_lab], beta_object=1.)
# Duration and initial position of the laser
ctau = 10. * lambda0
z0 = -2 * ctau
# Calculate laser amplitude
omega = 2 * np.pi * c / lambda0
E0 = a0 * m_e * c * omega / e
B0 = E0 / c
def laser_func(F, x, y, z, t, amplitude, length_scale):
"""
Function that describes a Gaussian laser with infinite waist
"""
return( F + amplitude * math.cos( 2*np.pi*(z-c*t)/lambda0 ) * \
math.exp( - (z - c*t - z0)**2/ctau**2 ) )
# Resolution and number of timesteps
dz = lambda0 / 16.
dt = dz / c
Nz = int((zmax - zmin) / dz) + 1
N_step = int(
(2. * 40. * lambda0 + zmax - zmin) / (dz * (1 + beta_boost))) + 1
# Get the speed of the plasma
uz_m, = boost.longitudinal_momentum([0.])
v_plasma, = boost.velocity([0.])
# The diagnostics
diag_period = N_step - 1 # Period of the diagnostics in number of timesteps
# Initial ionization level of the Nitrogen atoms
level_start = 2
# Initialize the simulation object, with the neutralizing electrons
# No particles are created because we do not pass the density
sim = Simulation(Nz,
zmax,
Nr,
rmax,
Nm,
dt,
zmin=zmin,
v_comoving=v_plasma,
use_galilean=False,
boundaries='open',
use_cuda=use_cuda)
# Add the charge-neutralizing electrons
elec = sim.add_new_species(q=-e,
m=m_e,
n=level_start * n_atoms,
p_nz=p_nz,
p_nr=p_nr,
p_nt=p_nt,
p_zmin=p_zmin,
p_zmax=p_zmax,
p_rmin=p_rmin,
p_rmax=p_rmax,
continuous_injection=False,
uz_m=uz_m)
# Add the N atoms
ions = sim.add_new_species(q=0,
m=14. * m_p,
n=n_atoms,
p_nz=p_nz,
p_nr=p_nr,
p_nt=p_nt,
p_zmin=p_zmin,
p_zmax=p_zmax,
p_rmin=p_rmin,
p_rmax=p_rmax,
continuous_injection=False,
uz_m=uz_m)
# Add the target electrons
if use_separate_electron_species:
# Use a dictionary of electron species: one per ionizable level
target_species = {}
level_max = 6 # N can go up to N7+, but here we stop at N6+
for i_level in range(level_start, level_max):
target_species[i_level] = sim.add_new_species(q=-e, m=m_e)
else:
# Use the pre-existing, charge-neutralizing electrons
target_species = elec
level_max = None # Default is going up to N7+
# Define ionization
ions.make_ionizable(element='N',
level_start=level_start,
level_max=level_max,
target_species=target_species)
# Set the moving window
sim.set_moving_window(v=v_plasma)
# Add a laser to the fields of the simulation (external fields)
sim.external_fields = [
ExternalField(laser_func, 'Ex', E0, 0.),
ExternalField(laser_func, 'By', B0, 0.)
]
# Add a particle diagnostic
sim.diags = [
ParticleDiagnostic(
diag_period,
{"ions": ions},
particle_data=["position", "gamma", "weighting", "E", "B"],
# Test output of fields and gamma for standard
# (non-boosted) particle diagnostics
write_dir='tests/diags',
comm=sim.comm)
]
if gamma_boost > 1:
T_sim_lab = (2. * 40. * lambda0_lab + zmax_lab - zmin_lab) / c
sim.diags.append(
BackTransformedParticleDiagnostic(zmin_lab,
zmax_lab,
v_lab=0.,
dt_snapshots_lab=T_sim_lab / 2.,
Ntot_snapshots_lab=3,
gamma_boost=gamma_boost,
period=diag_period,
fldobject=sim.fld,
species={"ions": ions},
comm=sim.comm,
write_dir='tests/lab_diags'))
# Run the simulation
sim.step(N_step, use_true_rho=True)
# Check the fraction of N5+ ions at the end of the simulation
w = ions.w
ioniz_level = ions.ionizer.ionization_level
# Get the total number of N atoms/ions (all ionization levels together)
ntot = w.sum()
# Get the total number of N5+ ions
n_N5 = w[ioniz_level == 5].sum()
# Get the fraction of N5+ ions, and check that it is close to 0.32
N5_fraction = n_N5 / ntot
print('N5+ fraction: %.4f' % N5_fraction)
assert ((N5_fraction > 0.30) and (N5_fraction < 0.34))
# When different electron species are created, check the fraction of
# each electron species
if use_separate_electron_species:
for i_level in range(level_start, level_max):
n_N = w[ioniz_level == i_level].sum()
assert np.allclose(target_species[i_level].w.sum(), n_N)
# Check consistency in the regular openPMD diagnostics
ts = OpenPMDTimeSeries('./tests/diags/hdf5/')
last_iteration = ts.iterations[-1]
w, q = ts.get_particle(['w', 'charge'],
species="ions",
iteration=last_iteration)
# Check that the openPMD file contains the same number of N5+ ions
n_N5_openpmd = np.sum(w[(4.5 * e < q) & (q < 5.5 * e)])
assert np.isclose(n_N5_openpmd, n_N5)
# Remove openPMD files
shutil.rmtree('./tests/diags/')
# Check consistency of the back-transformed openPMD diagnostics
if gamma_boost > 1.:
ts = OpenPMDTimeSeries('./tests/lab_diags/hdf5/')
last_iteration = ts.iterations[-1]
w, q = ts.get_particle(['w', 'charge'],
species="ions",
iteration=last_iteration)
# Check that the openPMD file contains the same number of N5+ ions
n_N5_openpmd = np.sum(w[(4.5 * e < q) & (q < 5.5 * e)])
assert np.isclose(n_N5_openpmd, n_N5)
# Remove openPMD files
shutil.rmtree('./tests/lab_diags/')
def add_elec_bunch_openPMD(sim,
ts_path,
z_off=0.,
species=None,
select=None,
iteration=None,
boost=None,
filter_currents=True):
"""
Introduce a relativistic electron 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.
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.
filter_currents : bool, optional
Whether to filter the currents in k space (True by default)
"""
# Import openPMD viewer
try:
from opmd_viewer import OpenPMDTimeSeries
except ImportError:
raise ImportError(
'The package `opmd_viewer` is required to restart from checkpoints.'
'\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)
# 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.amax(z) + np.amin(z)) / 2 + z_off
# Add the electrons to the simulation, and calculate the space charge
add_elec_bunch_from_arrays(sim,
x,
y,
z,
ux,
uy,
uz,
w,
boost=boost,
filter_currents=filter_currents)
plt.savefig('%s/rho_%09i.png' % (Path + '/rho', iter), bbox_inches='tight')
plt.close()
a0 = ts_laser.get_a0(iteration=iter, pol='x')
#print(a0)
a0PerSnapshot.append(a0)
waist = ts_laser.get_laser_waist(iteration=iter, pol='x')
#print(waist)
PulsewaistPerSnapshot.append(waist)
length = ts_laser.get_ctau(iteration=iter, pol='x')
#print(length)
PulselengthPerSnapshot.append(length)
ts_particles.get_particle(['z', 'uz'], species='electrons', iteration=iter)
z_selected, uz_selected = ts_particles.get_particle(
['z', 'uz'],
species='electrons',
iteration=iter,
select={'uz': [1, None]})
plt.plot(z_selected, uz_selected, 'b.', markersize=0.1)
plt.savefig('%s/uz_%09i.png' % (Path + '/uz', iter), bbox_inches='tight')
plt.close()
np.savetxt('TextFiles/' + 'a0PerSnapshot' + 'SimulationName' + '.txt',
a0PerSnapshot)
np.savetxt('TextFiles/' + 'PulsewaistPerSnapshot' + 'SimulationName' + '.txt',
PulsewaistPerSnapshot)
np.savetxt('TextFiles/' + 'PulselengthPerSnapshot' + 'SimulationName' + '.txt',
PulselengthPerSnapshot)
os.makedirs(Path + '/gamma_z')
if not os.path.exists(Path + '/y_z'):
os.makedirs(Path + '/y_z')
# The main loop will run for every i*factor,
# from 0 to ts_particles.iterations.size/factor
# Setting factor=2 will process every other file and so on.
factor = 2
a0PerSnapshot = []
PulsewaistPerSnapshot = []
PulselengthPerSnapshot = []
EzMinPerSnapshot = []
x, y, z, gamma = ts_particles.get_particle(
var_list=['x', 'y', 'z', 'gamma'],
iteration=ts_particles.iterations[0],
species='beam')
U0 = gamma.sum()
gamma_ave = gamma.mean()
sig_gamma = gamma.std()
gamma0 = gamma[0]
delta_gamma = gamma0 - gamma_ave
for i in range(0, int(ts_laser.iterations.size / factor)):
iter = ts_laser.iterations[factor * i]
print('Iteration:%d' % iter)
Elaser = ts_laser.get_laser_envelope(iteration=iter,
pol='x',
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 = [
BoostedParticleDiagnostic(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('../')
def run_simulation(gamma_boost):
"""
Run a simulation with a laser pulse going through a gas jet of ionizable
N5+ atoms, and check the fraction of atoms that are in the N5+ state.
Parameters
----------
gamma_boost: float
The Lorentz factor of the frame in which the simulation is carried out.
"""
# The simulation box
zmax_lab = 20.e-6 # Length of the box along z (meters)
zmin_lab = 0.e-6
Nr = 3 # Number of gridpoints along r
rmax = 10.e-6 # Length of the box along r (meters)
Nm = 2 # Number of modes used
# The particles of the plasma
p_zmin = 5.e-6 # Position of the beginning of the plasma (meters)
p_zmax = 15.e-6
p_rmin = 0. # Minimal radial position of the plasma (meters)
p_rmax = 100.e-6 # Maximal radial position of the plasma (meters)
n_e = 1. # The plasma density is chosen very low,
# to avoid collective effects
p_nz = 2 # Number of particles per cell along z
p_nr = 1 # Number of particles per cell along r
p_nt = 4 # Number of particles per cell along theta
# Boosted frame
boost = BoostConverter(gamma_boost)
# Boost the different quantities
beta_boost = np.sqrt(1. - 1. / gamma_boost**2)
zmin, zmax = boost.static_length([zmin_lab, zmax_lab])
p_zmin, p_zmax = boost.static_length([p_zmin, p_zmax])
n_e, = boost.static_density([n_e])
# Increase the number of particles per cell in order to keep sufficient
# statistics for the evaluation of the ionization fraction
if gamma_boost > 1:
p_nz = int(2 * gamma_boost * (1 + beta_boost) * p_nz)
# The laser
a0 = 1.8 # Laser amplitude
lambda0_lab = 0.8e-6 # Laser wavelength
# Boost the laser wavelength before calculating the laser amplitude
lambda0, = boost.copropag_length([lambda0_lab], beta_object=1.)
# Duration and initial position of the laser
ctau = 10. * lambda0
z0 = -2 * ctau
# Calculate laser amplitude
omega = 2 * np.pi * c / lambda0
E0 = a0 * m_e * c * omega / e
B0 = E0 / c
def laser_func(F, x, y, z, t, amplitude, length_scale):
"""
Function that describes a Gaussian laser with infinite waist
"""
return( F + amplitude * math.cos( 2*np.pi*(z-c*t)/lambda0 ) * \
math.exp( - (z - c*t - z0)**2/ctau**2 ) )
# Resolution and number of timesteps
dz = lambda0 / 16.
dt = dz / c
Nz = int((zmax - zmin) / dz) + 1
N_step = int(
(2. * 40. * lambda0 + zmax - zmin) / (dz * (1 + beta_boost))) + 1
# Get the speed of the plasma
uz_m, = boost.longitudinal_momentum([0.])
v_plasma, = boost.velocity([0.])
# The diagnostics
diag_period = N_step - 1 # Period of the diagnostics in number of timesteps
# Initialize the simulation object
sim = Simulation(
Nz,
zmax,
Nr,
rmax,
Nm,
dt,
p_zmax,
p_zmax, # No electrons get created because we pass p_zmin=p_zmax
p_rmin,
p_rmax,
p_nz,
p_nr,
p_nt,
n_e,
zmin=zmin,
initialize_ions=False,
v_comoving=v_plasma,
use_galilean=False,
boundaries='open',
use_cuda=use_cuda)
sim.set_moving_window(v=v_plasma)
# Add the N atoms
p_zmin, p_zmax, Npz = adapt_to_grid(sim.fld.interp[0].z, p_zmin, p_zmax,
p_nz)
p_rmin, p_rmax, Npr = adapt_to_grid(sim.fld.interp[0].r, p_rmin, p_rmax,
p_nr)
sim.ptcl.append(
Particles(q=e,
m=14. * m_p,
n=0.2 * n_e,
Npz=Npz,
zmin=p_zmin,
zmax=p_zmax,
Npr=Npr,
rmin=p_rmin,
rmax=p_rmax,
Nptheta=p_nt,
dt=dt,
use_cuda=use_cuda,
uz_m=uz_m,
grid_shape=sim.fld.interp[0].Ez.shape,
continuous_injection=False))
sim.ptcl[1].make_ionizable(element='N',
level_start=0,
target_species=sim.ptcl[0])
# Add a laser to the fields of the simulation (external fields)
sim.external_fields = [
ExternalField(laser_func, 'Ex', E0, 0.),
ExternalField(laser_func, 'By', B0, 0.)
]
# Add a field diagnostic
sim.diags = [
ParticleDiagnostic(diag_period, {"ions": sim.ptcl[1]},
write_dir='tests/diags',
comm=sim.comm)
]
if gamma_boost > 1:
T_sim_lab = (2. * 40. * lambda0_lab + zmax_lab - zmin_lab) / c
sim.diags.append(
BoostedParticleDiagnostic(zmin_lab,
zmax_lab,
v_lab=0.,
dt_snapshots_lab=T_sim_lab / 2.,
Ntot_snapshots_lab=3,
gamma_boost=gamma_boost,
period=diag_period,
fldobject=sim.fld,
species={"ions": sim.ptcl[1]},
comm=sim.comm,
write_dir='tests/lab_diags'))
# Run the simulation
sim.step(N_step, use_true_rho=True)
# Check the fraction of N5+ ions at the end of the simulation
w = sim.ptcl[1].w
ioniz_level = sim.ptcl[1].ionizer.ionization_level
# Get the total number of N atoms/ions (all ionization levels together)
ntot = w.sum()
# Get the total number of N5+ ions
n_N5 = w[ioniz_level == 5].sum()
# Get the fraction of N5+ ions, and check that it is close to 0.32
N5_fraction = n_N5 / ntot
print('N5+ fraction: %.4f' % N5_fraction)
assert ((N5_fraction > 0.30) and (N5_fraction < 0.34))
# Check consistency in the regular openPMD diagnostics
ts = OpenPMDTimeSeries('./tests/diags/hdf5/')
last_iteration = ts.iterations[-1]
w, q = ts.get_particle(['w', 'charge'],
species="ions",
iteration=last_iteration)
# Check that the openPMD file contains the same number of N5+ ions
n_N5_openpmd = np.sum(w[(4.5 * e < q) & (q < 5.5 * e)])
assert np.isclose(n_N5_openpmd, n_N5)
# Remove openPMD files
shutil.rmtree('./tests/diags/')
# Check consistency of the back-transformed openPMD diagnostics
if gamma_boost > 1.:
ts = OpenPMDTimeSeries('./tests/lab_diags/hdf5/')
last_iteration = ts.iterations[-1]
w, q = ts.get_particle(['w', 'charge'],
species="ions",
iteration=last_iteration)
# Check that the openPMD file contains the same number of N5+ ions
n_N5_openpmd = np.sum(w[(4.5 * e < q) & (q < 5.5 * e)])
assert np.isclose(n_N5_openpmd, n_N5)
# Remove openPMD files
shutil.rmtree('./tests/lab_diags/')
