[WIP] Use cupy, in which case all operations are performed on GPU

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

@@ -10,6 +10,8 @@
 )
 from lasy.utils.openpmd_output import write_to_openpmd_file
 
+from .backend import use_cupy, xp
+
 
 class Laser:
     """
@@ -87,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",
@@ -113,11 +115,11 @@ def __init__(
         # Get the spectral axis
         dt = self.grid.dx[time_axis_indx]
         Nt = self.grid.shape[time_axis_indx]
-        self.omega_1d = 2 * np.pi * np.fft.fftfreq(Nt, dt) + profile.omega0
+        self.omega_1d = 2 * xp.pi * xp.fft.fftfreq(Nt, dt) + profile.omega0
 
         # 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":
             if n_theta_evals is None:
@@ -126,17 +128,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)
@@ -179,7 +181,8 @@ def apply_optics(self, optical_element):
         # Apply optical element
         spectral_field = self.grid.get_spectral_field()
         if self.dim == "rt":
-            r, omega = np.meshgrid(self.grid.axes[0], self.omega_1d, indexing="ij")
+            r, omega = xp.meshgrid(self.grid.axes[0], self.omega_1d, indexing="ij")
+
             # The line below assumes that amplitude_multiplier
             # is cylindrically symmetric, hence we pass
             # `r` as `x` and 0 as `y`
@@ -194,15 +197,15 @@ 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], self.omega_1d, indexing="ij"
             )
             spectral_field *= optical_element.amplitude_multiplier(
                 x, y, omega, self.profile.omega0
             )
         self.grid.set_spectral_field(spectral_field)
 
-    def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True):
+    def propagate(self, distance, nr_boundary=None, show_progress=True):
         """
         Propagate the laser pulse by the distance specified.
 
@@ -216,16 +219,14 @@ def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True
             will be attenuated (to assert proper Hankel transform).
             Only used for ``'rt'``.
 
-        backend : string (optional)
-            Backend used by axiprop (see axiprop documentation).
         show_progress : bool (optional)
             Whether to show a progress bar when performing the computation
         """
         # apply boundary "absorption" if required
         if nr_boundary is not None:
             assert type(nr_boundary) is int and nr_boundary > 0
-            absorb_layer_axis = np.linspace(0, np.pi / 2, nr_boundary)
-            absorb_layer_shape = np.cos(absorb_layer_axis) ** 0.5
+            absorb_layer_axis = xp.linspace(0, xp.pi / 2, nr_boundary)
+            absorb_layer_shape = xp.cos(absorb_layer_axis) ** 0.5
             absorb_layer_shape[-1] = 0.0
             field = self.grid.get_temporal_field()
             if self.dim == "rt":
@@ -237,16 +238,27 @@ def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True
                 field[:, :nr_boundary, :] *= absorb_layer_shape[::-1][None, :, None]
             self.grid.set_temporal_field(field)
 
+        # Select backend
+        if use_cupy:
+            backend = "CU"
+        else:
+            backend = "NP"
+
+        k = self.omega_1d / c
         if self.dim == "rt":
             # Construct the propagator (check if exists)
             if not hasattr(self, "prop"):
                 spatial_axes = (self.grid.axes[0],)
                 self.prop = []
+                if use_cupy:
+                    # Move quantities to CPU to create propagator
+                    k = xp.asnumpy(k)
+                    spatial_axes = (xp.asnumpy(spatial_axes[0]),)
                 for m in self.grid.azimuthal_modes:
                     self.prop.append(
                         PropagatorResampling(
                             *spatial_axes,
-                            self.omega_1d / c,
+                            k,
                             mode=m,
                             backend=backend,
                             verbose=False,
@@ -255,14 +267,14 @@ def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True
             # Propagate the spectral image
             spectral_field = self.grid.get_spectral_field()
             for i_m in range(self.grid.azimuthal_modes.size):
-                transform_data = np.transpose(spectral_field[i_m]).copy()
+                transform_data = xp.transpose(spectral_field[i_m]).copy()
                 self.prop[i_m].step(
                     transform_data,
                     distance,
                     overwrite=True,
                     show_progress=show_progress,
                 )
-                spectral_field[i_m, :, :] = np.transpose(transform_data).copy()
+                spectral_field[i_m, :, :] = xp.transpose(transform_data).copy()
             self.grid.set_spectral_field(spectral_field)
         else:
             # Construct the propagator (check if exists)
@@ -279,11 +291,11 @@ def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True
                 )
             # Propagate the spectral image
             spectral_field = self.grid.get_spectral_field()
-            transform_data = np.moveaxis(spectral_field, -1, 0).copy()
+            transform_data = xp.moveaxis(spectral_field, -1, 0).copy()
             self.prop.step(
                 transform_data, distance, overwrite=True, show_progress=show_progress
             )
-            spectral_field = np.moveaxis(transform_data, 0, -1).copy()
+            spectral_field = xp.moveaxis(transform_data, 0, -1).copy()
 
         # Choose the time translation assuming propagation at v=c
         translate_time = distance / c
@@ -293,7 +305,7 @@ def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True
         # propagators, so it needs to be added by hand.
         # Note: subtracting by omega0 is only a global phase convention,
         # that derives from the definition of the envelope in lasy.
-        spectral_field *= np.exp(
+        spectral_field *= xp.exp(
             -1j * (self.omega_1d[None, None, :] - self.profile.omega0) * translate_time
         )
         self.grid.set_spectral_field(spectral_field)
@@ -349,7 +361,8 @@ def show(self, **kw):
         ----------
         **kw : additional arguments to be passed to matplotlib's imshow command
         """
-        temporal_field = self.grid.get_temporal_field()
+        # Get field on CPU
+        temporal_field = self.grid.get_temporal_field(to_cpu=True)
         if self.dim == "rt":
             # Show field in the plane y=0, above and below axis, with proper sign for each mode
             E = [

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):
@@ -48,4 +48,4 @@ def amplitude_multiplier(self, x, y, omega, omega0):
         """
         # 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
 
 
@@ -47,4 +48,4 @@ def amplitude_multiplier(self, x, y, omega, omega0):
             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/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
 
 
@@ -56,35 +57,35 @@ def __init__(self, wavelength, pol, array, dim, axes, axes_order=["x", "y", "t"]
             assert axes_order in [["x", "y", "t"], ["t", "y", "x"]]
 
             if axes_order == ["t", "y", "x"]:
-                self.array = np.swapaxes(array, 0, 2)
+                self.array = xp.swapaxes(array, 0, 2)
             else:
                 self.array = array
 
             self.combined_field_interp = RegularGridInterpolator(
                 (axes["x"], axes["y"], axes["t"]),
-                np.abs(array) + 1.0j * np.unwrap(np.angle(array), axis=-1),
+                xp.abs(array) + 1.0j * xp.unwrap(xp.angle(array), axis=-1),
                 bounds_error=False,
                 fill_value=0.0,
             )
         else:  # dim = "rt"
             assert axes_order in [["r", "t"], ["t", "r"]]
 
             if axes_order == ["t", "r"]:
-                self.array = np.swapaxes(array, 0, 2)
+                self.array = xp.swapaxes(array, 0, 2)
             else:
                 self.array = array
 
             # If the first point of radial axis is not 0, we "mirror" it,
             # to make correct interpolation within the first cell
             if axes["r"][0] != 0.0:
-                r = np.concatenate(([-axes["r"][0]], axes["r"]))
-                array = np.concatenate(([array[0]], array))
+                r = xp.concatenate(([-axes["r"][0]], axes["r"]))
+                array = xp.concatenate(([array[0]], array))
             else:
                 r = axes["r"]
 
             self.combined_field_interp = RegularGridInterpolator(
                 (r, axes["t"]),
-                np.abs(array) + 1.0j * np.unwrap(np.angle(array), axis=-1),
+                xp.abs(array) + 1.0j * xp.unwrap(xp.angle(array), axis=-1),
                 bounds_error=False,
                 fill_value=0.0,
             )
@@ -94,10 +95,10 @@ def evaluate(self, x, y, t):
         if self.dim == "xyt":
             combined_field = self.combined_field_interp((x, y, t))
         else:
-            combined_field = self.combined_field_interp((np.sqrt(x**2 + y**2), t))
+            combined_field = self.combined_field_interp((xp.sqrt(x**2 + y**2), t))
 
-        envelope = np.abs(np.real(combined_field)) * np.exp(
-            1.0j * np.imag(combined_field)
+        envelope = xp.abs(xp.real(combined_field)) * xp.exp(
+            1.0j * xp.imag(combined_field)
         )
 
         return envelope

lasy/profiles/from_insight_file.py

@@ -1,7 +1,8 @@
 import h5py
-import numpy as np
 from scipy.constants import c
 
+from lasy.backend import xp
+
 from .from_array_profile import FromArrayProfile
 
 
@@ -30,10 +31,10 @@ class FromInsightFile(FromArrayProfile):
     def __init__(self, file_path, pol, omega0="barycenter"):
         # read the data from H5 filed
         with h5py.File(file_path, "r") as hf:
-            data = np.asanyarray(hf["data/Exyt_0"][()], dtype=np.complex128, order="C")
-            t = np.asanyarray(hf["scales/t"][()], dtype=np.float64, order="C")
-            x = np.asanyarray(hf["scales/x"][()], dtype=np.float64, order="C")
-            y = np.asanyarray(hf["scales/y"][()], dtype=np.float64, order="C")
+            data = xp.asanyarray(hf["data/Exyt_0"][()], dtype=xp.complex128, order="C")
+            t = xp.asanyarray(hf["scales/t"][()], dtype=xp.float64, order="C")
+            x = xp.asanyarray(hf["scales/x"][()], dtype=xp.float64, order="C")
+            y = xp.asanyarray(hf["scales/y"][()], dtype=xp.float64, order="C")
 
         # convert data and axes to SI units
         t *= 1e-15
@@ -42,29 +43,29 @@ def __init__(self, file_path, pol, omega0="barycenter"):
 
         # get the field on axis and local frequencies
         field_onaxis = data[data.shape[0] // 2, data.shape[1] // 2, :]
-        omega_array = -np.gradient(np.unwrap(np.angle(field_onaxis)), t)
+        omega_array = -xp.gradient(xp.unwrap(xp.angle(field_onaxis)), t)
 
         # choose the central frequency
         if omega0 == "peak":
             # using peak field frequency
-            omega0 = omega_array[np.abs(field_onaxis).argmax()]
+            omega0 = omega_array[xp.abs(field_onaxis).argmax()]
         elif omega0 == "barycenter":
             # or "center of mass" frequency
-            omega0 = np.average(omega_array, weights=np.abs(field_onaxis) ** 2)
+            omega0 = xp.average(omega_array, weights=xp.abs(field_onaxis) ** 2)
         else:
             assert type(omega0) == float
 
         # check the complex field convention and correct if needed
         if omega0 < 0:
             omega0 *= -1
-            data = np.conj(data)
+            data = xp.conj(data)
             print("Warning: input field will be conjugated")
 
         # remove the envelope frequency
-        data *= np.exp(1j * omega0 * t[None, None, :])
+        data *= xp.exp(1j * omega0 * t[None, None, :])
 
         # created LASY profile using FromArrayProfile class
-        wavelength = 2 * np.pi * c / omega0
+        wavelength = 2 * xp.pi * c / omega0
         dim = "xyt"
         axes = {"x": x, "y": y, "t": t}
         super().__init__(

lasy/profiles/from_openpmd_profile.py

@@ -1,8 +1,8 @@
-import numpy as np
 import openpmd_api as io
 from openpmd_viewer import OpenPMDTimeSeries
 from scipy.constants import c
 
+from lasy.backend import xp
 from lasy.utils.laser_utils import create_grid, field_to_envelope
 from lasy.utils.openpmd_input import reorder_array
 
@@ -85,7 +85,7 @@ def __init__(
 
         # If `is_envelope` is not given, assume that complex arrays are envelopes.
         if is_envelope is None:
-            is_envelope = np.iscomplexobj(F)
+            is_envelope = xp.iscomplexobj(F)
 
         if theta is None:  # Envelope obtained from the full 3D array
             dim = "xyt"
@@ -109,7 +109,7 @@ def __init__(
             omg0 = it.meshes["laserEnvelope"].get_attribute("angularFrequency")
             array = F
 
-        wavelength = 2 * np.pi * c / omg0
+        wavelength = 2 * xp.pi * c / omg0
         if verbose:
             print(
                 "Wavelength used in the definition of the envelope (nm):",

lasy/profiles/gaussian_profile.py

@@ -89,7 +89,7 @@ class GaussianProfile(CombinedLongitudinalTransverseProfile):
     >>> extent[2:] *= 1e6
     >>> extent[:2] *= 1e15
     >>> tmin, tmax, rmin, rmax = extent
-    >>> vmax = np.abs(E_rt).max()
+    >>> vmax = xp.abs(E_rt).max()
     >>> plt.imshow(
     ...     E_rt,
     ...     origin="lower",