def init_fields(sim, w, ctau, k0, z0, zf, E0, m=1):
"""
Imprints the appropriate profile on the fields of the simulation.
Parameters
----------
sim: Simulation object from fbpic
w : float
The initial waist of the laser (in microns)
ctau : float
The initial temporal waist of the laser (in microns)
k0 : flat
The central wavevector of the laser (in microns^-1)
z0 : float
The position of the centroid on the z axis
zf : float
The position of the focal plane
E0 : float
The initial E0 of the pulse
m: int, optional
The mode on which to imprint the profile
For m = 1 : gaussian profile, linearly polarized beam
For m = 0 : annular profile, polarized in E_theta
"""
# Initialize the fields
a0 = E0 * e / (m_e * c**2 * k0)
tau = ctau / c
lambda0 = 2 * np.pi / k0
# Create the relevant laser profile
if m == 0:
# Build a radially-polarized pulse from 2 Laguerre-Gauss profiles
profile = LaguerreGaussLaser( 0, 1, 0.5*a0, w, tau, z0, zf=zf,
lambda0=lambda0, theta_pol=0., theta0=0. ) \
+ LaguerreGaussLaser( 0, 1, 0.5*a0, w, tau, z0, zf=zf,
lambda0=lambda0, theta_pol=np.pi/2, theta0=np.pi/2 )
elif m == 1:
profile = GaussianLaser(a0=a0,
waist=w,
tau=tau,
lambda0=lambda0,
z0=z0,
zf=zf)
elif m == 2:
profile = LaguerreGaussLaser(0,
1,
a0=a0,
waist=w,
tau=tau,
lambda0=lambda0,
z0=z0,
zf=zf)
# Add the profiles to the simulation
add_laser_pulse(sim, profile)
def add_laser(self, laser, injection_method):
# Call method of parent class
PICMI_Simulation.add_laser(self, laser, injection_method)
# Handle injection method
assert type(injection_method) == PICMI_LaserAntenna
# Handle laser profile method
if type(laser) == PICMI_GaussianLaser:
assert laser.propagation_direction[0] == 0.
assert laser.propagation_direction[1] == 0.
assert (laser.zeta is None) or (laser.zeta == 0)
assert (laser.beta is None) or (laser.beta == 0)
phi2_chirp = laser.phi2
if phi2_chirp is None:
phi2_chirp = 0
polarization_angle = np.arctan2(laser.polarization_direction[1],
laser.polarization_direction[0])
laser_profile = GaussianLaser(a0=laser.a0,
waist=laser.waist,
z0=laser.centroid_position[-1],
zf=laser.focal_position[-1],
tau=laser.duration,
theta_pol=polarization_angle,
phi2_chirp=phi2_chirp)
else:
raise ValueError('Unknown laser profile: %s' %
type(injection_method))
# Inject the laser
add_laser_pulse(self.fbpic_sim,
laser_profile,
method='antenna',
z0_antenna=injection_method.position[-1],
gamma_boost=self.gamma_boost)
def test_laser_periodic(show=False):
"""
Function that is run by py.test, when doing `python setup.py test`
Test the propagation of a laser in a periodic box.
"""
# Propagate the pulse in a single step
dt = zfoc * 1. / c
# Initialize the simulation object
sim = Simulation(Nz,
zmax,
Nr,
rmax,
Nm,
dt,
zmin=zmin,
boundaries={
'z': 'periodic',
'r': 'reflective'
})
# Initialize the laser fields
profile = FewCycleLaser(a0=a0, waist=w0, tau_fwhm=tau_fwhm, z0=0, zf=zfoc)
add_laser_pulse(sim, profile)
# Propagate the pulse
compare_fields(sim.fld.interp[1], sim.time, profile, show)
sim.step(1)
compare_fields(sim.fld.interp[1], sim.time, profile, show)
def test_laser_periodic(show=False):
"""
Function that is run by py.test, when doing `python setup.py test`
Test the propagation of a laser in a periodic box.
"""
# Propagate the pulse in a single step
dt = Lprop * 1. / c
# Initialize the simulation object
sim = Simulation(Nz,
zmax,
Nr,
rmax,
Nm,
dt,
n_order=n_order,
zmin=zmin,
boundaries={
'z': 'periodic',
'r': 'reflective'
})
# Initialize the laser fields
profile = FlattenedGaussianLaser(a0=a0,
w0=w0,
N=N,
tau=ctau / c,
z0=0,
zf=zfoc)
add_laser_pulse(sim, profile)
# Propagate the pulse
sim.step(1)
# Check the validity of the transverse field profile
# (Take the RMS field in order to suppress the laser oscillations)
trans_profile = np.sqrt(np.average(sim.fld.interp[1].Er.real**2, axis=0))
# Calculate the theortical profile out-of-focus
Zr = k0 * w0**2 / 2
w_th = w0 * (Lprop - zfoc) / Zr
r = sim.fld.interp[1].r
th_profile = trans_profile[0] * flat_gauss(r / w_th, N)
# Plot the profile, if requested by the user
if show:
import matplotlib.pyplot as plt
plt.plot(1.e6 * r, trans_profile, label='Simulated')
plt.plot(1.e6 * r, th_profile, label='Theoretical')
plt.legend(loc=0)
plt.xlabel('r (microns)')
plt.title('Transverse profile out-of-focus')
plt.tight_layout()
plt.show()
# Check the validity
assert np.allclose(th_profile, trans_profile, atol=rtol * th_profile[0])
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)
phi2_chirp=phi2_chirp)
trans_prof = FlattenedGaussianTransverseProfile(w0=w0, N=30, zf=zfoc)
elif case == 'donut_chirped':
long_prof = GaussianChirpedLongitudinalProfile(tau=ctau / c,
z0=0.,
phi2_chirp=phi2_chirp)
trans_prof = DonutLikeLaguerreGaussTransverseProfile(waist=w0,
zf=zfoc,
p=2,
m=1)
else:
raise ValueError('Unknown case')
# Construct Paraxial Approximation Laser
profile = ParaxialApproximationLaser(long_prof, trans_prof, E_laser)
# Initialize the laser fields
add_laser_pulse(sim, profile)
# Plot and compare with reference pulse
if reference_profile is not None:
r_2d, z_2d = np.meshgrid(sim.fld.interp[0].r,
sim.fld.interp[0].z,
indexing='ij')
Ex_reference = reference_profile.E_field(r_2d, 0, z_2d, 0)[0]
if plot:
import matplotlib.pyplot as plt
plt.subplot(211)
plt.imshow(2 * sim.fld.interp[1].Er.real.T)
plt.colorbar()
if reference_profile is not None:
plt.subplot(212)
plt.imshow(Ex_reference)
n_emit=0, #initial emittance is 0
gamma0=gamma0,
sig_gamma=sig_gamma,
n_physical_particles=Qtot / e,
n_macroparticles=5000,
zf=L0,
tf=(L0 - z0Beam) / c,
z_injection_plane=L0)
# Load initial fields
# Add a laser to the fields of the simulation
#add_laser( sim, a0, w0, ctau, z0 )
#add_laser_pulse(sim,PulseTrain_profile,method="antenna",z0_antenna=z0Antenna,v_antenna=0) #fw_propagating=True
add_laser_pulse(sim,
UniformPulseTrainN(NumberOfPulses, UniformPulseSpacing),
method="antenna",
z0_antenna=z0Antenna,
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
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()
# - Mode 0: Build a radially-polarized pulse from 2 Laguerre-Gauss profiles
profile0 = LaguerreGaussLaser( 0, 1, 0.5*a0, w0, tau, z0, zf=zf,
lambda0=lambda0, theta_pol=0., theta0=0. ) \
+ 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)
# NB: The code below is only executed when running the script,
# (`python lwfa_script.py`), but not when importing it (`import lwfa_script`).
if __name__ == '__main__':
# Initialize the simulation object
sim = Simulation( Nz, zmax, Nr, rmax, Nm, dt, n_e=None, zmin=zmin,
boundaries={"z":"open", "r":"reflective"}, n_order=n_order, use_cuda=use_cuda, verbose_level=2, )
# Create a Gaussian laser profile
laser_profile = GaussianLaser(a0=a0, waist=w0, tau=ctau / c, z0=z0,
zf=None, theta_pol=0., lambda0=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=-e, m=m_e, n=n_e,
dens_func=dens_func, p_zmin=p_zmin, p_zmax=p_zmax, p_rmax=p_rmax,
p_nz=p_nz, p_nr=p_nr, p_nt=p_nt )
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