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

lasy/__init__.py

@@ -1 +1,8 @@
 __version__ = "0.4.0"
+
+from .backend import use_cupy
+
+if use_cupy:
+    print("Lasy is using Cupy as a backend.")
+else:
+    print("Lasy is using Numpy as a backend.")

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 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",
@@ -112,7 +114,7 @@ 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":
             if n_theta_evals is None:
@@ -121,17 +123,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)
@@ -175,12 +177,12 @@ def apply_optics(self, optical_element):
         dt = self.grid.dx[time_axis_indx]
         omega0 = self.profile.omega0
         Nt = self.grid.shape[time_axis_indx]
-        omega_1d = 2 * np.pi * np.fft.fftfreq(Nt, dt) + omega0
+        omega_1d = 2 * xp.pi * xp.fft.fftfreq(Nt, dt) + omega0
 
         # Apply optical element
         spectral_field = self.grid.get_spectral_field()
         if self.dim == "rt":
-            r, omega = np.meshgrid(self.grid.axes[0], omega_1d, indexing="ij")
+            r, omega = xp.meshgrid(self.grid.axes[0], omega_1d, indexing="ij")
             # The line below assumes that amplitude_multiplier
             # is cylindrically symmetric, hence we pass
             # `r` as `x` and 0 as `y`
@@ -193,7 +195,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], omega_1d, indexing="ij"
             )
             spectral_field *= optical_element.amplitude_multiplier(x, y, omega)
@@ -221,8 +223,8 @@ def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True
         # 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":
@@ -238,7 +240,7 @@ def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True
         dt = self.grid.dx[time_axis_indx]
         omega0 = self.profile.omega0
         Nt = self.grid.shape[time_axis_indx]
-        omega = 2 * np.pi * np.fft.fftfreq(Nt, dt) + omega0
+        omega = 2 * xp.pi * xp.fft.fftfreq(Nt, dt) + omega0
 
         if self.dim == "rt":
             # Construct the propagator (check if exists)
@@ -258,14 +260,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)
@@ -282,11 +284,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
@@ -296,7 +298,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(-1j * (omega[None, None, :] - omega0) * translate_time)
+        spectral_field *= xp.exp(-1j * (omega[None, None, :] - omega0) * translate_time)
         self.grid.set_spectral_field(spectral_field)
 
         # Translate the domain
@@ -342,7 +344,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,4 +1,4 @@
-import numpy as np
+from lasy.backend import xp
 
 
 class OpticalElement(object):
@@ -38,4 +38,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 arrays x, y, 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
 
 
@@ -55,35 +56,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,
             )
@@ -93,10 +94,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",

lasy/profiles/longitudinal/cosine_profile.py

@@ -1,4 +1,4 @@
-import numpy as np
+from lasy.backend import xp
 
 from .longitudinal_profile import LongitudinalProfile
 
@@ -62,10 +62,10 @@ def evaluate(self, t):
         tn = (t - self.t_peak) / self.tau_fwhm
 
         envelope = (
-            np.cos(0.5 * np.pi * tn)
+            xp.cos(0.5 * xp.pi * tn)
             * (tn > -1)
             * (tn < 1)
-            * np.exp(+1.0j * (self.cep_phase + self.omega0 * self.t_peak))
+            * xp.exp(+1.0j * (self.cep_phase + self.omega0 * self.t_peak))
         )
 
         return envelope

lasy/profiles/longitudinal/gaussian_profile.py

@@ -1,4 +1,4 @@
-import numpy as np
+from lasy.backend import xp
 
 from .longitudinal_profile import LongitudinalProfile
 
@@ -55,7 +55,7 @@ def evaluate(self, t):
             Contains the value of the longitudinal envelope at the
             specified points. This array has the same shape as the array t.
         """
-        envelope = np.exp(
+        envelope = xp.exp(
             -((t - self.t_peak) ** 2) / self.tau**2
             + 1.0j * (self.cep_phase + self.omega0 * self.t_peak)
         )