How to use the fbpic.utils.threading.njit_parallel function in fbpic

@njit_parallel
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

@njit_parallel
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):

@njit_parallel
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):

@njit_parallel
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

@njit_parallel
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.

@njit_parallel
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],

@njit_parallel
def numba_correct_currents_curlfree_standard( rho_prev, rho_next, Jp, Jm, Jz,
                            kz, kr, inv_k2, inv_dt, Nz, Nr ):
    """
    Correct the currents in spectral space, using the curl-free correction
    which is adapted to the standard psatd
    """
    # 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] ) )

@njit_parallel
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] \

@njit_parallel
def numba_filter_scalar( field, 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)

@njit_parallel
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