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
