def numba_correct_currents_crossdeposition_comoving(rho_prev, rho_next, rho_next_z, rho_next_xy, Jp, Jm, Jz, kz, kr, j_corr_coef, T_eb, T_cc, inv_dt, Nz, Nr): """ Correct the currents in spectral space, using the cross-deposition algorithm adapted to the galilean/comoving-currents assumption. """ # Loop over the 2D grid for iz in prange(Nz): for ir in range(Nr): # Calculate the intermediate variable Dz and Dxy # (Such that Dz + Dxy is the error in the continuity equation) Dz = 1.j*kz[iz, ir]*Jz[iz, ir] \ + 0.5 * T_cc[iz, ir]*j_corr_coef[iz, ir] * \ ( rho_next[iz, ir] - T_eb[iz, ir] * rho_next_xy[iz, ir] \ + rho_next_z[iz, ir] - T_eb[iz, ir] * rho_prev[iz, ir] ) Dxy = kr[iz, ir]*( Jp[iz, ir] - Jm[iz, ir] ) \ + 0.5 * T_cc[iz, ir]*j_corr_coef[iz, ir] * \ ( rho_next[iz, ir] + T_eb[iz, ir] * rho_next_xy[iz, ir] \ - rho_next_z[iz, ir] - T_eb[iz, ir] * rho_prev[iz, ir] ) # Correct the currents accordingly if kr[iz, ir] != 0: inv_kr = 1. / kr[iz, ir] Jp[iz, ir] += -0.5 * Dxy * inv_kr Jm[iz, ir] += 0.5 * Dxy * inv_kr if kz[iz, ir] != 0: inv_kz = 1. / kz[iz, ir] Jz[iz, ir] += 1.j * Dz * inv_kz return
def push_p_ioniz_numba(ux, uy, uz, inv_gamma, Ex, Ey, Ez, Bx, By, Bz, m, Ntot, dt, ionization_level): """ Advance the particles' momenta, using numba """ # Set a few constants prefactor_econst = e * dt / (m * c) prefactor_bconst = 0.5 * e * dt / m # Loop over the particles (in parallel if threading is installed) for ip in prange(Ntot): # For neutral macroparticles, skip this step if ionization_level[ip] == 0: continue # Calculate the charge dependent constants econst = prefactor_econst * ionization_level[ip] bconst = prefactor_bconst * ionization_level[ip] # Perform the push ux[ip], uy[ip], uz[ip], inv_gamma[ip] = push_p_vay( ux[ip], uy[ip], uz[ip], inv_gamma[ip], Ex[ip], Ey[ip], Ez[ip], Bx[ip], By[ip], Bz[ip], econst, bconst) return ux, uy, uz, inv_gamma
def copy_ionized_electrons_numba( N_batch, batch_size, elec_old_Ntot, ion_Ntot, cumulative_n_ionized, ionized_from, i_level, store_electrons_per_level, elec_x, elec_y, elec_z, elec_inv_gamma, elec_ux, elec_uy, elec_uz, elec_w, elec_Ex, elec_Ey, elec_Ez, elec_Bx, elec_By, elec_Bz, ion_x, ion_y, ion_z, ion_inv_gamma, ion_ux, ion_uy, ion_uz, ion_w, ion_Ex, ion_Ey, ion_Ez, ion_Bx, ion_By, ion_Bz ): """ Create the new electrons by copying the properties (position, momentum, etc) of the ions that they originate from. """ # Loop over batches of particles (in parallel, if threading is enabled) for i_batch in prange( N_batch ): copy_ionized_electrons_batch( i_batch, batch_size, elec_old_Ntot, ion_Ntot, cumulative_n_ionized, ionized_from, i_level, store_electrons_per_level, elec_x, elec_y, elec_z, elec_inv_gamma, elec_ux, elec_uy, elec_uz, elec_w, elec_Ex, elec_Ey, elec_Ez, elec_Bx, elec_By, elec_Bz, ion_x, ion_y, ion_z, ion_inv_gamma, ion_ux, ion_uy, ion_uz, ion_w, ion_Ex, ion_Ey, ion_Ez, ion_Bx, ion_By, ion_Bz ) return( elec_x, elec_y, elec_z, elec_inv_gamma, elec_ux, elec_uy, elec_uz, elec_w, elec_Ex, elec_Ey, elec_Ez, elec_Bx, elec_By, elec_Bz )
def numba_filter_vector(fieldr, fieldt, fieldz, Nz, Nr, filter_array_z, filter_array_r): """ Multiply the input field by the filter_array Parameters : ------------ field : 2darray of complexs An array that represent the fields in spectral space filter_array_z, filter_array_r : 1darray of reals An array that damps the fields at high k, in z and r respectively Nz, Nr : ints Dimensions of the arrays """ # Loop over the 2D grid (parallel in z, if threading is installed) for iz in prange(Nz): for ir in range(Nr): fieldr[iz, ir] = filter_array_z[iz] * filter_array_r[ir] * fieldr[iz, ir] fieldt[iz, ir] = filter_array_z[iz] * filter_array_r[ir] * fieldt[iz, ir] fieldz[iz, ir] = filter_array_z[iz] * filter_array_r[ir] * fieldz[iz, ir]
def numba_push_eb_pml_standard(Ep_pml, Em_pml, Bp_pml, Bm_pml, Ez, Bz, C, S_w, kr, kz, Nz, Nr): """ Push the PML split fields over one timestep, using the standard psatd algorithm See the documentation of SpectralGrid.push_eb_with """ # Loop over the 2D grid for iz in prange(Nz): for ir in range(Nr): # Push the PML E field Ep_pml[iz, ir] = C[iz, ir]*Ep_pml[iz, ir] \ + c2*S_w[iz, ir]*( -1.j*0.5*kr[iz, ir]*Bz[iz, ir] ) Em_pml[iz, ir] = C[iz, ir]*Em_pml[iz, ir] \ + c2*S_w[iz, ir]*( -1.j*0.5*kr[iz, ir]*Bz[iz, ir] ) # Push the PML B field Bp_pml[iz, ir] = C[iz, ir]*Bp_pml[iz, ir] \ - S_w[iz, ir]*( -1.j*0.5*kr[iz, ir]*Ez[iz, ir] ) Bm_pml[iz, ir] = C[iz, ir]*Bm_pml[iz, ir] \ - S_w[iz, ir]*( -1.j*0.5*kr[iz, ir]*Ez[iz, ir] ) return
def sum_reduce_2d_array(global_array, reduced_array, m): """ Sum the array `global_array` along its first axis and add it into `reduced_array`, and fold the deposition guard cells of global_array into the regular cells of reduced_array. Parameters: ----------- global_array: 4darray of complexs Field array of shape (nthreads, Nm, 2+Nz+2, 2+Nr+2) where the additional 2's in z and r correspond to deposition guard cells that were used during the threaded deposition kernel. reduced array: 2darray of complex Field array of shape (Nz, Nr) m: int The azimuthal mode for which the reduction should be performed """ # Extract size of each dimension Nz = reduced_array.shape[0] # Parallel loop over z for iz in prange(Nz): # Get index inside reduced_array iz_global = iz + 2 reduce_slice(reduced_array, iz, global_array, iz_global, m) # Handle deposition guard cells in z reduce_slice(reduced_array, Nz - 2, global_array, 0, m) reduce_slice(reduced_array, Nz - 1, global_array, 1, m) reduce_slice(reduced_array, 0, global_array, Nz + 2, m) reduce_slice(reduced_array, 1, global_array, Nz + 3, m)
def ionize_ions_numba(N_batch, batch_size, Ntot, level_max, n_ionized, is_ionized, ionization_level, random_draw, adk_prefactor, adk_power, adk_exp_prefactor, ux, uy, uz, Ex, Ey, Ez, Bx, By, Bz, w, w_times_level): """ For each ion macroparticle, decide whether it is going to be further ionized during this timestep, based on the ADK rate. Increment the elements in `ionization_level` accordingly, and update `w_times_level` of the ions to take into account the change in level of the corresponding macroparticle. For the purpose of counting and creating the corresponding electrons, `is_ionized` (one element per macroparticle) is set to 1 at the position of the ionized ions, and `n_ionized` (one element per batch) counts the total number of ionized particles in the current batch. """ # Loop over batches of particles (in parallel, if threading is enabled) for i_batch in prange(N_batch): # Set the count of ionized particles in the batch to 0 n_ionized[i_batch] = 0 # Loop through the batch # (Note: a while loop is used here, because numba 0.34 does # not support nested prange and range loops) N_max = min((i_batch + 1) * batch_size, Ntot) ip = i_batch * batch_size while ip < N_max: # Skip the ionization routine, if the maximal ionization level # has already been reached for this macroparticle level = ionization_level[ip] if level >= level_max: is_ionized[ip] = 0 else: # Calculate the amplitude of the electric field, # in the frame of the electrons (device inline function) E, gamma = get_E_amplitude(ux[ip], uy[ip], uz[ip], Ex[ip], Ey[ip], Ez[ip], c * Bx[ip], c * By[ip], c * Bz[ip]) # Get ADK rate (device inline function) p = get_ionization_probability(E, gamma, adk_prefactor[level], adk_power[level], adk_exp_prefactor[level]) # Ionize particles if random_draw[ip] < p: # Set the corresponding flag and update particle count is_ionized[ip] = 1 n_ionized[i_batch] += 1 # Update the ionization level and the corresponding weight ionization_level[ip] += 1 w_times_level[ip] = w[ip] * ionization_level[ip] else: is_ionized[ip] = 0 # Increment ip ip = ip + 1 return (n_ionized, is_ionized, ionization_level, w_times_level)
def get_photon_density_gaussian_numba(photon_n, elec_Ntot, elec_x, elec_y, elec_z, ct, photon_n_lab_max, inv_laser_waist2, inv_laser_ctau2, laser_initial_z0, gamma_boost, beta_boost): """ Fill the array `photon_n` with the values of the photon density (in the simulation frame) in the scattering laser pulse, at the position of the electron macroparticles. Parameters ---------- elec_x, elec_y, elec_z: 1d arrays of floats The position of the electrons (in the frame of the simulation) ct: float Current time in the simulation frame (multiplied by c) photon_n_lab_max: float Peak photon density (in the lab frame) (i.e. at the peak of the Gaussian pulse) inv_laser_waist2, inv_laser_ctau2, laser_initial_z0: floats Properties of the Gaussian laser pulse (in the lab frame) gamma_boost, beta_boost: floats Properties of the Lorentz boost between the lab and simulation frame. """ # Loop over electrons (in parallel, if threading is enabled) for i_elec in prange(elec_Ntot): photon_n[i_elec] = get_photon_density_gaussian( elec_x[i_elec], elec_y[i_elec], elec_z[i_elec], ct, photon_n_lab_max, inv_laser_waist2, inv_laser_ctau2, laser_initial_z0, gamma_boost, beta_boost) return (photon_n)
def shift_spect_array_cpu(field_array, shift_factor, n_move): """ Shift the field 'field_array' by n_move cells on CPU. This is done in spectral space and corresponds to multiplying the fields with the factor exp(i*kz_true*dz)**n_move . Parameters ---------- field_array: 2darray of complexs Contains the value of the fields, and is modified by this function shift_factor: 1darray of complexs Contains the shift array, that is multiplied to the fields in spectral space to shift them by one cell in spatial space ( exp(i*kz_true*dz) ) n_move: int The number of cells by which the grid should be shifted """ Nz, Nr = field_array.shape # Loop over the 2D array (in parallel over z if threading is enabled) for iz in prange(Nz): power_shift = 1. + 0.j # Calculate the shift factor (raising to the power n_move ; # for negative n_move, we take the complex conjugate, since # shift_factor is of the form e^{i k dz}) for i in range(abs(n_move)): power_shift *= shift_factor[iz] if n_move < 0: power_shift = power_shift.conjugate() # Shift the fields for ir in range(Nr): field_array[iz, ir] *= power_shift
def numba_correct_currents_crossdeposition_standard(rho_prev, rho_next, rho_next_z, rho_next_xy, Jp, Jm, Jz, kz, kr, inv_dt, Nz, Nr): """ Correct the currents in spectral space, using the cross-deposition algorithm adapted to the standard psatd. """ # Loop over the 2D grid for iz in prange(Nz): for ir in range(Nr): # Calculate the intermediate variable Dz and Dxy # (Such that Dz + Dxy is the error in the continuity equation) Dz = 1.j*kz[iz, ir]*Jz[iz, ir] + 0.5 * inv_dt * \ ( rho_next[iz, ir] - rho_next_xy[iz, ir] + \ rho_next_z[iz, ir] - rho_prev[iz, ir] ) Dxy = kr[iz, ir]*( Jp[iz, ir] - Jm[iz, ir] ) + 0.5 * inv_dt * \ ( rho_next[iz, ir] - rho_next_z[iz, ir] + \ rho_next_xy[iz, ir] - rho_prev[iz, ir] ) # Correct the currents accordingly if kr[iz, ir] != 0: inv_kr = 1. / kr[iz, ir] Jp[iz, ir] += -0.5 * Dxy * inv_kr Jm[iz, ir] += 0.5 * Dxy * inv_kr if kz[iz, ir] != 0: inv_kz = 1. / kz[iz, ir] Jz[iz, ir] += 1.j * Dz * inv_kz return
def numba_push_eb_pml_comoving(Ep_pml, Em_pml, Bp_pml, Bm_pml, Ez, Bz, C, S_w, T_eb, kr, kz, Nz, Nr): """ Push the PML split fields over one timestep, using the galilean/comoving psatd algorithm See the documentation of SpectralGrid.push_eb_with """ # Loop over the 2D grid for iz in prange(Nz): for ir in range(Nr): # Push the E field Ep_pml[iz, ir] = T_eb[iz, ir]*C[iz, ir]*Ep_pml[iz, ir] \ + c2*T_eb[iz, ir]*S_w[iz, ir]*(-1.j*0.5*kr[iz, ir]*Bz[iz, ir]) Em_pml[iz, ir] = T_eb[iz, ir]*C[iz, ir]*Em_pml[iz, ir] \ + c2*T_eb[iz, ir]*S_w[iz, ir]*(-1.j*0.5*kr[iz, ir]*Bz[iz, ir]) # Push the B field Bp_pml[iz, ir] = T_eb[iz, ir]*C[iz, ir]*Bp_pml[iz, ir] \ - T_eb[iz, ir]*S_w[iz, ir]*( -1.j*0.5*kr[iz, ir]*Ez[iz, ir] ) Bm_pml[iz, ir] = T_eb[iz, ir]*C[iz, ir]*Bm_pml[iz, ir] \ - T_eb[iz, ir]*S_w[iz, ir]*( -1.j*0.5*kr[iz, ir]*Ez[iz, ir] ) return
def copy_particle_data_numba(Ntot, old_array, new_array): """ Copy the `Ntot` elements of `old_array` into `new_array`, on CPU """ # Loop over single particles (in parallel if threading is enabled) for ip in prange(Ntot): new_array[ip] = old_array[ip] return (new_array)
def determine_scatterings_numba(N_batch, batch_size, elec_Ntot, nscatter_per_elec, nscatter_per_batch, dt, elec_ux, elec_uy, elec_uz, elec_inv_gamma, ratio_w_electron_photon, photon_n, photon_p, photon_beta_x, photon_beta_y, photon_beta_z): """ For each electron macroparticle, decide how many photon macroparticles it will emit during `dt`, using the integrated Klein-Nishina formula. Note: this function uses a random generator within a `prange` loop. This implies that an indenpendent seed and random generator will be created for each thread. Electrons are processed in batches of size `batch_size`, with a parallel loop over batches. The batching allows quicker calculation of the total number of photons to be created. """ # Loop over batches of particles (in parallel, if threading is enabled) for i_batch in prange(N_batch): # Set the count of scattered particles in the batch to 0 nscatter_per_batch[i_batch] = 0 # Loop through the batch # (Note: a while loop is used here, because numba 0.34 does # not support nested prange and range loops) N_max = min((i_batch + 1) * batch_size, elec_Ntot) ip = i_batch * batch_size while ip < N_max: # Set the count of scattered photons for this electron to 0 nscatter_per_elec[ip] = 0 # For each electron, calculate the probability of scattering p = get_scattering_probability(dt, elec_ux[ip], elec_uy[ip], elec_uz[ip], elec_inv_gamma[ip], photon_n[ip], photon_p, photon_beta_x, photon_beta_y, photon_beta_z) # Determine the number of photons produced by this electron nscatter = int(p * ratio_w_electron_photon + random.random()) # Note: if p is 0, the above formula will return nscatter=0 # since random_draw is in [0, 1). Similarly, if p is very small, # nscatter will be 1 with probabiliy p * ratio_w_electron_photon, # and 0 otherwise. nscatter_per_elec[ip] = nscatter nscatter_per_batch[i_batch] += nscatter # Increment ip ip = ip + 1 return (nscatter_per_elec, nscatter_per_batch)
def push_x_numba(x, y, z, ux, uy, uz, inv_gamma, Ntot, dt, push_x, push_y, push_z): """ Advance the particles' positions over `dt` using the momenta ux, uy, uz, multiplied by the scalar coefficients x_push, y_push, z_push. """ # Half timestep, multiplied by c chdt = c * dt # Particle push (in parallel if threading is installed) for ip in prange(Ntot): x[ip] += chdt * inv_gamma[ip] * push_x * ux[ip] y[ip] += chdt * inv_gamma[ip] * push_y * uy[ip] z[ip] += chdt * inv_gamma[ip] * push_z * uz[ip] return x, y, z
def push_p_numba(ux, uy, uz, inv_gamma, Ex, Ey, Ez, Bx, By, Bz, q, m, Ntot, dt): """ Advance the particles' momenta, using numba """ # Set a few constants econst = q * dt / (m * c) bconst = 0.5 * q * dt / m # Loop over the particles (in parallel if threading is installed) for ip in prange(Ntot): ux[ip], uy[ip], uz[ip], inv_gamma[ip] = push_p_vay( ux[ip], uy[ip], uz[ip], inv_gamma[ip], Ex[ip], Ey[ip], Ez[ip], Bx[ip], By[ip], Bz[ip], econst, bconst) return ux, uy, uz, inv_gamma
def push_p_after_plane_numba(z, z_plane, ux, uy, uz, inv_gamma, Ex, Ey, Ez, Bx, By, Bz, q, m, Ntot, dt): """ Advance the particles' momenta, using numba. Only the particles that are located beyond the plane z=z_plane have their momentum modified ; the others particles move ballistically. """ # Set a few constants econst = q * dt / (m * c) bconst = 0.5 * q * dt / m # Loop over the particles (in parallel if threading is installed) for ip in prange(Ntot): if z[ip] > z_plane: ux[ip], uy[ip], uz[ip], inv_gamma[ip] = push_p_vay( ux[ip], uy[ip], uz[ip], inv_gamma[ip], Ex[ip], Ey[ip], Ez[ip], Bx[ip], By[ip], Bz[ip], econst, bconst)
def erase_eb_numba(Ex, Ey, Ez, Bx, By, Bz, Ntot): """ Reset the arrays of fields (i.e. set them to 0) Parameters ---------- Ex, Ey, Ez, Bx, By, Bz: 1d arrays of floats (One element per macroparticle) Represents the fields on the macroparticles """ for i in prange(Ntot): Ex[i] = 0 Ey[i] = 0 Ez[i] = 0 Bx[i] = 0 By[i] = 0 Bz[i] = 0 return Ex, Ey, Ez, Bx, By, Bz
def push_x_numba(x, y, z, ux, uy, uz, inv_gamma, Ntot, dt): """ Advance the particles' positions over one half-timestep This assumes that the positions (x, y, z) are initially either one half-timestep *behind* the momenta (ux, uy, uz), or at the same timestep as the momenta. """ # Half timestep, multiplied by c chdt = c * 0.5 * dt # Particle push (in parallel if threading is installed) for ip in prange(Ntot): x[ip] += chdt * inv_gamma[ip] * ux[ip] y[ip] += chdt * inv_gamma[ip] * uy[ip] z[ip] += chdt * inv_gamma[ip] * uz[ip] return x, y, z
def numba_erase_threading_buffer(global_array): """ Set the threading buffer `global_array` to 0 Parameter: ---------- global_array: 4darray of complexs An array that contains the duplicated charge/current for each thread """ nthreads, Nm, Nz, Nr = global_array.shape # Loop in parallel along nthreads for i_thread in prange(nthreads): # Loop through the modes and the grid for m in range(Nm): for iz in range(Nz): for ir in range(Nr): # Erase values global_array[i_thread, m, iz, ir] = 0.
def numba_correct_currents_crossdeposition_comoving(rho_prev, rho_next, rho_next_z, rho_next_xy, Jp, Jm, Jz, kz, kr, j_corr_coef, T_eb, T_cc, inv_dt, Nz, Nr): """ Correct the currents in spectral space, using the cross-deposition algorithm adapted to the galilean/comoving-currents assumption. """ # Loop over the 2D grid for iz in prange(Nz): # Loop through the radial points # (Note: a while loop is used here, because numba 0.34 does # not support nested prange and range loops) ir = 0 while ir < Nr: # Calculate the intermediate variable Dz and Dxy # (Such that Dz + Dxy is the error in the continuity equation) Dz = 1.j*kz[iz, ir]*Jz[iz, ir] \ + 0.5 * T_cc[iz, ir]*j_corr_coef[iz, ir] * \ ( rho_next[iz, ir] - T_eb[iz, ir] * rho_next_xy[iz, ir] \ + rho_next_z[iz, ir] - T_eb[iz, ir] * rho_prev[iz, ir] ) Dxy = kr[iz, ir]*( Jp[iz, ir] - Jm[iz, ir] ) \ + 0.5 * T_cc[iz, ir]*j_corr_coef[iz, ir] * \ ( rho_next[iz, ir] + T_eb[iz, ir] * rho_next_xy[iz, ir] \ - rho_next_z[iz, ir] - T_eb[iz, ir] * rho_prev[iz, ir] ) # Correct the currents accordingly if kr[iz, ir] != 0: inv_kr = 1. / kr[iz, ir] Jp[iz, ir] += -0.5 * Dxy * inv_kr Jm[iz, ir] += 0.5 * Dxy * inv_kr if kz[iz, ir] != 0: inv_kz = 1. / kz[iz, ir] Jz[iz, ir] += 1.j * Dz * inv_kz # Increment ir ir += 1 return
def numba_correct_currents_standard(rho_prev, rho_next, Jp, Jm, Jz, kz, kr, inv_k2, inv_dt, Nz, Nr): """ Correct the currents in spectral space, using the standard pstad """ # Loop over the 2D grid (parallel in z, if threading is installed) for iz in prange(Nz): for ir in range(Nr): # Calculate the intermediate variable F F = - inv_k2[iz, ir] * ( (rho_next[iz, ir] - rho_prev[iz, ir])*inv_dt \ + 1.j*kz[iz, ir]*Jz[iz, ir] \ + kr[iz, ir]*( Jp[iz, ir] - Jm[iz, ir] ) ) # Correct the currents accordingly Jp[iz, ir] += 0.5 * kr[iz, ir] * F Jm[iz, ir] += -0.5 * kr[iz, ir] * F Jz[iz, ir] += -1.j * kz[iz, ir] * F return
def numba_correct_currents_curlfree_comoving(rho_prev, rho_next, Jp, Jm, Jz, kz, kr, inv_k2, j_corr_coef, T_eb, T_cc, inv_dt, Nz, Nr): """ Correct the currents in spectral space, using the curl-free correction which is adapted to the galilean/comoving-currents assumption """ # Loop over the 2D grid (parallel in z, if threading is installed) for iz in prange(Nz): for ir in range(Nr): # Calculate the intermediate variable F F = - inv_k2[iz, ir] * ( T_cc[iz, ir]*j_corr_coef[iz, ir] \ * (rho_next[iz, ir] - rho_prev[iz, ir]*T_eb[iz, ir]) \ + 1.j*kz[iz, ir]*Jz[iz, ir] \ + kr[iz, ir]*( Jp[iz, ir] - Jm[iz, ir] ) ) # Correct the currents accordingly Jp[iz, ir] += 0.5 * kr[iz, ir] * F Jm[iz, ir] += -0.5 * kr[iz, ir] * F Jz[iz, ir] += -1.j * kz[iz, ir] * F return
def ionize_ions_numba( N_batch, batch_size, Ntot, level_start, level_max, n_levels, n_ionized, ionized_from, ionization_level, random_draw, adk_prefactor, adk_power, adk_exp_prefactor, ux, uy, uz, Ex, Ey, Ez, Bx, By, Bz, w, w_times_level ): """ For each ion macroparticle, decide whether it is going to be further ionized during this timestep, based on the ADK rate. Increment the elements in `ionization_level` accordingly, and update `w_times_level` of the ions to take into account the change in level of the corresponding macroparticle. For the purpose of counting and creating the corresponding electrons, `ionized_from` (one element per macroparticle) is set to -1 at the position of the unionized ions, and to the level (before ionization) otherwise `n_ionized` (one element per batch, and per ionizable level that needs to be distinguished) counts the total number of ionized particles in the current batch. """ # Loop over batches of particles (in parallel, if threading is enabled) for i_batch in prange( N_batch ): # Set the count of ionized particles in the batch to 0 for i_level in range(n_levels): n_ionized[i_level, i_batch] = 0 # Loop through the batch N_max = min( (i_batch+1)*batch_size, Ntot ) for ip in range(i_batch*batch_size, N_max): # Skip the ionization routine, if the maximal ionization level # has already been reached for this macroparticle level = ionization_level[ip] if level >= level_max: ionized_from[ip] = -1 else: # Calculate the amplitude of the electric field, # in the frame of the electrons (device inline function) E, gamma = get_E_amplitude( ux[ip], uy[ip], uz[ip], Ex[ip], Ey[ip], Ez[ip], c*Bx[ip], c*By[ip], c*Bz[ip] ) # Get ADK rate (device inline function) p = get_ionization_probability( E, gamma, adk_prefactor[level], adk_power[level], adk_exp_prefactor[level]) # Ionize particles if random_draw[ip] < p: # Set the corresponding flag and update particle count ionized_from[ip] = level-level_start if n_levels == 1: # No need to distinguish ionization levels n_ionized[0, i_batch] += 1 else: # Distinguish count for each ionizable level n_ionized[level-level_start, i_batch] += 1 # Update the ionization level and the corresponding weight ionization_level[ip] += 1 w_times_level[ip] = w[ip] * ionization_level[ip] else: ionized_from[ip] = -1 return( n_ionized, ionized_from, ionization_level, w_times_level )
def scatter_photons_electrons_numba( N_batch, batch_size, photon_old_Ntot, elec_Ntot, cumul_nscatter_per_batch, nscatter_per_elec, photon_p, photon_px, photon_py, photon_pz, photon_x, photon_y, photon_z, photon_inv_gamma, photon_ux, photon_uy, photon_uz, photon_w, elec_x, elec_y, elec_z, elec_inv_gamma, elec_ux, elec_uy, elec_uz, elec_w, inv_ratio_w_elec_photon): """ Given the number of photons that are emitted by each electron macroparticle, determine the properties (momentum, energy) of each scattered photon and fill the arrays `photon_*` accordingly. Also, apply a recoil on the electrons. Note: this function uses a random generator within a `prange` loop. This implies that an indenpendent seed and random generator will be created for each thread. """ # Loop over batches of particles for i_batch in prange(N_batch): # Photon index: this is incremented each time # a scattered photon is identified i_photon = photon_old_Ntot + cumul_nscatter_per_batch[i_batch] # Loop through the electrons in this batch N_max = min((i_batch + 1) * batch_size, elec_Ntot) for i_elec in range(i_batch * batch_size, N_max): # Prepare calculation of scattered photons from this electron if nscatter_per_elec[i_elec] > 0: # Prepare Lorentz transformation to the electron rest frame elec_gamma = 1. / elec_inv_gamma[i_elec] elec_u = math.sqrt(elec_ux[i_elec]**2 + elec_uy[i_elec]**2 + elec_uz[i_elec]**2) elec_beta = elec_u * elec_inv_gamma[i_elec] if elec_u != 0: elec_inv_u = 1. / elec_u elec_nx = elec_inv_u * elec_ux[i_elec] elec_ny = elec_inv_u * elec_uy[i_elec] elec_nz = elec_inv_u * elec_uz[i_elec] else: # Avoid division by 0; provide arbitrary direction # for the Lorentz transform (since beta=0 anyway) elec_nx = 0. elec_ny = 0. elec_nz = 1. # Transform momentum of photon to the electron rest frame photon_rest_p, photon_rest_px, \ photon_rest_py, photon_rest_pz = lorentz_transform( photon_p, photon_px, photon_py, photon_pz, elec_gamma, elec_beta, elec_nx, elec_ny, elec_nz ) # Find cos and sin of the spherical angle that represent # the direction of the incoming photon in the rest frame cos_theta = photon_rest_pz / photon_rest_p if cos_theta**2 < 1: sin_theta = math.sqrt(1 - cos_theta**2) inv_photon_rest_pxy = 1. / (sin_theta * photon_rest_p) cos_phi = photon_rest_px * inv_photon_rest_pxy sin_phi = photon_rest_py * inv_photon_rest_pxy else: sin_theta = 0 # Avoid division by 0; provide arbitrary direction # for the phi angle (since theta is 0 or pi anyway) cos_phi = 1. sin_phi = 0. # Loop through the number of scatterings for this electron for i_scat in range(nscatter_per_elec[i_elec]): # Draw scattering angle in the rest frame, from the # Klein-Nishina cross-section (See Ozmutl, E. N. # "Sampling of Angular Distribution in Compton Scattering" # Appl. Radiat. Isot. 43, 6, pp. 713-715 (1992)) k = photon_rest_p * INV_MC c0 = 2. * (2. * k**2 + 2. * k + 1.) / (2. * k + 1.)**3 b = (2. + c0) / (2. - c0) a = 2. * b - 1. # Use rejection method to draw x reject = True while reject: # - Draw x with an approximate probability distribution r1 = random.random() x = b - (b + 1.) * (0.5 * c0)**r1 # - Calculate approximate probability distribution h h = a / (b - x) # - Calculate expected (exact) probability distribution f factor = 1 + k * (1 - x) f = ((1 + x**2) * factor + k**2 * (1 - x)**2) / factor**3 # - Keep x according to rejection rule r2 = random.random() if r2 < f / h: reject = False # Get scattered momentum in the rest frame new_photon_rest_p = photon_rest_p / (1 + k * (1 - x)) # - First in a system of axes aligned with the incoming photon cos_theta_s = x sin_theta_s = math.sqrt(1 - x**2) phi_s = 2 * math.pi * random.random() cos_phi_s = math.cos(phi_s) sin_phi_s = math.sin(phi_s) new_photon_rest_pX = new_photon_rest_p * sin_theta_s * cos_phi_s new_photon_rest_pY = new_photon_rest_p * sin_theta_s * sin_phi_s new_photon_rest_pZ = new_photon_rest_p * cos_theta_s # - Then rotate it to the original system of axes new_photon_rest_px = sin_theta * cos_phi * new_photon_rest_pZ \ + cos_theta * cos_phi * new_photon_rest_pX \ - sin_phi * new_photon_rest_pY new_photon_rest_py = sin_theta * sin_phi * new_photon_rest_pZ \ + cos_theta * sin_phi * new_photon_rest_pX \ + cos_phi * new_photon_rest_pY new_photon_rest_pz = cos_theta * new_photon_rest_pZ \ - sin_theta * new_photon_rest_pX # Transform momentum of photon back to the simulation frame # (i.e. Lorentz transform with opposite direction) new_photon_p, new_photon_px, new_photon_py, new_photon_pz = \ lorentz_transform( new_photon_rest_p, new_photon_rest_px, new_photon_rest_py, new_photon_rest_pz, elec_gamma, elec_beta, -elec_nx, -elec_ny, -elec_nz) # Create the new photon by copying the electron position photon_x[i_photon] = elec_x[i_elec] photon_y[i_photon] = elec_y[i_elec] photon_z[i_photon] = elec_z[i_elec] photon_w[i_photon] = elec_w[i_elec] * inv_ratio_w_elec_photon # The photon's ux, uy, uz corresponds to the actual px, py, pz photon_ux[i_photon] = new_photon_px photon_uy[i_photon] = new_photon_py photon_uz[i_photon] = new_photon_pz # The photon's inv_gamma corresponds to 1./p (consistent # with the code for the particle pusher and for the # openPMD back-transformed diagnostics) photon_inv_gamma[i_photon] = 1. / new_photon_p # Update the photon index i_photon += 1 # Add recoil to electrons # Note: In order to reproduce the right distribution of electron # momentum, the electrons should recoil with the momentum # of *one single* photon, with a probability p (calculated by # get_scattering_probability). Here we reuse the momentum of # the last photon generated above. This requires that at least one # photon be created for this electron, which occurs with a # probability p*ratio_w_elec_photon. Thus, given that at least one # photon has been created, we should add recoil to the corresponding # electron only with a probability inv_ratio_w_elec_photon. if nscatter_per_elec[i_elec] > 0: if random.random() < inv_ratio_w_elec_photon: elec_ux[i_elec] += INV_MC * (photon_px - new_photon_px) elec_uy[i_elec] += INV_MC * (photon_py - new_photon_py) elec_uz[i_elec] += INV_MC * (photon_pz - new_photon_pz)
def gather_field_numba_cubic(x, y, z, invdz, zmin, Nz, invdr, rmin, Nr, Er_m0, Et_m0, Ez_m0, Er_m1, Et_m1, Ez_m1, Br_m0, Bt_m0, Bz_m0, Br_m1, Bt_m1, Bz_m1, Ex, Ey, Ez, Bx, By, Bz, nthreads, ptcl_chunk_indices): """ Gathering of the fields (E and B) using numba with multi-threading. Iterates over the particles, calculates the weighted amount of fields acting on each particle based on its shape (cubic). Fields are gathered in cylindrical coordinates and then transformed to cartesian coordinates. Supports only mode 0 and 1. Parameters ---------- x, y, z : 1darray of floats (in meters) The position of the particles invdz, invdr : float (in meters^-1) Inverse of the grid step along the considered direction zmin, rmin : float (in meters) Position of the edge of the simulation box along the direction considered Nz, Nr : int Number of gridpoints along the considered direction Er_m0, Et_m0, Ez_m0 : 2darray of complexs The electric fields on the interpolation grid for the mode 0 Er_m1, Et_m1, Ez_m1 : 2darray of complexs The electric fields on the interpolation grid for the mode 1 Br_m0, Bt_m0, Bz_m0 : 2darray of complexs The magnetic fields on the interpolation grid for the mode 0 Br_m1, Bt_m1, Bz_m1 : 2darray of complexs The magnetic fields on the interpolation grid for the mode 1 Ex, Ey, Ez : 1darray of floats The electric fields acting on the particles (is modified by this function) Bx, By, Bz : 1darray of floats The magnetic fields acting on the particles (is modified by this function) nthreads : int Number of CPU threads used with numba prange ptcl_chunk_indices : array of int, of size nthreads+1 The indices (of the particle array) between which each thread should loop. (i.e. divisions of particle array between threads) """ # Gather the field per cell in parallel for nt in prange( nthreads ): # Create private arrays for each thread # to store the particle index and shape Sr = np.empty( 4 ) Sz = np.empty( 4 ) # Loop over all particles in thread chunk for i in range( ptcl_chunk_indices[nt], ptcl_chunk_indices[nt+1] ): # Preliminary arrays for the cylindrical conversion # -------------------------------------------- # Position xj = x[i] yj = y[i] zj = z[i] # Cylindrical conversion rj = math.sqrt(xj**2 + yj**2) if (rj != 0.): invr = 1./rj cos = xj*invr # Cosine sin = yj*invr # Sine else: cos = 1. sin = 0. exptheta_m0 = 1. exptheta_m1 = cos - 1.j*sin # Get weights for the deposition # -------------------------------------------- # Positions of the particle, in the cell unit r_cell = invdr*(rj - rmin) - 0.5 z_cell = invdz*(zj - zmin) - 0.5 # Calculate the shape factors ir_lowest = int64(math.floor(r_cell)) - 1 r_local = r_cell-ir_lowest Sr[0] = -1./6. * (r_local-2.)**3 Sr[1] = 1./6. * (3.*(r_local-1.)**3 - 6.*(r_local-1.)**2 + 4.) Sr[2] = 1./6. * (3.*(2.-r_local)**3 - 6.*(2.-r_local)**2 + 4.) Sr[3] = -1./6. * (1.-r_local)**3 iz_lowest = int64(math.floor(z_cell)) - 1 z_local = z_cell-iz_lowest Sz[0] = -1./6. * (z_local-2.)**3 Sz[1] = 1./6. * (3.*(z_local-1.)**3 - 6.*(z_local-1.)**2 + 4.) Sz[2] = 1./6. * (3.*(2.-z_local)**3 - 6.*(2.-z_local)**2 + 4.) Sz[3] = -1./6. * (1.-z_local)**3 # E-Field # ------- Fr = 0. Ft = 0. Fz = 0. # Only perform gathering for particles that are inside the box radially if r_cell+0.5 < Nr: # Add contribution from mode 0 Fr, Ft, Fz = add_cubic_gather_for_mode( 0, Fr, Ft, Fz, exptheta_m0, Er_m0, Et_m0, Ez_m0, ir_lowest, iz_lowest, Sr, Sz, Nr, Nz ) # Add contribution from mode 1 Fr, Ft, Fz = add_cubic_gather_for_mode( 1, Fr, Ft, Fz, exptheta_m1, Er_m1, Et_m1, Ez_m1, ir_lowest, iz_lowest, Sr, Sz, Nr, Nz ) # Convert to Cartesian coordinates # and write to particle field arrays Ex[i] = cos*Fr - sin*Ft Ey[i] = sin*Fr + cos*Ft Ez[i] = Fz # B-Field # ------- # Clear the placeholders for the # gathered field for each coordinate Fr = 0. Ft = 0. Fz = 0. # Only perform gathering for particles that are inside the box radially if r_cell+0.5 < Nr: # Add contribution from mode 0 Fr, Ft, Fz = add_cubic_gather_for_mode( 0, Fr, Ft, Fz, exptheta_m0, Br_m0, Bt_m0, Bz_m0, ir_lowest, iz_lowest, Sr, Sz, Nr, Nz ) # Add contribution from mode 1 Fr, Ft, Fz = add_cubic_gather_for_mode( 1, Fr, Ft, Fz, exptheta_m1, Br_m1, Bt_m1, Bz_m1, ir_lowest, iz_lowest, Sr, Sz, Nr, Nz ) # Convert to Cartesian coordinates # and write to particle field arrays Bx[i] = cos*Fr - sin*Ft By[i] = sin*Fr + cos*Ft Bz[i] = Fz return Ex, Ey, Ez, Bx, By, Bz
def deposit_rho_numba_linear(x, y, z, w, q, invdz, zmin, Nz, invdr, rmin, Nr, rho_global, Nm, nthreads, ptcl_chunk_indices): """ Deposition of the charge density rho using numba prange on the CPU. Iterates over the threads in parallel, while each thread iterates over a batch of particles. Intermediate results for each threads are stored in copies of the global grid. At the end of the parallel loop, the thread-local field arrays are combined (summed) to a global array. (This final reduction is *not* done in this function) Calculates the weighted amount of rho that is deposited to the 4 cells surounding the particle based on its shape (linear). Parameters ---------- x, y, z : 1darray of floats (in meters) The position of the particles w : 1d array of floats The weights of the particles (For ionizable atoms: weight times the ionization level) q : float Charge of the species (For ionizable atoms: this is always the elementary charge e) rho_global : 4darrays of complexs Global helper arrays of shape (nthreads, Nm, 2+Nz+2, 2+Nr+2) where the additional 2's in z and r correspond to deposition guard cells. This array stores the thread local charge density on the interpolation grid for each mode. (is modified by this function) Nm : int The number of azimuthal modes invdz, invdr : float (in meters^-1) Inverse of the grid step along the considered direction zmin, rmin : float (in meters) Position of the edge of the simulation box, along the considered direction Nz, Nr : int Number of gridpoints along the considered direction nthreads : int Number of CPU threads used with numba prange ptcl_chunk_indices : array of int, of size nthreads+1 The indices (of the particle array) between which each thread should loop. (i.e. divisions of particle array between threads) """ # Deposit the field per cell in parallel (for threads < number of cells) for i_thread in prange(nthreads): # Allocate thread-local array rho_scal = np.zeros(Nm, dtype=np.complex128) # Loop over all particles in thread chunk for i_ptcl in range(ptcl_chunk_indices[i_thread], ptcl_chunk_indices[i_thread + 1]): # Position xj = x[i_ptcl] yj = y[i_ptcl] zj = z[i_ptcl] # Weights wj = q * w[i_ptcl] # Cylindrical conversion rj = math.sqrt(xj**2 + yj**2) # Avoid division by 0. if (rj != 0.): invr = 1. / rj cos = xj * invr # Cosine sin = yj * invr # Sine else: cos = 1. sin = 0. # Calculate contribution from this particle to each mode rho_scal[0] = wj for m in range(1, Nm): rho_scal[m] = (cos + 1.j * sin) * rho_scal[m - 1] # Positions of the particles, in the cell unit r_cell = invdr * (rj - rmin) - 0.5 z_cell = invdz * (zj - zmin) - 0.5 # Index of the lowest cell of `global_array` that gets modified # by this particle (note: `global_array` has 2 guard cells) # (`min` function avoids out-of-bounds access at high r) ir_cell = min(int(math.floor(r_cell)) + 2, Nr + 2) iz_cell = int(math.floor(z_cell)) + 2 # Add contribution of this particle to the global array for m in range(Nm): rho_global[i_thread, m, iz_cell + 0, ir_cell + 0] += Sz_linear( z_cell, 0) * Sr_linear(r_cell, 0) * rho_scal[m] rho_global[i_thread, m, iz_cell + 0, ir_cell + 1] += Sz_linear( z_cell, 0) * Sr_linear(r_cell, 1) * rho_scal[m] rho_global[i_thread, m, iz_cell + 1, ir_cell + 0] += Sz_linear( z_cell, 1) * Sr_linear(r_cell, 0) * rho_scal[m] rho_global[i_thread, m, iz_cell + 1, ir_cell + 1] += Sz_linear( z_cell, 1) * Sr_linear(r_cell, 1) * rho_scal[m] return
def gather_field_numba_linear(x, y, z, invdz, zmin, Nz, invdr, rmin, Nr, Er_m0, Et_m0, Ez_m0, Er_m1, Et_m1, Ez_m1, Br_m0, Bt_m0, Bz_m0, Br_m1, Bt_m1, Bz_m1, Ex, Ey, Ez, Bx, By, Bz ): """ Gathering of the fields (E and B) using numba with multi-threading. Iterates over the particles, calculates the weighted amount of fields acting on each particle based on its shape (linear). Fields are gathered in cylindrical coordinates and then transformed to cartesian coordinates. Supports only mode 0 and 1. Parameters ---------- x, y, z : 1darray of floats (in meters) The position of the particles invdz, invdr : float (in meters^-1) Inverse of the grid step along the considered direction zmin, rmin : float (in meters) Position of the edge of the simulation box along the direction considered Nz, Nr : int Number of gridpoints along the considered direction Er_m0, Et_m0, Ez_m0 : 2darray of complexs The electric fields on the interpolation grid for the mode 0 Er_m1, Et_m1, Ez_m1 : 2darray of complexs The electric fields on the interpolation grid for the mode 1 Br_m0, Bt_m0, Bz_m0 : 2darray of complexs The magnetic fields on the interpolation grid for the mode 0 Br_m1, Bt_m1, Bz_m1 : 2darray of complexs The magnetic fields on the interpolation grid for the mode 1 Ex, Ey, Ez : 1darray of floats The electric fields acting on the particles (is modified by this function) Bx, By, Bz : 1darray of floats The magnetic fields acting on the particles (is modified by this function) """ # Deposit the field per cell in parallel for i in prange(x.shape[0]): # Preliminary arrays for the cylindrical conversion # -------------------------------------------- # Position xj = x[i] yj = y[i] zj = z[i] # Cylindrical conversion rj = math.sqrt( xj**2 + yj**2 ) if (rj !=0. ) : invr = 1./rj cos = xj*invr # Cosine sin = yj*invr # Sine else : cos = 1. sin = 0. exptheta_m0 = 1. exptheta_m1 = cos - 1.j*sin # Get linear weights for the deposition # ------------------------------------- # Positions of the particles, in the cell unit r_cell = invdr*(rj - rmin) - 0.5 z_cell = invdz*(zj - zmin) - 0.5 # Original index of the uppper and lower cell ir_lower = int(math.floor( r_cell )) ir_upper = ir_lower + 1 iz_lower = int(math.floor( z_cell )) iz_upper = iz_lower + 1 # Linear weight Sr_lower = ir_upper - r_cell Sr_upper = r_cell - ir_lower Sz_lower = iz_upper - z_cell Sz_upper = z_cell - iz_lower # Set guard weights to zero Sr_guard = 0. # Treat the boundary conditions # ----------------------------- # guard cells in lower r if ir_lower < 0: Sr_guard = Sr_lower Sr_lower = 0. ir_lower = 0 # absorbing in upper r if ir_lower > Nr-1: ir_lower = Nr-1 if ir_upper > Nr-1: ir_upper = Nr-1 # periodic boundaries in z # lower z boundaries if iz_lower < 0: iz_lower += Nz if iz_upper < 0: iz_upper += Nz # upper z boundaries if iz_lower > Nz-1: iz_lower -= Nz if iz_upper > Nz-1: iz_upper -= Nz # Precalculate Shapes S_ll = Sz_lower*Sr_lower S_lu = Sz_lower*Sr_upper S_ul = Sz_upper*Sr_lower S_uu = Sz_upper*Sr_upper S_lg = Sz_lower*Sr_guard S_ug = Sz_upper*Sr_guard # E-Field # ------- Fr = 0. Ft = 0. Fz = 0. # Only perform gathering for particles that are inside the box radially if r_cell+0.5 < Nr: # Add contribution from mode 0 Fr, Ft, Fz = add_linear_gather_for_mode( 0, Fr, Ft, Fz, exptheta_m0, Er_m0, Et_m0, Ez_m0, iz_lower, iz_upper, ir_lower, ir_upper, S_ll, S_lu, S_lg, S_ul, S_uu, S_ug ) # Add contribution from mode 1 Fr, Ft, Fz = add_linear_gather_for_mode( 1, Fr, Ft, Fz, exptheta_m1, Er_m1, Et_m1, Ez_m1, iz_lower, iz_upper, ir_lower, ir_upper, S_ll, S_lu, S_lg, S_ul, S_uu, S_ug ) # Convert to Cartesian coordinates # and write to particle field arrays Ex[i] = cos*Fr - sin*Ft Ey[i] = sin*Fr + cos*Ft Ez[i] = Fz # B-Field # ------- # Clear the placeholders for the # gathered field for each coordinate Fr = 0. Ft = 0. Fz = 0. # Only perform gathering for particles that are inside the box radially if r_cell+0.5 < Nr: # Add contribution from mode 0 Fr, Ft, Fz = add_linear_gather_for_mode( 0, Fr, Ft, Fz, exptheta_m0, Br_m0, Bt_m0, Bz_m0, iz_lower, iz_upper, ir_lower, ir_upper, S_ll, S_lu, S_lg, S_ul, S_uu, S_ug ) # Add contribution from mode 1 Fr, Ft, Fz = add_linear_gather_for_mode( 1, Fr, Ft, Fz, exptheta_m1, Br_m1, Bt_m1, Bz_m1, iz_lower, iz_upper, ir_lower, ir_upper, S_ll, S_lu, S_lg, S_ul, S_uu, S_ug ) # Convert to Cartesian coordinates # and write to particle field arrays Bx[i] = cos*Fr - sin*Ft By[i] = sin*Fr + cos*Ft Bz[i] = Fz return Ex, Ey, Ez, Bx, By, Bz
def deposit_J_numba_cubic(x, y, z, w, q, ux, uy, uz, inv_gamma, invdz, zmin, Nz, invdr, rmin, Nr, j_r_global, j_t_global, j_z_global, Nm, nthreads, ptcl_chunk_indices, beta_n_m_0, beta_n_m_higher): """ Deposition of the current density J using numba prange on the CPU. Iterates over the threads in parallel, while each thread iterates over a batch of particles. Intermediate results for each threads are stored in copies of the global grid. At the end of the parallel loop, the thread-local field arrays are combined (summed) to the global array. (This final reduction is *not* done in this function) Calculates the weighted amount of J that is deposited to the 16 cells surounding the particle based on its shape (cubic). Parameters ---------- x, y, z : 1darray of floats (in meters) The position of the particles w : 1d array of floats The weights of the particles (For ionizable atoms: weight times the ionization level) q : float Charge of the species (For ionizable atoms: this is always the elementary charge e) ux, uy, uz : 1darray of floats (in meters * second^-1) The velocity of the particles inv_gamma : 1darray of floats The inverse of the relativistic gamma factor j_x_global : 4darrays of complexs Global helper arrays of shape (nthreads, Nm, 2+Nz+2, 2+Nr+2) where the additional 2's in z and r correspond to deposition guard cells. This array stores the thread local current component in each direction (r, t, z) on the interpolation grid for each mode. (is modified by this function) Nm : int The number of azimuthal modes invdz, invdr : float (in meters^-1) Inverse of the grid step along the considered direction zmin, rmin : float (in meters) Position of the edge of the simulation box, along the direction considered Nz, Nr : int Number of gridpoints along the considered direction nthreads : int Number of CPU threads used with numba prange ptcl_chunk_indices : array of int, of size nthreads+1 The indices (of the particle array) between which each thread should loop. (i.e. divisions of particle array between threads) beta_n_m_0 : 1darray of floats Ruyten-corrected particle shape factor coefficients for mode 0. beta_n_m_higher : 1darray of floats Ruyten-corrected particle shape factor coefficients for higher modes. Ignored when Nm == 1. """ # Deposit the field per cell in parallel (for threads < number of cells) for i_thread in prange(nthreads): # Allocate thread-local array jr_scal = np.zeros(Nm, dtype=np.complex128) jt_scal = np.zeros(Nm, dtype=np.complex128) jz_scal = np.zeros(Nm, dtype=np.complex128) # Loop over all particles in thread chunk for i_ptcl in range(ptcl_chunk_indices[i_thread], ptcl_chunk_indices[i_thread + 1]): # Position xj = x[i_ptcl] yj = y[i_ptcl] zj = z[i_ptcl] # Velocity uxj = ux[i_ptcl] uyj = uy[i_ptcl] uzj = uz[i_ptcl] # Inverse gamma inv_gammaj = inv_gamma[i_ptcl] # Weights wj = q * w[i_ptcl] # Cylindrical conversion rj = math.sqrt(xj**2 + yj**2) # Avoid division by 0. if (rj != 0.): invr = 1. / rj cos = xj * invr # Cosine sin = yj * invr # Sine else: cos = 1. sin = 0. # Calculate contribution from this particle to each mode jr_scal[0] = wj * c * inv_gammaj * (cos * uxj + sin * uyj) jt_scal[0] = wj * c * inv_gammaj * (cos * uyj - sin * uxj) jz_scal[0] = wj * c * inv_gammaj * uzj for m in range(1, Nm): jr_scal[m] = (cos + 1.j * sin) * jr_scal[m - 1] jt_scal[m] = (cos + 1.j * sin) * jt_scal[m - 1] jz_scal[m] = (cos + 1.j * sin) * jz_scal[m - 1] # Positions of the particles, in the cell unit r_cell = invdr * (rj - rmin) - 0.5 z_cell = invdz * (zj - zmin) - 0.5 # Index of the lowest cell of `global_array` that gets modified # by this particle (note: `global_array` has 2 guard cells) # (`min` function avoids out-of-bounds access at high r) ir_cell = min(int(math.ceil(r_cell)), Nr) iz_cell = int(math.ceil(z_cell)) ir = min(int(math.ceil(r_cell)), Nr) # Add contribution of this particle to the global array for m in range(Nm): # Ruyten-corrected shape factor coefficient if m == 0: bn = beta_n_m_0[ir] else: bn = beta_n_m_higher[ir] j_r_global[i_thread, m, iz_cell + 0, ir_cell + 0] += Sz_cubic( z_cell, 0) * Sr_cubic(r_cell, 0, -(-1)**m, bn) * jr_scal[m] j_r_global[i_thread, m, iz_cell + 0, ir_cell + 1] += Sz_cubic( z_cell, 0) * Sr_cubic(r_cell, 1, -(-1)**m, bn) * jr_scal[m] j_r_global[i_thread, m, iz_cell + 0, ir_cell + 2] += Sz_cubic( z_cell, 0) * Sr_cubic(r_cell, 2, -(-1)**m, bn) * jr_scal[m] j_r_global[i_thread, m, iz_cell + 0, ir_cell + 3] += Sz_cubic( z_cell, 0) * Sr_cubic(r_cell, 3, -(-1)**m, bn) * jr_scal[m] j_r_global[i_thread, m, iz_cell + 1, ir_cell + 0] += Sz_cubic( z_cell, 1) * Sr_cubic(r_cell, 0, -(-1)**m, bn) * jr_scal[m] j_r_global[i_thread, m, iz_cell + 1, ir_cell + 1] += Sz_cubic( z_cell, 1) * Sr_cubic(r_cell, 1, -(-1)**m, bn) * jr_scal[m] j_r_global[i_thread, m, iz_cell + 1, ir_cell + 2] += Sz_cubic( z_cell, 1) * Sr_cubic(r_cell, 2, -(-1)**m, bn) * jr_scal[m] j_r_global[i_thread, m, iz_cell + 1, ir_cell + 3] += Sz_cubic( z_cell, 1) * Sr_cubic(r_cell, 3, -(-1)**m, bn) * jr_scal[m] j_r_global[i_thread, m, iz_cell + 2, ir_cell + 0] += Sz_cubic( z_cell, 2) * Sr_cubic(r_cell, 0, -(-1)**m, bn) * jr_scal[m] j_r_global[i_thread, m, iz_cell + 2, ir_cell + 1] += Sz_cubic( z_cell, 2) * Sr_cubic(r_cell, 1, -(-1)**m, bn) * jr_scal[m] j_r_global[i_thread, m, iz_cell + 2, ir_cell + 2] += Sz_cubic( z_cell, 2) * Sr_cubic(r_cell, 2, -(-1)**m, bn) * jr_scal[m] j_r_global[i_thread, m, iz_cell + 2, ir_cell + 3] += Sz_cubic( z_cell, 2) * Sr_cubic(r_cell, 3, -(-1)**m, bn) * jr_scal[m] j_r_global[i_thread, m, iz_cell + 3, ir_cell + 0] += Sz_cubic( z_cell, 3) * Sr_cubic(r_cell, 0, -(-1)**m, bn) * jr_scal[m] j_r_global[i_thread, m, iz_cell + 3, ir_cell + 1] += Sz_cubic( z_cell, 3) * Sr_cubic(r_cell, 1, -(-1)**m, bn) * jr_scal[m] j_r_global[i_thread, m, iz_cell + 3, ir_cell + 2] += Sz_cubic( z_cell, 3) * Sr_cubic(r_cell, 2, -(-1)**m, bn) * jr_scal[m] j_r_global[i_thread, m, iz_cell + 3, ir_cell + 3] += Sz_cubic( z_cell, 3) * Sr_cubic(r_cell, 3, -(-1)**m, bn) * jr_scal[m] j_t_global[i_thread, m, iz_cell + 0, ir_cell + 0] += Sz_cubic( z_cell, 0) * Sr_cubic(r_cell, 0, -(-1)**m, bn) * jt_scal[m] j_t_global[i_thread, m, iz_cell + 0, ir_cell + 1] += Sz_cubic( z_cell, 0) * Sr_cubic(r_cell, 1, -(-1)**m, bn) * jt_scal[m] j_t_global[i_thread, m, iz_cell + 0, ir_cell + 2] += Sz_cubic( z_cell, 0) * Sr_cubic(r_cell, 2, -(-1)**m, bn) * jt_scal[m] j_t_global[i_thread, m, iz_cell + 0, ir_cell + 3] += Sz_cubic( z_cell, 0) * Sr_cubic(r_cell, 3, -(-1)**m, bn) * jt_scal[m] j_t_global[i_thread, m, iz_cell + 1, ir_cell + 0] += Sz_cubic( z_cell, 1) * Sr_cubic(r_cell, 0, -(-1)**m, bn) * jt_scal[m] j_t_global[i_thread, m, iz_cell + 1, ir_cell + 1] += Sz_cubic( z_cell, 1) * Sr_cubic(r_cell, 1, -(-1)**m, bn) * jt_scal[m] j_t_global[i_thread, m, iz_cell + 1, ir_cell + 2] += Sz_cubic( z_cell, 1) * Sr_cubic(r_cell, 2, -(-1)**m, bn) * jt_scal[m] j_t_global[i_thread, m, iz_cell + 1, ir_cell + 3] += Sz_cubic( z_cell, 1) * Sr_cubic(r_cell, 3, -(-1)**m, bn) * jt_scal[m] j_t_global[i_thread, m, iz_cell + 2, ir_cell + 0] += Sz_cubic( z_cell, 2) * Sr_cubic(r_cell, 0, -(-1)**m, bn) * jt_scal[m] j_t_global[i_thread, m, iz_cell + 2, ir_cell + 1] += Sz_cubic( z_cell, 2) * Sr_cubic(r_cell, 1, -(-1)**m, bn) * jt_scal[m] j_t_global[i_thread, m, iz_cell + 2, ir_cell + 2] += Sz_cubic( z_cell, 2) * Sr_cubic(r_cell, 2, -(-1)**m, bn) * jt_scal[m] j_t_global[i_thread, m, iz_cell + 2, ir_cell + 3] += Sz_cubic( z_cell, 2) * Sr_cubic(r_cell, 3, -(-1)**m, bn) * jt_scal[m] j_t_global[i_thread, m, iz_cell + 3, ir_cell + 0] += Sz_cubic( z_cell, 3) * Sr_cubic(r_cell, 0, -(-1)**m, bn) * jt_scal[m] j_t_global[i_thread, m, iz_cell + 3, ir_cell + 1] += Sz_cubic( z_cell, 3) * Sr_cubic(r_cell, 1, -(-1)**m, bn) * jt_scal[m] j_t_global[i_thread, m, iz_cell + 3, ir_cell + 2] += Sz_cubic( z_cell, 3) * Sr_cubic(r_cell, 2, -(-1)**m, bn) * jt_scal[m] j_t_global[i_thread, m, iz_cell + 3, ir_cell + 3] += Sz_cubic( z_cell, 3) * Sr_cubic(r_cell, 3, -(-1)**m, bn) * jt_scal[m] j_z_global[i_thread, m, iz_cell + 0, ir_cell + 0] += Sz_cubic( z_cell, 0) * Sr_cubic(r_cell, 0, (-1)**m, bn) * jz_scal[m] j_z_global[i_thread, m, iz_cell + 0, ir_cell + 1] += Sz_cubic( z_cell, 0) * Sr_cubic(r_cell, 1, (-1)**m, bn) * jz_scal[m] j_z_global[i_thread, m, iz_cell + 0, ir_cell + 2] += Sz_cubic( z_cell, 0) * Sr_cubic(r_cell, 2, (-1)**m, bn) * jz_scal[m] j_z_global[i_thread, m, iz_cell + 0, ir_cell + 3] += Sz_cubic( z_cell, 0) * Sr_cubic(r_cell, 3, (-1)**m, bn) * jz_scal[m] j_z_global[i_thread, m, iz_cell + 1, ir_cell + 0] += Sz_cubic( z_cell, 1) * Sr_cubic(r_cell, 0, (-1)**m, bn) * jz_scal[m] j_z_global[i_thread, m, iz_cell + 1, ir_cell + 1] += Sz_cubic( z_cell, 1) * Sr_cubic(r_cell, 1, (-1)**m, bn) * jz_scal[m] j_z_global[i_thread, m, iz_cell + 1, ir_cell + 2] += Sz_cubic( z_cell, 1) * Sr_cubic(r_cell, 2, (-1)**m, bn) * jz_scal[m] j_z_global[i_thread, m, iz_cell + 1, ir_cell + 3] += Sz_cubic( z_cell, 1) * Sr_cubic(r_cell, 3, (-1)**m, bn) * jz_scal[m] j_z_global[i_thread, m, iz_cell + 2, ir_cell + 0] += Sz_cubic( z_cell, 2) * Sr_cubic(r_cell, 0, (-1)**m, bn) * jz_scal[m] j_z_global[i_thread, m, iz_cell + 2, ir_cell + 1] += Sz_cubic( z_cell, 2) * Sr_cubic(r_cell, 1, (-1)**m, bn) * jz_scal[m] j_z_global[i_thread, m, iz_cell + 2, ir_cell + 2] += Sz_cubic( z_cell, 2) * Sr_cubic(r_cell, 2, (-1)**m, bn) * jz_scal[m] j_z_global[i_thread, m, iz_cell + 2, ir_cell + 3] += Sz_cubic( z_cell, 2) * Sr_cubic(r_cell, 3, (-1)**m, bn) * jz_scal[m] j_z_global[i_thread, m, iz_cell + 3, ir_cell + 0] += Sz_cubic( z_cell, 3) * Sr_cubic(r_cell, 0, (-1)**m, bn) * jz_scal[m] j_z_global[i_thread, m, iz_cell + 3, ir_cell + 1] += Sz_cubic( z_cell, 3) * Sr_cubic(r_cell, 1, (-1)**m, bn) * jz_scal[m] j_z_global[i_thread, m, iz_cell + 3, ir_cell + 2] += Sz_cubic( z_cell, 3) * Sr_cubic(r_cell, 2, (-1)**m, bn) * jz_scal[m] j_z_global[i_thread, m, iz_cell + 3, ir_cell + 3] += Sz_cubic( z_cell, 3) * Sr_cubic(r_cell, 3, (-1)**m, bn) * jz_scal[m] return
def numba_push_eb_comoving(Ep, Em, Ez, Bp, Bm, Bz, Jp, Jm, Jz, rho_prev, rho_next, rho_prev_coef, rho_next_coef, j_coef, C, S_w, T_eb, T_cc, T_rho, kr, kz, dt, V, use_true_rho, Nz, Nr): """ Push the fields over one timestep, using the psatd algorithm, with the assumptions of comoving currents (either with the galilean scheme or comoving scheme, depending on the values of the coefficients that are passed) See the documentation of SpectralGrid.push_eb_with """ # Loop over the grid (parallel in z, if threading is installed) for iz in prange(Nz): for ir in range(Nr): # Save the electric fields, since it is needed for the B push Ep_old = Ep[iz, ir] Em_old = Em[iz, ir] Ez_old = Ez[iz, ir] # Calculate useful auxiliary arrays if use_true_rho: # Evaluation using the rho projected on the grid rho_diff = rho_next_coef[iz, ir] * rho_next[iz, ir] \ - rho_prev_coef[iz, ir] * rho_prev[iz, ir] else: # Evaluation using div(E) and div(J) divE = kr[iz, ir]*( Ep[iz, ir] - Em[iz, ir] ) \ + 1.j*kz[iz, ir]*Ez[iz, ir] divJ = kr[iz, ir]*( Jp[iz, ir] - Jm[iz, ir] ) \ + 1.j*kz[iz, ir]*Jz[iz, ir] rho_diff = ( T_eb[iz,ir] * rho_next_coef[iz, ir] \ - rho_prev_coef[iz, ir] ) \ * epsilon_0 * divE + T_rho[iz, ir] \ * rho_next_coef[iz, ir] * divJ # Push the E field Ep[iz, ir] = \ T_eb[iz, ir]*C[iz, ir]*Ep[iz, ir] + 0.5*kr[iz, ir]*rho_diff \ + j_coef[iz, ir]*1.j*kz[iz, ir]*V*Jp[iz, ir] \ + c2*T_eb[iz, ir]*S_w[iz, ir]*( -1.j*0.5*kr[iz, ir]*Bz[iz, ir] \ + kz[iz, ir]*Bp[iz, ir] - mu_0*T_cc[iz, ir]*Jp[iz, ir] ) Em[iz, ir] = \ T_eb[iz, ir]*C[iz, ir]*Em[iz, ir] - 0.5*kr[iz, ir]*rho_diff \ + j_coef[iz, ir]*1.j*kz[iz, ir]*V*Jm[iz, ir] \ + c2*T_eb[iz, ir]*S_w[iz, ir]*( -1.j*0.5*kr[iz, ir]*Bz[iz, ir] \ - kz[iz, ir]*Bm[iz, ir] - mu_0*T_cc[iz, ir]*Jm[iz, ir] ) Ez[iz, ir] = \ T_eb[iz, ir]*C[iz, ir]*Ez[iz, ir] - 1.j*kz[iz, ir]*rho_diff \ + j_coef[iz, ir]*1.j*kz[iz, ir]*V*Jz[iz, ir] \ + c2*T_eb[iz, ir]*S_w[iz, ir]*( 1.j*kr[iz, ir]*Bp[iz, ir] \ + 1.j*kr[iz, ir]*Bm[iz, ir] - mu_0*T_cc[iz, ir]*Jz[iz, ir] ) # Push the B field Bp[iz, ir] = T_eb[iz, ir]*C[iz, ir]*Bp[iz, ir] \ - T_eb[iz, ir]*S_w[iz, ir]*( -1.j*0.5*kr[iz, ir]*Ez_old \ + kz[iz, ir]*Ep_old ) \ + j_coef[iz, ir]*( -1.j*0.5*kr[iz, ir]*Jz[iz, ir] \ + kz[iz, ir]*Jp[iz, ir] ) Bm[iz, ir] = T_eb[iz, ir]*C[iz, ir]*Bm[iz, ir] \ - T_eb[iz, ir]*S_w[iz, ir]*( -1.j*0.5*kr[iz, ir]*Ez_old \ - kz[iz, ir]*Em_old ) \ + j_coef[iz, ir]*( -1.j*0.5*kr[iz, ir]*Jz[iz, ir] \ - kz[iz, ir]*Jm[iz, ir] ) Bz[iz, ir] = T_eb[iz, ir]*C[iz, ir]*Bz[iz, ir] \ - T_eb[iz, ir]*S_w[iz, ir]*( 1.j*kr[iz, ir]*Ep_old \ + 1.j*kr[iz, ir]*Em_old ) \ + j_coef[iz, ir]*( 1.j*kr[iz, ir]*Jp[iz, ir] \ + 1.j*kr[iz, ir]*Jm[iz, ir] ) return
def numba_push_eb_standard(Ep, Em, Ez, Bp, Bm, Bz, Jp, Jm, Jz, rho_prev, rho_next, rho_prev_coef, rho_next_coef, j_coef, C, S_w, kr, kz, dt, use_true_rho, Nz, Nr): """ Push the fields over one timestep, using the standard psatd algorithm See the documentation of SpectralGrid.push_eb_with """ # Loop over the 2D grid (parallel in z, if threading is installed) for iz in prange(Nz): for ir in range(Nr): # Save the electric fields, since it is needed for the B push Ep_old = Ep[iz, ir] Em_old = Em[iz, ir] Ez_old = Ez[iz, ir] # Calculate useful auxiliary arrays if use_true_rho: # Evaluation using the rho projected on the grid rho_diff = rho_next_coef[iz, ir] * rho_next[iz, ir] \ - rho_prev_coef[iz, ir] * rho_prev[iz, ir] else: # Evaluation using div(E) and div(J) divE = kr[iz, ir]*( Ep[iz, ir] - Em[iz, ir] ) \ + 1.j*kz[iz, ir]*Ez[iz, ir] divJ = kr[iz, ir]*( Jp[iz, ir] - Jm[iz, ir] ) \ + 1.j*kz[iz, ir]*Jz[iz, ir] rho_diff = (rho_next_coef[iz, ir] - rho_prev_coef[iz, ir]) \ * epsilon_0 * divE - rho_next_coef[iz, ir] * dt * divJ # Push the E field Ep[iz, ir] = C[iz, ir]*Ep[iz, ir] + 0.5*kr[iz, ir]*rho_diff \ + c2*S_w[iz, ir]*( -1.j*0.5*kr[iz, ir]*Bz[iz, ir] \ + kz[iz, ir]*Bp[iz, ir] - mu_0*Jp[iz, ir] ) Em[iz, ir] = C[iz, ir]*Em[iz, ir] - 0.5*kr[iz, ir]*rho_diff \ + c2*S_w[iz, ir]*( -1.j*0.5*kr[iz, ir]*Bz[iz, ir] \ - kz[iz, ir]*Bm[iz, ir] - mu_0*Jm[iz, ir] ) Ez[iz, ir] = C[iz, ir]*Ez[iz, ir] - 1.j*kz[iz, ir]*rho_diff \ + c2*S_w[iz, ir]*( 1.j*kr[iz, ir]*Bp[iz, ir] \ + 1.j*kr[iz, ir]*Bm[iz, ir] - mu_0*Jz[iz, ir] ) # Push the B field Bp[iz, ir] = C[iz, ir]*Bp[iz, ir] \ - S_w[iz, ir]*( -1.j*0.5*kr[iz, ir]*Ez_old \ + kz[iz, ir]*Ep_old ) \ + j_coef[iz, ir]*( -1.j*0.5*kr[iz, ir]*Jz[iz, ir] \ + kz[iz, ir]*Jp[iz, ir] ) Bm[iz, ir] = C[iz, ir]*Bm[iz, ir] \ - S_w[iz, ir]*( -1.j*0.5*kr[iz, ir]*Ez_old \ - kz[iz, ir]*Em_old ) \ + j_coef[iz, ir]*( -1.j*0.5*kr[iz, ir]*Jz[iz, ir] \ - kz[iz, ir]*Jm[iz, ir] ) Bz[iz, ir] = C[iz, ir]*Bz[iz, ir] \ - S_w[iz, ir]*( 1.j*kr[iz, ir]*Ep_old \ + 1.j*kr[iz, ir]*Em_old ) \ + j_coef[iz, ir]*( 1.j*kr[iz, ir]*Jp[iz, ir] \ + 1.j*kr[iz, ir]*Jm[iz, ir] ) return