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 run_continuous_injection( gamma_boost, dens_func,
p_zmin, p_zmax, show, N_check=3 ):
# Chose the time step
dt = (zmax-zmin)/Nz/c
# Initialize the different structures
sim = Simulation( Nz, zmax, Nr, rmax, Nm, dt,
p_zmin, p_zmax, 0, p_rmax, p_nz, p_nr, p_nt, 0.5*n,
dens_func=dens_func, initialize_ions=False, zmin=zmin,
use_cuda=use_cuda, gamma_boost=gamma_boost, boundaries='open' )
# Add another species with a different number of particles per cell
sim.add_new_species( -e, m_e, 0.5*n, dens_func,
2*p_nz, 2*p_nr, 2*p_nt,
p_zmin, p_zmax, 0, p_rmax )
# Set the moving window, which handles the continuous injection
# The moving window has an arbitrary velocity (0.7*c) so as to check
# that the injection is correct in this case also
sim.set_moving_window( v=c )
# Check that the density is correct after different timesteps
N_step = int( Nz/N_check/2 )
for i in range( N_check ):
sim.step( N_step, move_momenta=False )
check_density( sim, gamma_boost, dens_func, show )
def uniform_electron_plasma(shape, show=False):
# Initialize the different structures
sim = Simulation( Nz, zmax, Nr, rmax, Nm, zmax/Nz/c,
0, zmax, 0, p_rmax, p_nz, p_nr, p_nt, n,
initialize_ions=False, particle_shape=shape )
# Deposit the charge
sim.fld.erase('rho')
for species in sim.ptcl :
species.deposit( sim.fld, 'rho')
sim.fld.divide_by_volume('rho')
# Check that the density has the correct value
if show is False:
# Integer index at which the plasma stops
Nrmax = int( Nr * p_rmax * 1./rmax )
# Check that the density is correct in mode 0, below this index
assert np.allclose( -n*e,
sim.fld.interp[0].rho[:,:Nrmax-2], 1.e-10 )
# Check that the density is correct in mode 0, above this index
assert np.allclose( 0,
sim.fld.interp[0].rho[:,Nrmax+2:], 1.e-10 )
# Check the density in the mode 1
assert np.allclose( 0,
sim.fld.interp[1].rho[:,:], 1.e-10 )
# Show the results
else:
plt.title('Uniform plasma, mode 0')
sim.fld.interp[0].show('rho')
plt.show()
plt.title('Uniform plasma, mode 1')
sim.fld.interp[1].show('rho')
plt.show()
def neutral_plasma_shifted(shape, show=False):
# Initialize the different structures
sim = Simulation(Nz,
zmax,
Nr,
rmax,
Nm,
zmax / Nz / c,
0,
zmax,
0,
p_rmax,
p_nz,
p_nr,
p_nt,
n,
initialize_ions=True,
particle_shape=shape)
# Shift the electrons
dr = rmax / Nr
sim.ptcl[0].x += frac_shift * dr
# Deposit the charge
sim.fld.erase('rho')
for species in sim.ptcl:
species.deposit(sim.fld, 'rho')
sim.fld.sum_reduce_deposition_array('rho')
sim.fld.divide_by_volume('rho')
# Check that the density has the correct value
if show is False:
# Integer index at which the plasma stops
Nrmax = int(Nr * p_rmax * 1. / rmax)
# Check that the density is correct in mode 0, below this index
assert np.allclose(0,
sim.fld.interp[0].rho[:, :Nrmax - 2],
atol=n * e * 1.e-3)
assert np.allclose(0,
sim.fld.interp[1].rho[:, :Nrmax - 2],
atol=n * e * 1.e-3)
# Check that the density is correct in mode 0, above this index
assert np.allclose(0, sim.fld.interp[0].rho[:, Nrmax + 2:], 1.e-10)
# Check the density in the mode 1
assert np.allclose(0,
sim.fld.interp[1].rho[:, Nrmax + 2:],
atol=n * e * 1.e-10)
# Show the results
else:
import matplotlib.pyplot as plt
plt.title('Shifted plasma, mode 0')
plt.imshow(sim.fld.interp[0].rho.real, aspect='auto')
plt.colorbar()
plt.show()
plt.title('Shifted plasma, mode 1')
plt.imshow(sim.fld.interp[1].rho.real, aspect='auto')
plt.colorbar()
plt.show()
def init(self, kw):
# Get the grid
grid = self.solver.grid
if not type(grid) == PICMI_CylindricalGrid:
raise ValueError('When using fbpic with PICMI, '
'the grid needs to be a CylindricalGrid object.')
# Check rmin and boundary conditions
assert grid.rmin == 0.
assert grid.bc_zmin == grid.bc_zmax
assert grid.bc_zmax in ['periodic', 'open']
assert grid.bc_rmax in ['reflective', 'open']
# Determine timestep
if self.solver.cfl is not None:
dz = (grid.zmax - grid.zmin) / grid.nz
dt = self.solver.cfl * dz / c
elif self.time_step_size is not None:
dt = self.time_step_size
else:
raise ValueError('You need to either set the `cfl` of the solver\n'
'or the `timestep_size` of the `Simulation`.')
# Convert API for the smoother
if self.solver.source_smoother is None:
smoother = BinomialSmoother()
else:
smoother = BinomialSmoother(
n_passes=self.solver.source_smoother.n_pass,
compensator=self.solver.source_smoother.compensation)
# Order of the stencil for z derivatives in the Maxwell solver
if self.solver.stencil_order is None:
n_order = -1
else:
n_order = self.solver.stencil_order[-1]
# Initialize and store the FBPIC simulation object
self.fbpic_sim = FBPICSimulation(Nz=int(grid.nz),
zmin=grid.zmin,
zmax=grid.zmax,
Nr=int(grid.nr),
rmax=grid.rmax,
Nm=grid.n_azimuthal_modes,
dt=dt,
use_cuda=True,
smoother=smoother,
n_order=n_order,
boundaries={
'z': grid.bc_zmax,
'r': grid.bc_rmax
})
# Set the moving window
if grid.moving_window_velocity is not None:
self.fbpic_sim.set_moving_window(grid.moving_window_velocity[-1])
def simulate_periodic_plasma_wave( particle_shape, show=False ):
"Simulate a periodic plasma wave and check its fields"
# Initialization of 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, n_order=n_order, use_cuda=use_cuda,
particle_shape=particle_shape )
# Save the initial density in spectral space, and consider it
# to be the density of the (uninitialized) ions
sim.deposit('rho_prev', exchange=True)
sim.fld.spect2interp('rho_prev')
rho_ions = [ ]
for m in range(len(sim.fld.interp)):
rho_ions.append( -sim.fld.interp[m].rho.copy() )
# Impart velocities to the electrons
# (The electrons are initially homogeneous, but have an
# intial non-zero velocity that develops into a plasma wave)
impart_momenta( sim.ptcl[0], epsilons, k0, w0, wp )
# Run the simulation
sim.step( N_step, correct_currents=True )
# Plot the results and compare with analytical theory
compare_fields( sim, show )
# Test check that div(E) - rho = 0 (directly in spectral space)
check_charge_conservation( sim, rho_ions )
def uniform_electron_plasma(shape, show=False):
# Initialize the different structures
sim = Simulation(Nz,
zmax,
Nr,
rmax,
Nm,
zmax / Nz / c,
0,
zmax,
0,
p_rmax,
p_nz,
p_nr,
p_nt,
n,
initialize_ions=False,
particle_shape=shape)
# Deposit the charge
# (Move the simulation to GPU if needed, for this step)
with GpuMemoryManager(sim):
sim.fld.erase('rho')
for species in sim.ptcl:
species.deposit(sim.fld, 'rho')
sim.fld.sum_reduce_deposition_array('rho')
sim.fld.divide_by_volume('rho')
# Check that the density has the correct value
if show is False:
# Integer index at which the plasma stops
Nrmax = int(Nr * p_rmax * 1. / rmax)
# Check that the density is correct in mode 0, below this index
assert np.allclose(-n * e, sim.fld.interp[0].rho[:, :Nrmax - 2], 2.e-3)
# Check that the density is correct in mode 0, above this index
assert np.allclose(0, sim.fld.interp[0].rho[:, Nrmax + 2:], 1.e-10)
# Check the density in the mode 1
assert np.allclose(0, sim.fld.interp[1].rho[:, :], 1.e-10)
# Show the results
else:
import matplotlib.pyplot as plt
plt.title('Uniform electron plasma, mode 0')
plt.imshow(sim.fld.interp[0].rho.real, aspect='auto')
plt.colorbar()
plt.show()
plt.title('Uniform electron plasma, mode 1')
plt.imshow(sim.fld.interp[1].rho.real, aspect='auto')
plt.colorbar()
plt.show()
def test_periodic_plasma_wave(show=False):
"Function that is run by py.test, when doing `python setup.py test`"
# Initialization of 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,
n_order=n_order,
use_cuda=use_cuda)
# Impart velocities to the electrons
# (The electrons are initially homogeneous, but have an
# intial non-zero velocity that develops into a plasma wave)
impart_momenta(sim.ptcl[0], epsilon, k0, w0, wp)
# Choose whether to correct the currents
if sim.comm.size == 1:
correct_currents = True
else:
correct_currents = False
# Run the simulation
sim.step(N_step, correct_currents=correct_currents)
# Plot the results
compare_fields(sim, show)
def charge_cylinder(shape, show=False):
"On-axis cylinder of charge for different radii"
# Initialize the simulation object
sim = Simulation( Nz, zmax, Nr, rmax, Nm, (zmax-zmin)/Nz/c,
p_zmin, p_zmax, p_rmin, p_rmax, p_nz, p_nr, p_nt, n_e,
zmin=zmin, boundaries={'z':'periodic', 'r':'reflective'}, verbose_level=0,
smoother=BinomialSmoother(1, False), particle_shape=shape)
# store results in dict
res = {}
# Scale the radius of the cylinder and calculate the space charge field
for i, scale in enumerate(scales):
res[scale] = calculate_fields(sim, scale)
if show is False:
for i, scale in enumerate(scales):
r, Er, Er_theory, rho_n = res[scale]
# Check that the density is correct in mode 0, below this index
assert np.allclose( (-Er*r)[-5:], (Er_theory*r)[-5:], 1.e-3 )
else:
import matplotlib.pyplot as plt
import matplotlib
# Segmented colormap
cmap = matplotlib.cm.get_cmap('YlGnBu')
colors = np.array([co for co in cmap(np.linspace(0.2,0.8,7))])[::-1]
# Plot charge density
plt.title('Cylinder charge density')
for i, scale in enumerate(scales):
r, Er, Er_theory, rho_n = res[scale]
plt.plot(r*1.e6, rho_n/(n_e*e), label=str(scale), color=colors[i])
plt.legend()
plt.xlim(0, 1.25)
plt.ylabel(r'$\rho_n$')
plt.xlabel(r'$r$')
plt.show()
# Plot simulated and analytically calculated field
plt.title('Cylinder space charge field')
for i, scale in enumerate(scales):
r, Er, Er_theory, rho_n = res[scale]
E0 = (n_e*e*p_rmax**2/(2*epsilon_0))
plt.plot(r*1.e6, -Er*r/E0, label=str(scale), color=colors[i])
plt.plot(r*1.e6, Er_theory*r/E0, color=colors[i], ls='--')
plt.legend()
plt.xlim(0, 1.25)
plt.ylabel(r'$-E_{r,norm} \times r$')
plt.xlabel(r'$r$')
plt.show()
def simulate_beam_focusing( z_injection_plane, write_dir ):
"""
Simulate a focusing beam in the boosted frame
Parameters
----------
z_injection_plane: float or None
when this is not None, the injection through a plane is
activated.
write_dir: string
The directory where the boosted diagnostics are written.
"""
# Initialize the simulation object
sim = Simulation( Nz, zmax, Nr, rmax, Nm, dt, zmin=zmin,
gamma_boost=gamma_boost, boundaries='open',
use_cuda=use_cuda, v_comoving=v_comoving )
# Note: no macroparticles get created because we do not pass
# the density and number of particle per cell
# Remove the plasma particles
sim.ptcl = []
# Initialize the bunch, along with its space charge
add_elec_bunch_gaussian( sim, sigma_r, sigma_z, n_emit, gamma0,
sigma_gamma, Q, N, tf=(z_focus-z0)/c, zf=z_focus, boost=boost,
z_injection_plane=z_injection_plane )
# Configure the moving window
sim.set_moving_window( v=c )
# Add a field diagnostic
sim.diags = [
BackTransformedParticleDiagnostic( zmin, zmax, c, dt_snapshot_lab,
Ntot_snapshot_lab, gamma_boost, period=100, fldobject=sim.fld,
species={'bunch':sim.ptcl[0]}, comm=sim.comm, write_dir=write_dir)
]
# Run the simulation
sim.step( N_step )
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/')
p_rmax = 5.e-6 # Maximal radial position of the bunch (meters)
n_e = 4.e18 * 1.e6 # Density (electrons.meters^-3)
p_nz = 2 # Number of particles per cell along z
p_nr = 2 # Number of particles per cell along r
p_nt = 4 # Number of particles per cell along theta
# 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,
n_order=n_order,
v_comoving=-0.999999 * c,
use_galilean=True)
# Configure the moving window
#sim.moving_win = MovingWindow( sim.fld.interp[0],
# ncells_damp=2,
# ncells_zero=2 )
# Suppress the particles that were intialized by default and add the bunch
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')
n = np.where(z < 0, 0., n)
return (n)
# ---------------------------
# Carrying out the simulation
# ---------------------------
if __name__ == '__main__':
# Initialize the simulation object
sim = Simulation(Nz,
zmax,
Nr,
rmax,
Nm,
dt,
zmin=zmin,
boundaries='open',
initialize_ions=False,
n_order=n_order,
use_cuda=use_cuda)
# By default the simulation initializes an electron species (sim.ptcl[0])
# Because we did not pass the arguments `n`, `p_nz`, `p_nr`, `p_nz`,
# this electron species does not contain any macroparticles.
# It is okay to just remove it from the list of species.
sim.ptcl = []
# Add the Helium ions (pre-ionized up to level 1),
# the Nitrogen ions (pre-ionized up to level 5)
# and the associated electrons (from the pre-ionized levels)
atoms_He = sim.add_new_species(q=e,
n = np.where(z < 0, 0., n)
return (n)
# ---------------------------
# Carrying out the simulation
# ---------------------------
if __name__ == '__main__':
# Initialize the simulation object
sim = Simulation(Nz,
zmax,
Nr,
rmax,
Nm,
dt,
zmin=zmin,
boundaries='open',
initialize_ions=False,
n_order=n_order,
use_cuda=use_cuda)
# By default the simulation initializes an electron species (sim.ptcl[0])
# Because we did not pass the arguments `n`, `p_nz`, `p_nr`, `p_nz`,
# this electron species does not contain any macroparticles.
# It is okay to just remove it from the list of species.
sim.ptcl = []
# Add the Helium ions (pre-ionized up to level 1),
# the Nitrogen ions (pre-ionized up to level 5)
# and the associated electrons (from the pre-ionized levels)
atoms_He = sim.add_new_species(q=e,
def propagate_pulse(Nz,
Nr,
Nm,
zmin,
zmax,
Lr,
L_prop,
zf,
dt,
N_diag,
w0,
ctau,
k0,
E0,
m,
N_show,
n_order,
rtol,
boundaries,
v_window=0,
use_galilean=False,
v_comoving=0,
show=False):
"""
Propagate the beam over a distance L_prop in Nt steps,
and extracts the waist and a0 at each step.
Parameters
----------
show : bool
Wether to show the fields, so that the user can manually check
the agreement with the theory.
If True, this will periodically show the map of the fields (with
a period N_show), as well as (eventually) the evoluation of a0 and w0.
If False, this
N_diag : int
Number of diagnostic points (i.e. measure of waist and a0)
along the propagation
Nz, Nr : int
The number of points on the grid in z and r respectively
Nm : int
The number of modes in the azimuthal direction
zmin, zmax : float
The limits of the box in z
Lr : float
The size of the box in the r direction
(In the case of Lr, this is the distance from the *axis*
to the outer boundary)
L_prop : float
The total propagation distance (in meters)
zf : float
The position of the focal plane of the laser (only works for m=1)
dt : float
The timestep of the simulation
w0 : float
The initial waist of the laser (in meters)
ctau : float
The initial temporal waist of the laser (in meters)
k0 : flat
The central wavevector of the laser (in meters^-1)
E0 : float
The initial E0 of the pulse
m : int
Index of the mode to be tested
For m = 1 : test with a gaussian, linearly polarized beam
For m = 0 : test with an annular beam, polarized in E_theta
n_order : int
Order of the stencil
rtol : float
Relative precision with which the results are tested
boundaries : string
Type of boundary condition
Either 'open' or 'periodic'
v_window : float
Speed of the moving window
v_comoving : float
Velocity at which the currents are assumed to move
use_galilean: bool
Whether to use a galilean frame that moves at the speed v_comoving
Returns
-------
A dictionary containing :
- 'E' : 1d array containing the values of the amplitude
- 'w' : 1d array containing the values of waist
- 'fld' : the Fields object at the end of the simulation.
"""
# Initialize the simulation object
sim = Simulation(Nz,
zmax,
Nr,
Lr,
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.,
n_order=n_order,
zmin=zmin,
use_cuda=use_cuda,
boundaries=boundaries,
v_comoving=v_comoving,
exchange_period=1,
use_galilean=use_galilean)
# Remove the particles
sim.ptcl = []
# Set the moving window object
if v_window != 0:
sim.set_moving_window(v=v_window)
# Initialize the laser fields
z0 = (zmax + zmin) / 2
init_fields(sim, w0, ctau, k0, z0, zf, E0, m)
# Create the arrays to get the waist and amplitude
w = np.zeros(N_diag)
E = np.zeros(N_diag)
# Calculate the number of steps to run between each diagnostic
Ntot_step = int(round(L_prop / (c * dt)))
N_step = int(round(Ntot_step / N_diag))
# Loop over the iterations
print('Running the simulation...')
for it in range(N_diag):
print('Diagnostic point %d/%d' % (it, N_diag))
# Fit the fields to find the waist and a0
w[it], E[it] = fit_fields(sim.fld, m)
# Plot the fields during the simulation
if show == True and it % N_show == 0:
import matplotlib.pyplot as plt
plt.clf()
sim.fld.interp[m].show('Er')
plt.show()
# Advance the Maxwell equations
sim.step(N_step, show_progress=False)
# Get the analytical solution
z_prop = c * dt * N_step * np.arange(N_diag)
ZR = 0.5 * k0 * w0**2
w_analytic = w0 * np.sqrt(1 + (z_prop - zf)**2 / ZR**2)
E_analytic = E0 / (1 + (z_prop - zf)**2 / ZR**2)**(1. / 2)
# Either plot the results and check them manually
if show is True:
import matplotlib.pyplot as plt
plt.suptitle('Diffraction of a pulse in the mode %d' % m)
plt.subplot(121)
plt.plot(1.e6 * z_prop, 1.e6 * w, 'o', label='Simulation')
plt.plot(1.e6 * z_prop, 1.e6 * w_analytic, '--', label='Theory')
plt.xlabel('z (microns)')
plt.ylabel('w (microns)')
plt.title('Waist')
plt.legend(loc=0)
plt.subplot(122)
plt.plot(1.e6 * z_prop, E, 'o', label='Simulation')
plt.plot(1.e6 * z_prop, E_analytic, '--', label='Theory')
plt.xlabel('z (microns)')
plt.ylabel('E')
plt.legend(loc=0)
plt.title('Amplitude')
plt.show()
# or automatically check that the theoretical and simulated curves
# of w and E are close
else:
assert np.allclose(w, w_analytic, rtol=rtol)
assert np.allclose(E, E_analytic, rtol=5.e-3)
print('The simulation results agree with the theory to %e.' % rtol)
# Return a dictionary of the results
return ({'E': E, 'w': w, 'fld': sim.fld})
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()
# NB: The code below is only executed when running the script,
# (`python boosted_frame_sim.py`), but not when importing it.
if __name__ == '__main__':
# 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,
dens_func=dens_func,
zmin=zmin,
initialize_ions=True,
v_comoving=v_comoving,
gamma_boost=gamma_boost,
n_order=n_order,
boundaries='open',
use_cuda=use_cuda)
# Add an electron bunch
add_elec_bunch(sim,
bunch_gamma,
bunch_n,
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')
zf = -20.e-6
# The simulation timestep
dt = (zmax - zmin) / Nz / c # Timestep (seconds)
# Initialize the simulation object
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)
# ---------------------------
# NB: The code below is only executed when running the script,
# (`python -i lpa_sim.py`), but not when importing it (`import lpa_sim`).
if __name__ == '__main__':
# Initialize the simulation object
# Parametric scan: use the flag `use_all_mpi_ranks=False` to
# have each MPI rank run an independent simulation
sim = Simulation(Nz,
zmax,
Nr,
rmax,
Nm,
dt,
zmin=zmin,
boundaries={
'z': 'open',
'r': 'reflective'
},
n_order=n_order,
use_cuda=use_cuda,
use_all_mpi_ranks=False)
# 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,
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')
class Simulation(PICMI_Simulation):
# Redefine the `init` method, as required by the picmi `_ClassWithInit`
def init(self, kw):
self.sim_kw = {}
for argname in ['use_ruyten_shapes', 'use_modified_volume']:
if f'fbpic_{argname}' in kw:
self.sim_kw[argname] = kw.pop(f'fbpic_{argname}')
self.step_kw = {}
for argname in [
'correct_currents', 'correct_divE', 'use_true_rho',
'move_positions', 'move_momenta', 'show_progress'
]:
if f'fbpic_{argname}' in kw:
self.step_kw[argname] = kw.pop(f'fbpic_{argname}')
# Get the grid
grid = self.solver.grid
if not type(grid) == PICMI_CylindricalGrid:
raise ValueError('When using fbpic with PICMI, '
'the grid needs to be a CylindricalGrid object.')
# Check rmin and boundary conditions
assert grid.lower_bound[0] == 0.
assert grid.lower_boundary_conditions[
1] == grid.upper_boundary_conditions[1]
if grid.lower_boundary_conditions[1] == 'reflective':
warnings.warn(
"FBPIC does not support reflective boundary condition in z.\n"
"The z boundary condition was automatically converted to 'open'."
)
grid.lower_boundary_conditions[1] = 'open'
grid.upper_boundary_conditions[1] = 'open'
assert grid.upper_boundary_conditions[1] in ['periodic', 'open']
assert grid.upper_boundary_conditions[0] in ['reflective', 'open']
# Determine timestep
if self.solver.cfl is not None:
dz = (grid.upper_bound[1] -
grid.lower_bound[1]) / grid.number_of_cells[1]
dr = (grid.upper_bound[0] -
grid.lower_bound[0]) / grid.number_of_cells[0]
if self.gamma_boost is not None:
beta = np.sqrt(1. - 1. / self.gamma_boost**2)
dr = dr / ((1 + beta) * self.gamma_boost)
dt = self.solver.cfl * min(dz, dr) / c
elif self.time_step_size is not None:
dt = self.time_step_size
else:
raise ValueError('You need to either set the `cfl` of the solver\n'
'or the `timestep_size` of the `Simulation`.')
# Convert API for the smoother
if self.solver.source_smoother is None:
smoother = BinomialSmoother()
else:
if self.solver.source_smoother.n_pass is None:
n_passes = 1
else:
n_passes = {
'r': self.solver.source_smoother.n_pass[0],
'z': self.solver.source_smoother.n_pass[1]
}
if self.solver.source_smoother.compensation is None:
compensator = False
else:
compensator = all(self.solver.source_smoother.compensation)
smoother = BinomialSmoother(n_passes=n_passes,
compensator=compensator)
# Convert verbose level:
verbose_level = self.verbose
if verbose_level is None:
verbose_level = 1
# Order of the stencil for z derivatives in the Maxwell solver
if self.solver.stencil_order is None:
n_order = -1
else:
n_order = self.solver.stencil_order[-1]
# Number of guard cells
if grid.guard_cells is None:
n_guard = None
else:
n_guard = grid.guard_cells[-1]
if self.solver.galilean_velocity is None:
v_comoving = None
else:
v_comoving = self.solver.galilean_velocity[-1]
# Initialize and store the FBPIC simulation object
self.fbpic_sim = FBPICSimulation(Nz=int(grid.number_of_cells[1]),
zmin=grid.lower_bound[1],
zmax=grid.upper_bound[1],
Nr=int(grid.number_of_cells[0]),
rmax=grid.upper_bound[0],
Nm=grid.n_azimuthal_modes,
dt=dt,
use_cuda=True,
smoother=smoother,
n_order=n_order,
boundaries={
'z':
grid.upper_boundary_conditions[1],
'r':
grid.upper_boundary_conditions[0]
},
n_guard=n_guard,
verbose_level=verbose_level,
particle_shape=self.particle_shape,
v_comoving=v_comoving,
gamma_boost=self.gamma_boost,
**self.sim_kw)
# Set the moving window
if grid.moving_window_velocity is not None:
self.fbpic_sim.set_moving_window(grid.moving_window_velocity[-1])
# Redefine the method `add_laser` from the PICMI Simulation class
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)
# Redefine the method `add_species` from the PICMI Simulation class
def add_species(self, species, layout, initialize_self_field=False):
# Call method of parent class
PICMI_Simulation.add_species(self, species, layout,
initialize_self_field)
# Call generic method internally
self._add_species_generic(species,
layout,
injection_plane_position=None,
injection_plane_normal_vector=None,
initialize_self_field=initialize_self_field)
def add_species_through_plane(self,
species,
layout,
injection_plane_position,
injection_plane_normal_vector,
initialize_self_field=False):
# Call method of parent class
PICMI_Simulation.add_species_through_plane(
self,
species,
layout,
injection_plane_position,
injection_plane_normal_vector,
initialize_self_field=initialize_self_field)
# Call generic method internally
self._add_species_generic(
species,
layout,
injection_plane_position=injection_plane_position,
injection_plane_normal_vector=injection_plane_normal_vector,
initialize_self_field=initialize_self_field)
def _add_species_generic(self, species, layout, injection_plane_position,
injection_plane_normal_vector,
initialize_self_field):
# Extract list of species
if type(species) == PICMI_Species:
species_instances_list = [species]
elif type(species) == PICMI_MultiSpecies:
species_instances_list = species.species_instances_list
else:
raise ValueError('Unknown type: %s' % type(species))
# Loop over species and create FBPIC species
for s in species_instances_list:
# Get their charge and mass
if s.particle_type is not None:
s.charge = particle_charge[s.particle_type]
s.mass = particle_mass[s.particle_type]
# If `charge_state` is set, redefine the charge and mass
if s.charge_state is not None:
s.charge = s.charge_state * e
s.mass -= s.charge_state * m_e
# Add the species to the FBPIC simulation
fbpic_species = self._create_new_fbpic_species(
s, layout, injection_plane_position,
injection_plane_normal_vector, initialize_self_field)
# Register a pointer to the FBPIC species in the PICMI species itself
# (Useful for particle diagnostics later on)
s.fbpic_species = fbpic_species
# Loop over species and handle ionization
for s in species_instances_list:
for interaction in s.interactions:
assert interaction[0] == 'ionization'
assert interaction[1] == 'ADK'
picmi_target = interaction[2]
if not hasattr(picmi_target, 'fbpic_species'):
raise RuntimeError(
'For ionization with PICMI+FBPIC:\n'
'You need to add the target species to the simulation,'
' before the other species.')
fbpic_target = picmi_target.fbpic_species
fbpic_source = s.fbpic_species
fbpic_source.make_ionizable(element=s.particle_type,
level_start=s.charge_state,
target_species=fbpic_target)
def _create_new_fbpic_species(self, s, layout, injection_plane_position,
injection_plane_normal_vector,
initialize_self_field):
# - For the case of a plasma/beam defined in a gridded layout
if type(layout) == PICMI_GriddedLayout:
# - Uniform distribution
if type(s.initial_distribution) == PICMI_UniformDistribution:
n0 = s.initial_distribution.density
dens_func = None
# - Analytic distribution
elif type(s.initial_distribution) == PICMI_AnalyticDistribution:
import numexpr
density_expression = s.initial_distribution.density_expression
if s.density_scale is not None:
n0 = s.density_scale
else:
n0 = 1.
def dens_func(x, y, z):
d = locals()
d.update(s.initial_distribution.user_defined_kw)
n = numexpr.evaluate(density_expression, local_dict=d)
return n
else:
raise ValueError(
'Unknown combination of layout and distribution')
p_nr = layout.n_macroparticle_per_cell[0]
p_nt = layout.n_macroparticle_per_cell[1]
p_nz = layout.n_macroparticle_per_cell[2]
if initialize_self_field or (injection_plane_position is not None):
assert s.initial_distribution.fill_in != True
if injection_plane_position is None:
z_injection_plane = None
else:
z_injection_plane = injection_plane_position[-1]
gamma0_beta0 = s.initial_distribution.directed_velocity[-1] / c
gamma0 = (1 + gamma0_beta0**2)**.5
dist = s.initial_distribution
fbpic_species = add_particle_bunch(
self.fbpic_sim,
q=s.charge,
m=s.mass,
gamma0=gamma0,
n=n0,
dens_func=dens_func,
p_nz=p_nz,
p_nr=p_nr,
p_nt=p_nt,
p_zmin=dist.lower_bound[-1]
if dist.lower_bound[-1] is not None else -np.inf,
p_zmax=dist.upper_bound[-1]
if dist.upper_bound[-1] is not None else +np.inf,
p_rmin=0,
p_rmax=dist.upper_bound[0]
if dist.upper_bound[0] is not None else +np.inf,
boost=self.fbpic_sim.boost,
z_injection_plane=z_injection_plane,
initialize_self_field=initialize_self_field,
boost_positions_in_dens_func=True)
else:
dist = s.initial_distribution
fbpic_species = self.fbpic_sim.add_new_species(
q=s.charge,
m=s.mass,
n=n0,
dens_func=dens_func,
p_nz=p_nz,
p_nr=p_nr,
p_nt=p_nt,
p_zmin=dist.lower_bound[-1]
if dist.lower_bound[-1] is not None else -np.inf,
p_zmax=dist.upper_bound[-1]
if dist.upper_bound[-1] is not None else +np.inf,
p_rmax=dist.upper_bound[0]
if dist.upper_bound[0] is not None else +np.inf,
continuous_injection=s.initial_distribution.fill_in,
boost_positions_in_dens_func=True)
# - For the case of a Gaussian beam
elif (type(s.initial_distribution)==PICMI_GaussianBunchDistribution) \
and (type(layout) == PICMI_PseudoRandomLayout):
dist = s.initial_distribution
gamma0_beta0 = dist.centroid_velocity[-1] / c
gamma0 = (1 + gamma0_beta0**2)**.5
sig_r = dist.rms_bunch_size[0]
sig_z = dist.rms_bunch_size[-1]
sig_gamma = dist.rms_velocity[-1] / c
sig_vr = dist.rms_velocity[0] / gamma0
if sig_vr != 0:
tf = -sig_r**2 / sig_vr**2 * dist.velocity_divergence[0]
else:
tf = 0.
zf = dist.centroid_position[-1] + \
dist.centroid_velocity[-1]/gamma0 * tf
# Calculate size at focus and emittance
sig_r0 = (sig_r**2 - (sig_vr * tf)**2)**0.5
n_emit = gamma0 * sig_r0 * sig_vr / c
# Get the number of physical particles
n_physical_particles = dist.n_physical_particles
if s.density_scale is not None:
n_physical_particles *= s.density_scale
fbpic_species = add_particle_bunch_gaussian(
self.fbpic_sim,
q=s.charge,
m=s.mass,
gamma0=gamma0,
sig_gamma=sig_gamma,
sig_r=sig_r0,
sig_z=sig_z,
n_emit=n_emit,
n_physical_particles=n_physical_particles,
n_macroparticles=layout.n_macroparticles,
zf=zf,
tf=tf,
boost=self.fbpic_sim.boost,
initialize_self_field=initialize_self_field)
# - For the case of an empty species
elif (s.initial_distribution is None) and (layout is None):
fbpic_species = self.fbpic_sim.add_new_species(q=s.charge,
m=s.mass)
else:
raise ValueError('Unknown combination of layout and distribution')
return fbpic_species
# Redefine the method `add_diagnostic` of the parent class
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)
# Redefine the method `add_diagnostic` of the parent class
def add_applied_field(self, applied_field):
# Call method of parent class
PICMI_Simulation.add_applied_field(self, applied_field)
if type(applied_field) == PICMI_Mirror:
assert applied_field.z_front_location is not None
mirror = Mirror(z_lab=applied_field.z_front_location,
n_cells=applied_field.number_of_cells,
gamma_boost=self.fbpic_sim.boost.gamma0)
self.fbpic_sim.mirrors.append(mirror)
elif type(applied_field) == PICMI_ConstantAppliedField:
# TODO: Handle bounds
for field_name in ['Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz']:
field_value = getattr(applied_field, field_name)
if field_value is None:
continue
def field_func(F, x, y, z, t, amplitude, length_scale):
return (F + amplitude * field_value)
# Pass it to FBPIC
self.fbpic_sim.external_fields.append(
ExternalField(field_func, field_name, 1., 0.))
elif type(applied_field) == PICMI_AnalyticAppliedField:
# TODO: Handle bounds
for field_name in ['Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz']:
# Extract expression and execute it inside a function definition
expression = getattr(applied_field, field_name + '_expression')
if expression is None:
continue
fieldfunc = None
define_function_code = \
"""def fieldfunc( F, x, y, z, t, amplitude, length_scale ):\n return( F + amplitude * %s )""" %expression
# Take into account user-defined variables
for k in applied_field.user_defined_kw:
define_function_code = \
"%s = %s\n" %(k,applied_field.user_defined_kw[k]) \
+ define_function_code
exec(define_function_code, globals())
# Pass it to FBPIC
self.fbpic_sim.external_fields.append(
ExternalField(fieldfunc, field_name, 1., 0.))
else:
raise ValueError("Unrecognized `applied_field` type.")
# Redefine the method `step` of the parent class
def step(self, nsteps=None):
if nsteps is None:
nsteps = self.max_steps
self.fbpic_sim.step(nsteps, **self.step_kw)
# NB: The code below is only executed when running the script,
# (`python boosted_frame_sim.py`), but not when importing it.
if __name__ == '__main__':
# 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,
dens_func=dens_func,
zmin=zmin,
initialize_ions=True,
v_comoving=v_comoving,
gamma_boost=boost.gamma0,
n_order=n_order,
boundaries='open',
use_cuda=use_cuda)
# Add an electron bunch
add_particle_bunch(sim,
-e,
m_e,
def init(self, kw):
self.sim_kw = {}
for argname in ['use_ruyten_shapes', 'use_modified_volume']:
if f'fbpic_{argname}' in kw:
self.sim_kw[argname] = kw.pop(f'fbpic_{argname}')
self.step_kw = {}
for argname in [
'correct_currents', 'correct_divE', 'use_true_rho',
'move_positions', 'move_momenta', 'show_progress'
]:
if f'fbpic_{argname}' in kw:
self.step_kw[argname] = kw.pop(f'fbpic_{argname}')
# Get the grid
grid = self.solver.grid
if not type(grid) == PICMI_CylindricalGrid:
raise ValueError('When using fbpic with PICMI, '
'the grid needs to be a CylindricalGrid object.')
# Check rmin and boundary conditions
assert grid.lower_bound[0] == 0.
assert grid.lower_boundary_conditions[
1] == grid.upper_boundary_conditions[1]
if grid.lower_boundary_conditions[1] == 'reflective':
warnings.warn(
"FBPIC does not support reflective boundary condition in z.\n"
"The z boundary condition was automatically converted to 'open'."
)
grid.lower_boundary_conditions[1] = 'open'
grid.upper_boundary_conditions[1] = 'open'
assert grid.upper_boundary_conditions[1] in ['periodic', 'open']
assert grid.upper_boundary_conditions[0] in ['reflective', 'open']
# Determine timestep
if self.solver.cfl is not None:
dz = (grid.upper_bound[1] -
grid.lower_bound[1]) / grid.number_of_cells[1]
dr = (grid.upper_bound[0] -
grid.lower_bound[0]) / grid.number_of_cells[0]
if self.gamma_boost is not None:
beta = np.sqrt(1. - 1. / self.gamma_boost**2)
dr = dr / ((1 + beta) * self.gamma_boost)
dt = self.solver.cfl * min(dz, dr) / c
elif self.time_step_size is not None:
dt = self.time_step_size
else:
raise ValueError('You need to either set the `cfl` of the solver\n'
'or the `timestep_size` of the `Simulation`.')
# Convert API for the smoother
if self.solver.source_smoother is None:
smoother = BinomialSmoother()
else:
if self.solver.source_smoother.n_pass is None:
n_passes = 1
else:
n_passes = {
'r': self.solver.source_smoother.n_pass[0],
'z': self.solver.source_smoother.n_pass[1]
}
if self.solver.source_smoother.compensation is None:
compensator = False
else:
compensator = all(self.solver.source_smoother.compensation)
smoother = BinomialSmoother(n_passes=n_passes,
compensator=compensator)
# Convert verbose level:
verbose_level = self.verbose
if verbose_level is None:
verbose_level = 1
# Order of the stencil for z derivatives in the Maxwell solver
if self.solver.stencil_order is None:
n_order = -1
else:
n_order = self.solver.stencil_order[-1]
# Number of guard cells
if grid.guard_cells is None:
n_guard = None
else:
n_guard = grid.guard_cells[-1]
if self.solver.galilean_velocity is None:
v_comoving = None
else:
v_comoving = self.solver.galilean_velocity[-1]
# Initialize and store the FBPIC simulation object
self.fbpic_sim = FBPICSimulation(Nz=int(grid.number_of_cells[1]),
zmin=grid.lower_bound[1],
zmax=grid.upper_bound[1],
Nr=int(grid.number_of_cells[0]),
rmax=grid.upper_bound[0],
Nm=grid.n_azimuthal_modes,
dt=dt,
use_cuda=True,
smoother=smoother,
n_order=n_order,
boundaries={
'z':
grid.upper_boundary_conditions[1],
'r':
grid.upper_boundary_conditions[0]
},
n_guard=n_guard,
verbose_level=verbose_level,
particle_shape=self.particle_shape,
v_comoving=v_comoving,
gamma_boost=self.gamma_boost,
**self.sim_kw)
# Set the moving window
if grid.moving_window_velocity is not None:
self.fbpic_sim.set_moving_window(grid.moving_window_velocity[-1])
class Simulation(PICMI_Simulation):
# Redefine the `init` method, as required by the picmi `_ClassWithInit`
def init(self, kw):
# Get the grid
grid = self.solver.grid
if not type(grid) == PICMI_CylindricalGrid:
raise ValueError('When using fbpic with PICMI, '
'the grid needs to be a CylindricalGrid object.')
# Check rmin and boundary conditions
assert grid.rmin == 0.
assert grid.bc_zmin == grid.bc_zmax
assert grid.bc_zmax in ['periodic', 'open']
assert grid.bc_rmax in ['reflective', 'open']
# Determine timestep
if self.solver.cfl is not None:
dz = (grid.zmax - grid.zmin) / grid.nz
dt = self.solver.cfl * dz / c
elif self.time_step_size is not None:
dt = self.time_step_size
else:
raise ValueError('You need to either set the `cfl` of the solver\n'
'or the `timestep_size` of the `Simulation`.')
# Convert API for the smoother
if self.solver.source_smoother is None:
smoother = BinomialSmoother()
else:
smoother = BinomialSmoother(
n_passes=self.solver.source_smoother.n_pass,
compensator=self.solver.source_smoother.compensation)
# Initialize and store the FBPIC simulation object
self.fbpic_sim = FBPICSimulation(Nz=int(grid.nz),
zmin=grid.zmin,
zmax=grid.zmax,
Nr=int(grid.nr),
rmax=grid.rmax,
Nm=grid.n_azimuthal_modes,
dt=dt,
use_cuda=True,
smoother=smoother,
n_order=32,
boundaries={
'z': grid.bc_zmax,
'r': grid.bc_rmax
})
# Set the moving window
if grid.moving_window_zvelocity is not None:
self.fbpic_sim.set_moving_window(grid.moving_window_zvelocity)
# Redefine the method `add_laser` from the PICMI Simulation class
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])
# Redefine the method `add_species` from the PICMI Simulation class
def add_species(self, species, layout, initialize_self_field=False):
# Call method of parent class
PICMI_Simulation.add_species(self, species, layout,
initialize_self_field)
# Extract list of species
if type(species) == PICMI_Species:
species_instances_list = [species]
elif type(species) == PICMI_MultiSpecies:
species_instances_list = species.species_instances_list
else:
raise ValueError('Unknown type: %s' % type(species))
# Loop over species and create FBPIC species
for s in species_instances_list:
# Get their charge and mass
if s.particle_type is not None:
s.charge = particle_charge[s.particle_type]
s.mass = particle_mass[s.particle_type]
# If `charge_state` is set, redefine the charge and mass
if s.charge_state is not None:
s.charge = s.charge_state * e
s.mass -= s.charge_state * m_e
# Add the species to the FBPIC simulation
fbpic_species = self._create_new_fbpic_species(
s, layout, initialize_self_field)
# Register a pointer to the FBPIC species in the PICMI species itself
# (Useful for particle diagnostics later on)
s.fbpic_species = fbpic_species
# Loop over species and handle ionization
for s in species_instances_list:
for interaction in s.interactions:
assert interaction[0] == 'ionization'
assert interaction[1] == 'ADK'
picmi_target = interaction[2]
if not hasattr(picmi_target, 'fbpic_species'):
raise RuntimeError(
'For ionization with PICMI+FBPIC:\n'
'You need to add the target species to the simulation,'
' before the other species.')
fbpic_target = picmi_target.fbpic_species
fbpic_source = s.fbpic_species
fbpic_source.make_ionizable(element=s.particle_type,
level_start=s.charge_state,
target_species=fbpic_target)
def _create_new_fbpic_species(self, s, layout, initialize_self_field):
# - For the case of a plasma defined in a gridded layout
if (type(s.initial_distribution)==PICMI_AnalyticDistribution) and \
(type(layout) == PICMI_GriddedLayout):
assert initialize_self_field == False
import numexpr
density_expression = s.initial_distribution.density_expression
if s.density_scale is not None:
density_expression = "%f*(%s)" \
%(s.density_scale, density_expression)
def dens_func(z, r):
n = numexpr.evaluate(density_expression)
return n
p_nr = layout.n_macroparticle_per_cell[0]
p_nt = layout.n_macroparticle_per_cell[1]
p_nz = layout.n_macroparticle_per_cell[2]
fbpic_species = self.fbpic_sim.add_new_species(
q=s.charge,
m=s.mass,
n=1.,
dens_func=dens_func,
p_nz=p_nz,
p_nr=p_nr,
p_nt=p_nt,
p_zmin=s.initial_distribution.lower_bound[-1],
p_zmax=s.initial_distribution.upper_bound[-1],
continuous_injection=s.initial_distribution.fill_in)
# - For the case of a Gaussian beam
elif (type(s.initial_distribution)==PICMI_GaussianBunchDistribution) \
and (type(layout) == PICMI_PseudoRandomLayout):
dist = s.initial_distribution
gamma0_beta0 = dist.centroid_velocity[-1] / c
gamma0 = (1 + gamma0_beta0**2)**.5
sig_r = dist.rms_bunch_size[0]
sig_z = dist.rms_bunch_size[-1]
sig_gamma = dist.rms_velocity[-1] / c
sig_vr = dist.rms_velocity[0] / gamma0
if sig_vr != 0:
tf = -sig_r**2 / sig_vr**2 * dist.velocity_divergence[0]
else:
tf = 0.
zf = dist.centroid_position[-1] + \
dist.centroid_velocity[-1]/gamma0 * tf
# Calculate size at focus and emittance
sig_r0 = (sig_r**2 - (sig_vr * tf)**2)**0.5
n_emit = gamma0 * sig_r0 * sig_vr / c
# Get the number of physical particles
n_physical_particles = dist.n_physical_particles
if s.density_scale is not None:
n_physical_particles *= s.density_scale
fbpic_species = add_particle_bunch_gaussian(
self.fbpic_sim,
q=s.charge,
m=s.mass,
gamma0=gamma0,
sig_gamma=sig_gamma,
sig_r=sig_r0,
sig_z=sig_z,
n_emit=n_emit,
n_physical_particles=n_physical_particles,
n_macroparticles=layout.n_macroparticles,
zf=zf,
tf=tf,
initialize_self_field=initialize_self_field)
# - For the case of an empty species
elif (s.initial_distribution is None) and (layout is None):
fbpic_species = self.fbpic_sim.add_new_species(q=s.charge,
m=s.mass)
else:
raise ValueError('Unknown combination of layout and distribution')
return fbpic_species
# Redefine the method `add_diagnostic` of the parent class
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:
diag = FieldDiagnostic(period=diagnostic.period,
fldobject=self.fbpic_sim.fld,
comm=self.fbpic_sim.comm,
fieldtypes=diagnostic.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
diag = ParticleDiagnostic(period=diagnostic.period,
species=species_dict,
comm=self.fbpic_sim.comm,
particle_data=diagnostic.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)
# Redefine the method `step` of the parent class
def step(self, nsteps):
self.fbpic_sim.step(nsteps)
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)
# ---------------------------
# Carrying out the simulation
# ---------------------------
# NB: The code below is only executed when running the script,
# (`python boosted_frame_sim.py`), but not when importing it.
if __name__ == '__main__':
# Initialize the simulation object
sim = Simulation(Nz,
zmax,
Nr,
rmax,
Nm,
dt,
zmin=zmin,
v_comoving=v_comoving,
gamma_boost=boost.gamma0,
n_order=n_order,
use_cuda=use_cuda,
boundaries={
'z': 'open',
'r': 'reflective'
})
# 'r': 'open' can also be used, but is more computationally expensive
# Add the plasma electron and plasma ions
plasma_elec = sim.add_new_species(q=-e,
m=m_e,
n=n_e,
dens_func=dens_func,
boost_positions_in_dens_func=True,
# 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,
p_zmin,
p_zmax,
p_rmin,
p_rmax,
p_nz,
p_nr,
p_nt,
n_e,
dens_func=dens_func,
zmin=zmin,
boundaries='open',
n_order=n_order,
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:
