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

To help you get started, we’ve selected a few fbpic examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github fbpic / fbpic / fbpic / fields / numba_methods.py View on Github external
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
github fbpic / fbpic / fbpic / particles / gathering / threading_methods_one_mode.py View on Github external
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
github fbpic / fbpic / fbpic / fields / numba_methods.py View on Github external
"""
    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]
github fbpic / fbpic / fbpic / particles / elementary_process / cuda_numba_utils.py View on Github external
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 )
github fbpic / fbpic / fbpic / particles / push / numba_methods.py View on Github external
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
github fbpic / fbpic / fbpic / particles / push / numba_methods.py View on Github external
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
github fbpic / fbpic / fbpic / particles / deposition / threading_methods.py View on Github external
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
github fbpic / fbpic / fbpic / fields / numba_methods.py View on Github external
"""
    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]
github fbpic / fbpic / fbpic / particles / gathering / threading_methods.py View on Github external
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.