Only perform spectral-to-temporal and temporal-to-spectral transforms…

… when needed (#256)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

examples/example_gerchberg_saxton_algo.py

@@ -38,7 +38,9 @@
 
 # NOW ADD THE PHASE TO EACH SLICE OF THE FOCUS
 phase3D = np.repeat(phase[:, :, np.newaxis], npoints[2], axis=2)
-laser.grid.field = np.abs(laser.grid.field) * np.exp(1j * phase3D)
+laser.grid.set_temporal_field(
+    np.abs(laser.grid.get_temporal_field()) * np.exp(1j * phase3D)
+)
 
 # PROPAGATE THE FIELD FIELD FOWARDS AND BACKWARDS BY 1 MM
 propDist = 2e-3
@@ -72,17 +74,16 @@ def addColorbar(im, ax, label=None):
     cax.set_ylabel(label)
 
 
-im0 = ax[0, 0].imshow(
-    np.abs(laser.grid.field[:, :, tIndx]) ** 2, extent=extent, cmap="PuRd"
-)
+field = laser.grid.get_temporal_field()
+im0 = ax[0, 0].imshow(np.abs(field[:, :, tIndx]) ** 2, extent=extent, cmap="PuRd")
 addColorbar(im0, ax[0, 0], "Intensity (norm.)")
 ax[0, 0].set_title("Inten. z = 0.0 mm")
 ax[0, 0].set_xlabel("x ($\mu m$)")
 ax[0, 0].set_ylabel("y ($\mu m$)")
 
 
 im1 = ax[0, 1].imshow(
-    np.angle(laser.grid.field[:, :, tIndx]) * phaseMask, extent=extent, cmap="coolwarm"
+    np.angle(field[:, :, tIndx]) * phaseMask, extent=extent, cmap="coolwarm"
 )
 addColorbar(im1, ax[0, 1], "Phase (rad.)")
 ax[0, 1].set_title("Phase z = 0.0 mm")
@@ -91,13 +92,14 @@ def addColorbar(im, ax, label=None):
 
 
 laser_calc = copy.deepcopy(laserBackward)
-laser_calc.grid.field = np.abs(laser_calc.grid.field) * np.exp(1j * phaseBackward)
+laser_calc.grid.set_temporal_field(
+    np.abs(laser_calc.grid.get_temporal_field()) * np.exp(1j * phaseBackward)
+)
 laser_calc.propagate(propDist)
+
+field_calc = laser_calc.grid.get_temporal_field()
 im2 = ax[1, 0].imshow(
-    np.abs(
-        np.abs(laser.grid.field[:, :, tIndx]) ** 2
-        - np.abs(laser_calc.grid.field[:, :, tIndx]) ** 2
-    ),
+    np.abs(np.abs(field[:, :, tIndx]) ** 2 - np.abs(field_calc[:, :, tIndx]) ** 2),
     extent=extent,
     cmap="PuRd",
 )
@@ -106,9 +108,7 @@ def addColorbar(im, ax, label=None):
 ax[1, 0].set_ylabel("y ($\mu m$)")
 addColorbar(im2, ax[1, 0], "Intensity (norm.)")
 
-phaseResidual = np.angle(laser_calc.grid.field[:, :, tIndx]) - np.angle(
-    laser.grid.field[:, :, tIndx]
-)
+phaseResidual = np.angle(field_calc[:, :, tIndx]) - np.angle(field[:, :, tIndx])
 phaseResidual -= phaseResidual[int(npoints[1] / 2), int(npoints[0] / 2)]
 maxPhaseRes = np.max(np.abs(phaseResidual) * phaseMask)
 im3 = ax[1, 1].imshow(
@@ -123,26 +123,20 @@ def addColorbar(im, ax, label=None):
 ax[1, 1].set_ylabel("y ($\mu m$)")
 addColorbar(im3, ax[1, 1], "Phase (rad.)")
 
-
-im4 = ax[0, 2].imshow(
-    np.abs(laserBackward.grid.field[:, :, tIndx]) ** 2, extent=extent, cmap="PuRd"
-)
+field_bw = laserBackward.grid.get_temporal_field()
+im4 = ax[0, 2].imshow(np.abs(field_bw[:, :, tIndx]) ** 2, extent=extent, cmap="PuRd")
 addColorbar(im4, ax[0, 2], "Intensity (norm.)")
 ax[0, 2].set_title("Inten. z = %.1f mm" % (-propDist * 1e3))
 ax[0, 2].set_xlabel("x ($\mu m$)")
 ax[0, 2].set_ylabel("y ($\mu m$)")
 
-im5 = ax[0, 3].imshow(
-    np.angle(laserBackward.grid.field[:, :, tIndx]), extent=extent, cmap="coolwarm"
-)
+im5 = ax[0, 3].imshow(np.angle(field_bw[:, :, tIndx]), extent=extent, cmap="coolwarm")
 addColorbar(im5, ax[0, 3], "Phase (rad.)")
 ax[0, 3].set_title("Phase z = %.1f mm" % (-propDist * 1e3))
 ax[0, 3].set_xlabel("x ($\mu m$)")
 ax[0, 3].set_ylabel("y ($\mu m$)")
 
-phaseResidual = (
-    np.angle(laserBackward.grid.field[:, :, tIndx]) - phaseBackward[:, :, tIndx]
-)
+phaseResidual = np.angle(field_bw[:, :, tIndx]) - phaseBackward[:, :, tIndx]
 phaseResidual -= phaseResidual[int(npoints[1] / 2), int(npoints[0] / 2)]
 maxPhaseRes = np.max(np.abs(phaseResidual) * phaseMask)
 im6 = ax[0, 4].imshow(
@@ -157,24 +151,19 @@ def addColorbar(im, ax, label=None):
 ax[0, 4].set_xlabel("x ($\mu m$)")
 ax[0, 4].set_ylabel("y ($\mu m$)")
 
-im7 = ax[1, 2].imshow(
-    np.abs(laserForward.grid.field[:, :, tIndx]) ** 2, extent=extent, cmap="PuRd"
-)
+field_fw = laserForward.grid.get_temporal_field()
+im7 = ax[1, 2].imshow(np.abs(field_fw[:, :, tIndx]) ** 2, extent=extent, cmap="PuRd")
 addColorbar(im7, ax[1, 2], "Intensity (norm.)")
 ax[1, 2].set_title("Inten. z = %.1f mm" % (propDist * 1e3))
 ax[1, 2].set_xlabel("x ($\mu m$)")
 ax[1, 2].set_ylabel("y ($\mu m$)")
-im8 = ax[1, 3].imshow(
-    np.angle(laserForward.grid.field[:, :, tIndx]), extent=extent, cmap="coolwarm"
-)
+im8 = ax[1, 3].imshow(np.angle(field_fw[:, :, tIndx]), extent=extent, cmap="coolwarm")
 addColorbar(im8, ax[1, 3], "Phase (rad.)")
 ax[1, 3].set_title("Phase z = %.1f mm" % (propDist * 1e3))
 ax[1, 3].set_xlabel("x ($\mu m$)")
 ax[1, 3].set_ylabel("y ($\mu m$)")
 
-phaseResidual = (
-    np.angle(laserForward.grid.field[:, :, tIndx]) - phaseForward[:, :, tIndx]
-)
+phaseResidual = np.angle(field_fw[:, :, tIndx]) - phaseForward[:, :, tIndx]
 phaseResidual -= phaseResidual[int(npoints[1] / 2), int(npoints[0] / 2)]
 maxPhaseRes = np.max(np.abs(phaseResidual) * phaseMask)
 im9 = ax[1, 4].imshow(

lasy/laser.py

@@ -2,7 +2,7 @@
 from axiprop.lib import PropagatorFFT2, PropagatorResampling
 from scipy.constants import c
 
-from lasy.utils.grid import Grid
+from lasy.utils.grid import Grid, time_axis_indx
 from lasy.utils.laser_utils import (
     normalize_energy,
     normalize_peak_field_amplitude,
@@ -113,7 +113,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")
-            self.grid.field[...] = profile.evaluate(x, y, t)
+            self.grid.set_temporal_field(profile.evaluate(x, y, t))
         elif self.dim == "rt":
             if n_theta_evals is None:
                 # Generate 2*n_azimuthal_modes - 1 evenly-spaced values of
@@ -129,11 +129,12 @@ def __init__(
             envelope = profile.evaluate(x, y, t)
             # Perform the azimuthal decomposition
             azimuthal_modes = np.fft.ifft(envelope, axis=0)
-            self.grid.field[:n_azimuthal_modes] = azimuthal_modes[:n_azimuthal_modes]
+            field = azimuthal_modes[:n_azimuthal_modes]
             if n_azimuthal_modes > 1:
-                self.grid.field[-n_azimuthal_modes + 1 :] = azimuthal_modes[
-                    -n_azimuthal_modes + 1 :
-                ]
+                field = np.concatenate(
+                    (field, azimuthal_modes[-n_azimuthal_modes + 1 :])
+                )
+            self.grid.set_temporal_field(field)
 
         # For profiles that define the energy, normalize the amplitude
         if hasattr(profile, "laser_energy"):
@@ -170,17 +171,14 @@ def apply_optics(self, optical_element):
             Represents a thin optical element, through which the laser
             propagates.
         """
-        # Transform the field from temporal to frequency domain
-        time_axis_indx = -1
-        field_fft = np.fft.ifft(self.grid.field, axis=time_axis_indx, norm="backward")
-
         # Create the frequency axis
         dt = self.grid.dx[time_axis_indx]
         omega0 = self.profile.omega0
-        Nt = self.grid.field.shape[time_axis_indx]
+        Nt = self.grid.shape[time_axis_indx]
         omega_1d = 2 * np.pi * np.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")
             # The line below assumes that amplitude_multiplier
@@ -193,17 +191,13 @@ def apply_optics(self, optical_element):
             # FT_theta[ multiplier * field ] = multiplier * FT_theta[ field ]
             # Thus, we can simply multiply each azimuthal mode by the multiplier.
             for i_m in range(self.grid.azimuthal_modes.size):
-                field_fft[i_m, :, :] *= multiplier
+                spectral_field[i_m, :, :] *= multiplier
         else:
             x, y, omega = np.meshgrid(
                 self.grid.axes[0], self.grid.axes[1], omega_1d, indexing="ij"
             )
-            field_fft *= optical_element.amplitude_multiplier(x, y, omega)
-
-        # Transform field from frequency to temporal domain
-        self.grid.field[:, :, :] = np.fft.fft(
-            field_fft, axis=time_axis_indx, norm="backward"
-        )
+            spectral_field *= optical_element.amplitude_multiplier(x, y, omega)
+        self.grid.set_spectral_field(spectral_field)
 
     def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True):
         """
@@ -224,44 +218,28 @@ def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True
         show_progress : bool (optional)
             Whether to show a progress bar when performing the computation
         """
-        time_axis_indx = -1
-
         # 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_shape[-1] = 0.0
+            field = self.grid.get_temporal_field()
             if self.dim == "rt":
-                self.grid.field[:, -nr_boundary:, :] *= absorb_layer_shape[
-                    None, :, None
-                ]
+                field[:, -nr_boundary:, :] *= absorb_layer_shape[None, :, None]
             else:
-                self.grid.field[-nr_boundary:, :, :] *= absorb_layer_shape[
-                    :, None, None
-                ]
-                self.grid.field[:nr_boundary, :, :] *= absorb_layer_shape[::-1][
-                    :, None, None
-                ]
-                self.grid.field[:, -nr_boundary:, :] *= absorb_layer_shape[
-                    None, :, None
-                ]
-                self.grid.field[:, :nr_boundary, :] *= absorb_layer_shape[::-1][
-                    None, :, None
-                ]
-
-        # Transform the field from temporal to frequency domain
-        field_fft = np.fft.ifft(self.grid.field, axis=time_axis_indx, norm="backward")
+                field[-nr_boundary:, :, :] *= absorb_layer_shape[:, None, None]
+                field[:nr_boundary, :, :] *= absorb_layer_shape[::-1][:, None, None]
+                field[:, -nr_boundary:, :] *= absorb_layer_shape[None, :, None]
+                field[:, :nr_boundary, :] *= absorb_layer_shape[::-1][None, :, None]
+            self.grid.set_temporal_field(field)
 
         # Create the frequency axis
         dt = self.grid.dx[time_axis_indx]
         omega0 = self.profile.omega0
-        Nt = self.grid.field.shape[time_axis_indx]
+        Nt = self.grid.shape[time_axis_indx]
         omega = 2 * np.pi * np.fft.fftfreq(Nt, dt) + omega0
 
-        # make 3D shape for the frequency axis
-        omega_shape = (1, 1, self.grid.field.shape[time_axis_indx])
-
         if self.dim == "rt":
             # Construct the propagator (check if exists)
             if not hasattr(self, "prop"):
@@ -278,19 +256,21 @@ 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(field_fft[i_m]).copy()
+                transform_data = np.transpose(spectral_field[i_m]).copy()
                 self.prop[i_m].step(
                     transform_data,
                     distance,
                     overwrite=True,
                     show_progress=show_progress,
                 )
-                field_fft[i_m, :, :] = np.transpose(transform_data).copy()
+                spectral_field[i_m, :, :] = np.transpose(transform_data).copy()
+            self.grid.set_spectral_field(spectral_field)
         else:
             # Construct the propagator (check if exists)
             if not hasattr(self, "prop"):
-                Nx, Ny, Nt = self.grid.field.shape
+                Nx, Ny, Nt = self.grid.shape
                 Lx = self.grid.hi[0] - self.grid.lo[0]
                 Ly = self.grid.hi[1] - self.grid.lo[1]
                 spatial_axes = ((Lx, Nx), (Ly, Ny))
@@ -301,31 +281,29 @@ def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True
                     verbose=False,
                 )
             # Propagate the spectral image
-            transform_data = np.moveaxis(field_fft, -1, 0).copy()
+            spectral_field = self.grid.get_spectral_field()
+            transform_data = np.moveaxis(spectral_field, -1, 0).copy()
             self.prop.step(
                 transform_data, distance, overwrite=True, show_progress=show_progress
             )
-            field_fft[:, :, :] = np.moveaxis(transform_data, 0, -1).copy()
+            spectral_field = np.moveaxis(transform_data, 0, -1).copy()
 
         # Choose the time translation assuming propagation at v=c
         translate_time = distance / c
 
+        # This translation (e.g. delay in time, compared to t=0, associated
+        # with the propagation) is not automatically handled by the above
+        # 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)
+        self.grid.set_spectral_field(spectral_field)
+
         # Translate the domain
         self.grid.lo[time_axis_indx] += translate_time
         self.grid.hi[time_axis_indx] += translate_time
         self.grid.axes[time_axis_indx] += translate_time
 
-        # Translate the phase of spectral image
-        field_fft *= np.exp(-1j * translate_time * omega.reshape(omega_shape))
-
-        # Transform field from frequency to temporal domain
-        self.grid.field[:, :, :] = np.fft.fft(
-            field_fft, axis=time_axis_indx, norm="backward"
-        )
-
-        # Translate phase of the retrieved envelope by the distance
-        self.grid.field *= np.exp(1j * omega0 * distance / c)
-
     def write_to_file(
         self, file_prefix="laser", file_format="h5", save_as_vector_potential=False
     ):
@@ -364,11 +342,12 @@ def show(self, **kw):
         ----------
         **kw: additional arguments to be passed to matplotlib's imshow command
         """
+        temporal_field = self.grid.get_temporal_field()
         if self.dim == "rt":
             # Show field in the plane y=0, above and below axis, with proper sign for each mode
             E = [
                 np.concatenate(
-                    ((-1.0) ** m * self.grid.field[m, ::-1], self.grid.field[m])
+                    ((-1.0) ** m * temporal_field[m, ::-1], temporal_field[m])
                 )
                 for m in self.grid.azimuthal_modes
             ]
@@ -382,8 +361,8 @@ def show(self, **kw):
 
         else:
             # In 3D show an image in the xt plane
-            i_slice = int(self.grid.field.shape[1] // 2)
-            E = self.grid.field[:, i_slice, :]
+            i_slice = int(temporal_field.shape[1] // 2)
+            E = temporal_field[:, i_slice, :]
             extent = [
                 self.grid.lo[-1],
                 self.grid.hi[-1],

lasy/profiles/from_openpmd_profile.py

@@ -102,7 +102,7 @@ def __init__(
         if not is_envelope:
             grid = create_grid(F, axes, dim)
             grid, omg0 = field_to_envelope(grid, dim, phase_unwrap_nd)
-            array = grid.field[0]
+            array = grid.get_temporal_field()[0]
         else:
             s = io.Series(path + "/" + prefix + "_%T.h5", io.Access.read_only)
             it = s.iterations[iteration]