[WIP] Lasy GPU implementation

.github/badges/coverage-badge.json

@@ -0,0 +1,6 @@
+{
+  "schemaVersion": 1,
+  "label": "coverage",
+  "message": "84%",
+  "color": "green"
+}

.github/workflows/coverage.yml

@@ -1,7 +1,7 @@
 name: Weekly Coverage Update
 on:
   schedule:
-    - cron: "0 0 * * 1"
+    - cron: "30 11 * * 4"
   workflow_dispatch:
 
 jobs:
@@ -99,3 +99,4 @@ jobs:
         branch: update-coverage-badge
         base: development
         delete-branch: true
+        push-to-fork: LASY-org/LASY

lasy/backend.py

@@ -0,0 +1,10 @@
+try:
+    import cupy as xp
+
+    use_cupy = True
+except ImportError:
+    import numpy as xp
+
+    use_cupy = False
+
+__all__ = ["use_cupy", "xp"]

lasy/laser.py

@@ -1,5 +1,3 @@
-import numpy as np
-
 from lasy.utils.grid import Grid
 from lasy.utils.laser_utils import (
     normalize_average_intensity,
@@ -12,6 +10,8 @@
 from lasy.utils.openpmd_helper import write_to_openpmd_file
 from lasy.utils.plotting import show_laser
 
+from .backend import xp
+
 
 class Laser:
     """
@@ -89,7 +89,7 @@ class Laser:
     >>>     extent[2:] *= 1e6
     >>>     extent[:2] *= 1e12
     >>>     tmin, tmax, rmin, rmax = extent
-    >>>     vmax = np.abs(E_rt).max()
+    >>>     vmax = xp.abs(E_rt).max()
     >>>     axes[step].imshow(
     ...         E_rt,
     ...         origin="lower",
@@ -122,13 +122,13 @@ def __init__(
 
         # Create the grid on which to evaluate the laser, evaluate it
         if self.dim == "xyt":
-            x, y, t = np.meshgrid(*self.grid.axes, indexing="ij")
+            x, y, t = xp.meshgrid(*self.grid.axes, indexing="ij")
             self.grid.set_temporal_field(profile.evaluate(x, y, t))
         elif self.dim == "rt":
             profile_rt = profile.dim == "rt" if hasattr(profile, "dim") else False
             if profile_rt:
-                r, t = np.meshgrid(*self.grid.axes, indexing="ij")
-                field = np.zeros(
+                r, t = xp.meshgrid(*self.grid.axes, indexing="ij")
+                field = xp.zeros(
                     (2 * self.grid.n_azimuthal_modes - 1, *r.shape), dtype="complex128"
                 )
                 for mode in range(2 * self.grid.n_azimuthal_modes - 1):
@@ -140,17 +140,17 @@ def __init__(
                     n_theta_evals = 2 * self.grid.n_azimuthal_modes - 1
                 # Make sure that there are enough points to resolve the azimuthal modes
                 assert n_theta_evals >= 2 * self.grid.n_azimuthal_modes - 1
-                theta1d = 2 * np.pi / n_theta_evals * np.arange(n_theta_evals)
-                theta, r, t = np.meshgrid(theta1d, *self.grid.axes, indexing="ij")
-                x = r * np.cos(theta)
-                y = r * np.sin(theta)
+                theta1d = 2 * xp.pi / n_theta_evals * xp.arange(n_theta_evals)
+                theta, r, t = xp.meshgrid(theta1d, *self.grid.axes, indexing="ij")
+                x = r * xp.cos(theta)
+                y = r * xp.sin(theta)
                 # Evaluate the profile on the generated grid
                 envelope = profile.evaluate(x, y, t)
                 # Perform the azimuthal decomposition
-                azimuthal_modes = np.fft.ifft(envelope, axis=0)
+                azimuthal_modes = xp.fft.ifft(envelope, axis=0)
                 field = azimuthal_modes[:n_azimuthal_modes]
                 if n_azimuthal_modes > 1:
-                    field = np.concatenate(
+                    field = xp.concatenate(
                         (field, azimuthal_modes[-n_azimuthal_modes + 1 :])
                     )
             self.grid.set_temporal_field(field)
@@ -202,14 +202,14 @@ def apply_optics(self, optical_element):
         # Apply optical element
         spectral_field, spectral_axis = self.grid.get_spectral_field()
         if self.dim == "rt":
-            r, omega = np.meshgrid(
+            r, omega = xp.meshgrid(
                 self.grid.axes[0], spectral_axis + self.profile.omega0, indexing="ij"
             )
             # The line below assumes that amplitude_multiplier
             # is cylindrically symmetric, hence we pass
             # `r` as `x` and an array of 0s as `y`
             multiplier = optical_element.amplitude_multiplier(
-                r, np.zeros_like(r), omega
+                r, xp.zeros_like(r), omega
             )
             # The azimuthal modes are the components of the Fourier transform
             # along theta (FT_theta). Because the multiplier is assumed to be
@@ -219,7 +219,7 @@ def apply_optics(self, optical_element):
             for i_m in range(self.grid.azimuthal_modes.size):
                 spectral_field[i_m, :, :] *= multiplier
         else:
-            x, y, omega = np.meshgrid(
+            x, y, omega = xp.meshgrid(
                 self.grid.axes[0],
                 self.grid.axes[1],
                 spectral_axis + self.profile.omega0,

lasy/optical_elements/axicon.py

@@ -1,6 +1,7 @@
-import numpy as np
 from scipy.constants import c
 
+from lasy.backend import xp
+
 from .optical_element import OpticalElement
 
 
@@ -48,6 +49,6 @@ def amplitude_multiplier(self, x, y, omega):
             Contains the value of the multiplier at the specified points.
             This array has the same shape as the array omega.
         """
-        return np.exp(
-            -2j * (omega / c) * np.sqrt(x**2 + y**2) * np.tan(0.5 * self.gamma)
+        return xp.exp(
+            -2j * (omega / c) * xp.sqrt(x**2 + y**2) * xp.tan(0.5 * self.gamma)
         )

lasy/optical_elements/axiparabola.py

@@ -1,6 +1,7 @@
-import numpy as np
 from scipy.constants import c
 
+from lasy.backend import xp
+
 from .optical_element import OpticalElement
 
 
@@ -61,7 +62,7 @@ def amplitude_multiplier(self, x, y, omega):
         )
 
         # Calculate phase shift
-        T = np.exp(-2j * (omega / c) * sag)
+        T = xp.exp(-2j * (omega / c) * sag)
         # Remove intensity beyond R
         T[x**2 + y**2 > self.R**2] = 0
 

lasy/optical_elements/optical_element.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 
-import numpy as np
+from lasy.backend import xp
 
 
 class OpticalElement(ABC):
@@ -45,4 +45,4 @@ def amplitude_multiplier(self, x, y, omega):
         """
         # The base class only defines dummy multiplier
         # (This should be replaced by any class that inherits from this one.)
-        return np.ones_like(x, dtype="complex128")
+        return xp.ones_like(x, dtype="complex128")

lasy/optical_elements/parabolic_mirror.py

@@ -1,6 +1,7 @@
-import numpy as np
 from scipy.constants import c
 
+from lasy.backend import xp
+
 from .optical_element import OpticalElement
 
 
@@ -44,4 +45,4 @@ def amplitude_multiplier(self, x, y, omega):
             Contains the value of the multiplier at the specified points.
             This array has the same shape as the array omega.
         """
-        return np.exp(-1j * omega * (x**2 + y**2) / (2 * c * self.f))
+        return xp.exp(-1j * omega * (x**2 + y**2) / (2 * c * self.f))

lasy/optical_elements/polynomial_spectral_phase.py

@@ -1,4 +1,4 @@
-import numpy as np
+from lasy.backend import xp
 
 from .optical_element import OpticalElement
 
@@ -70,4 +70,4 @@ def amplitude_multiplier(self, x, y, omega):
             + self.fod / 24 * (omega - self.omega0) ** 4
         )
 
-        return np.exp(1j * spectral_phase)
+        return xp.exp(1j * spectral_phase)

lasy/optical_elements/spectral_filter.py

@@ -1,4 +1,4 @@
-import numpy as np
+from lasy.backend import xp
 
 from .optical_element import OpticalElement
 
@@ -44,12 +44,12 @@ def amplitude_multiplier(self, x, y, omega):
         """
         # Sort the transmission function and angular frequencies to
         # eliminate problems with interpolation (e.g. due to fftshifted data)
-        order = np.argsort(self.omega_in)
+        order = xp.argsort(self.omega_in)
         self.omega_in = self.omega_in[order]
         self.transmission = self.transmission[order]
 
         # interpolate transmission function to omega axis of the laser
-        transmission = np.interp(omega, self.omega_in, self.transmission)
+        transmission = xp.interp(omega, self.omega_in, self.transmission)
 
         # return the square root of the energy/intensity transmission
-        return np.sqrt(transmission)
+        return xp.sqrt(transmission)

lasy/optical_elements/spectral_phase.py

@@ -1,4 +1,4 @@
-import numpy as np
+from lasy.backend import xp
 
 from .optical_element import OpticalElement
 
@@ -44,11 +44,11 @@ def amplitude_multiplier(self, x, y, omega):
         """
         # Sort the transmission function and angular frequencies to
         # eliminate problems with interpolation (e.g. due to fftshifted data)
-        order = np.argsort(self.omega_in)
+        order = xp.argsort(self.omega_in)
         self.omega_in = self.omega_in[order]
         self.phase = self.phase[order]
 
         # interpolate transmission function to omega axis of the laser
-        phase = np.interp(omega, self.omega_in, self.phase)
+        phase = xp.interp(omega, self.omega_in, self.phase)
 
-        return np.exp(1j * phase)
+        return xp.exp(1j * phase)

lasy/optical_elements/zernike_aberrations.py

@@ -1,5 +1,4 @@
-import numpy as np
-
+from lasy.backend import xp
 from lasy.utils.zernike import zernike
 
 from .optical_element import OpticalElement
@@ -59,8 +58,8 @@ def amplitude_multiplier(self, x, y, omega):
             Contains the value of the multiplier at the specified points.
             This array has the same shape as the array omega.
         """
-        rr = np.sqrt(x**2 + y**2)
-        phase = np.zeros_like(rr)
+        rr = xp.sqrt(x**2 + y**2)
+        phase = xp.zeros_like(rr)
 
         for j in list(self.zernike_amplitudes):
             # Create the zernike phase and ensure it has the same number of dimensions as phase
@@ -70,8 +69,8 @@ def amplitude_multiplier(self, x, y, omega):
 
             # Increase the length of the frequency dimension such that the shape is suitable to be added
             # to the phase array, then add it
-            phase += self.zernike_amplitudes[j] * np.broadcast_to(
+            phase += self.zernike_amplitudes[j] * xp.broadcast_to(
                 zernike_phase, phase.shape
             )
 
-        return np.exp(1j * phase)
+        return xp.exp(1j * phase)

lasy/profiles/from_array_profile.py

@@ -1,6 +1,7 @@
-import numpy as np
 from scipy.interpolate import RegularGridInterpolator
 
+from lasy.backend import xp
+
 from .profile import Profile
 
 
@@ -62,7 +63,7 @@ def __init__(self, wavelength, pol, array, dim, axes, axes_order=["x", "y", "t"]
 
             self.combined_field_interp = RegularGridInterpolator(
                 (axes["x"], axes["y"], axes["t"]),
-                np.abs(self.array) + 1.0j * np.unwrap(np.angle(self.array), axis=-1),
+                xp.abs(self.array) + 1.0j * xp.unwrap(xp.angle(self.array), axis=-1),
                 bounds_error=False,
                 fill_value=0.0,
             )
@@ -73,12 +74,12 @@ def __init__(self, wavelength, pol, array, dim, axes, axes_order=["x", "y", "t"]
             # to make correct interpolation within the first cell
             if axes["r"][0] != 0.0:
                 # add mirrored point to the axis
-                r = np.concatenate(([-axes["r"][0]], axes["r"]))
+                r = xp.concatenate(([-axes["r"][0]], axes["r"]))
                 # takes first element of the array in the radial dimension
                 subarray = self.array[:, 0, :]
                 # add it at the beginning to be the value at the mirrored point
-                self.array = np.concatenate(
-                    (subarray[:, np.newaxis, :], self.array), axis=1
+                self.array = xp.concatenate(
+                    (subarray[:, xp.newaxis, :], self.array), axis=1
                 )
             else:
                 r = axes["r"]
@@ -94,8 +95,8 @@ def __init__(self, wavelength, pol, array, dim, axes, axes_order=["x", "y", "t"]
                 self.field_interp_modes.append(
                     RegularGridInterpolator(
                         (r, axes["t"]),
-                        np.abs(self.array[imode, :, :])
-                        + 1.0j * np.unwrap(np.angle(self.array[imode, :, :]), axis=-1),
+                        xp.abs(self.array[imode, :, :])
+                        + 1.0j * xp.unwrap(xp.angle(self.array[imode, :, :]), axis=-1),
                         bounds_error=False,
                         fill_value=0.0,
                     )
@@ -106,19 +107,19 @@ def evaluate(self, x, y, t):
         if self.dim == "xyt":
             combined_field = self.combined_field_interp((x, y, t))
         else:
-            r = np.sqrt(x**2 + y**2)
-            theta = np.angle(x + 1j * y)
-            combined_field = np.zeros_like(x, dtype="complex128")
+            r = xp.sqrt(x**2 + y**2)
+            theta = xp.angle(x + 1j * y)
+            combined_field = xp.zeros_like(x, dtype="complex128")
             nmodes = (len(self.field_interp_modes) + 1) // 2
             for imode in range(-nmodes + 1, nmodes):
-                combined_field += self.field_interp_modes[imode]((r, t)) * np.exp(
+                combined_field += self.field_interp_modes[imode]((r, t)) * xp.exp(
                     -1j * imode * theta
                 )
 
-        return np.abs(np.real(combined_field)) * np.exp(1.0j * np.imag(combined_field))
+        return xp.abs(xp.real(combined_field)) * xp.exp(1.0j * xp.imag(combined_field))
 
     def evaluate_mrt(self, mode, r, t):
         """Return the envelope field of the scaled profile."""
         assert self.dim == "rt"
         combined_field = self.field_interp_modes[mode]((r, t))
-        return np.abs(np.real(combined_field)) * np.exp(1.0j * np.imag(combined_field))
+        return xp.abs(xp.real(combined_field)) * xp.exp(1.0j * xp.imag(combined_field))