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

    @cuda.jit()
    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]

@cuda.jit()
def extract_array_from_gpu( part_idx_start, array, selected ):
    """
    Extract a selection of particles from the GPU and
    store them in a 1D array (N_part,)
    Selection goes from starting index (part_idx_start)
    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)

@cuda.jit()
def extract_particles_from_gpu( part_idx_start, x, y, z, ux, uy, uz, w,
                                inv_gamma, selected ):
    """
    Extract a selection of particles from the GPU and
    store them in a 2D array (8, N_part) in the following
    order: x, y, z, ux, uy, uz, w, inv_gamma.
    Selection goes from starting index (part_idx_start)
    to (part_idx_start + N_part-1), where N_part is derived
    from the shape of the 2D array (selected).

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

    @cuda.jit()
    def add_rho_to_gpu_array( iz_min, rho_buffer, rho, m ):
        """
        Add the small-size array rho_buffer into the full-size array rho
        on the GPU

        Parameters
        ----------
        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

    @cuda.jit
    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

    @cuda.jit
    def shift_spect_array_gpu( field_array, shift_factor, n_move ):
        """
        Shift the field 'field_array' by n_move cells on the GPU.
        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) )

    @cuda.jit
    def copy_particles( N_elements, source_array, source_start,
                                    target_array, target_start ):
        """
        Copy `N_elements` elements from `source_array` to `target_array`

        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

    @cuda.jit
    def cuda_damp_pml_EB( Et, Et_pml, Ez, Bt, Bt_pml, Bz,
                      damp_array, n_pml ) :
        """
        Damp the E and B fields in the PML cells (i.e. the last n_pml cells
        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

    @cuda.jit
    def extract_slice_cuda( Nr, iz, Sz, slice_arr,
        Er, Et, Ez, Br, Bt, Bz, Jr, Jt, Jz, rho, m ):
        """
        Extract a slice of the fields at iz and iz+1, and interpolated
        between those two points using Sz and (1-Sz)

        Parameters
        ----------
        Nr: int
            Number of cells transversally

        iz: int
            Index at which to extract the fields

        Sz: float
            Interpolation shape factor used at iz