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