Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
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):
# 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 * 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
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
"""
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):
field[iz,ir] = filter_array_z[iz]*filter_array_r[ir]*field[iz,ir]
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 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_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
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
"""
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]
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.