def read_osiris_beam(file_path, plasma_dens):
"""Reads particle data from an OSIRIS paricle file and returns it in the
unis used by APtools.
Parameters
----------
file_path : str
Path to the file with particle data
plasma_dens : float
Plasma density in units od cm^{-3} used to convert the beam data to
non-normalized units
Returns
-------
A tuple with 7 arrays containing the 6D phase space and charge of the
particles.
"""
s_d = plasma_skin_depth(plasma_dens)
file_content = H5File(file_path)
# get data
q = np.array(file_content.get('q')) * ct.e
x = np.array(file_content.get('x2')) * s_d
y = np.array(file_content.get('x3')) * s_d
z = np.array(file_content.get('x1')) * s_d
px = np.array(file_content.get('p2'))
py = np.array(file_content.get('p3'))
pz = np.array(file_content.get('p1'))
return x, y, z, px, py, pz, q
def read_hipace_beam(file_path, plasma_dens):
"""Reads particle data from an HiPACE paricle file and returns it in the
unis used by APtools.
Parameters
----------
file_path : str
Path to the file with particle data
plasma_dens : float
Plasma density in units od cm^{-3} used to convert the beam data to
non-normalized units
Returns
-------
A tuple with 7 arrays containing the 6D phase space and charge of the
particles.
"""
s_d = plasma_skin_depth(plasma_dens)
file_content = H5File(file_path)
# sim parameters
n_cells = file_content.attrs['NX']
sim_size = (file_content.attrs['XMAX'] - file_content.attrs['XMIN'])
cell_vol = np.prod(sim_size / n_cells)
q_norm = cell_vol * plasma_dens * 1e6 * s_d**3 * ct.e
# get data
q = np.array(file_content.get('q')) * q_norm
x = np.array(file_content.get('x2')) * s_d
y = np.array(file_content.get('x3')) * s_d
z = np.array(file_content.get('x1')) * s_d
px = np.array(file_content.get('p2'))
py = np.array(file_content.get('p3'))
pz = np.array(file_content.get('p1'))
return x, y, z, px, py, pz, q
def get_beam_function(beam_part, n_r, n_xi, n_p, r_arr, xi_arr, p_shape):
"""
Return a function of r and xi which gives the azimuthal magnetic field
from a particle distribution. This is Eq. (18) in the original paper.
For details about input parameters see method 'calculate_wakefields'.
"""
# Plasma skin depth.
s_d = ge.plasma_skin_depth(n_p / 1e6)
# Grid parameters.
dr = r_arr[1] - r_arr[0]
dxi = xi_arr[1] - xi_arr[0]
r_min = r_arr[0]
xi_min = xi_arr[0]
# Grid arrays with guard cells.
r_grid_g = (0.5 + np.arange(-2, n_r + 2)) * dr
xi_grid_g = np.arange(-2, n_xi + 2) * dxi + xi_min
# Get and normalize particle coordinate arrays.
x, y, xi, q = beam_part
xi_n = xi / s_d
x_n = x / s_d
y_n = y / s_d
# Calculate particle weights.
w = q / ct.e / (2 * np.pi * dr * dxi * s_d**3 * n_p)
# Obtain charge distribution (using cubic particle shape by default).
q_dist = charge_distribution_cyl(xi_n,
x_n,
y_n,
w,
xi_min,
r_min,
n_xi,
n_r,
dxi,
dr,
p_shape=p_shape)
# Calculate radial integral (Eq. (18)).
r_int = np.cumsum(q_dist, axis=1) / np.abs(r_grid_g) * dr
# Create and return interpolator.
return scint.interp2d(r_grid_g, xi_grid_g, -r_int)
def __calculate_wakefields(self, x, y, xi, px, py, pz, q, t):
self.current_t = t
if self.current_t_wf is None:
self.current_t_wf = t
elif (self.current_t_wf != t
and t >= self.current_t_wf + self.dz_fields / ct.c):
self.current_t_wf = t
else:
return
z_beam = t * ct.c + np.average(xi) # z postion of beam center
n_p = self.density_function(z_beam)
# calculate distance to laser focus
if self.laser_evolution:
dz_foc = self.laser_z_foc - ct.c * t
else:
dz_foc = 0
flds = calculate_wakefields(self.laser, [x, y, xi, q],
self.r_max,
self.xi_min,
self.xi_max,
self.n_r,
self.n_xi,
self.ppc,
n_p,
dz_foc,
p_shape=self.p_shape)
n_p_mesh, W_r, E_z, E_z_p, K_r, psi_mesh, xi_arr, r_arr = flds
E_0 = ge.plasma_cold_non_relativisct_wave_breaking_field(n_p * 1e-6)
s_d = ge.plasma_skin_depth(n_p * 1e-6)
self.rho = n_p_mesh
self.E_z = E_z.T * E_0
self.W_x = W_r.T * E_0
self.K_x = K_r.T * E_0 / s_d / ct.c
self.E_z_p = E_z_p.T * E_0 / s_d
self.xi_fld = xi_arr * s_d
self.r_fld = r_arr * s_d
def get_beam_function(beam_part, r_max, xi_min, xi_max, n_r, n_xi, n_p):
"""
Return a function of r and xi which gives the azimuthal magnetic field
from a particle distribution. This is Eq. (18) in the original paper.
For details about input parameters see method 'calculate_wakefields'.
"""
x, y, xi, q = beam_part
s_d = ge.plasma_skin_depth(n_p / 1e6)
xi_n = xi / s_d
r_n = np.sqrt(x**2 + y**2) / s_d
x_edges = np.linspace(xi_min, xi_max, n_xi)
y_edges = np.linspace(0, r_max, n_r)
bins = [x_edges, y_edges]
dr = y_edges[1] - y_edges[0]
dxi = x_edges[1] - x_edges[0]
hist_weights = q / ct.e / (2 * np.pi * dr * dxi * s_d**3 * n_p)
bunch_hist, *_ = np.histogram2d(xi_n, r_n, bins=bins, weights=hist_weights)
r_b = y_edges[1:] - dr / 2
xi_b = x_edges[1:] - dxi / 2
bunch_rint = np.cumsum(bunch_hist, axis=1) / r_b * dr
return scint.interp2d(r_b, xi_b, -bunch_rint)
def __calculate_wakefields(self, x, y, xi, px, py, pz, q, t):
if self.current_t is None:
self.current_t = t
elif self.current_t != t and t >= self.current_t + self.dz_fields / ct.c:
self.current_t = t
else:
return
z_beam = t * ct.c + np.average(xi) # z postion of beam center
n_p = self.density_function(z_beam)
# calculate distance to laser focus
if self.laser_evolution:
dz_foc = self.laser_z_foc - ct.c * t
else:
dz_foc = 0
flds = calculate_wakefields(self.laser, [x, y, xi, q], self.r_max,
self.xi_min, self.xi_max, self.n_r,
self.n_xi, self.ppc, n_p, dz_foc)
n_p_mesh, W_r, E_z, E_z_p, K_r, psi_mesh, xi_arr, r_arr = flds
# For debugging
# plt.plot(E_z[:,0])
# plt.plot(K_r[:,0])
# plt.plot(a0_0[:,0])
# plt.show()
# plt.subplot(411)
# plt.imshow(E_z.T*E_0, aspect='auto',
# extent=(self.xi_min, self.xi_max, 0, self.r_max))
# plt.plot(E_z[:,0]*E_0)
# plt.subplot(412)
# plt.imshow(K_r.T*E_0/s_d, aspect='auto',
# extent=(self.xi_min, self.xi_max, 0, self.r_max))
# plt.subplot(413)
# plt.imshow(E_z_p.T*E_0/s_d, aspect='auto',
# extent=(self.xi_min, self.xi_max, 0, self.r_max))
# plt.subplot(414)
# plt.imshow(beam_hist.T, aspect='auto',
# extent=(self.xi_min, self.xi_max, 0, self.r_max))
# plt.show()
E_0 = ge.plasma_cold_non_relativisct_wave_breaking_field(n_p * 1e-6)
s_d = ge.plasma_skin_depth(n_p * 1e-6)
self.E_z = RectBivariateSpline(xi_arr * s_d,
r_arr * s_d,
E_z.T * E_0,
kx=2,
ky=2)
self.W_x = RectBivariateSpline(xi_arr * s_d,
r_arr * s_d,
W_r.T * E_0,
kx=2,
ky=2)
self.K_x = RectBivariateSpline(xi_arr * s_d,
r_arr * s_d,
K_r.T * E_0 / s_d / ct.c,
kx=2,
ky=2)
self.E_z_p = RectBivariateSpline(xi_arr * s_d,
r_arr * s_d,
E_z_p.T * E_0 / s_d,
kx=2,
ky=2)
def __calculate_wakefields(self, x, y, xi, px, py, pz, q, t):
if self.current_t != t:
self.current_t = t
else:
return
z_beam = t * ct.c + np.average(xi) # z postion of beam center
n_p = self.density_function(z_beam)
if n_p == self.current_n_p and not self.laser_evolution:
# If density has not changed and driver does not evolve, it is
# not necessary to recompute fields.
return
self.current_n_p = n_p
s_d = ge.plasma_skin_depth(n_p * 1e-6)
r = np.linspace(0, self.r_max, self.n_r)
dz = (self.xi_max - self.xi_min) / self.n_xi / s_d
dr = self.r_max / self.n_r / s_d
r_part = np.sqrt(x**2 + y**2)
beam_hist, *_ = np.histogram2d(xi,
r_part,
bins=[self.n_xi, self.n_r],
range=[[self.xi_min, self.xi_max],
[0, self.r_max]],
weights=q / ct.e)
# ,
# weights=1/(ct.pi*dr*2*dz))
n = np.arange(self.n_r)
disc_area = np.pi * dr**2 * (1 + 2 * n)
beam_hist *= 1 / (disc_area * dz * n_p) / s_d**3
n_iter = self.n_xi - 1
u_1 = np.zeros((n_iter + 1, len(r)))
u_2 = np.zeros((n_iter + 1, len(r)))
z_arr = np.zeros(n_iter + 1)
z_arr[-1] = self.xi_max / s_d
# calculate distance to laser focus
if self.laser_evolution:
dz_foc = self.laser_z_foc - ct.c * t
else:
dz_foc = 0
for i in np.arange(n_iter):
z = z_arr[-1] - i * dz
# get laser a0 at z, z+dz/2 and z+dz
if self.driver is not None:
a0_0 = self.driver.get_a0_profile(r, z * s_d, dz_foc)
a0_1 = self.driver.get_a0_profile(r, (z - dz / 2) * s_d,
dz_foc)
a0_2 = self.driver.get_a0_profile(r, (z - dz) * s_d, dz_foc)
else:
a0_0 = np.zeros(r.shape[0])
a0_1 = np.zeros(r.shape[0])
a0_2 = np.zeros(r.shape[0])
# perform runge-kutta
A = dz * self.__wakefield_ode_system(u_1[-1 - i], u_2[-1 - i], r,
z * s_d, a0_0,
beam_hist[-(i + 1)])
B = dz * self.__wakefield_ode_system(
u_1[-1 - i] + A[0] / 2, u_2[-1 - i] + A[1] / 2, r,
(z - dz / 2) * s_d, a0_1, beam_hist[-(i + 1)])
C = dz * self.__wakefield_ode_system(
u_1[-1 - i] + B[0] / 2, u_2[-1 - i] + B[1] / 2, r,
(z - dz / 2) * s_d, a0_1, beam_hist[-(i + 1)])
D = dz * self.__wakefield_ode_system(
u_1[-1 - i] + C[0], u_2[-1 - i] + C[1], r,
(z - dz) * s_d, a0_2, beam_hist[-(i + 1)])
u_1[-2 -
i] = u_1[-1 - i] + 1 / 6 * (A[0] + 2 * B[0] + 2 * C[0] + D[0])
u_2[-2 -
i] = u_2[-1 - i] + 1 / 6 * (A[1] + 2 * B[1] + 2 * C[1] + D[1])
z_arr[-2 - i] = z - dz
E_z = -np.gradient(u_1, dz, axis=0, edge_order=2)
E_z_p = np.gradient(E_z, dz, axis=0, edge_order=2)
W_r = -np.gradient(u_1, dr, axis=1, edge_order=2)
K_r = np.gradient(W_r, dr, axis=1, edge_order=2)
E_0 = ge.plasma_cold_non_relativisct_wave_breaking_field(n_p * 1e-6)
# For debugging
# plt.plot(E_z[:,0])
# plt.plot(K_r[:,0])
# plt.plot(a0_0[:,0])
# plt.show()
# plt.subplot(411)
# plt.imshow(E_z.T*E_0, aspect='auto',
# extent=(self.xi_min, self.xi_max, 0, self.r_max))
# plt.plot(E_z[:,0]*E_0)
# plt.subplot(412)
# plt.imshow(W_r.T*E_0, aspect='auto',
# extent=(self.xi_min, self.xi_max, 0, self.r_max))
# plt.subplot(413)
# plt.imshow(E_z_p.T*E_0/s_d, aspect='auto',
# extent=(self.xi_min, self.xi_max, 0, self.r_max))
# plt.subplot(414)
# plt.imshow(beam_hist.T, aspect='auto',
# extent=(self.xi_min, self.xi_max, 0, self.r_max))
# plt.show()
self.E_z = RectBivariateSpline(z_arr * s_d, r, E_z * E_0, kx=2, ky=2)
self.W_x = RectBivariateSpline(z_arr * s_d, r, W_r * E_0, kx=2, ky=2)
self.K_x = RectBivariateSpline(z_arr * s_d,
r,
K_r * E_0 / s_d / ct.c,
kx=2,
ky=2)
self.E_z_p = RectBivariateSpline(z_arr * s_d,
r,
E_z_p * E_0 / s_d,
kx=2,
ky=2)
def calculate_wakefields(laser, beam_part, r_max, xi_min, xi_max, n_r, n_xi,
ppc, n_p, laser_z_foc):
"""
Calculate the plasma wakefields generated by the given laser pulse and
electron beam in the specified grid points.
Parameters:
-----------
laser : LaserPulse (optional)
Laser driver of the plasma stage.
beam_part : list
List of numpy arrays containing the spatial coordinates and charge of
all beam particles, i.e [x, y, xi, q].
r_max : float
Maximum radial position up to which plasma wakefield will be
calculated.
xi_min : float
Minimum longitudinal (speed of light frame) position up to which
plasma wakefield will be calculated.
xi_max : float
Maximum longitudinal (speed of light frame) position up to which
plasma wakefield will be calculated.
n_r : int
Number of grid elements along r in which to calculate the wakefields.
n_xi : int
Number of grid elements along xi in which to calculate the wakefields.
ppc : int (optional)
Number of plasma particles per 1d cell along the radial direction.
n_p : float
Plasma density in units of m^{-3}.
laser_z_foc : float
Focal position of the laser along z in meters. It is measured as
the distance from the beginning of the PlasmaStage. A negative
value implies that the focal point is located before the
PlasmaStage.
"""
s_d = ge.plasma_skin_depth(n_p * 1e-6)
r_max = r_max / s_d
xi_min = xi_min / s_d
xi_max = xi_max / s_d
# Laser parameters.
if laser is not None:
laser_params = [
laser.a_0, laser.l_0, laser.w_0, laser.tau, laser.xi_c,
laser.polarization, laser_z_foc
]
else:
laser_params = None
# Initialize plasma particles.
dr = r_max / n_r
dr_p = dr / ppc
n_part = n_r * ppc
r = np.linspace(dr_p / 2, r_max - dr_p / 2, n_part)
pr = np.zeros_like(r)
gamma = np.ones_like(r)
q = dr_p * r
# Iteration steps.
dxi = (xi_max - xi_min) / n_xi
# Initialize field arrays.
psi_mesh = np.zeros((n_r, n_xi))
dr_psi_mesh = np.zeros((n_r, n_xi))
dxi_psi_mesh = np.zeros((n_r, n_xi))
b_theta_bar_mesh = np.zeros((n_r, n_xi))
b_theta_0_mesh = np.zeros((n_r, n_xi))
r_arr = np.linspace(dr / 2, r_max - dr / 2, n_r)
xi_arr = np.linspace(xi_min, xi_max, n_xi)
# Calculate beam source term (b_theta_0) from particle distribution.
if beam_part is not None:
beam_source = get_beam_function(beam_part, r_max, xi_min, xi_max, n_r,
n_xi, n_p)
else:
beam_source = None
# Main loop.
for step in np.arange(n_xi):
xi = xi_max - dxi * step
# Evolve plasma to next xi step.
r, pr = evolve_plasma(r, pr, q, xi, dxi, laser_params, beam_source,
s_d)
# Remove plasma particles leaving simulation boundaries (plus margin).
idx_keep = np.where(r <= r_max + 0.1)
r = r[idx_keep]
pr = pr[idx_keep]
gamma = gamma[idx_keep]
q = q[idx_keep]
# Calculate fields at specified r locations.
fields = calculate_fields(r_arr, xi, r, pr, q, laser_params,
beam_source, s_d)
i = -1 - step
# Unpack fields.
(psi_mesh[:, i], dr_psi_mesh[:, i], dxi_psi_mesh[:, i],
b_theta_bar_mesh[:, i], b_theta_0_mesh[:, i]) = fields
# Calculate derived fields (E_r, n_p, K_r and E_z').
dr_psi_mesh, dxi_psi_mesh = np.gradient(psi_mesh, dr, dxi, edge_order=2)
dxi_psi_mesh *= -1
dr_psi_mesh *= -1
e_r_mesh = b_theta_bar_mesh + b_theta_0_mesh - dr_psi_mesh
r_arr_v = np.vstack(r_arr)
n_p = (
np.gradient(r_arr_v * e_r_mesh, dr, axis=0, edge_order=2) / r_arr_v -
np.gradient(dxi_psi_mesh, dxi, axis=1, edge_order=2) - 1)
k_r_mesh = np.gradient(dr_psi_mesh, dr, axis=0, edge_order=2)
e_z_p_mesh = np.gradient(dxi_psi_mesh, dxi, axis=1, edge_order=2)
return (n_p, dr_psi_mesh, dxi_psi_mesh, e_z_p_mesh, k_r_mesh, psi_mesh,
xi_arr, r_arr)
def __calculate_wakefields(self, x, y, xi, px, py, pz, q, t):
if self.current_t != t:
self.current_t = t
else:
return
z_beam = t * ct.c + np.average(xi) # z postion of beam center
n_p = self.density_function(z_beam)
if n_p == self.current_n_p and not self.laser_evolution:
# If density has not changed and driver does not evolve, it is
# not necessary to recompute fields.
return
self.current_n_p = n_p
s_d = ge.plasma_skin_depth(n_p * 1e-6)
dz = (self.xi_max - self.xi_min) / self.n_xi / s_d
dr = self.r_max / self.n_r / s_d
r = np.linspace(dr / 2, self.r_max / s_d - dr / 2, self.n_r)
# Get charge distribution and remove guard cells.
beam_hist = charge_distribution_cyl(xi / s_d,
x / s_d,
y / s_d,
q / ct.e,
self.xi_min / s_d,
r[0],
self.n_xi,
self.n_r,
dz,
dr,
p_shape=self.p_shape)
beam_hist = beam_hist[2:-2, 2:-2]
n = np.arange(self.n_r)
disc_area = np.pi * dr**2 * (1 + 2 * n)
beam_hist *= 1 / (disc_area * dz * n_p) / s_d**3
n_iter = self.n_xi - 1
u_1 = np.zeros((n_iter + 1, len(r)))
u_2 = np.zeros((n_iter + 1, len(r)))
z_arr = np.zeros(n_iter + 1)
z_arr[-1] = self.xi_max / s_d
# calculate distance to laser focus
if self.laser_evolution:
dz_foc = self.laser_z_foc - ct.c * t
else:
dz_foc = 0
for i in np.arange(n_iter):
z = z_arr[-1] - i * dz
# get laser a0 at z, z+dz/2 and z+dz
if self.driver is not None:
a0_0 = self.driver.get_a0_profile(r * s_d, z * s_d, dz_foc)
a0_1 = self.driver.get_a0_profile(r * s_d, (z - dz / 2) * s_d,
dz_foc)
a0_2 = self.driver.get_a0_profile(r * s_d, (z - dz) * s_d,
dz_foc)
else:
a0_0 = np.zeros(r.shape[0])
a0_1 = np.zeros(r.shape[0])
a0_2 = np.zeros(r.shape[0])
# perform runge-kutta
A = dz * self.__wakefield_ode_system(u_1[-1 - i], u_2[-1 - i],
r * s_d, z * s_d, a0_0,
beam_hist[-(i + 1)])
B = dz * self.__wakefield_ode_system(
u_1[-1 - i] + A[0] / 2, u_2[-1 - i] + A[1] / 2, r * s_d,
(z - dz / 2) * s_d, a0_1, beam_hist[-(i + 1)])
C = dz * self.__wakefield_ode_system(
u_1[-1 - i] + B[0] / 2, u_2[-1 - i] + B[1] / 2, r * s_d,
(z - dz / 2) * s_d, a0_1, beam_hist[-(i + 1)])
D = dz * self.__wakefield_ode_system(
u_1[-1 - i] + C[0], u_2[-1 - i] + C[1], r * s_d,
(z - dz) * s_d, a0_2, beam_hist[-(i + 1)])
u_1[-2 -
i] = u_1[-1 - i] + 1 / 6 * (A[0] + 2 * B[0] + 2 * C[0] + D[0])
u_2[-2 -
i] = u_2[-1 - i] + 1 / 6 * (A[1] + 2 * B[1] + 2 * C[1] + D[1])
z_arr[-2 - i] = z - dz
E_z = -np.gradient(u_1, dz, axis=0, edge_order=2)
E_z_p = np.gradient(E_z, dz, axis=0, edge_order=2)
W_r = -np.gradient(u_1, dr, axis=1, edge_order=2)
K_r = np.gradient(W_r, dr, axis=1, edge_order=2)
E_0 = ge.plasma_cold_non_relativisct_wave_breaking_field(n_p * 1e-6)
self.E_z = E_z * E_0
self.W_x = W_r * E_0
self.K_x = K_r * E_0 / s_d / ct.c
self.E_z_p = E_z_p * E_0 / s_d
self.xi_fld = z_arr * s_d
self.r_fld = r * s_d
import matplotlib.pyplot as plt
from laser_profiles import GaussianLaser
import aptools.plasma_accel.general_equations as ge
from envelope_solver import *
import scipy.constants as ct
import time
from spotsize_evolution import pulse_width, spotsize
# plasma parameters
n_p = 1e18 # cm^-3
s_d = ge.plasma_skin_depth(n_p)
k_p = 1 / s_d
w_p = ge.plasma_frequency(n_p)
r_0 = 10e-6
# laser parameters in SI units
tau = 25e-15 # s
w_0 = 20e-6 # m
l_0 = 0.8e-6 # m
z_c = 0. # m
a_0 = 1.
k_0 = 2 * np.pi / l_0
k0p = k_0 / k_p
zR = ge.laser_rayleigh_length(w_0, l_0)
zR_norm = zR / s_d
print('Normalized Rayleigh length = {}'.format(zR_norm))
# create laser pulse
laser_fb = GaussianLaser(a_0, w_0, tau, z_c, zf=0., lambda0=l_0)
# grid