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)
ctau,
z0,
lambda0=lambda0,
zf=zfoc,
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'],
# Add a laser to the fields of the simulation
add_laser(sim,
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,
def run_and_check_laser_antenna(gamma_b,
show,
write_files,
z0,
v=0,
forward_propagating=True):
"""
Generic function, which runs and check the laser antenna for
both boosted frame and lab frame
Parameters
----------
gamma_b: float or None
The Lorentz factor of the boosted frame
show: bool
Whether to show the images of the laser as pop-up windows
write_files: bool
Whether to output openPMD data of the laser
v: float (m/s)
Speed of the laser antenna
"""
# Initialize the simulation object
sim = Simulation(Nz,
zmax,
Nr,
rmax,
Nm,
dt,
p_zmin=0,
p_zmax=0,
p_rmin=0,
p_rmax=0,
p_nz=2,
p_nr=2,
p_nt=2,
n_e=0.,
zmin=zmin,
use_cuda=use_cuda,
boundaries='open',
gamma_boost=gamma_b)
# Remove the particles
sim.ptcl = []
# Add the laser
add_laser(sim,
a0,
w0,
ctau,
z0,
zf=zf,
method='antenna',
z0_antenna=z0_antenna,
v_antenna=v,
gamma_boost=gamma_b,
fw_propagating=forward_propagating)
# Calculate the number of steps between each output
N_step = int(round(Ntot_step / N_show))
# Add diagnostic
if write_files:
sim.diags = [
FieldDiagnostic(N_step,
sim.fld,
comm=None,
fieldtypes=["rho", "E", "B", "J"])
]
# Loop over the iterations
print('Running the simulation...')
for it in range(N_show):
print('Diagnostic point %d/%d' % (it, N_show))
# Advance the Maxwell equations
sim.step(N_step, show_progress=False)
# Plot the fields during the simulation
if show == True:
show_fields(sim.fld.interp[1], 'Er')
# Finish the remaining iterations
sim.step(Ntot_step - N_show * N_step, show_progress=False)
# Check the transverse E and B field
Nz_half = int(sim.fld.interp[1].Nz / 2) + 2
z = sim.fld.interp[1].z[Nz_half:-(sim.comm.n_guard+sim.comm.n_damp+\
sim.comm.n_inject)]
r = sim.fld.interp[1].r
# Loop through the different fields
for fieldtype, info_in_real_part, factor in [ ('Er', True, 2.), \
('Et', False, 2.), ('Br', False, 2.*c), ('Bt', True, 2.*c) ]:
# factor correspond to the factor that has to be applied
# in order to get a value which is comparable to an electric field
# (Because of the definition of the interpolation grid, the )
field = getattr(sim.fld.interp[1], fieldtype)\
[Nz_half:-(sim.comm.n_guard+sim.comm.n_damp+\
sim.comm.n_inject)]
print('Checking %s' % fieldtype)
check_fields(factor * field, z, r, info_in_real_part, z0, gamma_b,
forward_propagating)
print('OK')
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')
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 test_cpu_gpu_deposition(show=False):
"Function that is run by py.test, when doing `python setup.py test`"
# 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,
p_zmin,
p_zmax,
p_rmin,
p_rmax,
p_nz,
p_nr,
p_nt,
n_e,
zmin=zmin,
use_cuda=use_cuda)
# Tweak the velocity of the electron bunch
gamma0 = 10.
sim.ptcl[0].ux = sim.ptcl[0].x
sim.ptcl[0].uy = sim.ptcl[0].y
sim.ptcl[0].uz = np.sqrt(gamma0**2 - sim.ptcl[0].ux**2 -
sim.ptcl[0].uy**2 - 1)
sim.ptcl[0].inv_gamma = 1. / gamma0 * np.ones_like(sim.ptcl[0].x)
sim.ptcl[0].m = 1e10
# 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')
p_zmin, p_zmax, p_rmin, p_rmax, p_nz, p_nr, p_nt, n_e,
dens_func=dens_func, zmin=zmin, boundaries='open',
use_cuda=use_cuda )
# 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 a field diagnostic
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('')
sim = Simulation(Nz,
zmax,
Nr,
rmax,
Nm,
dt,
0,
0,
0,
0,
2,
2,
4,
0.,
zmin=zmin,
n_order=n_order,
boundaries={
'z': 'open',
'r': 'reflective'
})
# Configure the moving window
sim.set_moving_window(v=c)
# Suppress the particles that were intialized by default and add the bunch
sim.ptcl = []
add_elec_bunch_gaussian(sim, sig_r, sig_z, n_emit, gamma0, sig_gamma, Q, N, tf,
zf)
# Set the diagnostics
sim.diags = [FieldDiagnostic(10, sim.fld, comm=sim.comm)]
# Perform one simulation step (essentially in order to write the diags)
sim.step(1)
add_laser(sim, a0, w0, ctau, z0)
if use_restart is True:
# 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 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
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)
v_antenna=0) #fw_propagating=True
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={
"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:
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()
+ LaguerreGaussLaser( 0, 1, 0.5*a0, w0, tau, z0, zf=zf,
lambda0=lambda0, theta_pol=np.pi/2, theta0=np.pi/2 )
# - Mode 1: Use a regular linearly-polarized pulse
profile1 = GaussianLaser(a0=a0,
waist=w0,
tau=tau,
lambda0=lambda0,
z0=z0,
zf=zf)
if not restart:
# Add the profiles to the simulation
add_laser_pulse(sim, profile0)
add_laser_pulse(sim, profile1)
else:
restart_from_checkpoint(sim)
# Calculate the total number of steps
N_step = int(round(L_prop / (c * dt)))
diag_period = int(round(N_step / N_diag))
# Add openPMD diagnostics
sim.diags = [
FieldDiagnostic(diag_period, sim.fld, fieldtypes=["E"], comm=sim.comm)
]
set_periodic_checkpoint(sim, N_step // 2)
# Do only half the steps
sim.step(N_step // 2 + 1)
p_nr,
p_nt,
n_e,
use_cuda=use_cuda,
n_guard=n_guard,
boundaries='open')
# 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