How to use the fbpic.utils.cuda.cuda.grid function in fbpic

Index at which to extract the fields

        Sz: float
            Interpolation shape factor used at iz

        slice_arr: cuda.device_array
            Array of floats of shape (10, 2*Nm-1, Nr)

        Er, Et, etc...: cuda.device_array
            Array of complexs of shape (Nz, Nr), for the azimuthal mode m

        m: int
            Index of the azimuthal mode involved
        """
        # One thread per radial position
        ir = cuda.grid(1)
        # Intermediate variables
        izp = iz+1
        Szp = 1. - Sz

        if ir < Nr:
            # Interpolate the field in the longitudinal direction
            # and store it into pre-packed arrays

            # For the higher modes:
            # There is a factor 2 here so as to comply with the convention in
            # Lifschitz et al., which is also the convention of FBPIC
            # For performance, this is included in the shape factor.
            if m > 0:
                Sz = 2*Sz
                Szp = 2*Szp
                # Index at which the mode should be added

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
        """
        # Get a 2D CUDA grid
        iz, ir = cuda.grid(2)

        # Only access values that are actually in the array
        if ir < field_array.shape[1] and iz < field_array.shape[0]:
            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 fields
            field_array[iz, ir] *= power_shift

to (part_idx_start + N_part-1), where N_part is derived
    from the shape of the array `selected`.

    Parameters
    ----------
    part_idx_start : int
        The starting index needed for the extraction process.
        ( minimum particle index to be extracted )

    array : 1D arrays of ints or floats
        The GPU particle arrays for a given species. (e.g. particle id)

    selected : 1D array of ints or floats
        An empty GPU array to store the particles that are extracted.
    """
    i = cuda.grid(1)
    N_part = selected.shape[0]

    if i < N_part:
        selected[i] = array[part_idx_start+i]

Parameters:
        ------------
        N_elements: int
            The number of elements to copy

        source_array, target_array: 1d device arrays of floats
            The arrays from/to which the data should be copied
            (represents *one* of the particle quantities)

        source_start, target_start: ints
            The indices at which to start the copy, in both the
            source and the target.
        """
        # Get a 1D CUDA grid (the index corresponds to a particle index)
        i = cuda.grid(1)

        # Copy the particles into the right buffer
        if i < N_elements:
            target_array[i+target_start] = source_array[i+source_start]

def copy_particle_data_cuda( Ntot, old_array, new_array ):
        """
        Copy the `Ntot` elements of `old_array` into `new_array`, on GPU
        """
        # Loop over single particles
        ip = cuda.grid(1)
        if ip < Ntot:
            new_array[ip] = old_array[ip]

def shift_particles_periodic_cuda( z, zmin, zmax ):
        """
        Shift the particle positions by an integer number of box length,
        so that outside particle are back inside the physical domain

        Parameters:
        -----------
        z: 1darray of floats
            The z position of the particles (one element per particle)
        zmin, zmax: floats
            Positions of the edges of the periodic box
        """
        # Get a 1D CUDA grid (the index corresponds to a particle index)
        i = cuda.grid(1)
        # Get box length
        l_box = zmax - zmin
        # Shift particle position
        if i < z.shape[0]:
            while z[i] >= zmax:
                z[i] -= l_box
            while z[i] < zmin:
                z[i] += l_box

----------
        iz_min: int
            The index of the lowest cell in z that surrounds the antenna

        rho_buffer: 3darray of complexs
            Array of shape (Nm, 2, Nr) that stores the values of rho
            in the 2 cells that surround the antenna (for each mode).

        rho: 2darray of complexs
            Array of shape (Nz, Nr) that contains rho in the mode m

        m: int
           The index of the azimuthal mode involved
        """
        # Use one thread per radial cell
        ir = cuda.grid(1)

        # Add the values
        if ir < rho.shape[1]:
            rho[iz_min, ir] += rho_buffer[m, 0, ir]
            rho[iz_min+1, ir] += rho_buffer[m, 1, ir]

Parameters:
        -----------
        iz_min: int

        Jr_buffer, Jt_buffer, Jz_buffer: 3darrays of complexs
            Arrays of shape (Nm, 2, Nr) that store the values of rho
            in the 2 cells that surround the antenna (for each mode).

        Jr, Jt, Jz: 2darrays of complexs
            Arrays of shape (Nz, Nr) that contain rho in the mode m

        m: int
           The index of the azimuthal mode involved
        """
        # Use one thread per radial cell
        ir = cuda.grid(1)

        # Add the values
        if ir < Jr.shape[1]:
            Jr[iz_min, ir] += Jr_buffer[m, 0, ir]
            Jr[iz_min+1, ir] += Jr_buffer[m, 1, ir]

            Jt[iz_min, ir] += Jt_buffer[m, 0, ir]
            Jt[iz_min+1, ir] += Jt_buffer[m, 1, ir]

            Jz[iz_min, ir] += Jz_buffer[m, 0, ir]
            Jz[iz_min+1, ir] += Jz_buffer[m, 1, ir]

in r), in an anisotropic manner which is given by the PML principles

        Parameters :
        ------------
        Et, Et_pml, Ez, Bt, Bt_pml, Bz : 2darrays of complexs
            Contain the fields to be damped
            The first axis corresponds to z and the second to r

        damp_array: 1darray of floats
            An array of length n_guards, which contains the damping factors

        n_pml: int
            Number of PML cells
        """
        # Obtain Cuda grid
        iz, i_pml = cuda.grid(2)

        # Obtain the size of the array along z and r
        Nz, Nr = Et.shape

        # Modify the fields
        if i_pml < n_pml:
            # Apply the damping arrays
            if iz < Nz:
                # Get the damping factor
                damp_factor= damp_array[i_pml]
                # Get the index in the bigger field array
                ir = Nr - n_pml + i_pml
                # Substract the theta PML fields to the regular theta fields
                Et[iz,ir] -= Et_pml[iz,ir]
                Bt[iz,ir] -= Bt_pml[iz,ir]
                # Damp the theta PML fields