From 9e2d5e6eedcd223b3c243301d93bb8161ac0f174 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= Date: Tue, 18 Jun 2024 22:50:21 +0200 Subject: [PATCH] Formatting --- .../qs_rz_baxevanis_ion/__init__.py | 2 +- .../qs_rz_baxevanis_ion/adaptive_grid.py | 177 ++++++---- .../qs_rz_baxevanis_ion/b_theta.py | 179 ++++++---- .../qs_rz_baxevanis_ion/b_theta_bunch.py | 45 ++- .../qs_rz_baxevanis_ion/deposition.py | 43 +-- .../qs_rz_baxevanis_ion/gather.py | 14 +- .../qs_rz_baxevanis_ion/plasma_particles.py | 310 ++++++++++++------ .../qs_rz_baxevanis_ion/plasma_push/ab2.py | 60 ++-- .../psi_and_derivatives.py | 117 ++++--- .../qs_rz_baxevanis_ion/solver.py | 164 ++++++--- .../qs_rz_baxevanis_ion/utils.py | 53 +-- .../qs_rz_baxevanis_ion/wakefield.py | 279 +++++++++------- 12 files changed, 941 insertions(+), 502 deletions(-) diff --git a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/__init__.py b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/__init__.py index 1748cc7..46e1fb5 100644 --- a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/__init__.py +++ b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/__init__.py @@ -1,3 +1,3 @@ from .wakefield import Quasistatic2DWakefieldIon -__all__ = ['Quasistatic2DWakefieldIon'] +__all__ = ["Quasistatic2DWakefieldIon"] diff --git a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/adaptive_grid.py b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/adaptive_grid.py index a645832..9eb91a6 100644 --- a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/adaptive_grid.py +++ b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/adaptive_grid.py @@ -14,7 +14,7 @@ from .utils import longitudinal_gradient, radial_gradient -class AdaptiveGrid(): +class AdaptiveGrid: """Grid whose size dynamically adapts to the extent of a particle bunch. The number of radial cells is fixed, but its transverse extent is @@ -49,6 +49,7 @@ class AdaptiveGrid(): deposit to and gather from the base grid (if they haven't escaped from it too). """ + def __init__( self, x: np.ndarray, @@ -59,7 +60,7 @@ def __init__( nxi: int, xi_plasma: np.ndarray, r_max: Optional[float] = None, - r_lim: Optional[float] = None + r_lim: Optional[float] = None, ): self.bunch_name = bunch_name self.xi_plasma = xi_plasma @@ -72,7 +73,7 @@ def __init__( self._r_max_hist = [] self._update(x, y, xi) - + @property def r_min_cell(self): """Radial position of first cell (ignoring guard cells).""" @@ -111,9 +112,8 @@ def update_if_needed(self, x, y, xi, n_p, pp_hist): # Only trigger radial update if the radial size is not fixed. if self._r_max is None: r_max_beam = np.max(np.sqrt(x**2 + y**2)) - update_r = ( - (r_max_beam > self.r_max_cell) or - (r_max_beam < self.r_max_cell * 0.9) + update_r = (r_max_beam > self.r_max_cell) or ( + r_max_beam < self.r_max_cell * 0.9 ) # It a radial limit is set, update only if limit has not been # reached. @@ -127,9 +127,8 @@ def update_if_needed(self, x, y, xi, n_p, pp_hist): xi_min_beam = np.min(xi) xi_max_beam = np.max(xi) - update_xi = ( - (xi_min_beam < self.xi_grid[0 + self.nxi_border]) or - (xi_max_beam > self.xi_grid[-1 - self.nxi_border]) + update_xi = (xi_min_beam < self.xi_grid[0 + self.nxi_border]) or ( + xi_max_beam > self.xi_grid[-1 - self.nxi_border] ) if update_r or update_xi: self._update(x, y, xi) @@ -152,19 +151,30 @@ def calculate_fields(self, n_p, pp_hist, reset_fields=True): self._reset_fields() s_d = ge.plasma_skin_depth(n_p * 1e-6) calculate_fields_on_grid( - self.i_grid, self.r_grid, s_d, - self.psi_grid, self.b_t, pp_hist['r_hist'], pp_hist['log_r_hist'], - pp_hist['sum_1_hist'], pp_hist['sum_2_hist'], - pp_hist['a_0_hist'], pp_hist['a_i_hist'], pp_hist['b_i_hist']) + self.i_grid, + self.r_grid, + s_d, + self.psi_grid, + self.b_t, + pp_hist["r_hist"], + pp_hist["log_r_hist"], + pp_hist["sum_1_hist"], + pp_hist["sum_2_hist"], + pp_hist["a_0_hist"], + pp_hist["a_i_hist"], + pp_hist["b_i_hist"], + ) E_0 = ge.plasma_cold_non_relativisct_wave_breaking_field(n_p * 1e-6) longitudinal_gradient( - self.psi_grid[2:-2, 2:-2], self.dxi/s_d, self.e_z[2:-2, 2:-2]) + self.psi_grid[2:-2, 2:-2], self.dxi / s_d, self.e_z[2:-2, 2:-2] + ) radial_gradient( - self.psi_grid[2:-2, 2:-2], self.dr/s_d, self.e_r[2:-2, 2:-2]) + self.psi_grid[2:-2, 2:-2], self.dr / s_d, self.e_r[2:-2, 2:-2] + ) self.e_r -= self.b_t - self.e_z *= - E_0 - self.e_r *= - E_0 + self.e_z *= -E_0 + self.e_r *= -E_0 self.b_t *= E_0 / ct.c def calculate_bunch_source(self, bunch, n_p, p_shape): @@ -179,12 +189,23 @@ def calculate_bunch_source(self, bunch, n_p, p_shape): p_shape : str The particle shape. """ - self.b_t_bunch[:] = 0. - self.q_bunch[:] = 0. + self.b_t_bunch[:] = 0.0 + self.q_bunch[:] = 0.0 all_deposited = deposit_bunch_charge( - bunch.x, bunch.y, bunch.xi, bunch.q, n_p, - self.nr-self.nr_border, self.nxi, self.r_grid, self.xi_grid, - self.dr, self.dxi, p_shape, self.q_bunch) + bunch.x, + bunch.y, + bunch.xi, + bunch.q, + n_p, + self.nr - self.nr_border, + self.nxi, + self.r_grid, + self.xi_grid, + self.dr, + self.dxi, + p_shape, + self.q_bunch, + ) calculate_bunch_source(self.q_bunch, self.nr, self.nxi, self.b_t_bunch) return all_deposited @@ -199,9 +220,25 @@ def gather_fields(self, x, y, z, ex, ey, ez, bx, by, bz): The arrays where the gathered field components will be stored. """ return gather_main_fields_cyl_linear( - self.e_r, self.e_z, self.b_t, self.xi_min, self.xi_max, - self.r_min_cell, self.r_max_cell, self.dxi, self.dr, x, y, z, - ex, ey, ez, bx, by, bz) + self.e_r, + self.e_z, + self.b_t, + self.xi_min, + self.xi_max, + self.r_min_cell, + self.r_max_cell, + self.dxi, + self.dr, + x, + y, + z, + ex, + ey, + ez, + bx, + by, + bz, + ) def get_openpmd_data(self, global_time, diags): """Get the field data at the grid to store in the openpmd diagnostics. @@ -221,8 +258,8 @@ def get_openpmd_data(self, global_time, diags): # Grid parameters. grid_spacing = [self.dr, self.dxi] - grid_labels = ['r', 'z'] - grid_global_offset = [0., global_time*ct.c + self.xi_min] + grid_labels = ["r", "z"] + grid_global_offset = [0.0, global_time * ct.c + self.xi_min] # Initialize field diags lists. names = [] @@ -231,45 +268,45 @@ def get_openpmd_data(self, global_time, diags): arrays = [] # Add requested fields to lists. - if 'E' in diags: - names += ['E'] - comps += [['r', 'z']] + if "E" in diags: + names += ["E"] + comps += [["r", "z"]] attrs += [{}] arrays += [ - [np.ascontiguousarray(self.e_r.T[2:-2, 2:-2]), - np.ascontiguousarray(self.e_z.T[2:-2, 2:-2])] + [ + np.ascontiguousarray(self.e_r.T[2:-2, 2:-2]), + np.ascontiguousarray(self.e_z.T[2:-2, 2:-2]), + ] ] - if 'B' in diags: - names += ['B'] - comps += [['t']] + if "B" in diags: + names += ["B"] + comps += [["t"]] attrs += [{}] - arrays += [ - [np.ascontiguousarray(self.b_t.T[2:-2, 2:-2])] - ] + arrays += [[np.ascontiguousarray(self.b_t.T[2:-2, 2:-2])]] # Create dictionary with all diagnostics data. - comp_pos = [[0.5, 0.]] * len(names) + comp_pos = [[0.5, 0.0]] * len(names) fld_zip = zip(names, comps, attrs, arrays, comp_pos) diag_data = {} - diag_data['fields'] = [] + diag_data["fields"] = [] for fld, comps, attrs, arrays, pos in fld_zip: - fld += '_' + self.bunch_name - diag_data['fields'].append(fld) + fld += "_" + self.bunch_name + diag_data["fields"].append(fld) diag_data[fld] = {} if comps is not None: - diag_data[fld]['comps'] = {} + diag_data[fld]["comps"] = {} for comp, arr in zip(comps, arrays): - diag_data[fld]['comps'][comp] = {} - diag_data[fld]['comps'][comp]['array'] = arr - diag_data[fld]['comps'][comp]['position'] = pos + diag_data[fld]["comps"][comp] = {} + diag_data[fld]["comps"][comp]["array"] = arr + diag_data[fld]["comps"][comp]["position"] = pos else: - diag_data[fld]['array'] = arrays[0] - diag_data[fld]['position'] = pos - diag_data[fld]['grid'] = {} - diag_data[fld]['grid']['spacing'] = grid_spacing - diag_data[fld]['grid']['labels'] = grid_labels - diag_data[fld]['grid']['global_offset'] = grid_global_offset - diag_data[fld]['attributes'] = attrs + diag_data[fld]["array"] = arrays[0] + diag_data[fld]["position"] = pos + diag_data[fld]["grid"] = {} + diag_data[fld]["grid"]["spacing"] = grid_spacing + diag_data[fld]["grid"]["labels"] = grid_labels + diag_data[fld]["grid"]["global_offset"] = grid_global_offset + diag_data[fld]["attributes"] = attrs return diag_data @@ -285,14 +322,14 @@ def _update(self, x, y, xi): self._r_max_hist.append(r_max) self.dr = r_max / (self.nr - self.nr_border) r_max += self.nr_border * self.dr - self.r_grid = np.linspace(self.dr/2, r_max - self.dr/2, self.nr) + self.r_grid = np.linspace(self.dr / 2, r_max - self.dr / 2, self.nr) # Create grid in xi xi_min_beam = np.min(xi) xi_max_beam = np.max(xi) self.i_grid = np.where( - (self.xi_plasma > xi_min_beam - self.dxi * (1 + self.nxi_border)) & - (self.xi_plasma < xi_max_beam + self.dxi * (1 + self.nxi_border)) + (self.xi_plasma > xi_min_beam - self.dxi * (1 + self.nxi_border)) + & (self.xi_plasma < xi_max_beam + self.dxi * (1 + self.nxi_border)) )[0] self.xi_grid = self.xi_plasma[self.i_grid] self.xi_max = self.xi_grid[-1] @@ -309,17 +346,27 @@ def _update(self, x, y, xi): def _reset_fields(self): """Reset value of the fields at the grid.""" - self.psi_grid[:] = 0. - self.b_t[:] = 0. - self.e_r[:] = 0. - self.e_z[:] = 0. + self.psi_grid[:] = 0.0 + self.b_t[:] = 0.0 + self.e_r[:] = 0.0 + self.e_z[:] = 0.0 @njit_serial() def calculate_fields_on_grid( - i_grid, r_grid, s_d, - psi_grid, bt_grid, r_hist, log_r_hist, sum_1_hist, sum_2_hist, - a_0_hist, a_i_hist, b_i_hist): + i_grid, + r_grid, + s_d, + psi_grid, + bt_grid, + r_hist, + log_r_hist, + sum_1_hist, + sum_2_hist, + a_0_hist, + a_i_hist, + b_i_hist, +): """Compute the plasma fields on the grid. Compiling this method in numba avoids significant overhead. @@ -336,7 +383,7 @@ def calculate_fields_on_grid( log_r=log_r_hist[j, :n_elec], sum_1_arr=sum_1_hist[j, :n_elec + 1], sum_2_arr=sum_2_hist[j, :n_elec + 1], - psi=psi + psi=psi, ) calculate_psi_with_interpolation( r_eval=r_grid / s_d, @@ -353,5 +400,5 @@ def calculate_fields_on_grid( a=a_i_hist[j], b=b_i_hist[j], r=r_hist[j, :n_elec], - b_theta=b_theta + b_theta=b_theta, ) diff --git a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/b_theta.py b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/b_theta.py index 6a2ff11..b293eea 100644 --- a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/b_theta.py +++ b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/b_theta.py @@ -9,15 +9,29 @@ @njit_serial() def calculate_b_theta_at_particles( - r_e, pr_e, w_e, w_center_e, gamma_e, q_e, - r_i, + r_e, + pr_e, + w_e, + w_center_e, + gamma_e, + q_e, + r_i, ion_motion, - psi_e, dr_psi_e, dxi_psi_e, - b_t_0_e, nabla_a2_e, - A, B, C, - K, U, - a_0, a, b, - b_t_e, b_t_i + psi_e, + dr_psi_e, + dxi_psi_e, + b_t_0_e, + nabla_a2_e, + A, + B, + C, + K, + U, + a_0, + a, + b, + b_t_e, + b_t_i, ): """Calculate the azimuthal magnetic field at the plasma particles. @@ -27,7 +41,7 @@ def calculate_b_theta_at_particles( slower than the electrons. The value of b_theta at a a radial position r is calculated as - + b_theta = a_i * r + b_i / r This requires determining the values of the a_i and b_i coefficients, @@ -107,26 +121,34 @@ def calculate_b_theta_at_particles( # Calculate the A_i, B_i, C_i coefficients in Eq. (26). calculate_ABC( - r_e, pr_e, gamma_e, - psi_e, dr_psi_e, dxi_psi_e, b_t_0_e, - nabla_a2_e, A, B, C + r_e, + pr_e, + gamma_e, + psi_e, + dr_psi_e, + dxi_psi_e, + b_t_0_e, + nabla_a2_e, + A, + B, + C, ) # Calculate the a_i, b_i coefficients in Eq. (27). calculate_KU(r_e, q_e, w_e, w_center_e, A, K, U) - calculate_ai_bi_from_axis(r_e, q_e, w_e, w_center_e, A, B, C, K, U, a_0, a, b) + calculate_ai_bi_from_axis( + r_e, q_e, w_e, w_center_e, A, B, C, K, U, a_0, a, b + ) # Calculate b_theta at plasma particles. calculate_b_theta_at_particle_centers(a, b, r_e, b_t_e) check_b_theta(b_t_e) if ion_motion: - calculate_b_theta_with_interpolation( - r_i, a_0[0], a, b, r_e, b_t_i - ) + calculate_b_theta_with_interpolation(r_i, a_0[0], a, b, r_e, b_t_i) check_b_theta(b_t_i) -@njit_serial(error_model='numpy') +@njit_serial(error_model="numpy") def calculate_b_theta_with_interpolation(r_fld, a_0, a, b, r, b_theta): """ Calculate the azimuthal magnetic field from the plasma at the radial @@ -138,9 +160,9 @@ def calculate_b_theta_with_interpolation(r_fld, a_0, a, b, r, b_theta): n_points = r_fld.shape[0] i_last = 0 a_i = a_0 - b_i = 0. - b_theta_left = 0. - r_left = 0. + b_i = 0.0 + b_theta_left = 0.0 + r_left = 0.0 for j in range(n_points): r_j = r_fld[j] # Get index of last plasma particle with r_i < r_j, continuing from @@ -166,7 +188,8 @@ def calculate_b_theta_with_interpolation(r_fld, a_0, a, b, r, b_theta): b_theta_j = a_i * r_j + b_i / r_j b_theta[j] = b_theta_j -@njit_serial(error_model='numpy') + +@njit_serial(error_model="numpy") def calculate_b_theta_at_particle_centers(a, b, r, b_theta): """ Calculate the azimuthal magnetic field from the plasma at the radial @@ -182,7 +205,7 @@ def calculate_b_theta_at_particle_centers(a, b, r, b_theta): b_theta[i] = a_i * r_i + b_i / r_i -@njit_serial(error_model='numpy') +@njit_serial(error_model="numpy") def calculate_ai_bi_from_axis(r, q, w, w_center, A, B, C, K, U, a_0, a, b): """ Calculate the values of a_i and b_i which are needed to determine @@ -195,10 +218,10 @@ def calculate_ai_bi_from_axis(r, q, w, w_center, A, B, C, K, U, a_0, a, b): n_part = r.shape[0] # Establish initial conditions (T_0 = 0, P_0 = 0) - T_im1 = 0. - P_im1 = 0. + T_im1 = 0.0 + P_im1 = 0.0 - a_0[:] = 0. + a_0[:] = 0.0 i_start = 0 @@ -214,30 +237,56 @@ def calculate_ai_bi_from_axis(r, q, w, w_center, A, B, C, K, U, a_0, a, b): C_i = C[i] A_inv_r_i = A_i / r_i A_r_i = A_i * r_i - A_r_i_3 = A_r_i * r_i * r_i + A_r_i_3 = A_r_i * r_i * r_i # Calculate value of coefficients at the center of the particles. - l_i = 1. + 0.5 * q_center_i * A_r_i + l_i = 1.0 + 0.5 * q_center_i * A_r_i m_i = 0.5 * q_center_i * A_inv_r_i n_i = -0.5 * q_center_i * A_r_i_3 - o_i = 1. - 0.5 * q_center_i * A_r_i - a[i] = l_i * T_im1 + m_i * P_im1 + 0.5 * q_center_i * B_i + 0.25 * q_center_i * q_center_i * A_i * C_i - b[i] = n_i * T_im1 + o_i * P_im1 + r_i * ( - q_center_i * C_i - 0.5 * q_center_i * B_i * r_i - 0.25 * q_center_i * q_center_i * A_i * C_i * r_i) + o_i = 1.0 - 0.5 * q_center_i * A_r_i + a[i] = ( + l_i * T_im1 + + m_i * P_im1 + + 0.5 * q_center_i * B_i + + 0.25 * q_center_i * q_center_i * A_i * C_i + ) + b[i] = ( + n_i * T_im1 + + o_i * P_im1 + + r_i + * ( + q_center_i * C_i + - 0.5 * q_center_i * B_i * r_i + - 0.25 * q_center_i * q_center_i * A_i * C_i * r_i + ) + ) # But add total charge for next iteration. - l_i = 1. + 0.5 * q_i * A_r_i + l_i = 1.0 + 0.5 * q_i * A_r_i m_i = 0.5 * q_i * A_inv_r_i n_i = -0.5 * q_i * A_r_i_3 - o_i = 1. - 0.5 * q_i * A_r_i - T_i = l_i * T_im1 + m_i * P_im1 + 0.5 * q_i * B_i + 0.25 * q_i * q_i * A_i * C_i - P_i = n_i * T_im1 + o_i * P_im1 + r_i * ( - q_i * C_i - 0.5 * q_i * B_i * r_i - 0.25 * q_i * q_i * A_i * C_i * r_i) + o_i = 1.0 - 0.5 * q_i * A_r_i + T_i = ( + l_i * T_im1 + + m_i * P_im1 + + 0.5 * q_i * B_i + + 0.25 * q_i * q_i * A_i * C_i + ) + P_i = ( + n_i * T_im1 + + o_i * P_im1 + + r_i + * ( + q_i * C_i + - 0.5 * q_i * B_i * r_i + - 0.25 * q_i * q_i * A_i * C_i * r_i + ) + ) T_im1 = T_i P_im1 = P_i # Calculate a_0_diff. - a_0_diff = - a[i] / K[i] + a_0_diff = -a[i] / K[i] a_0 += a_0_diff # Calculate a_i (in T_i) and b_i (in P_i) as functions of a_0_diff. @@ -257,9 +306,12 @@ def calculate_ai_bi_from_axis(r, q, w, w_center, A, B, C, K, U, a_0, a, b): # Angel: if T_old + K_old (small number) is less than 10 orders # of magnitude smaller than T_old - K_old (big number), then we # have enough precision (from simulation tests). - if (i == i_start or i == (n_part-1) or - abs(T_old + K_old) >= 1e-10 * abs(T_old - K_old) and - abs(P_old + U_old) >= 1e-10 * abs(P_old - U_old)): + if ( + i == i_start + or i == (n_part - 1) + or abs(T_old + K_old) >= 1e-10 * abs(T_old - K_old) + and abs(P_old + U_old) >= 1e-10 * abs(P_old - U_old) + ): # Calculate a_i and b_i as functions of a_0_diff. # Store the result in T and P a[i] = T_old + K_old @@ -277,11 +329,12 @@ def calculate_ai_bi_from_axis(r, q, w, w_center, A, B, C, K, U, a_0, a, b): i_start = i_stop -@njit_serial(error_model='numpy') -def calculate_ABC(r, pr, gamma, psi, dr_psi, dxi_psi, b_theta_0, - nabla_a2, A, B, C): +@njit_serial(error_model="numpy") +def calculate_ABC( + r, pr, gamma, psi, dr_psi, dxi_psi, b_theta_0, nabla_a2, A, B, C +): """Calculate the A_i, B_i and C_i coefficients of the linear system. - + The coefficients are missing the q_i * w_i term. They are multiplied by it in following functions. """ @@ -297,33 +350,35 @@ def calculate_ABC(r, pr, gamma, psi, dr_psi, dxi_psi, b_theta_0, b_theta_0_i = b_theta_0[i] nabla_a2_i = nabla_a2[i] - a = 1. + psi_i - inv_a = 1. / a + a = 1.0 + psi_i + inv_a = 1.0 / a inv_a2 = inv_a * inv_a inv_a3 = inv_a2 * inv_a - inv_r_i = 1. / r_i + inv_r_i = 1.0 / r_i b = inv_a * inv_r_i c = inv_a2 * inv_r_i pr_i2 = pr_i * pr_i A[i] = b - B[i] = (- (gamma_i * dr_psi_i) * c - + (pr_i2 * dr_psi_i) * inv_r_i * inv_a3 - + (pr_i * dxi_psi_i) * c - + pr_i2 * inv_r_i * inv_r_i * inv_a2 - + b_theta_0_i * b - + nabla_a2_i * c * 0.5) - C[i] = (pr_i2 * c - (gamma_i * inv_a - 1.) * inv_r_i) + B[i] = ( + -(gamma_i * dr_psi_i) * c + + (pr_i2 * dr_psi_i) * inv_r_i * inv_a3 + + (pr_i * dxi_psi_i) * c + + pr_i2 * inv_r_i * inv_r_i * inv_a2 + + b_theta_0_i * b + + nabla_a2_i * c * 0.5 + ) + C[i] = pr_i2 * c - (gamma_i * inv_a - 1.0) * inv_r_i -@njit_serial(error_model='numpy') +@njit_serial(error_model="numpy") def calculate_KU(r, q, w, w_center, A, K, U): """Calculate the K_i and U_i values of the linear system.""" n_part = r.shape[0] # Establish initial conditions (K_0 = 1, U_0 = 0) - K_im1 = 1. - U_im1 = 0. + K_im1 = 1.0 + U_im1 = 0.0 for i in range(n_part): r_i = r[i] @@ -332,21 +387,21 @@ def calculate_KU(r, q, w, w_center, A, K, U): A_i = A[i] A_inv_r_i = A_i / r_i A_r_i = A_i * r_i - A_r_i_3 = A_r_i * r_i * r_i + A_r_i_3 = A_r_i * r_i * r_i # Calculate value of coefficients at the center of the particles. - l_i = (1. + 0.5 * q_center_i * A_r_i) + l_i = 1.0 + 0.5 * q_center_i * A_r_i m_i = 0.5 * q_center_i * A_inv_r_i n_i = -0.5 * q_center_i * A_r_i_3 - o_i = (1. - 0.5 * q_center_i * A_r_i) + o_i = 1.0 - 0.5 * q_center_i * A_r_i K[i] = l_i * K_im1 + m_i * U_im1 U[i] = n_i * K_im1 + o_i * U_im1 # But add total charge for next iteration. - l_i = (1. + 0.5 * q_i * A_r_i) + l_i = 1.0 + 0.5 * q_i * A_r_i m_i = 0.5 * q_i * A_inv_r_i n_i = -0.5 * q_i * A_r_i_3 - o_i = (1. - 0.5 * q_i * A_r_i) + o_i = 1.0 - 0.5 * q_i * A_r_i K_i = l_i * K_im1 + m_i * U_im1 U_i = n_i * K_im1 + o_i * U_im1 K_im1 = K_i diff --git a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/b_theta_bunch.py b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/b_theta_bunch.py index 71b3b9c..a78b3b9 100644 --- a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/b_theta_bunch.py +++ b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/b_theta_bunch.py @@ -1,4 +1,5 @@ """Defines a method to compute the source terms from a particle bunch.""" + import numpy as np import scipy.constants as ct @@ -24,7 +25,7 @@ def calculate_bunch_source(q_bunch, n_r, n_xi, b_t_bunch): A (nz+4, nr+4) array where the magnetic field will be stored. """ for i in range(n_xi): - cumsum = 0. + cumsum = 0.0 # Iterate until n_r + 2 to get bunch source also in guard cells. for j in range(n_r + 2): q_ij = q_bunch[2 + i, 2 + j] @@ -32,17 +33,29 @@ def calculate_bunch_source(q_bunch, n_r, n_xi, b_t_bunch): # At each grid cell, calculate integral only until cell center by # assuming that half the charge is evenly distributed within the # cell (i.e., subtract half the charge) - b_t_bunch[2 + i, 2 + j] = (cumsum - 0.5 * q_ij) * 1. / (0.5 + j) + b_t_bunch[2 + i, 2 + j] = (cumsum - 0.5 * q_ij) * 1.0 / (0.5 + j) # At the first grid point along r, subtract an additional 1/4 of the # charge. This comes from assuming that the density has to be zero on # axis. - b_t_bunch[2 + i, 2] -= 0.25 * q_bunch[2 + i, 2] * 2. + b_t_bunch[2 + i, 2] -= 0.25 * q_bunch[2 + i, 2] * 2.0 @njit_serial() def deposit_bunch_charge( - x, y, z, q, n_p, n_r, n_xi, r_grid, xi_grid, dr, dxi, p_shape, q_bunch, - r_min_deposit=0. + x, + y, + z, + q, + n_p, + n_r, + n_xi, + r_grid, + xi_grid, + dr, + dxi, + p_shape, + q_bunch, + r_min_deposit=0.0, ): """ Deposit the charge of particle bunch in a 2D grid. @@ -73,13 +86,25 @@ def deposit_bunch_charge( Whether all particles managed to deposit on the grid. """ n_part = x.shape[0] - s_d = ct.c / np.sqrt(ct.e**2 * n_p / (ct.m_e*ct.epsilon_0)) - k = 1. / (2 * np.pi * ct.e * dr * dxi * s_d * n_p) + s_d = ct.c / np.sqrt(ct.e**2 * n_p / (ct.m_e * ct.epsilon_0)) + k = 1.0 / (2 * np.pi * ct.e * dr * dxi * s_d * n_p) w = np.empty(n_part) for i in range(n_part): w[i] = q[i] * k return deposit_3d_distribution( - z, x, y, w, xi_grid[0], r_grid[0], n_xi, - n_r, dxi, dr, q_bunch, p_shape=p_shape, - use_ruyten=True, r_min_deposit=r_min_deposit) + z, + x, + y, + w, + xi_grid[0], + r_grid[0], + n_xi, + n_r, + dxi, + dr, + q_bunch, + p_shape=p_shape, + use_ruyten=True, + r_min_deposit=r_min_deposit, + ) diff --git a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/deposition.py b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/deposition.py index 392c83e..c60f10e 100644 --- a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/deposition.py +++ b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/deposition.py @@ -9,8 +9,9 @@ @njit_serial() -def deposit_plasma_particles(r, w, r_min, nr, dr, deposition_array, - p_shape='cubic'): +def deposit_plasma_particles( + r, w, r_min, nr, dr, deposition_array, p_shape="cubic" +): """ Deposit the the weight of a 1D slice of plasma particles into a 2D r-z grid. @@ -37,17 +38,19 @@ def deposit_plasma_particles(r, w, r_min, nr, dr, deposition_array, Particle shape to be used. Possible values are 'linear' or 'cubic'. """ - if p_shape == 'linear': + if p_shape == "linear": return deposit_plasma_particles_linear( - r, w, r_min, nr, dr, deposition_array) - elif p_shape == 'cubic': + r, w, r_min, nr, dr, deposition_array + ) + elif p_shape == "cubic": return deposit_plasma_particles_cubic( - r, w, r_min, nr, dr, deposition_array) + r, w, r_min, nr, dr, deposition_array + ) @njit_serial(fastmath=True, error_model="numpy") def deposit_plasma_particles_linear(r, q, r_min, nr, dr, deposition_array): - """ Calculate charge distribution assuming linear particle shape. """ + """Calculate charge distribution assuming linear particle shape.""" r_max = nr * dr @@ -69,7 +72,7 @@ def deposit_plasma_particles_linear(r, q, r_min, nr, dr, deposition_array): u_r = r_cell + 2 - ir_cell # Precalculate quantities. - rsl_0 = 1. - u_r + rsl_0 = 1.0 - u_r rsl_1 = u_r # Add contribution of particle to density array. @@ -79,15 +82,15 @@ def deposit_plasma_particles_linear(r, q, r_min, nr, dr, deposition_array): # Apply correction on axis (ensures uniform density in a uniform # plasma) deposition_array[2] -= deposition_array[1] - deposition_array[1] = 0. + deposition_array[1] = 0.0 for i in range(nr): - deposition_array[i+2] /= (r_min + i * dr) * dr + deposition_array[i + 2] /= (r_min + i * dr) * dr @njit_serial(fastmath=True, error_model="numpy") def deposit_plasma_particles_cubic(r, q, r_min, nr, dr, deposition_array): - """ Calculate charge distribution assuming cubic particle shape. """ + """Calculate charge distribution assuming cubic particle shape.""" r_max = nr * dr @@ -109,14 +112,14 @@ def deposit_plasma_particles_cubic(r, q, r_min, nr, dr, deposition_array): u_r = r_cell - ir_cell + 1 # Precalculate quantities for shape coefficients. - inv_6 = 1. / 6. - v_r = 1. - u_r + inv_6 = 1.0 / 6.0 + v_r = 1.0 - u_r # Cubic particle shape coefficients in z and r. - rsc_0 = inv_6 * v_r ** 3 - rsc_1 = inv_6 * (3. * u_r**3 - 6. * u_r**2 + 4.) - rsc_2 = inv_6 * (3. * v_r**3 - 6. * v_r**2 + 4.) - rsc_3 = inv_6 * u_r ** 3 + rsc_0 = inv_6 * v_r**3 + rsc_1 = inv_6 * (3.0 * u_r**3 - 6.0 * u_r**2 + 4.0) + rsc_2 = inv_6 * (3.0 * v_r**3 - 6.0 * v_r**2 + 4.0) + rsc_3 = inv_6 * u_r**3 # Add contribution of particle to density array. deposition_array[ir_cell + 0] += rsc_0 * w_i @@ -128,8 +131,8 @@ def deposit_plasma_particles_cubic(r, q, r_min, nr, dr, deposition_array): # plasma) deposition_array[2] -= deposition_array[1] deposition_array[3] -= deposition_array[0] - deposition_array[0] = 0. - deposition_array[1] = 0. + deposition_array[0] = 0.0 + deposition_array[1] = 0.0 for i in range(nr): - deposition_array[i+2] /= (r_min + i * dr) * dr + deposition_array[i + 2] /= (r_min + i * dr) * dr diff --git a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/gather.py b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/gather.py index e94e022..495238f 100644 --- a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/gather.py +++ b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/gather.py @@ -61,11 +61,11 @@ def gather_laser_sources(a2, nabla_a, r_min, r_max, dr, r, a2_pp, nabla_a_pp): # Interpolate in r dr_u = ir_upper - r_i_cell dr_l = 1 - dr_u - a2_pp[i] = dr_u*fld_1_l + dr_l*fld_1_u - nabla_a_pp[i] = dr_u*fld_2_l + dr_l*fld_2_u + a2_pp[i] = dr_u * fld_1_l + dr_l * fld_1_u + nabla_a_pp[i] = dr_u * fld_2_l + dr_l * fld_2_u else: - a2_pp[i] = 0. - nabla_a_pp[i] = 0. + a2_pp[i] = 0.0 + nabla_a_pp[i] = 0.0 @njit_serial(error_model="numpy") @@ -131,12 +131,12 @@ def gather_bunch_sources(b_t, r_min, r_max, dr, r, b_t_pp): fld_l = b_t[ir_lower] * sign fld_u = b_t[ir_upper] else: - r_lower = (0.5 + ir_lower-2) * dr - r_upper = (0.5 + ir_upper-2) * dr + r_lower = (0.5 + ir_lower - 2) * dr + r_upper = (0.5 + ir_upper - 2) * dr fld_l = b_t[-1] * r_max / r_lower * sign fld_u = b_t[-1] * r_max / r_upper # Interpolate in r dr_u = ir_upper - r_i_cell dr_l = 1 - dr_u - b_t_pp[i] += dr_u*fld_l + dr_l*fld_u + b_t_pp[i] += dr_u * fld_l + dr_l * fld_u diff --git a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/plasma_particles.py b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/plasma_particles.py index d7ce275..5a01b2d 100644 --- a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/plasma_particles.py +++ b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/plasma_particles.py @@ -1,22 +1,32 @@ """Contains the definition of the `PlasmaParticles` class.""" + from typing import Optional, List, Callable import numpy as np import scipy.constants as ct -from .psi_and_derivatives import (calculate_psi_with_interpolation, - calculate_psi_and_derivatives_at_particles) +from .psi_and_derivatives import ( + calculate_psi_with_interpolation, + calculate_psi_and_derivatives_at_particles, +) from .deposition import deposit_plasma_particles from .gather import gather_bunch_sources, gather_laser_sources -from .b_theta import (calculate_b_theta_at_particles, - calculate_b_theta_with_interpolation) +from .b_theta import ( + calculate_b_theta_at_particles, + calculate_b_theta_with_interpolation, +) from .plasma_push.ab2 import evolve_plasma_ab2 from .utils import ( - calculate_chi, calculate_rho, update_gamma_and_pz, sort_particle_arrays, - check_gamma, log) + calculate_chi, + calculate_rho, + update_gamma_and_pz, + sort_particle_arrays, + check_gamma, + log, +) -class PlasmaParticles(): +class PlasmaParticles: """ Class containing a 1D slice of plasma particles. @@ -75,14 +85,14 @@ def __init__( nr: int, nz: int, radial_density: Callable[[float], float], - max_gamma: Optional[float] = 10., + max_gamma: Optional[float] = 10.0, ion_motion: Optional[bool] = True, ion_mass: Optional[float] = ct.m_p, free_electrons_per_ion: Optional[int] = 1, - pusher: Optional[str] = 'ab2', - shape: Optional[str] = 'linear', + pusher: Optional[str] = "ab2", + shape: Optional[str] = "linear", store_history: Optional[bool] = False, - diags: Optional[List[str]] = [] + diags: Optional[List[str]] = [], ): # Store parameters. @@ -106,7 +116,7 @@ def initialize(self): """Initialize column of plasma particles.""" # Create radial distribution of plasma particles. - rmin = 0. + rmin = 0.0 for i in range(self.ppc.shape[0]): rmax = self.ppc[i, 0] ppc = self.ppc[i, 1] @@ -138,13 +148,13 @@ def initialize(self): gamma = np.ones(self.n_elec) tag = np.arange(self.n_elec, dtype=np.int32) w = dr_p * r * self.radial_density(r) - w_center = w / 2 - dr_p ** 2 / 8 + w_center = w / 2 - dr_p**2 / 8 # Charge and mass of the macroparticles of each species. self.m_elec = self.free_electrons_per_ion self.m_ion = self.ion_mass / ct.m_e self.q_species_elec = self.free_electrons_per_ion - self.q_species_ion = - self.free_electrons_per_ion + self.q_species_ion = -self.free_electrons_per_ion # Combine arrays of both species. self.r = np.concatenate((r, r)) @@ -173,7 +183,7 @@ def initialize(self): self.b_i_hist = np.zeros((self.nz, self.n_elec)) self.a_0_hist = np.zeros(self.nz) self.i_push = 0 - self.xi_current = 0. + self.xi_current = 0.0 self.ions_computed = False @@ -183,16 +193,16 @@ def initialize(self): self._make_species_views() # Allocate arrays needed for the particle pusher. - if self.pusher == 'ab2': + if self.pusher == "ab2": self._allocate_ab2_arrays() def sort(self): """Sort plasma particles radially. - + The `q_species` and `m` arrays do not need to be sorted because all particles have the same value. """ - i_sort_e = np.argsort(self.r_elec, kind='stable') + i_sort_e = np.argsort(self.r_elec, kind="stable") sort_particle_arrays( self.r_elec, self.dr_p_elec, @@ -208,7 +218,7 @@ def sort(self): i_sort_e, ) if self.ion_motion: - i_sort_i = np.argsort(self.r_ion, kind='stable') + i_sort_i = np.argsort(self.r_ion, kind="stable") sort_particle_arrays( self.r_ion, self.dr_p_ion, @@ -228,19 +238,32 @@ def gather_laser_sources(self, a2, nabla_a2, r_min, r_max, dr): """Gather the source terms (a^2 and nabla(a)^2) from the laser.""" if self.ion_motion: gather_laser_sources( - a2, nabla_a2, r_min, r_max, dr, - self.r, self._a2, self._nabla_a2 + a2, + nabla_a2, + r_min, + r_max, + dr, + self.r, + self._a2, + self._nabla_a2, ) else: gather_laser_sources( - a2, nabla_a2, r_min, r_max, dr, - self.r_elec, self._a2_e, self._nabla_a2_e + a2, + nabla_a2, + r_min, + r_max, + dr, + self.r_elec, + self._a2_e, + self._nabla_a2_e, ) - def gather_bunch_sources(self, source_arrays, source_xi_indices, - source_metadata, slice_i): + def gather_bunch_sources( + self, source_arrays, source_xi_indices, source_metadata, slice_i + ): """Gather the source terms (b_theta) from the particle bunches.""" - self._b_t_0[:] = 0. + self._b_t_0[:] = 0.0 for i in range(len(source_arrays)): array = source_arrays[i] idx = source_xi_indices[i] @@ -251,11 +274,18 @@ def gather_bunch_sources(self, source_arrays, source_xi_indices, if slice_i in idx: xi_index = slice_i + 2 - idx[0] if self.ion_motion: - gather_bunch_sources(array[xi_index], r_min, r_max, dr, - self.r, self._b_t_0) + gather_bunch_sources( + array[xi_index], r_min, r_max, dr, self.r, self._b_t_0 + ) else: - gather_bunch_sources(array[xi_index], r_min, r_max, dr, - self.r_elec, self._b_t_0_e) + gather_bunch_sources( + array[xi_index], + r_min, + r_max, + dr, + self.r_elec, + self._b_t_0_e, + ) def calculate_fields(self): """Calculate the fields at the plasma particles.""" @@ -265,73 +295,144 @@ def calculate_fields(self): log(self.r_ion, self.log_r_ion) calculate_psi_and_derivatives_at_particles( - self.r_elec, self.log_r_elec, self.pr_elec, self.w_elec, - self.w_center_elec, self.q_species_elec, - self.r_ion, self.log_r_ion, self.pr_ion, self.w_ion, - self.w_center_ion, self.q_species_ion, - self.ion_motion, self.ions_computed, - self._sum_1_e, self._sum_2_e, self._sum_3_e, - self._sum_1_i, self._sum_2_i, self._sum_3_i, - self._psi_e, self._dr_psi_e, self._dxi_psi_e, - self._psi_i, self._dr_psi_i, self._dxi_psi_i, - self._psi, self._dr_psi, self._dxi_psi + self.r_elec, + self.log_r_elec, + self.pr_elec, + self.w_elec, + self.w_center_elec, + self.q_species_elec, + self.r_ion, + self.log_r_ion, + self.pr_ion, + self.w_ion, + self.w_center_ion, + self.q_species_ion, + self.ion_motion, + self.ions_computed, + self._sum_1_e, + self._sum_2_e, + self._sum_3_e, + self._sum_1_i, + self._sum_2_i, + self._sum_3_i, + self._psi_e, + self._dr_psi_e, + self._dxi_psi_e, + self._psi_i, + self._dr_psi_i, + self._dxi_psi_i, + self._psi, + self._dr_psi, + self._dxi_psi, ) update_gamma_and_pz( - self.gamma_elec, self.pz_elec, self.pr_elec, - self._a2_e, self._psi_e, self.q_species_elec, self.m_elec - ) + self.gamma_elec, + self.pz_elec, + self.pr_elec, + self._a2_e, + self._psi_e, + self.q_species_elec, + self.m_elec, + ) if self.ion_motion: update_gamma_and_pz( - self.gamma_ion, self.pz_ion, self.pr_ion, - self._a2_i, self._psi_i, self.q_species_ion, self.m_ion + self.gamma_ion, + self.pz_ion, + self.pr_ion, + self._a2_i, + self._psi_i, + self.q_species_ion, + self.m_ion, ) - check_gamma(self.gamma_elec, self.pz_elec, self.pr_elec, - self.max_gamma) + check_gamma( + self.gamma_elec, self.pz_elec, self.pr_elec, self.max_gamma + ) calculate_b_theta_at_particles( - self.r_elec, self.pr_elec, self.w_elec, self.w_center_elec, - self.gamma_elec, self.q_species_elec, + self.r_elec, + self.pr_elec, + self.w_elec, + self.w_center_elec, + self.gamma_elec, + self.q_species_elec, self.r_ion, self.ion_motion, - self._psi_e, self._dr_psi_e, self._dxi_psi_e, - self._b_t_0_e, self._nabla_a2_e, - self._A, self._B, self._C, - self._K, self._U, - self._a_0, self._a_i, self._b_i, - self._b_t_e, self._b_t_i + self._psi_e, + self._dr_psi_e, + self._dxi_psi_e, + self._b_t_0_e, + self._nabla_a2_e, + self._A, + self._B, + self._C, + self._K, + self._U, + self._a_0, + self._a_i, + self._b_i, + self._b_t_e, + self._b_t_i, ) def calculate_psi_at_grid(self, r_eval, psi): """Calculate psi on the current grid slice.""" calculate_psi_with_interpolation( - r_eval, self.r_elec, self.log_r_elec, self._sum_1_e, self._sum_2_e, - psi + r_eval, + self.r_elec, + self.log_r_elec, + self._sum_1_e, + self._sum_2_e, + psi, ) calculate_psi_with_interpolation( - r_eval, self.r_ion, self.log_r_ion, self._sum_1_i, self._sum_2_i, - psi, add=True + r_eval, + self.r_ion, + self.log_r_ion, + self._sum_1_i, + self._sum_2_i, + psi, + add=True, ) def calculate_b_theta_at_grid(self, r_eval, b_theta): """Calculate b_theta on the current grid slice.""" calculate_b_theta_with_interpolation( - r_eval, self._a_0[0], self._a_i, self._b_i, self.r_elec, - b_theta + r_eval, self._a_0[0], self._a_i, self._b_i, self.r_elec, b_theta ) def evolve(self, dxi): """Evolve plasma particles to next longitudinal slice.""" evolve_plasma_ab2( - dxi, self.r_elec, self.pr_elec, self.gamma_elec, self.m_elec, - self.q_species_elec, self.r_to_x_elec, - self._nabla_a2_e, self._b_t_0_e, - self._b_t_e, self._psi_e, self._dr_psi_e, self._dr_e, self._dpr_e + dxi, + self.r_elec, + self.pr_elec, + self.gamma_elec, + self.m_elec, + self.q_species_elec, + self.r_to_x_elec, + self._nabla_a2_e, + self._b_t_0_e, + self._b_t_e, + self._psi_e, + self._dr_psi_e, + self._dr_e, + self._dpr_e, ) if self.ion_motion: evolve_plasma_ab2( - dxi, self.r_ion, self.pr_ion, self.gamma_ion, self.m_ion, - self.q_species_ion, self.r_to_x_ion, - self._nabla_a2_i, self._b_t_0_i, - self._b_t_i, self._psi_i, self._dr_psi_i, self._dr_i, self._dpr_i + dxi, + self.r_ion, + self.pr_ion, + self.gamma_ion, + self.m_ion, + self.q_species_ion, + self.r_to_x_ion, + self._nabla_a2_i, + self._b_t_0_i, + self._b_t_i, + self._psi_i, + self._dr_psi_i, + self._dr_i, + self._dpr_i, ) if self.store_history: @@ -342,13 +443,20 @@ def evolve(self, dxi): def calculate_weights(self): """Calculate the plasma density weights of each particle.""" calculate_rho( - self.q_species_elec, self.w_elec, self.pz_elec, self.gamma_elec, - self._rho_e) + self.q_species_elec, + self.w_elec, + self.pz_elec, + self.gamma_elec, + self._rho_e, + ) if self.ion_motion or not self.ions_computed: calculate_rho( - self.q_species_ion, self.w_ion, self.pz_ion, self.gamma_ion, - self._rho_i) - + self.q_species_ion, + self.w_ion, + self.pz_ion, + self.gamma_ion, + self._rho_i, + ) def deposit_rho(self, rho, rho_e, rho_i, r_fld, nr, dr): """Deposit plasma density on a grid slice.""" @@ -368,8 +476,12 @@ def deposit_rho(self, rho, rho_e, rho_i, r_fld, nr, dr): def deposit_chi(self, chi, r_fld, nr, dr): """Deposit plasma susceptibility on a grid slice.""" calculate_chi( - self.q_species_elec, self.w_elec, self.pz_elec, self.gamma_elec, - self._chi_e) + self.q_species_elec, + self.w_elec, + self.pz_elec, + self.gamma_elec, + self._chi_e, + ) deposit_plasma_particles( self.r_elec, self._chi_e, r_fld[0], nr, dr, chi, self.shape ) @@ -384,37 +496,37 @@ def get_history(self): """ if self.store_history: history = { - 'r_hist': self.r_hist, - 'log_r_hist': self.log_r_hist, - 'xi_hist': self.xi_hist, - 'pr_hist': self.pr_hist, - 'pz_hist': self.pz_hist, - 'w_hist': self.w_hist, - 'r_to_x_hist': self.r_to_x_hist, - 'tag_hist': self.tag_hist, - 'sum_1_hist': self.sum_1_hist, - 'sum_2_hist': self.sum_2_hist, - 'a_i_hist': self.a_i_hist, - 'b_i_hist': self.b_i_hist, - 'a_0_hist': self.a_0_hist, + "r_hist": self.r_hist, + "log_r_hist": self.log_r_hist, + "xi_hist": self.xi_hist, + "pr_hist": self.pr_hist, + "pz_hist": self.pz_hist, + "w_hist": self.w_hist, + "r_to_x_hist": self.r_to_x_hist, + "tag_hist": self.tag_hist, + "sum_1_hist": self.sum_1_hist, + "sum_2_hist": self.sum_2_hist, + "a_i_hist": self.a_i_hist, + "b_i_hist": self.b_i_hist, + "a_0_hist": self.a_0_hist, } return history def store_current_step(self): """Store current particle properties in the history arrays.""" - if 'r' in self.diags or self.store_history: + if "r" in self.diags or self.store_history: self.r_hist[-1 - self.i_push] = self.r - if 'z' in self.diags: + if "z" in self.diags: self.xi_hist[-1 - self.i_push] = self.xi_current - if 'pr' in self.diags: + if "pr" in self.diags: self.pr_hist[-1 - self.i_push] = self.pr - if 'pz' in self.diags: + if "pz" in self.diags: self.pz_hist[-1 - self.i_push] = self.pz - if 'w' in self.diags: + if "w" in self.diags: self.w_hist[-1 - self.i_push] = self._rho - if 'r_to_x' in self.diags: + if "r_to_x" in self.diags: self.r_to_x_hist[-1 - self.i_push] = self.r_to_x - if 'tag' in self.diags: + if "tag" in self.diags: self.tag_hist[-1 - self.i_push] = self.tag if self.store_history: self.a_0_hist[-1 - self.i_push] = self._a_0[0] @@ -521,8 +633,8 @@ def _allocate_ab2_arrays(self): size = self.n_elec self._dr = np.zeros((2, size)) self._dpr = np.zeros((2, size)) - self._dr_e = self._dr[:, :self.n_elec] - self._dpr_e = self._dpr[:, :self.n_elec] + self._dr_e = self._dr[:, : self.n_elec] + self._dpr_e = self._dpr[:, : self.n_elec] self._dr_i = self._dr[:, self.n_elec:] self._dpr_i = self._dpr[:, self.n_elec:] @@ -531,7 +643,7 @@ def _move_auxiliary_arrays_to_next_slice(self): When storing the particle history, some auxiliary arrays (e.g., those storing the cumulative sums, the a_i, b_i coefficients, ...) have to be - stored at every longitudinal step. In principle, this used to be done + stored at every longitudinal step. In principle, this used to be done by writing the 1D auxiliary arrays into the corresponding slice of the 2D history arrays. However, this is time consuming as it leads to copying data at every step. In order to avoid this, the auxiliary diff --git a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/plasma_push/ab2.py b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/plasma_push/ab2.py index 62d3309..9899ea6 100644 --- a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/plasma_push/ab2.py +++ b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/plasma_push/ab2.py @@ -1,17 +1,25 @@ """ Contains the 5th order Adams–Bashforth pusher for the plasma particles. """ - -import numpy as np - from wake_t.utilities.numba import njit_serial @njit_serial() def evolve_plasma_ab2( - dxi, r, pr, gamma, m, q, r_to_x, - nabla_a2, b_theta_0, b_theta, psi, dr_psi, - dr, dpr - ): + dxi, + r, + pr, + gamma, + m, + q, + r_to_x, + nabla_a2, + b_theta_0, + b_theta, + psi, + dr_psi, + dr, + dpr, +): """ Evolve the r and pr coordinates of plasma particles to the next xi step using an Adams-Bashforth method of 2nd order. @@ -33,8 +41,17 @@ def evolve_plasma_ab2( """ calculate_derivatives( - pr, gamma, m, q, b_theta_0, nabla_a2, b_theta, - psi, dr_psi, dr[0], dpr[0] + pr, + gamma, + m, + q, + b_theta_0, + nabla_a2, + b_theta, + psi, + dr_psi, + dr[0], + dpr[0], ) # Push radial position. @@ -54,7 +71,8 @@ def evolve_plasma_ab2( @njit_serial(fastmath=True, error_model="numpy") def calculate_derivatives( - pr, gamma, m, q, b_theta_0, nabla_a2, b_theta_bar, psi, dr_psi, dr, dpr): + pr, gamma, m, q, b_theta_0, nabla_a2, b_theta_bar, psi, dr_psi, dr, dpr +): """ Calculate the derivative of the radial position and the radial momentum of the plasma particles at the current slice. @@ -85,11 +103,13 @@ def calculate_derivatives( # Calculate derivatives of r and pr. q_over_m = q / m for i in range(pr.shape[0]): - inv_psi_i = 1. / (1. + psi[i] * q_over_m) - dpr[i] = (gamma[i] * dr_psi[i] * inv_psi_i - - b_theta_bar[i] - - b_theta_0[i] - - nabla_a2[i] * 0.5 * inv_psi_i * q_over_m) * q_over_m + inv_psi_i = 1.0 / (1.0 + psi[i] * q_over_m) + dpr[i] = ( + gamma[i] * dr_psi[i] * inv_psi_i + - b_theta_bar[i] + - b_theta_0[i] + - nabla_a2[i] * 0.5 * inv_psi_i * q_over_m + ) * q_over_m dr[i] = pr[i] * inv_psi_i @@ -114,9 +134,9 @@ def apply_ab2(x, dt, dx): def check_axis_crossing(r, pr, dr, dpr, r_to_x): """Check for particles with r < 0 and invert them.""" for i in range(r.shape[0]): - if r[i] < 0.: - r[i] *= -1. - pr[i] *= -1. - dr[i] *= -1. - dpr[i] *= -1. + if r[i] < 0.0: + r[i] *= -1.0 + pr[i] *= -1.0 + dr[i] *= -1.0 + dpr[i] *= -1.0 r_to_x[i] *= -1 diff --git a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/psi_and_derivatives.py b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/psi_and_derivatives.py index 36c0376..e9fbc43 100644 --- a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/psi_and_derivatives.py +++ b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/psi_and_derivatives.py @@ -11,14 +11,35 @@ @njit_serial(fastmath=True, error_model="numpy") def calculate_psi_and_derivatives_at_particles( - r_e, log_r_e, pr_e, w_e, w_center_e, q_e, - r_i, log_r_i, pr_i, w_i, w_center_i, q_i, - ion_motion, calculate_ion_sums, - sum_1_e, sum_2_e, sum_3_e, - sum_1_i, sum_2_i, sum_3_i, - psi_e, dr_psi_e, dxi_psi_e, - psi_i, dr_psi_i, dxi_psi_i, - psi, dr_psi, dxi_psi, + r_e, + log_r_e, + pr_e, + w_e, + w_center_e, + q_e, + r_i, + log_r_i, + pr_i, + w_i, + w_center_i, + q_i, + ion_motion, + calculate_ion_sums, + sum_1_e, + sum_2_e, + sum_3_e, + sum_1_i, + sum_2_i, + sum_3_i, + psi_e, + dr_psi_e, + dxi_psi_e, + psi_i, + dr_psi_i, + dxi_psi_i, + psi, + dr_psi, + dxi_psi, ): """Calculate wakefield potential and derivatives at the plasma particles. @@ -88,7 +109,9 @@ def calculate_psi_and_derivatives_at_particles( # Calculate cumulative sum 3 (Eq. (32)). calculate_cumulative_sum_3(q_e, r_e, pr_e, w_e, w_center_e, psi_e, sum_3_e) if ion_motion or not calculate_ion_sums: - calculate_cumulative_sum_3(q_i, r_i, pr_i, w_i, w_center_i, psi_i, sum_3_i) + calculate_cumulative_sum_3( + q_i, r_i, pr_i, w_i, w_center_i, psi_i, sum_3_i + ) # Calculate the dxi_psi background at the neighboring points. # For the electrons, compute the psi and dr_psi due to the ions at @@ -96,10 +119,13 @@ def calculate_psi_and_derivatives_at_particles( # electrons at r_neighbor_i. calculate_dxi_psi_at_particle_centers(r_e, sum_3_e, dxi_psi_e) if ion_motion: - calculate_dxi_psi_with_interpolation(r_e, r_i, sum_3_i, dxi_psi_e, add=True) + calculate_dxi_psi_with_interpolation( + r_e, r_i, sum_3_i, dxi_psi_e, add=True + ) calculate_dxi_psi_at_particle_centers(r_i, sum_3_i, dxi_psi_i) - calculate_dxi_psi_with_interpolation(r_i, r_e, sum_3_e, dxi_psi_i, add=True) - + calculate_dxi_psi_with_interpolation( + r_i, r_e, sum_3_e, dxi_psi_i, add=True + ) # Check that the values of dxi_psi are within a reasonable range (prevents # issues at the peak of a blowout wake, for example). @@ -109,7 +135,7 @@ def calculate_psi_and_derivatives_at_particles( @njit_serial(fastmath=True) def calculate_cumulative_sum_1(q, w, w_center, sum_1_arr): """Calculate the cumulative sum in Eq. (29).""" - sum_1 = 0. + sum_1 = 0.0 for i in range(w.shape[0]): w_i = w[i] w_center_i = w_center[i] @@ -124,7 +150,7 @@ def calculate_cumulative_sum_1(q, w, w_center, sum_1_arr): @njit_serial(fastmath=True) def calculate_cumulative_sum_2(q, log_r, w, w_center, sum_2_arr): """Calculate the cumulative sum in Eq. (31).""" - sum_2 = 0. + sum_2 = 0.0 for i in range(log_r.shape[0]): log_r_i = log_r[i] w_i = w[i] @@ -140,7 +166,7 @@ def calculate_cumulative_sum_2(q, log_r, w, w_center, sum_2_arr): @njit_serial(fastmath=True, error_model="numpy") def calculate_cumulative_sum_3(q, r, pr, w, w_center, psi, sum_3_arr): """Calculate the cumulative sum in Eq. (32).""" - sum_3 = 0. + sum_3 = 0.0 for i in range(r.shape[0]): r_i = r[i] pr_i = pr[i] @@ -157,7 +183,8 @@ def calculate_cumulative_sum_3(q, r, pr, w, w_center, psi, sum_3_arr): @njit_serial(fastmath=True, error_model="numpy") def calculate_psi_with_interpolation( - r_eval, r, log_r, sum_1_arr, sum_2_arr, psi, add=False): + r_eval, r, log_r, sum_1_arr, sum_2_arr, psi, add=False +): """Calculate psi at the radial positions given in `r_eval`.""" # Get number of plasma particles. n_part = r.shape[0] @@ -172,10 +199,10 @@ def calculate_psi_with_interpolation( # Calculate fields at r_eval. i_last = 0 - r_left = 0. - sum_1_left = 0. - sum_2_left = 0. - psi_left = 0. + r_left = 0.0 + sum_1_left = 0.0 + sum_2_left = 0.0 + psi_left = 0.0 for j in range(n_points): r_j = r_eval[j] # Get index of last plasma particle with r_i < r_j, continuing from @@ -198,7 +225,7 @@ def calculate_psi_with_interpolation( psi_right = sum_1_right * log_r_right - sum_2_right # Interpolate sums. - inv_dr = 1. / (r_right - r_left) + inv_dr = 1.0 / (r_right - r_left) slope_2 = (psi_right - psi_left) * inv_dr psi_j = psi_left + slope_2 * (r_j - r_left) + sum_2_max else: @@ -215,7 +242,8 @@ def calculate_psi_with_interpolation( @njit_serial(fastmath=True, error_model="numpy") def calculate_psi_and_dr_psi_with_interpolation( - r_eval, r, log_r, sum_1_arr, sum_2_arr, psi, dr_psi, add=False): + r_eval, r, log_r, sum_1_arr, sum_2_arr, psi, dr_psi, add=False +): """Calculate psi and dr_psi at the radial positions given in `r_eval`.""" # Get number of plasma particles. n_part = r.shape[0] @@ -230,11 +258,11 @@ def calculate_psi_and_dr_psi_with_interpolation( # Calculate fields at r_eval. i_last = 0 - r_left = 0. - sum_1_left = 0. - sum_2_left = 0. - psi_left = 0. - dr_psi_left = 0. + r_left = 0.0 + sum_1_left = 0.0 + sum_2_left = 0.0 + psi_left = 0.0 + dr_psi_left = 0.0 for j in range(n_points): r_j = r_eval[j] # Get index of last plasma particle with r_i < r_j, continuing from @@ -259,7 +287,7 @@ def calculate_psi_and_dr_psi_with_interpolation( psi_right = sum_1_right * log_r_right - sum_2_right # Interpolate sums. - inv_dr = 1. / (r_right - r_left) + inv_dr = 1.0 / (r_right - r_left) slope_1 = (dr_psi_right - dr_psi_left) * inv_dr slope_2 = (psi_right - psi_left) * inv_dr dr_psi_j = dr_psi_left + slope_1 * (r_j - r_left) @@ -281,8 +309,13 @@ def calculate_psi_and_dr_psi_with_interpolation( @njit_serial(fastmath=True, error_model="numpy") def calculate_psi_and_dr_psi_at_particle_centers( - r, log_r, sum_1_arr, sum_2_arr, psi, dr_psi, - ): + r, + log_r, + sum_1_arr, + sum_2_arr, + psi, + dr_psi, +): # Get number of particles. n_part = r.shape[0] @@ -302,14 +335,16 @@ def calculate_psi_and_dr_psi_at_particle_centers( @njit_serial() -def calculate_dxi_psi_with_interpolation(r_eval, r, sum_3_arr, dxi_psi, add=False): +def calculate_dxi_psi_with_interpolation( + r_eval, r, sum_3_arr, dxi_psi, add=False +): """Calculate dxi_psi at the radial position given in `r_eval`.""" # Get number of plasma particles. n_part = r.shape[0] # Get number of points to evaluate. n_points = r_eval.shape[0] - + # Calculate dxi_psi after the last plasma plasma particle # This is used to ensure the boundary condition dxi_psi=0, which also # assumes that the total electron and ion charge are the same. @@ -317,8 +352,8 @@ def calculate_dxi_psi_with_interpolation(r_eval, r, sum_3_arr, dxi_psi, add=Fals # Calculate fields at r_eval. i_last = 0 - r_left = 0. - dxi_psi_left = 0. + r_left = 0.0 + dxi_psi_left = 0.0 for j in range(n_points): r_j = r_eval[j] # Get index of last plasma particle with r_i < r_j, continuing from @@ -331,13 +366,13 @@ def calculate_dxi_psi_with_interpolation(r_eval, r, sum_3_arr, dxi_psi, add=Fals if i_last < n_part: if i_last > 0: r_left = r[i_last - 1] - dxi_psi_left = - sum_3_arr[i_last - 1] + dxi_psi_left = -sum_3_arr[i_last - 1] r_right = r[i_last] - dxi_psi_right = - sum_3_arr[i_last] + dxi_psi_right = -sum_3_arr[i_last] slope = (dxi_psi_right - dxi_psi_left) / (r_right - r_left) dxi_psi_j = dxi_psi_left + slope * (r_j - r_left) + sum_3_max else: - dxi_psi_j = - sum_3_arr[-1] + sum_3_max + dxi_psi_j = -sum_3_arr[-1] + sum_3_max if add: dxi_psi[j] += dxi_psi_j else: @@ -346,8 +381,10 @@ def calculate_dxi_psi_with_interpolation(r_eval, r, sum_3_arr, dxi_psi, add=Fals @njit_serial(fastmath=True, error_model="numpy") def calculate_dxi_psi_at_particle_centers( - r, sum_3_arr, dxi_psi, - ): + r, + sum_3_arr, + dxi_psi, +): # Get number of particles. n_part = r.shape[0] @@ -358,7 +395,7 @@ def calculate_dxi_psi_at_particle_centers( # Calculate fields. for i in range(n_part): - dxi_psi[i] = - sum_3_arr[i] + sum_3_max + dxi_psi[i] = -sum_3_arr[i] + sum_3_max @njit_serial() diff --git a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/solver.py b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/solver.py index e8f4fb7..2834f0f 100644 --- a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/solver.py +++ b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/solver.py @@ -5,24 +5,40 @@ See https://journals.aps.org/prab/abstract/10.1103/PhysRevAccelBeams.21.071301 for the full details about this model. """ + import numpy as np import scipy.constants as ct import aptools.plasma_accel.general_equations as ge from .plasma_particles import PlasmaParticles -from wake_t.utilities.numba import njit_serial from .utils import longitudinal_gradient, radial_gradient -def calculate_wakefields(laser_a2, r_max, xi_min, xi_max, - n_r, n_xi, ppc, n_p, r_max_plasma=None, - radial_density=None, p_shape='cubic', - max_gamma=10., plasma_pusher='ab2', - ion_motion=False, ion_mass=ct.m_p, - free_electrons_per_ion=1, - bunch_source_arrays=[], bunch_source_xi_indices=[], - bunch_source_metadata=[], store_plasma_history=False, - calculate_rho=True, particle_diags=[], fld_arrays=[]): +def calculate_wakefields( + laser_a2, + r_max, + xi_min, + xi_max, + n_r, + n_xi, + ppc, + n_p, + r_max_plasma=None, + radial_density=None, + p_shape="cubic", + max_gamma=10.0, + plasma_pusher="ab2", + ion_motion=False, + ion_mass=ct.m_p, + free_electrons_per_ion=1, + bunch_source_arrays=[], + bunch_source_xi_indices=[], + bunch_source_metadata=[], + store_plasma_history=False, + calculate_rho=True, + particle_diags=[], + fld_arrays=[], +): """ Calculate the plasma wakefields generated by the given laser pulse and electron beam in the specified grid points. @@ -116,8 +132,8 @@ def radial_density_normalized(r): xi_fld = xi_fld / s_d # Initialize field arrays, including guard cells. - nabla_a2 = np.zeros((n_xi+4, n_r+4)) - psi = np.zeros((n_xi+4, n_r+4)) + nabla_a2 = np.zeros((n_xi + 4, n_r + 4)) + psi = np.zeros((n_xi + 4, n_r + 4)) # Laser source. laser_source = laser_a2 is not None @@ -127,41 +143,99 @@ def radial_density_normalized(r): # Calculate plasma response (including density, susceptibility, potential # and magnetic field) pp_hist = calculate_plasma_response( - r_max, r_max_plasma, radial_density_normalized, dr, ppc, n_r, - plasma_pusher, p_shape, max_gamma, ion_motion, ion_mass, - free_electrons_per_ion, n_xi, laser_a2, nabla_a2, laser_source, - bunch_source_arrays, bunch_source_xi_indices, bunch_source_metadata, - r_fld, psi, B_t, rho, rho_e, rho_i, chi, dxi, + r_max, + r_max_plasma, + radial_density_normalized, + dr, + ppc, + n_r, + plasma_pusher, + p_shape, + max_gamma, + ion_motion, + ion_mass, + free_electrons_per_ion, + n_xi, + laser_a2, + nabla_a2, + laser_source, + bunch_source_arrays, + bunch_source_xi_indices, + bunch_source_metadata, + r_fld, + psi, + B_t, + rho, + rho_e, + rho_i, + chi, + dxi, store_plasma_history=store_plasma_history, - calculate_rho=calculate_rho, particle_diags=particle_diags + calculate_rho=calculate_rho, + particle_diags=particle_diags, ) # Calculate derived fields (E_z, W_r, and E_r). - E_0 = ge.plasma_cold_non_relativisct_wave_breaking_field(n_p*1e-6) + E_0 = ge.plasma_cold_non_relativisct_wave_breaking_field(n_p * 1e-6) longitudinal_gradient(psi[2:-2, 2:-2], dxi, E_z[2:-2, 2:-2]) radial_gradient(psi[2:-2, 2:-2], dr, E_r[2:-2, 2:-2]) E_r -= B_t - E_z *= - E_0 - E_r *= - E_0 + E_z *= -E_0 + E_r *= -E_0 # B_t[:] = (b_t_bar + b_t_beam) * E_0 / ct.c B_t *= E_0 / ct.c return pp_hist def calculate_plasma_response( - r_max, r_max_plasma, radial_density_normalized, dr, ppc, n_r, - plasma_pusher, p_shape, max_gamma, ion_motion, ion_mass, - free_electrons_per_ion, n_xi, a2, nabla_a2, laser_source, - bunch_source_arrays, bunch_source_xi_indices, bunch_source_metadata, - r_fld, psi, b_t_bar, rho, - rho_e, rho_i, chi, dxi, store_plasma_history, calculate_rho, - particle_diags + r_max, + r_max_plasma, + radial_density_normalized, + dr, + ppc, + n_r, + plasma_pusher, + p_shape, + max_gamma, + ion_motion, + ion_mass, + free_electrons_per_ion, + n_xi, + a2, + nabla_a2, + laser_source, + bunch_source_arrays, + bunch_source_xi_indices, + bunch_source_metadata, + r_fld, + psi, + b_t_bar, + rho, + rho_e, + rho_i, + chi, + dxi, + store_plasma_history, + calculate_rho, + particle_diags, ): # Initialize plasma particles. pp = PlasmaParticles( - r_max, r_max_plasma, dr, ppc, n_r, n_xi, radial_density_normalized, - max_gamma, ion_motion, ion_mass, free_electrons_per_ion, - plasma_pusher, p_shape, store_plasma_history, particle_diags + r_max, + r_max_plasma, + dr, + ppc, + n_r, + n_xi, + radial_density_normalized, + max_gamma, + ion_motion, + ion_mass, + free_electrons_per_ion, + plasma_pusher, + p_shape, + store_plasma_history, + particle_diags, ) pp.initialize() @@ -174,23 +248,33 @@ def calculate_plasma_response( if laser_source: pp.gather_laser_sources( - a2[slice_i+2], nabla_a2[slice_i+2], r_fld[0], r_fld[-1], dr + a2[slice_i + 2], nabla_a2[slice_i + 2], r_fld[0], r_fld[-1], dr ) - pp.gather_bunch_sources(bunch_source_arrays, bunch_source_xi_indices, - bunch_source_metadata, slice_i) + pp.gather_bunch_sources( + bunch_source_arrays, + bunch_source_xi_indices, + bunch_source_metadata, + slice_i, + ) pp.calculate_fields() - pp.calculate_psi_at_grid(r_fld, psi[slice_i+2, 2:-2]) - pp.calculate_b_theta_at_grid(r_fld, b_t_bar[slice_i+2, 2:-2]) + pp.calculate_psi_at_grid(r_fld, psi[slice_i + 2, 2:-2]) + pp.calculate_b_theta_at_grid(r_fld, b_t_bar[slice_i + 2, 2:-2]) if calculate_rho: - pp.deposit_rho(rho[slice_i+2], rho_e[slice_i+2], rho_i[slice_i+2], - r_fld, n_r, dr) - elif 'w' in particle_diags: + pp.deposit_rho( + rho[slice_i + 2], + rho_e[slice_i + 2], + rho_i[slice_i + 2], + r_fld, + n_r, + dr, + ) + elif "w" in particle_diags: pp.calculate_weights() if laser_source: - pp.deposit_chi(chi[slice_i+2], r_fld, n_r, dr) + pp.deposit_chi(chi[slice_i + 2], r_fld, n_r, dr) pp.ions_computed = True diff --git a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/utils.py b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/utils.py index 9ff62e4..c0dd6aa 100644 --- a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/utils.py +++ b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/utils.py @@ -20,8 +20,8 @@ def calculate_chi(q, w, pz, gamma, chi): for i in range(w.shape[0]): w_i = w[i] pz_i = pz[i] - inv_gamma_i = 1. / gamma[i] - chi[i] = q * w_i / (1. - pz_i * inv_gamma_i) * inv_gamma_i + inv_gamma_i = 1.0 / gamma[i] + chi[i] = q * w_i / (1.0 - pz_i * inv_gamma_i) * inv_gamma_i @njit_serial(error_model="numpy") @@ -30,8 +30,8 @@ def calculate_rho(q, w, pz, gamma, rho): for i in range(w.shape[0]): w_i = w[i] pz_i = pz[i] - inv_gamma_i = 1. / gamma[i] - rho[i] = q * w_i / (1. - pz_i * inv_gamma_i) + inv_gamma_i = 1.0 / gamma[i] + rho[i] = q * w_i / (1.0 - pz_i * inv_gamma_i) @njit_serial() @@ -52,19 +52,19 @@ def longitudinal_gradient(f, dz, dz_f): Array where the longitudinal gradient will be stored. """ nz, nr = f.shape - inv_dz = 1. / dz + inv_dz = 1.0 / dz inv_h = 0.5 * inv_dz - a = - 1.5 * inv_dz - b = 2. * inv_dz - c = - 0.5 * inv_dz + a = -1.5 * inv_dz + b = 2.0 * inv_dz + c = -0.5 * inv_dz for j in range(nr): - f_right = f[nz-1, j] + f_right = f[nz - 1, j] for i in range(1, nz - 1): f_left = f[i - 1, j] - f_right = f[i + 1, j] + f_right = f[i + 1, j] dz_f[i, j] = (f_right - f_left) * inv_h dz_f[0, j] = a * f[0, j] + b * f[1, j] + c * f[2, j] - dz_f[-1, j] = - a * f[-1, j] - b * f[-2, j] - c * f[-3, j] + dz_f[-1, j] = -a * f[-1, j] - b * f[-2, j] - c * f[-3, j] @njit_serial() @@ -86,18 +86,18 @@ def radial_gradient(f, dr, dr_f): Array where the radial gradient will be stored. """ nz, nr = f.shape - inv_dr = 1. / dr + inv_dr = 1.0 / dr inv_h = 0.5 * inv_dr - a = - 1.5 * inv_dr - b = 2. * inv_dr - c = - 0.5 * inv_dr + a = -1.5 * inv_dr + b = 2.0 * inv_dr + c = -0.5 * inv_dr for i in range(nz): for j in range(1, nr - 1): f_left = f[i, j - 1] f_right = f[i, j + 1] dr_f[i, j] = (f_right - f_left) * inv_h dr_f[i, 0] = (f[i, 1] - f[i, 0]) * inv_h - dr_f[i, -1] = - a * f[i, -1] - b * f[i, -2] - c * f[i, -3] + dr_f[i, -1] = -a * f[i, -1] - b * f[i, -2] - c * f[i, -3] @njit_serial() @@ -121,7 +121,7 @@ def calculate_laser_a2(a_complex, a2): a2[2 + i, 2 + j] = ar * ar + ai * ai -@njit_serial(error_model='numpy') +@njit_serial(error_model="numpy") def update_gamma_and_pz(gamma, pz, pr, a2, psi, q, m): """ Update the gamma factor and longitudinal momentum of the plasma particles. @@ -141,12 +141,11 @@ def update_gamma_and_pz(gamma, pz, pr, a2, psi, q, m): q_over_m = q / m for i in range(pr.shape[0]): psi_i = psi[i] * q_over_m - pz_i = ( - (1 + pr[i] ** 2 + q_over_m ** 2 * a2[i] - (1 + psi_i) ** 2) / - (2 * (1 + psi_i)) + pz_i = (1 + pr[i] ** 2 + q_over_m**2 * a2[i] - (1 + psi_i) ** 2) / ( + 2 * (1 + psi_i) ) pz[i] = pz_i - gamma[i] = 1. + pz_i + psi_i + gamma[i] = 1.0 + pz_i + psi_i @njit_serial() @@ -154,15 +153,17 @@ def check_gamma(gamma, pz, pr, max_gamma): """Check that the gamma of particles does not exceed `max_gamma`""" for i in range(gamma.shape[0]): if gamma[i] > max_gamma: - gamma[i] = 1. - pz[i] = 0. - pr[i] = 0. + gamma[i] = 1.0 + pz[i] = 0.0 + pr[i] = 0.0 @njit_serial() -def sort_particle_arrays(a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11, indices): +def sort_particle_arrays( + a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11, indices +): """Sort all the particle arrays with a given sorting order. - + Implementing it like this looks very ugly, but it is much faster than repeating `array = array[indices]` for each array. It is also much faster than implementing a `sort_array` function that is called on each array. diff --git a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/wakefield.py b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/wakefield.py index 106d02e..7e74243 100644 --- a/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/wakefield.py +++ b/wake_t/physics_models/plasma_wakefields/qs_rz_baxevanis_ion/wakefield.py @@ -186,9 +186,9 @@ def __init__( dz_fields: Optional[float] = None, r_max_plasma: Optional[float] = None, parabolic_coefficient: Optional[float] = None, - p_shape: Optional[str] = 'cubic', + p_shape: Optional[str] = "cubic", max_gamma: Optional[float] = 10, - plasma_pusher: Optional[str] = 'ab2', + plasma_pusher: Optional[str] = "ab2", ion_motion: Optional[bool] = False, ion_mass: Optional[float] = ct.m_p, free_electrons_per_ion: Optional[int] = 1, @@ -198,14 +198,19 @@ def __init__( laser_envelope_nxi: Optional[int] = None, laser_envelope_nr: Optional[int] = None, laser_envelope_use_phase: Optional[bool] = True, - field_diags: Optional[List[str]] = ['rho', 'E', 'B', 'a_mod', - 'a_phase'], + field_diags: Optional[List[str]] = [ + "rho", + "E", + "B", + "a_mod", + "a_phase", + ], particle_diags: Optional[List[str]] = [], use_adaptive_grids: Optional[bool] = False, adaptive_grid_nr: Optional[Union[int, List[int]]] = 16, adaptive_grid_r_max: Optional[Union[float, List[float]]] = None, adaptive_grid_r_lim: Optional[Union[float, List[float]]] = None, - adaptive_grid_diags: Optional[List[str]] = ['E', 'B'], + adaptive_grid_diags: Optional[List[str]] = ["E", "B"], ) -> None: # Checks for backward compatibility. if plasma_pusher not in ["ab2"]: @@ -214,7 +219,7 @@ def __init__( warn( "Plasma pusher {plasma_pusher} has been deprecated. " "Using 'ab2' pusher instead.", - DeprecationWarning + DeprecationWarning, ) else: raise ValueError( @@ -228,7 +233,7 @@ def __init__( "compatibility, but the " "new recommended approach is to provide a density function " "that takes `z` and `r` as input parameters.", - DeprecationWarning + DeprecationWarning, ) parabolic_coefficient = self._get_parabolic_coefficient_fn( parabolic_coefficient @@ -249,11 +254,11 @@ def __init__( self.adaptive_grid_r_lim = adaptive_grid_r_lim self.adaptive_grid_diags = adaptive_grid_diags self.bunch_grids: Dict[str, AdaptiveGrid] = {} - self._t_reset_bunch_arrays = -1. + self._t_reset_bunch_arrays = -1.0 if len(self.ppc.shape) in [0, 1]: self.ppc = np.array([[self.r_max_plasma, self.ppc.flatten()[0]]]) self.parabolic_coefficient = parabolic_coefficient - + super().__init__( density_function=density_function, r_max=r_max, @@ -271,28 +276,37 @@ def __init__( laser_envelope_use_phase=laser_envelope_use_phase, field_diags=field_diags, particle_diags=particle_diags, - model_name='quasistatic_2d_ion' + model_name="quasistatic_2d_ion", ) def _initialize_properties(self, bunches): super()._initialize_properties(bunches) # Add bunch source array (needed if not using adaptive grids). - self.b_t_bunch = np.zeros((self.n_xi+4, self.n_r+4)) - self.q_bunch = np.zeros((self.n_xi+4, self.n_r+4)) - self.laser_a2 = np.zeros((self.n_xi+4, self.n_r+4)) - self.fld_arrays = [self.rho, self.rho_e, self.rho_i, self.chi, self.e_r, - self.e_z, self.b_t, self.xi_fld, self.r_fld] + self.b_t_bunch = np.zeros((self.n_xi + 4, self.n_r + 4)) + self.q_bunch = np.zeros((self.n_xi + 4, self.n_r + 4)) + self.laser_a2 = np.zeros((self.n_xi + 4, self.n_r + 4)) + self.fld_arrays = [ + self.rho, + self.rho_e, + self.rho_i, + self.chi, + self.e_r, + self.e_z, + self.b_t, + self.xi_fld, + self.r_fld, + ] def _calculate_wakefield(self, bunches: List[ParticleBunch]): - radial_density = self._get_radial_density(self.t*ct.c) + radial_density = self._get_radial_density(self.t * ct.c) # Get square of laser envelope if self.laser is not None: calculate_laser_a2(self.laser.get_envelope(), self.laser_a2) # If linearly polarized, divide by 2 so that the ponderomotive # force on the plasma particles is correct. - if self.laser.polarization == 'linear': - self.laser_a2 /= 2. + if self.laser.polarization == "linear": + self.laser_a2 /= 2.0 laser_a2 = self.laser_a2 else: laser_a2 = None @@ -314,9 +328,9 @@ def _calculate_wakefield(self, bunches: List[ParticleBunch]): # Get radial grid resolution. if isinstance(self.adaptive_grid_nr, list): assert len(self.adaptive_grid_nr) == len(bunches), ( - 'Several resolutions for the adaptive grids have been ' - 'given, but they do not match the number of tracked ' - 'bunches' + "Several resolutions for the adaptive grids have been " + "given, but they do not match the number of tracked " + "bunches" ) nr_grids = self.adaptive_grid_nr else: @@ -324,9 +338,9 @@ def _calculate_wakefield(self, bunches: List[ParticleBunch]): # Get radial extent. if isinstance(self.adaptive_grid_r_max, list): assert len(self.adaptive_grid_r_max) == len(bunches), ( - 'Several `r_max` for the adaptive grids have been ' - 'given, but they do not match the number of tracked ' - 'bunches' + "Several `r_max` for the adaptive grids have been " + "given, but they do not match the number of tracked " + "bunches" ) r_max_grids = self.adaptive_grid_r_max else: @@ -334,9 +348,9 @@ def _calculate_wakefield(self, bunches: List[ParticleBunch]): # Get radial extent limit. if isinstance(self.adaptive_grid_r_lim, list): assert len(self.adaptive_grid_r_lim) == len(bunches), ( - 'Several `r_lim` for the adaptive grids have been ' - 'given, but they do not match the number of tracked ' - 'bunches' + "Several `r_lim` for the adaptive grids have been " + "given, but they do not match the number of tracked " + "bunches" ) r_lim_grids = self.adaptive_grid_r_lim else: @@ -352,7 +366,7 @@ def _calculate_wakefield(self, bunches: List[ParticleBunch]): bunches_with_grid: List[ParticleBunch] = [] bunches_without_grid: List[ParticleBunch] = [] for i, bunch in enumerate(bunches): - if nr_grids[i] is not None: + if nr_grids[i] is not None: bunches_with_grid.append(bunch) if bunch.name not in self.bunch_grids: self.bunch_grids[bunch.name] = AdaptiveGrid( @@ -364,7 +378,7 @@ def _calculate_wakefield(self, bunches: List[ParticleBunch]): self.n_xi, self.xi_fld, r_max_grids[i], - r_lim_grids[i] + r_lim_grids[i], ) else: bunches_without_grid.append(bunch) @@ -377,9 +391,8 @@ def _calculate_wakefield(self, bunches: List[ParticleBunch]): bunch_source_arrays.append(grid.b_t_bunch) bunch_source_xi_indices.append(grid.i_grid) bunch_source_metadata.append( - np.array( - [grid.r_min_cell, grid.r_max_cell_guard, grid.dr] - ) / s_d + np.array([grid.r_min_cell, grid.r_max_cell_guard, grid.dr]) + / s_d ) if not all_deposited: self._reset_bunch_arrays() @@ -397,7 +410,7 @@ def _calculate_wakefield(self, bunches: List[ParticleBunch]): self.dxi, self.p_shape, self.q_bunch, - r_min_deposit=grid.r_max + r_min_deposit=grid.r_max, ) deposit_outliers_on_base_grid = True @@ -408,9 +421,19 @@ def _calculate_wakefield(self, bunches: List[ParticleBunch]): self._reset_bunch_arrays() for bunch in bunches_without_grid: deposit_bunch_charge( - bunch.x, bunch.y, bunch.xi, bunch.q, - self.n_p, self.n_r, self.n_xi, self.r_fld, self.xi_fld, - self.dr, self.dxi, self.p_shape, self.q_bunch + bunch.x, + bunch.y, + bunch.xi, + bunch.q, + self.n_p, + self.n_r, + self.n_xi, + self.r_fld, + self.xi_fld, + self.dr, + self.dxi, + self.p_shape, + self.q_bunch, ) calculate_bunch_source( self.q_bunch, self.n_r, self.n_xi, self.b_t_bunch @@ -422,22 +445,31 @@ def _calculate_wakefield(self, bunches: List[ParticleBunch]): [ self.r_fld[0], self.r_fld[-1] + 2 * self.dr, # r of last guard cell. - self.dr + self.dr, ] - ) / s_d + ) + / s_d ) # Calculate rho only if requested in the diagnostics. - calculate_rho = any('rho' in diag for diag in self.field_diags) + calculate_rho = any("rho" in diag for diag in self.field_diags) # Calculate plasma wakefields self.pp = calculate_wakefields( - laser_a2, self.r_max, self.xi_min, self.xi_max, - self.n_r, self.n_xi, self.ppc, self.n_p, + laser_a2, + self.r_max, + self.xi_min, + self.xi_max, + self.n_r, + self.n_xi, + self.ppc, + self.n_p, r_max_plasma=self.r_max_plasma, radial_density=radial_density, - p_shape=self.p_shape, max_gamma=self.max_gamma, - plasma_pusher=self.plasma_pusher, ion_motion=self.ion_motion, + p_shape=self.p_shape, + max_gamma=self.max_gamma, + plasma_pusher=self.plasma_pusher, + ion_motion=self.ion_motion, ion_mass=self.ion_mass, free_electrons_per_ion=self.free_electrons_per_ion, fld_arrays=self.fld_arrays, @@ -446,14 +478,14 @@ def _calculate_wakefield(self, bunches: List[ParticleBunch]): bunch_source_metadata=bunch_source_metadata, store_plasma_history=store_plasma_history, calculate_rho=calculate_rho, - particle_diags=self.particle_diags + particle_diags=self.particle_diags, ) # Add bunch density to total density. if calculate_rho: rho_bunch = -self.q_bunch[2:-2, 2:-2] / (self.r_fld / s_d) self.rho[2:-2, 2:-2] += rho_bunch - + # Calculate fields on adaptive grids. if self.use_adaptive_grids: for _, grid in self.bunch_grids.items(): @@ -462,32 +494,38 @@ def _calculate_wakefield(self, bunches: List[ParticleBunch]): def _reset_bunch_arrays(self): """Reset to zero the bunch arrays of the base grid.""" if self.t > self._t_reset_bunch_arrays: - self.b_t_bunch[:] = 0. - self.q_bunch[:] = 0. + self.b_t_bunch[:] = 0.0 + self.q_bunch[:] = 0.0 self._t_reset_bunch_arrays = self.t def _get_radial_density(self, z_current): - """ Get radial density profile function """ + """Get radial density profile function""" + def radial_density(r): n_p = self.density_function(z_current, r) if self.parabolic_coefficient is not None: pc = self.parabolic_coefficient(z_current) - n_p = n_p + n_p * pc * r ** 2 + n_p = n_p + n_p * pc * r**2 return n_p + return radial_density - + def _get_parabolic_coefficient_fn(self, parabolic_coefficient): - """ Get parabolic_coefficient profile function """ + """Get parabolic_coefficient profile function""" if isinstance(parabolic_coefficient, float): + def uniform_parabolic_coefficient(z): return np.ones_like(z) * parabolic_coefficient + return uniform_parabolic_coefficient elif callable(parabolic_coefficient): return parabolic_coefficient else: raise ValueError( - 'Type {} not supported for parabolic_coefficient.'.format( - type(parabolic_coefficient))) + "Type {} not supported for parabolic_coefficient.".format( + type(parabolic_coefficient) + ) + ) def _gather(self, x, y, z, t, ex, ey, ez, bx, by, bz, bunch_name): # If using adaptive grids, gather fields from them. @@ -500,10 +538,25 @@ def _gather(self, x, y, z, t, ex, ey, ez, bx, by, bz, bunch_name): # base grid. if not all_gathered: gather_main_fields_cyl_linear( - self.e_r, self.e_z, self.b_t, self.xi_fld[0], - self.xi_fld[-1], self.r_fld[0], self.r_fld[-1], - self.dxi, self.dr, x, y, z, - ex, ey, ez, bx, by, bz, r_min_gather=grid.r_max_cell + self.e_r, + self.e_z, + self.b_t, + self.xi_fld[0], + self.xi_fld[-1], + self.r_fld[0], + self.r_fld[-1], + self.dxi, + self.dr, + x, + y, + z, + ex, + ey, + ez, + bx, + by, + bz, + r_min_gather=grid.r_max_cell, ) # Otherwise, use base implementation. @@ -515,73 +568,75 @@ def _get_openpmd_diagnostics_data(self, global_time): # Add fields from adaptive grids to openpmd diagnostics. if self.use_adaptive_grids: for _, grid in self.bunch_grids.items(): - grid_data = grid.get_openpmd_data(global_time, - self.adaptive_grid_diags) - diag_data['fields'] += grid_data['fields'] - for field in grid_data['fields']: + grid_data = grid.get_openpmd_data( + global_time, self.adaptive_grid_diags + ) + diag_data["fields"] += grid_data["fields"] + for field in grid_data["fields"]: diag_data[field] = grid_data[field] # Add plasma particles to openpmd diagnostics. particle_diags = self._get_plasma_particle_diagnostics(global_time) diag_data = {**diag_data, **particle_diags} - diag_data['species'] = list(particle_diags.keys()) + diag_data["species"] = list(particle_diags.keys()) return diag_data def _get_plasma_particle_diagnostics(self, global_time): """Return dict with plasma particle diagnostics.""" diag_dict = {} - elec_name = 'plasma_electrons' - ions_name = 'plasma_ions' + elec_name = "plasma_electrons" + ions_name = "plasma_ions" if len(self.particle_diags) > 0: - n_elec = int(self.pp['r_hist'].shape[-1] / 2) - s_d = ge.plasma_skin_depth(self.n_p * 1e-6) + n_elec = int(self.pp["r_hist"].shape[-1] / 2) + s_d = ge.plasma_skin_depth(self.n_p * 1e-6) diag_dict[elec_name] = { - 'q': - ct.e, - 'm': ct.m_e, - 'name': elec_name, - 'geometry': 'rz' + "q": -ct.e, + "m": ct.m_e, + "name": elec_name, + "geometry": "rz", } diag_dict[ions_name] = { - 'q': ct.e * self.free_electrons_per_ion, - 'm': self.ion_mass, - 'name': ions_name, - 'geometry': 'rz' + "q": ct.e * self.free_electrons_per_ion, + "m": self.ion_mass, + "name": ions_name, + "geometry": "rz", } - if 'r' in self.particle_diags: - r_e = self.pp['r_hist'][:, :n_elec] * s_d - r_i = self.pp['r_hist'][:, n_elec:] * s_d - diag_dict[elec_name]['r'] = r_e - diag_dict[ions_name]['r'] = r_i - if 'z' in self.particle_diags: - z_e = self.pp['xi_hist'][:, :n_elec] * s_d + self.xi_max - z_i = self.pp['xi_hist'][:, n_elec:] * s_d + self.xi_max - diag_dict[elec_name]['z'] = z_e - diag_dict[ions_name]['z'] = z_i - diag_dict[elec_name]['z_off'] = global_time * ct.c - diag_dict[ions_name]['z_off'] = global_time * ct.c - if 'pr' in self.particle_diags: - pr_e = self.pp['pr_hist'][:, :n_elec] * ct.m_e * ct.c - pr_i = self.pp['pr_hist'][:, n_elec:] * self.ion_mass * ct.c - diag_dict[elec_name]['pr'] = pr_e - diag_dict[ions_name]['pr'] = pr_i - if 'pz' in self.particle_diags: - pz_e = self.pp['pz_hist'][:, :n_elec] * ct.m_e * ct.c - pz_i = self.pp['pz_hist'][:, n_elec:] * self.ion_mass * ct.c - diag_dict[elec_name]['pz'] = pz_e - diag_dict[ions_name]['pz'] = pz_i - if 'w' in self.particle_diags: - w_e = self.pp['w_hist'][:, :n_elec] * self.n_p - w_i = self.pp['w_hist'][:, n_elec:] * ( - self.n_p / self.free_electrons_per_ion) - diag_dict[elec_name]['w'] = w_e - diag_dict[ions_name]['w'] = w_i - if 'r_to_x' in self.particle_diags: - nc_e = self.pp['r_to_x_hist'][:, :n_elec] - nc_i = self.pp['r_to_x_hist'][:, n_elec:] - diag_dict[elec_name]['r_to_x'] = nc_e - diag_dict[ions_name]['r_to_x'] = nc_i - if 'tag' in self.particle_diags: - tag_e = self.pp['tag_hist'][:, :n_elec] - tag_i = self.pp['tag_hist'][:, n_elec:] - diag_dict[elec_name]['tag'] = tag_e - diag_dict[ions_name]['tag'] = tag_i + if "r" in self.particle_diags: + r_e = self.pp["r_hist"][:, :n_elec] * s_d + r_i = self.pp["r_hist"][:, n_elec:] * s_d + diag_dict[elec_name]["r"] = r_e + diag_dict[ions_name]["r"] = r_i + if "z" in self.particle_diags: + z_e = self.pp["xi_hist"][:, :n_elec] * s_d + self.xi_max + z_i = self.pp["xi_hist"][:, n_elec:] * s_d + self.xi_max + diag_dict[elec_name]["z"] = z_e + diag_dict[ions_name]["z"] = z_i + diag_dict[elec_name]["z_off"] = global_time * ct.c + diag_dict[ions_name]["z_off"] = global_time * ct.c + if "pr" in self.particle_diags: + pr_e = self.pp["pr_hist"][:, :n_elec] * ct.m_e * ct.c + pr_i = self.pp["pr_hist"][:, n_elec:] * self.ion_mass * ct.c + diag_dict[elec_name]["pr"] = pr_e + diag_dict[ions_name]["pr"] = pr_i + if "pz" in self.particle_diags: + pz_e = self.pp["pz_hist"][:, :n_elec] * ct.m_e * ct.c + pz_i = self.pp["pz_hist"][:, n_elec:] * self.ion_mass * ct.c + diag_dict[elec_name]["pz"] = pz_e + diag_dict[ions_name]["pz"] = pz_i + if "w" in self.particle_diags: + w_e = self.pp["w_hist"][:, :n_elec] * self.n_p + w_i = self.pp["w_hist"][:, n_elec:] * ( + self.n_p / self.free_electrons_per_ion + ) + diag_dict[elec_name]["w"] = w_e + diag_dict[ions_name]["w"] = w_i + if "r_to_x" in self.particle_diags: + nc_e = self.pp["r_to_x_hist"][:, :n_elec] + nc_i = self.pp["r_to_x_hist"][:, n_elec:] + diag_dict[elec_name]["r_to_x"] = nc_e + diag_dict[ions_name]["r_to_x"] = nc_i + if "tag" in self.particle_diags: + tag_e = self.pp["tag_hist"][:, :n_elec] + tag_i = self.pp["tag_hist"][:, n_elec:] + diag_dict[elec_name]["tag"] = tag_e + diag_dict[ions_name]["tag"] = tag_i return diag_dict