Add GPU support using CuPy

[AlexanderSinn]

Based on #437

This PR adds GPU support using CuPy. Similar to how it is done in pywarpx, the lasy.backend file defines xp which is then used instead of numpy np everywhere. xp is either cupy or numpy depending on if the cupy module was found. Additionally lasy.backend defines the helper functions to_cpu and to_gpu which can be used to convert between numpy and cupy arrays. If numpy is used as a backend, the functions are still available but do nothing.

lasy.backend:

try:
    import cupy as xp
    from cupyx.scipy.interpolate import RegularGridInterpolator
    from cupyx.scipy.signal import hilbert, zoom_fft
    from cupyx.scipy.special import j0

    use_cupy = True

    def to_cpu(arr):
        """Convert array from cupy to numpy."""
        if isinstance(arr, xp.ndarray):
            return xp.asnumpy(arr)
        elif isinstance(arr, list):
            return [to_cpu(a) for a in arr]
        elif isinstance(arr, tuple):
            return (to_cpu(a) for a in arr)
        else:
            return arr

    def to_gpu(arr):
        """Convert array from numpy to cupy."""
        if not isinstance(arr, xp.ndarray):
            return xp.asarray(arr)
        else:
            return arr

except ImportError:
    import numpy as xp
    from scipy.interpolate import RegularGridInterpolator
    from scipy.signal import hilbert, zoom_fft
    from scipy.special import j0

    use_cupy = False

    def to_cpu(arr):
        """Convert array from cupy to numpy."""
        return arr

    def to_gpu(arr):
        """Convert array from numpy to cupy."""
        return arr


__all__ = [
    "use_cupy",
    "xp",
    "to_cpu",
    "to_gpu",
    "RegularGridInterpolator",
    "hilbert",
    "zoom_fft",
    "j0",
]

Changes in this PR:

For every single occurrence in lasy of:

import numpy as np
np.call_some_function()

This PR replaces it with:

from lasy.backend import xp
xp.call_some_function()

Except in refractive_index.py where numpy is needed to work with numdifftools.

IMPORTANT: All future additions to lasy should also only use xp instead of np.

In case a lasy array is passed to a different Python library that does not support cupy, such as matplotlib, numdifftools, scipy (the functions not provided in cupyx.scipy), openpmd-api and axiprop (has CuPy backend, but the interface mostly uses CPU arrays), the to_cpu helper function has to be used. For example:

import matplotlib.pyplot as plt
from lasy.backend import to_cpu, xp
plt.plot(to_cpu(laser.grid.axes[-1] * 1e15), to_cpu(get_laser_power(laser.dim, laser.grid)/ 1e12))

When indexing into a CuPy array or using xp.sum(...), a dimension zero array is returned instead of a Python float to avoid copies to CPU and synchronization. This makes it necessary to explicitly cast to float() in a few places where scalars are needed.

[AlexanderSinn]

Now the CI passes when I run it on a GPU (A100) with CuPy. It is twice as fast compated to running it on the CPU.

[asinn@max-mpag003]~/lasy% python -m pytest -W ignore::FutureWarning tests/
=================================== test session starts ====================================
platform linux -- Python 3.9.19, pytest-8.4.2, pluggy-1.6.0
rootdir: /home/asinn/lasy
configfile: pyproject.toml
plugins: anyio-4.4.0
collected 67 items

tests/test_angular_spectrum_propagator.py ..                                         [  2%]
tests/test_axicon.py ..                                                              [  5%]
tests/test_axiparabola.py .                                                          [  7%]
tests/test_collins.py .                                                              [  8%]
tests/test_denoise.py .                                                              [ 10%]
tests/test_fresnel_chirpztransform_propagator.py ....                                [ 16%]
tests/test_gaussian_propagator.py ...                                                [ 20%]
tests/test_gsalgo.py .                                                               [ 22%]
tests/test_laser_profiles.py ..................                                      [ 49%]
tests/test_laser_utils.py ..                                                         [ 52%]
tests/test_parabolic_mirror.py ..                                                    [ 55%]
tests/test_polynomial_spectral_phase.py ...                                          [ 59%]
tests/test_refractive_index.py ...                                                   [ 64%]
tests/test_resampling.py ...                                                         [ 68%]
tests/test_special_laser.py ..                                                       [ 71%]
tests/test_speckles.py ................                                              [ 95%]
tests/test_spm.py .                                                                  [ 97%]
tests/test_t2z2t.py ..                                                               [100%]

============================== 67 passed in 114.48s (0:01:54) ==============================

CPU: 67 passed in 235.59s (0:03:55)