def add_diagnostic(self, diagnostic):
# Call method of parent class
PICMI_Simulation.add_diagnostic(self, diagnostic)
# Handle diagnostic
if diagnostic.step_min is None:
iteration_min = 0
else:
iteration_min = diagnostic.step_min
if diagnostic.step_max is None:
iteration_max = np.inf
else:
iteration_max = diagnostic.step_max
# Register field diagnostic
if type(diagnostic) == PICMI_FieldDiagnostic:
if diagnostic.data_list is None:
data_list = ['rho', 'E', 'B', 'J']
else:
data_list = diagnostic.data_list
diag = FieldDiagnostic(period=diagnostic.period,
fldobject=self.fbpic_sim.fld,
comm=self.fbpic_sim.comm,
fieldtypes=data_list,
write_dir=diagnostic.write_dir,
iteration_min=iteration_min,
iteration_max=iteration_max)
# Register particle diagnostic
elif type(diagnostic) == PICMI_ParticleDiagnostic:
species_dict = {}
for s in diagnostic.species:
if s.name is None:
raise ValueError('When using a species in a diagnostic, '
'its name must be set.')
species_dict[s.name] = s.fbpic_species
if diagnostic.data_list is None:
data_list = ['position', 'momentum', 'weighting']
else:
data_list = diagnostic.data_list
diag = ParticleDiagnostic(period=diagnostic.period,
species=species_dict,
comm=self.fbpic_sim.comm,
particle_data=data_list,
write_dir=diagnostic.write_dir,
iteration_min=iteration_min,
iteration_max=iteration_max)
# Add it to the FBPIC simulation
self.fbpic_sim.diags.append(diag)
gamma_boost=boost.gamma0)
# Convert parameter to boosted frame
v_window_boosted, = boost.velocity([v_window])
# Configure the moving window
sim.set_moving_window(v=v_window_boosted)
# Add a field diagnostic
sim.diags = [
# Diagnostics in the boosted frame
FieldDiagnostic(dt_period=dt_boosted_diag_period,
fldobject=sim.fld,
comm=sim.comm),
ParticleDiagnostic(dt_period=dt_boosted_diag_period,
species={
"electrons": sim.ptcl[0],
"bunch": sim.ptcl[2]
},
comm=sim.comm),
# Diagnostics in the lab frame (back-transformed)
BackTransformedFieldDiagnostic(zmin,
zmax,
v_window,
dt_lab_diag_period,
N_lab_diag,
boost.gamma0,
fieldtypes=['rho', 'E', 'B'],
period=write_period,
fldobject=sim.fld,
comm=sim.comm),
BackTransformedParticleDiagnostic(zmin,
zmax,
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/')
if use_restart is False:
# Track electrons if required (species 0 correspond to the electrons)
if track_electrons:
elec.track(sim.comm)
else:
# Load the fields and particles from the latest checkpoint file
restart_from_checkpoint(sim)
# Configure the moving window
sim.set_moving_window(v=v_window)
# Add diagnostics
sim.diags = [
FieldDiagnostic(diag_period, sim.fld, comm=sim.comm),
ParticleDiagnostic(diag_period, {
"electrons from N": elec_from_N,
"electrons": elec
},
comm=sim.comm),
# Since rho from `FieldDiagnostic` is 0 almost everywhere
# (neutral plasma), it is useful to see the charge density
# of individual particles
ParticleChargeDensityDiagnostic(diag_period, sim, {"electrons": elec})
]
# Add checkpoints
if save_checkpoints:
set_periodic_checkpoint(sim, checkpoint_period)
### Run the simulation
sim.step(N_step)
p_zmin=p_zmin)
# Activate ionization of He ions (for levels above 1).
# Store the created electrons in the species `elec`
atoms_He.make_ionizable('He', target_species=elec, level_start=1)
# Activate ionization of N ions (for levels above 5).
# Store the created electrons in a new dedicated electron species that
# does not contain any macroparticles initially
elec_from_N = sim.add_new_species(q=-e, m=m_e)
atoms_N.make_ionizable('N', target_species=elec_from_N, level_start=5)
# Add a laser to the fields of the simulation
add_laser(sim, a0, w0, ctau, z0, zf=z_foc)
# Configure the moving window
sim.set_moving_window(v=v_window)
# Add a diagnostics
sim.diags = [
FieldDiagnostic(diag_period, sim.fld, comm=sim.comm),
ParticleDiagnostic(diag_period, {
"e- from N": elec_from_N,
"e-": elec
},
comm=sim.comm)
]
### Run the simulation
sim.step(N_step)
def add_diagnostic(self, diagnostic):
# Call method of parent class
PICMI_Simulation.add_diagnostic(self, diagnostic)
# Handle iteration_min/max in regular diagnostic
if type(diagnostic) in [
PICMI_FieldDiagnostic, PICMI_ParticleDiagnostic
]:
if diagnostic.step_min is None:
iteration_min = 0
else:
iteration_min = diagnostic.step_min
if diagnostic.step_max is None:
iteration_max = np.inf
else:
iteration_max = diagnostic.step_max
# Register field diagnostic
if type(diagnostic) in [
PICMI_FieldDiagnostic, PICMI_LabFrameFieldDiagnostic
]:
if diagnostic.data_list is None:
data_list = ['rho', 'E', 'B', 'J']
else:
data_list = set() # Use set to avoid redundancy
for data in diagnostic.data_list:
if data in ['Ex', 'Ey', 'Ez', 'E']:
data_list.add('E')
elif data in ['Bx', 'By', 'Bz', 'B']:
data_list.add('B')
elif data in ['Jx', 'Jy', 'Jz', 'J']:
data_list.add('J')
elif data == 'rho':
data_list.add('rho')
data_list = list(data_list)
if type(diagnostic) == PICMI_FieldDiagnostic:
diag = FieldDiagnostic(period=diagnostic.period,
fldobject=self.fbpic_sim.fld,
comm=self.fbpic_sim.comm,
fieldtypes=data_list,
write_dir=diagnostic.write_dir,
iteration_min=iteration_min,
iteration_max=iteration_max)
# Register particle density diagnostic
rho_density_list = []
if diagnostic.data_list is not None:
for data in diagnostic.data_list:
if data.startswith('rho_'):
# particle density diagnostics, rho_speciesname
rho_density_list.append(data)
if rho_density_list:
species_dict = {}
for data in rho_density_list:
sname = data[4:]
for s in self.species:
if s.name == sname:
species_dict[s.name] = s.fbpic_species
pdd_diag = ParticleChargeDensityDiagnostic(
period=diagnostic.period,
sim=self.fbpic_sim,
species=species_dict,
write_dir=diagnostic.write_dir,
iteration_min=iteration_min,
iteration_max=iteration_max)
self.fbpic_sim.diags.append(pdd_diag)
elif type(diagnostic) == PICMI_LabFrameFieldDiagnostic:
diag = BackTransformedFieldDiagnostic(
zmin_lab=diagnostic.grid.lower_bound[1],
zmax_lab=diagnostic.grid.upper_bound[1],
v_lab=c,
dt_snapshots_lab=diagnostic.dt_snapshots,
Ntot_snapshots_lab=diagnostic.num_snapshots,
gamma_boost=self.gamma_boost,
period=100,
fldobject=self.fbpic_sim.fld,
comm=self.fbpic_sim.comm,
fieldtypes=diagnostic.data_list,
write_dir=diagnostic.write_dir)
# Register particle diagnostic
elif type(diagnostic) in [
PICMI_ParticleDiagnostic, PICMI_LabFrameParticleDiagnostic
]:
species_dict = {}
for s in diagnostic.species:
if s.name is None:
raise ValueError('When using a species in a diagnostic, '
'its name must be set.')
species_dict[s.name] = s.fbpic_species
if diagnostic.data_list is None:
data_list = ['position', 'momentum', 'weighting']
else:
data_list = diagnostic.data_list
if type(diagnostic) == PICMI_ParticleDiagnostic:
diag = ParticleDiagnostic(period=diagnostic.period,
species=species_dict,
comm=self.fbpic_sim.comm,
particle_data=data_list,
write_dir=diagnostic.write_dir,
iteration_min=iteration_min,
iteration_max=iteration_max)
else:
diag = BackTransformedParticleDiagnostic(
zmin_lab=diagnostic.grid.lower_bound[1],
zmax_lab=diagnostic.grid.upper_bound[1],
v_lab=c,
dt_snapshots_lab=diagnostic.dt_snapshots,
Ntot_snapshots_lab=diagnostic.num_snapshots,
gamma_boost=self.gamma_boost,
period=100,
fldobject=self.fbpic_sim.fld,
species=species_dict,
comm=self.fbpic_sim.comm,
particle_data=data_list,
write_dir=diagnostic.write_dir)
# Add it to the FBPIC simulation
self.fbpic_sim.diags.append(diag)
def run_simulation(gamma_boost, show):
"""
Run a simulation with a relativistic electron bunch crosses a laser
Parameters
----------
gamma_boost: float
The Lorentz factor of the frame in which the simulation is carried out.
show: bool
Whether to show a plot of the angular distribution
"""
# Boosted frame
boost = BoostConverter(gamma_boost)
# The simulation timestep
diag_period = 100
N_step = 101 # Number of iterations to perform
# Calculate timestep to resolve the interaction with enough points
laser_duration_boosted, = boost.copropag_length([laser_duration],
beta_object=-1)
bunch_sigma_z_boosted, = boost.copropag_length([bunch_sigma_z],
beta_object=1)
dt = (4 * laser_duration_boosted + bunch_sigma_z_boosted / c) / N_step
# Initialize the simulation object
zmax, zmin = boost.copropag_length([zmax_lab, zmin_lab], beta_object=1.)
sim = Simulation(Nz,
zmax,
Nr,
rmax,
Nm,
dt,
p_zmin=0,
p_zmax=0,
p_rmin=0,
p_rmax=0,
p_nz=1,
p_nr=1,
p_nt=1,
n_e=1,
dens_func=None,
zmin=zmin,
boundaries='periodic',
use_cuda=use_cuda)
# Remove particles that were previously created
sim.ptcl = []
print('Initialized simulation')
# Add electron bunch (automatically converted to boosted-frame)
add_elec_bunch_gaussian(sim,
sig_r=1.e-6,
sig_z=bunch_sigma_z,
n_emit=0.,
gamma0=gamma_bunch_mean,
sig_gamma=gamma_bunch_rms,
Q=Q_bunch,
N=N_bunch,
tf=0.0,
zf=0.5 * (zmax + zmin),
boost=boost)
elec = sim.ptcl[0]
print('Initialized electron bunch')
# Add a photon species
photons = Particles(q=0,
m=0,
n=0,
Npz=1,
zmin=0,
zmax=0,
Npr=1,
rmin=0,
rmax=0,
Nptheta=1,
dt=sim.dt,
ux_m=0.,
uy_m=0.,
uz_m=0.,
ux_th=0.,
uy_th=0.,
uz_th=0.,
dens_func=None,
continuous_injection=False,
grid_shape=sim.fld.interp[0].Ez.shape,
particle_shape='linear',
use_cuda=sim.use_cuda)
sim.ptcl.append(photons)
print('Initialized photons')
# Activate Compton scattering for electrons of the bunch
elec.activate_compton(target_species=photons,
laser_energy=laser_energy,
laser_wavelength=laser_wavelength,
laser_waist=laser_waist,
laser_ctau=laser_ctau,
laser_initial_z0=laser_initial_z0,
ratio_w_electron_photon=50,
boost=boost)
print('Activated Compton')
# Add diagnostics
if write_hdf5:
sim.diags = [
ParticleDiagnostic(diag_period,
species={
'electrons': elec,
'photons': photons
},
comm=sim.comm)
]
# Get initial total momentum
initial_total_elec_px = (elec.w * elec.ux).sum() * m_e * c
initial_total_elec_py = (elec.w * elec.uy).sum() * m_e * c
initial_total_elec_pz = (elec.w * elec.uz).sum() * m_e * c
### Run the simulation
for species in sim.ptcl:
species.send_particles_to_gpu()
for i_step in range(N_step):
for species in sim.ptcl:
species.halfpush_x()
elec.handle_elementary_processes(sim.time + 0.5 * sim.dt)
for species in sim.ptcl:
species.halfpush_x()
# Increment time and run diagnostics
sim.time += sim.dt
sim.iteration += 1
for diag in sim.diags:
diag.write(sim.iteration)
# Print fraction of photons produced
if i_step % 10 == 0:
for species in sim.ptcl:
species.receive_particles_from_gpu()
simulated_frac = photons.w.sum() / elec.w.sum()
for species in sim.ptcl:
species.send_particles_to_gpu()
print( 'Iteration %d: Photon fraction per electron = %f' \
%(i_step, simulated_frac) )
for species in sim.ptcl:
species.receive_particles_from_gpu()
# Check estimation of photon fraction
check_photon_fraction(simulated_frac)
# Check conservation of momentum (is only conserved )
if elec.compton_scatterer.ratio_w_electron_photon == 1:
check_momentum_conservation(gamma_boost, photons, elec,
initial_total_elec_px,
initial_total_elec_py,
initial_total_elec_pz)
# Transform the photon momenta back into the lab frame
photon_u = 1. / photons.inv_gamma
photon_lab_pz = boost.gamma0 * (photons.uz + boost.beta0 * photon_u)
photon_lab_p = boost.gamma0 * (photon_u + boost.beta0 * photons.uz)
# Plot the scaled angle and frequency
if show:
import matplotlib.pyplot as plt
# Bin the photons on a grid in frequency and angle
freq_min = 0.5
freq_max = 1.2
N_freq = 500
gammatheta_min = 0.
gammatheta_max = 1.
N_gammatheta = 100
hist_range = [[freq_min, freq_max], [gammatheta_min, gammatheta_max]]
extent = [freq_min, freq_max, gammatheta_min, gammatheta_max]
fundamental_frequency = 4 * gamma_bunch_mean**2 * c / laser_wavelength
photon_scaled_freq = photon_lab_p * c / (h * fundamental_frequency)
gamma_theta = gamma_bunch_mean * np.arccos(
photon_lab_pz / photon_lab_p)
grid, freq_bins, gammatheta_bins = np.histogram2d(
photon_scaled_freq,
gamma_theta,
weights=photons.w,
range=hist_range,
bins=[N_freq, N_gammatheta])
# Normalize by solid angle, frequency and number of photons
dw = (freq_bins[1] - freq_bins[0]) * 2 * np.pi * fundamental_frequency
dtheta = (gammatheta_bins[1] - gammatheta_bins[0]) / gamma_bunch_mean
domega = 2. * np.pi * np.sin(
gammatheta_bins / gamma_bunch_mean) * dtheta
grid /= dw * domega[np.newaxis, 1:] * elec.w.sum()
grid = np.where(grid == 0, np.nan, grid)
plt.imshow(grid.T,
origin='lower',
extent=extent,
cmap='gist_earth',
aspect='auto',
vmax=1.8e-16)
plt.title('Particles, $d^2N/d\omega \,d\Omega$')
plt.xlabel('Scaled energy ($\omega/4\gamma^2\omega_\ell$)')
plt.ylabel(r'$\gamma \theta$')
plt.colorbar()
# Plot theory
plt.plot(1. / (1 + gammatheta_bins**2), gammatheta_bins, color='r')
plt.show()
plt.clf()
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/')
restart_from_checkpoint(sim)
# Configure the moving window
sim.set_moving_window(v=v_window)
# Add a field diagnostic
# Parametric scan: each MPI rank should output its data to a
# different directory
write_dir = 'diags_a0_%.2f' % a0
sim.diags = [
FieldDiagnostic(diag_period,
sim.fld,
comm=sim.comm,
write_dir=write_dir),
ParticleDiagnostic(diag_period, {"electrons": elec},
select={"uz": [1., None]},
comm=sim.comm,
write_dir=write_dir)
]
# Add checkpoints
if save_checkpoints:
set_periodic_checkpoint(sim, checkpoint_period)
# Number of iterations to perform
N_step = int(T_interact / sim.dt)
### Run the simulation
sim.step(N_step)
print('')
def test_linear_wakefield(Nm=1, show=False):
"""
Run a simulation of linear laser-wakefield and compare the fields
with the analytical solution.
Parameters
----------
Nm: int
The number of azimuthal modes used in the simulation (Use 1, 2 or 3)
This also determines the profile of the driving laser:
- Nm=1: azimuthally-polarized annular laser
(laser in mode m=0, wakefield in mode m=0)
- Nm=2: linearly-polarized Gaussian laser
(laser in mode m=1, wakefield in mode m=0)
- Nm=3: linearly-polarized Laguerre-Gauss laser
(laser in mode m=0 and m=2, wakefield in mode m=0 and m=2,
show: bool
Whether to have pop-up windows show the comparison between
analytical and simulated results
"""
# Automatically choose higher number of macroparticles along theta
p_nt = 2 * Nm
# Initialize the simulation object
sim = Simulation(Nz,
zmax,
Nr,
rmax,
Nm,
dt,
p_zmin,
p_zmax,
p_rmin,
p_rmax,
p_nz,
p_nr,
p_nt,
n_e,
use_cuda=use_cuda,
boundaries={
'z': 'open',
'r': 'reflective'
})
# Create the relevant laser profile
if Nm == 1:
# Build an azimuthally-polarized pulse from 2 Laguerre-Gauss profiles
profile = LaguerreGaussLaser( 0, 1, a0=a0, waist=w0, tau=tau, z0=z0,
theta_pol=np.pi/2, theta0=0. ) \
+ LaguerreGaussLaser( 0, 1, a0=a0, waist=w0, tau=tau, z0=z0,
theta_pol=0., theta0=-np.pi/2 )
elif Nm == 2:
profile = GaussianLaser(a0=a0,
waist=w0,
tau=tau,
z0=z0,
theta_pol=np.pi / 2)
elif Nm == 3:
profile = LaguerreGaussLaser(0,
1,
a0=a0,
waist=w0,
tau=tau,
z0=z0,
theta_pol=np.pi / 2)
add_laser_pulse(sim, profile)
# Configure the moving window
sim.set_moving_window(v=c)
# Add diagnostics
if write_fields:
sim.diags.append(FieldDiagnostic(diag_period, sim.fld, sim.comm))
if write_particles:
sim.diags.append(
ParticleDiagnostic(diag_period, {'electrons': sim.ptcl[0]},
sim.comm))
# Prevent current correction for MPI simulation
if sim.comm.size > 1:
correct_currents = False
else:
correct_currents = True
# Run the simulation
sim.step(N_step, correct_currents=correct_currents)
# Compare the fields
compare_fields(sim, Nm, show)
# Track electrons if required (species 0 correspond to the electrons)
if track_electrons:
elec.track(sim.comm)
else:
# Load the fields and particles from the latest checkpoint file
restart_from_checkpoint(sim)
# Configure the moving window
sim.set_moving_window(v=v_window)
# Add diagnostics
sim.diags = [
FieldDiagnostic(dt_period / sim.dt, sim.fld, comm=sim.comm),
ParticleDiagnostic(dt_period / sim.dt, {"electrons": elec},
select={
"ux": [0, None],
"uy": [0, None]
},
comm=sim.comm),
ParticleDiagnostic(dt_period / sim.dt, {"beam": beam},
particle_data=[
"position", "momentum", "weighting", "E", "B",
"gamma"
],
comm=sim.comm)
]
# Add checkpoints
if save_checkpoints:
set_periodic_checkpoint(sim, 50000)
# Number of iterations to perform
N_step = int(T_interact / sim.dt)
add_laser(sim, a0, w0, ctau, z0, lambda0=lambda0, zf=zfoc)
if use_restart is False:
# Track electrons if required (species 0 correspond to the electrons)
if track_electrons:
elec.track(sim.comm)
else:
# Load the fields and particles from the latest checkpoint file
restart_from_checkpoint(sim)
# Configure the moving window
sim.set_moving_window(v=v_window)
# Add diagnostics
sim.diags = [
FieldDiagnostic(dt_period / sim.dt, sim.fld, comm=sim.comm),
ParticleDiagnostic(dt_period / sim.dt, {"electrons": elec},
select={"uz": [1., None]},
comm=sim.comm)
]
# Add checkpoints
if save_checkpoints:
set_periodic_checkpoint(sim, 500000)
# Number of iterations to perform
N_step = int(T_interact / sim.dt)
### Run the simulation
sim.step(N_step)
print('')
def run_fbpic(job: Job) -> None:
"""
This ``signac-flow`` operation runs a ``fbpic`` simulation.
:param job: the job instance is a handle to the data of a unique statepoint
"""
from fbpic.main import Simulation
from fbpic.lpa_utils.laser import add_laser_pulse, GaussianLaser
from fbpic.openpmd_diag import FieldDiagnostic, ParticleDiagnostic
# The density profile
def dens_func(z: np.ndarray, r: np.ndarray) -> np.ndarray:
"""Returns relative density at position z and r.
:param z: longitudinal positions, 1d array
:param r: radial positions, 1d array
:return: a 1d array ``n`` containing the density (between 0 and 1) at the given positions (z, r)
"""
# Allocate relative density
n = np.ones_like(z)
# Make linear ramp
n = np.where(
z < job.sp.ramp_start + job.sp.ramp_length,
(z - job.sp.ramp_start) / job.sp.ramp_length,
n,
)
# Supress density before the ramp
n = np.where(z < job.sp.ramp_start, 0.0, n)
return n
# plot density profile for checking
all_z = np.linspace(job.sp.zmin, job.sp.p_zmax, 1000)
dens = dens_func(all_z, 0.0)
width_inch = job.sp.p_zmax / 1e-5
major_locator = pyplot.MultipleLocator(10)
minor_locator = pyplot.MultipleLocator(5)
major_locator.MAXTICKS = 10000
minor_locator.MAXTICKS = 10000
def mark_on_plot(*, ax, parameter: str, y=1.1):
ax.annotate(s=parameter,
xy=(job.sp[parameter] * 1e6, y),
xycoords="data")
ax.axvline(x=job.sp[parameter] * 1e6, linestyle="--", color="red")
return ax
fig, ax = pyplot.subplots(figsize=(width_inch, 4.8))
ax.plot(all_z * 1e6, dens)
ax.set_xlabel(r"$%s \;(\mu m)$" % "z")
ax.set_ylim(-0.1, 1.2)
ax.set_xlim(job.sp.zmin * 1e6 - 20, job.sp.p_zmax * 1e6 + 20)
ax.set_ylabel("Density profile $n$")
ax.xaxis.set_major_locator(major_locator)
ax.xaxis.set_minor_locator(minor_locator)
mark_on_plot(ax=ax, parameter="zmin")
mark_on_plot(ax=ax, parameter="zmax")
mark_on_plot(ax=ax, parameter="p_zmin", y=0.9)
mark_on_plot(ax=ax, parameter="z0", y=0.8)
mark_on_plot(ax=ax, parameter="zf", y=0.6)
mark_on_plot(ax=ax, parameter="ramp_start", y=0.7)
mark_on_plot(ax=ax, parameter="L_interact")
mark_on_plot(ax=ax, parameter="p_zmax")
ax.annotate(s="ramp_start + ramp_length",
xy=(job.sp.ramp_start * 1e6 + job.sp.ramp_length * 1e6, 1.1),
xycoords="data")
ax.axvline(x=job.sp.ramp_start * 1e6 + job.sp.ramp_length * 1e6,
linestyle="--",
color="red")
ax.fill_between(all_z * 1e6, dens, alpha=0.5)
fig.savefig(job.fn("check_density.png"))
# redirect stdout to "stdout.txt"
orig_stdout = sys.stdout
f = open(job.fn("stdout.txt"), "w")
sys.stdout = f
# Initialize the simulation object
sim = Simulation(
job.sp.Nz,
job.sp.zmax,
job.sp.Nr,
job.sp.rmax,
job.sp.Nm,
job.sp.dt,
n_e=None, # no electrons
zmin=job.sp.zmin,
boundaries={
"z": "open",
"r": "reflective"
},
n_order=-1,
use_cuda=True,
verbose_level=2,
)
# Create a Gaussian laser profile
laser_profile = GaussianLaser(a0=job.sp.a0,
waist=job.sp.w0,
tau=job.sp.ctau / c_light,
z0=job.sp.z0,
zf=job.sp.zf,
theta_pol=0.,
lambda0=job.sp.lambda0,
cep_phase=0.,
phi2_chirp=0.,
propagation_direction=1)
# Add it to the simulation
add_laser_pulse(sim,
laser_profile,
gamma_boost=None,
method='direct',
z0_antenna=None,
v_antenna=0.)
# Create the plasma electrons
elec = sim.add_new_species(q=-q_e,
m=m_e,
n=job.sp.n_e,
dens_func=dens_func,
p_zmin=job.sp.p_zmin,
p_zmax=job.sp.p_zmax,
p_rmax=job.sp.p_rmax,
p_nz=job.sp.p_nz,
p_nr=job.sp.p_nr,
p_nt=job.sp.p_nt)
# Track electrons, useful for betatron radiation
# elec.track(sim.comm)
# Configure the moving window
sim.set_moving_window(v=c_light)
# Add diagnostics
write_dir = os.path.join(job.ws, "diags")
sim.diags = [
FieldDiagnostic(job.sp.diag_period,
sim.fld,
comm=sim.comm,
write_dir=write_dir,
fieldtypes=["rho", "E"]),
ParticleDiagnostic(job.sp.diag_period, {"electrons": elec},
select={"uz": [1., None]},
comm=sim.comm,
write_dir=write_dir,
particle_data=["momentum", "weighting"]),
]
# Plot the Ex component of the laser
# Get the fields in the half-plane theta=0 (Sum mode 0 and mode 1)
gathered_grids = [
sim.comm.gather_grid(sim.fld.interp[m]) for m in range(job.sp.Nm)
]
rgrid = gathered_grids[0].r
zgrid = gathered_grids[0].z
# construct the Er field for theta=0
Er = gathered_grids[0].Er.T.real
for m in range(1, job.sp.Nm):
# There is a factor 2 here so as to comply with the convention in
# Lifschitz et al., which is also the convention adopted in Warp Circ
Er += 2 * gathered_grids[m].Er.T.real
e0 = electric_field_amplitude_norm(lambda0=job.sp.lambda0)
fig = pyplot.figure(figsize=(8, 8))
sliceplots.Plot2D(
fig=fig,
arr2d=Er / e0,
h_axis=zgrid * 1e6,
v_axis=rgrid * 1e6,
zlabel=r"$E_r/E_0$",
xlabel=r"$z \;(\mu m)$",
ylabel=r"$r \;(\mu m)$",
extent=(
zgrid[0] * 1e6, # + 40
zgrid[-1] * 1e6, # - 20
rgrid[0] * 1e6,
rgrid[-1] * 1e6, # - 15,
),
cbar=True,
vmin=-3,
vmax=3,
hslice_val=0.0, # do a 1D slice through the middle of the simulation box
)
fig.savefig(job.fn('check_laser.png'))
# set deterministic random seed
np.random.seed(0)
# Run the simulation
sim.step(job.sp.N_step, show_progress=False)
# redirect stdout back and close "stdout.txt"
sys.stdout = orig_stdout
f.close()
# Track electrons if required (species 0 correspond to the electrons)
if track_electrons:
elec.track(sim.comm)
else:
# Load the fields and particles from the latest checkpoint file
restart_from_checkpoint(sim)
# Configure the moving window
sim.set_moving_window(v=v_window)
# Add diagnostics
sim.diags = [
FieldDiagnostic(dt_period / sim.dt, sim.fld, comm=sim.comm),
ParticleDiagnostic(dt_period / sim.dt, {"electrons": elec},
select={
"ux": [0, None],
"uy": [0, None]
},
comm=sim.comm)
]
# Add checkpoints
if save_checkpoints:
set_periodic_checkpoint(sim, 50000)
# Number of iterations to perform
N_step = int(T_interact / sim.dt)
### Run the simulation
sim.step(N_step)
print('')
a0,
w0,
ctau,
z0,
lambda0=lambda0,
zf=zfoc,
gamma_boost=gamma_boost)
# Configure the moving window
sim.set_moving_window(v=v_window)
# Add a field diagnostic
sim.diags = [
FieldDiagnostic(diag_period, sim.fld, sim.comm),
ParticleDiagnostic(diag_period, {
"electrons": sim.ptcl[0],
"bunch": sim.ptcl[2]
}, sim.comm),
BoostedFieldDiagnostic(zmin,
zmax,
c,
dt_snapshot_lab,
Ntot_snapshot_lab,
gamma_boost,
period=diag_period,
fldobject=sim.fld,
comm=sim.comm),
BoostedParticleDiagnostic(zmin,
zmax,
c,
dt_snapshot_lab,
Ntot_snapshot_lab,
# Load initial fields
# Add a laser to the fields of the simulation
add_laser(sim, a0, w0, ctau, z0)
if use_restart is False:
# Track electrons if required (species 0 correspond to the electrons)
if track_electrons:
sim.ptcl[0].track(sim.comm)
else:
# Load the fields and particles from the latest checkpoint file
restart_from_checkpoint(sim)
# Configure the moving window
sim.set_moving_window(v=v_window)
# Add diagnostics
sim.diags = [
FieldDiagnostic(diag_period, sim.fld, comm=sim.comm),
ParticleDiagnostic(diag_period, {"electrons": sim.ptcl[0]},
select={"uz": [1., None]},
comm=sim.comm)
]
# Add checkpoints
if save_checkpoints:
set_periodic_checkpoint(sim, checkpoint_period)
### Run the simulation
sim.step(N_step)
print('')
gamma_boost=boost.gamma0)
# Convert parameter to boosted frame
v_window_boosted, = boost.velocity([v_window])
# Configure the moving window
sim.set_moving_window(v=v_window_boosted)
# Add a field diagnostic
sim.diags = [
# Diagnostics in the boosted frame
FieldDiagnostic(dt_period=dt_boosted_diag_period,
fldobject=sim.fld,
comm=sim.comm),
ParticleDiagnostic(dt_period=dt_boosted_diag_period,
species={
"electrons": plasma_elec,
"bunch": bunch
},
comm=sim.comm),
# Diagnostics in the lab frame (back-transformed)
BackTransformedFieldDiagnostic(zmin,
zmax,
v_window,
dt_lab_diag_period,
N_lab_diag,
boost.gamma0,
fieldtypes=['rho', 'E', 'B'],
period=write_period,
fldobject=sim.fld,
comm=sim.comm),
BackTransformedParticleDiagnostic(zmin,
zmax,
# Add a laser to the fields of the simulation
add_laser(sim, a0, w0, ctau, z0)
# Configure the moving window
sim.set_moving_window(v=c)
# Add diagnostics
if write_fields:
sim.diags.append(FieldDiagnostic(diag_period, sim.fld, sim.comm))
sim.diags.append(
FieldDiagnostic(diag_period,
sim.fld,
None,
write_dir='proc%d' % sim.comm.rank))
if write_particles:
sim.diags.append(
ParticleDiagnostic(diag_period, {'electrons': sim.ptcl[0]}, sim.comm))
if __name__ == '__main__':
# Prevent current correction for MPI simulation
if sim.comm.size > 1:
correct_currents = False
else:
correct_currents = True
# Run the simulation
sim.step(N_step, correct_currents=correct_currents)
# Plot the fields
compare_fields(sim)