From 350059b7b55520f98b2824c6b190618b42a8a870 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Tue, 16 Jul 2024 14:34:09 +0200
Subject: [PATCH] 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 |  53 ++++------
 lasy/laser.py                             | 101 ++++++++-----------
 lasy/profiles/from_openpmd_profile.py     |   2 +-
 lasy/utils/grid.py                        | 112 +++++++++++++++++++++-
 lasy/utils/laser_utils.py                 |  97 +++++++++++--------
 lasy/utils/openpmd_output.py              |   2 +-
 lasy/utils/phase_retrieval.py             |  15 +--
 tests/test_gaussian_propagator.py         |   7 +-
 tests/test_gsalgo.py                      |   3 +-
 tests/test_parabolic_mirror.py            |   7 +-
 tests/test_speckles.py                    |   8 +-
 tests/test_t2z2t.py                       |   7 +-
 12 files changed, 252 insertions(+), 162 deletions(-)

diff --git a/examples/example_gerchberg_saxton_algo.py b/examples/example_gerchberg_saxton_algo.py
index d1724103..a4e21430 100644
--- a/examples/example_gerchberg_saxton_algo.py
+++ b/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,9 +74,8 @@ 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$)")
@@ -82,7 +83,7 @@ def addColorbar(im, ax, label=None):
 
 
 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(
diff --git a/lasy/laser.py b/lasy/laser.py
index fd03dee5..93fe456f 100644
--- a/lasy/laser.py
+++ b/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],
diff --git a/lasy/profiles/from_openpmd_profile.py b/lasy/profiles/from_openpmd_profile.py
index 01d6cbf3..a44b39b3 100644
--- a/lasy/profiles/from_openpmd_profile.py
+++ b/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]
diff --git a/lasy/utils/grid.py b/lasy/utils/grid.py
index 1d6719dd..557fb50a 100644
--- a/lasy/utils/grid.py
+++ b/lasy/utils/grid.py
@@ -1,9 +1,11 @@
 import numpy as np
 
+time_axis_indx = -1
+
 
 class Grid:
     """
-    Store an array (typically the envelope) and corresponding metadata.
+    Store the envelope in temporal and spectral space and corresponding metadata.
 
     Parameters
     ----------
@@ -53,11 +55,111 @@ def __init__(self, dim, lo, hi, npoints, n_azimuthal_modes=None):
 
         # Data
         if dim == "xyt":
-            self.field = np.zeros(self.npoints, dtype="complex128")
+            self.shape = self.npoints
         elif dim == "rt":
             # Azimuthal modes are arranged in the following order:
             # 0, 1, 2, ..., n_azimuthal_modes-1, -n_azimuthal_modes+1, ..., -1
             ncomp = 2 * self.n_azimuthal_modes - 1
-            self.field = np.zeros(
-                (ncomp, self.npoints[0], self.npoints[1]), dtype="complex128"
-            )
+            self.shape = (ncomp, self.npoints[0], self.npoints[1])
+        self.temporal_field = np.zeros(self.shape, dtype="complex128")
+        self.temporal_field_valid = False
+        self.spectral_field = np.zeros(self.shape, dtype="complex128")
+        self.spectral_field_valid = False
+
+    def set_temporal_field(self, field):
+        """
+        Set the temporal field.
+
+        Parameters
+        ----------
+        field : ndarray of complexs
+            The temporal field.
+        """
+        assert field.shape == self.temporal_field.shape
+        assert field.dtype == "complex128"
+        self.temporal_field[:, :, :] = field
+        self.temporal_field_valid = True
+        self.spectral_field_valid = False  # Invalidates the spectral field
+
+    def set_spectral_field(self, field):
+        """
+        Set the spectral field.
+
+        Parameters
+        ----------
+        field : ndarray of complexs
+            The spectral field.
+        """
+        assert field.shape == self.spectral_field.shape
+        assert field.dtype == "complex128"
+        self.spectral_field[:, :, :] = field
+        self.spectral_field_valid = True
+        self.temporal_field_valid = False  # Invalidates the temporal field
+
+    def get_temporal_field(self):
+        """
+        Return a copy of the temporal field.
+
+        (Modifying the returned object will not modify the original field stored
+        in the Grid object ; one must use set_temporal_field to do so.)
+
+        Returns
+        -------
+        field : ndarray of complexs
+            The temporal field.
+        """
+        # We return a copy, so that the user cannot modify
+        # the original field, unless get_temporal_field is called
+        if self.temporal_field_valid:
+            return self.temporal_field.copy()
+        elif self.spectral_field_valid:
+            self.spectral2temporal_fft()
+            return self.temporal_field.copy()
+        else:
+            raise ValueError("Both temporal and spectral fields are invalid")
+
+    def get_spectral_field(self):
+        """
+        Return a copy of the spectral field.
+
+        (Modifying the returned object will not modify the original field stored
+        in the Grid object ; one must use set_spectral_field to do so.)
+
+        Returns
+        -------
+        field : ndarray of complexs
+            The spectral field.
+        """
+        # We return a copy, so that the user cannot modify
+        # the original field, unless set_spectral_field is called
+        if self.spectral_field_valid:
+            return self.spectral_field.copy()
+        elif self.temporal_field_valid:
+            self.temporal2spectral_fft()
+            return self.spectral_field.copy()
+        else:
+            raise ValueError("Both temporal and spectral fields are invalid")
+
+    def temporal2spectral_fft(self):
+        """
+        Perform the Fourier transform of field from temporal to spectral space.
+
+        (Only along the time axis, not along the transverse spatial coordinates.)
+        """
+        assert self.temporal_field_valid
+        self.spectral_field = np.fft.ifft(
+            self.temporal_field, axis=time_axis_indx, norm="backward"
+        )
+        self.spectral_field_valid = True
+
+    def spectral2temporal_fft(self):
+        """
+        Perform the Fourier transform of field from spectral to temporal space.
+
+        (Only along the time axis, not along the transverse spatial coordinates.)
+        """
+        assert self.spectral_field_valid
+        self.temporal_field = np.fft.fft(
+            self.spectral_field, axis=time_axis_indx, norm="backward"
+        )
+        self.temporal_field_valid = True
diff --git a/lasy/utils/laser_utils.py b/lasy/utils/laser_utils.py
index 19625e96..76e8fef3 100644
--- a/lasy/utils/laser_utils.py
+++ b/lasy/utils/laser_utils.py
@@ -39,7 +39,7 @@ def compute_laser_energy(dim, grid):
     # specified laser wavelength.
     # This probably needs to be generalized for few-cycle laser pulses.
 
-    envelope = grid.field
+    envelope = grid.get_temporal_field()
 
     dV = get_grid_cell_volume(grid, dim)
 
@@ -84,7 +84,9 @@ def normalize_energy(dim, energy, grid):
         print("Field is zero everywhere, normalization will be skipped")
     else:
         norm_factor = (energy / current_energy) ** 0.5
-        grid.field *= norm_factor
+        field = grid.get_temporal_field()
+        field *= norm_factor
+        grid.set_temporal_field(field)
 
 
 def normalize_peak_field_amplitude(amplitude, grid):
@@ -100,10 +102,13 @@ def normalize_peak_field_amplitude(amplitude, grid):
         Contains value of the laser envelope and metadata.
     """
     if amplitude is not None:
-        if np.abs(grid.field).max() == 0.0:
+        field = grid.get_temporal_field()
+        field_max = np.abs(field).max()
+        if field_max == 0.0:
             print("Field is zero everywhere, normalization will be skipped")
         else:
-            grid.field *= amplitude / np.abs(grid.field).max()
+            field *= amplitude / field_max
+            grid.set_temporal_field(field)
 
 
 def normalize_peak_intensity(peak_intensity, grid):
@@ -119,12 +124,13 @@ def normalize_peak_intensity(peak_intensity, grid):
         Contains value of the laser envelope and metadata.
     """
     if peak_intensity is not None:
-        intensity = np.abs(epsilon_0 * grid.field**2 / 2 * c)
+        field = grid.get_temporal_field()
+        intensity = np.abs(epsilon_0 * field**2 / 2 * c)
         input_peak_intensity = intensity.max()
         if input_peak_intensity == 0.0:
             print("Field is zero everywhere, normalization will be skipped")
         else:
-            grid.field *= np.sqrt(peak_intensity / input_peak_intensity)
+            grid.set_temporal_field(np.sqrt(peak_intensity / input_peak_intensity))
 
 
 def get_full_field(laser, theta=0, slice=0, slice_axis="x", Nt=None):
@@ -151,7 +157,7 @@ def get_full_field(laser, theta=0, slice=0, slice_axis="x", Nt=None):
             Physical extent of the reconstructed field.
     """
     omega0 = laser.profile.omega0
-    env = laser.grid.field.copy()
+    env = laser.grid.get_temporal_field()
     time_axis = laser.grid.axes[-1]
 
     if laser.dim == "rt":
@@ -295,18 +301,17 @@ def get_spectrum(
         Array with the angular frequencies of the spectrum.
     """
     # Get the frequencies of the fft output.
-    freq = np.fft.fftfreq(grid.field.shape[-1], d=(grid.axes[-1][1] - grid.axes[-1][0]))
+    freq = np.fft.fftfreq(grid.shape[-1], d=(grid.axes[-1][1] - grid.axes[-1][0]))
     omega = 2 * np.pi * freq
 
     # Get on axis or full field.
+    field = grid.get_temporal_field()
     if method == "on_axis":
         if dim == "xyt":
-            nx, ny, nt = grid.field.shape
-            field = grid.field[nx // 2, ny // 2]
+            nx, ny, _ = field.shape
+            field = field[nx // 2, ny // 2]
         else:
-            field = grid.field[0, 0]
-    else:
-        field = grid.field
+            field = field[0, 0]
 
     # Get spectrum.
     if is_envelope:
@@ -421,19 +426,21 @@ def get_frequency(
         Central angular frequency (averaged omega, weighted by the local
         envelope amplitude).
     """
+    field = grid.get_temporal_field()
+
     # Assumes t is last dimension!
     if is_envelope:
         assert omega0 is not None
-        phase = np.unwrap(np.angle(grid.field))
+        phase = np.unwrap(np.angle(field))
         omega = omega0 + np.gradient(-phase, grid.axes[-1], axis=-1, edge_order=2)
-        central_omega = np.average(omega, weights=np.abs(grid.field))
+        central_omega = np.average(omega, weights=np.abs(field))
     else:
         assert dim in ["xyt", "rt"]
         if dim == "xyt" and phase_unwrap_nd:
             print("WARNING: using 3D phase unwrapping, this can be expensive")
 
-        h = grid.field if is_hilbert else hilbert_transform(grid)
-        h = np.squeeze(grid.field)
+        h = field if is_hilbert else hilbert_transform(grid)
+        h = np.squeeze(field)
         if phase_unwrap_nd:
             try:
                 from skimage.restoration import unwrap_phase
@@ -482,10 +489,11 @@ def get_duration(grid, dim):
     """
     # Calculate weights of each grid cell (amplitude of the field).
     dV = get_grid_cell_volume(grid, dim)
+    field = grid.get_temporal_field()
     if dim == "xyt":
-        weights = np.abs(grid.field) ** 2 * dV
+        weights = np.abs(field) ** 2 * dV
     else:  # dim == "rt":
-        weights = np.abs(grid.field) ** 2 * dV[np.newaxis, :, np.newaxis]
+        weights = np.abs(field) ** 2 * dV[np.newaxis, :, np.newaxis]
     # project weights to longitudinal axes
     weights = np.sum(weights, axis=(0, 1))
     return weighted_std(grid.axes[-1], weights)
@@ -513,7 +521,7 @@ def field_to_vector_potential(grid, omega0):
     # term in: E = -dA/dt + 1j * omega0 * A where E and A are the field and
     # vector potential envelopes, respectively
     omega, _ = get_frequency(grid, is_envelope=True, omega0=omega0)
-    return -1j * e * grid.field / (m_e * omega * c)
+    return -1j * e * grid.get_temporal_field() / (m_e * omega * c)
 
 
 def vector_potential_to_field(grid, omega0, direct=True):
@@ -538,15 +546,16 @@ def vector_potential_to_field(grid, omega0, direct=True):
     -------
     Envelope of the electric field (V/m).
     """
+    field = grid.get_temporal_field()
     if direct:
         A = (
-            -np.gradient(grid.field, grid.axes[-1], axis=-1, edge_order=2)
-            + 1j * omega0 * grid.field
+            -np.gradient(field, grid.axes[-1], axis=-1, edge_order=2)
+            + 1j * omega0 * field
         )
         return m_e * c / e * A
     else:
         omega, _ = get_frequency(grid, is_envelope=True, omega0=omega0)
-        return 1j * m_e * omega * c * grid.field / e
+        return 1j * m_e * omega * c * field / e
 
 
 def field_to_envelope(grid, dim, phase_unwrap_nd=False):
@@ -571,8 +580,10 @@ def field_to_envelope(grid, dim, phase_unwrap_nd=False):
     tuple
         A tuple with the envelope array and the central wavelength.
     """
+    field = grid.get_temporal_field()
+
     # hilbert transform needs inverted time axis.
-    grid.field = hilbert_transform(grid)
+    field = hilbert_transform(field)
 
     # Get central wavelength from array
     omg_h, omg0_h = get_frequency(
@@ -582,12 +593,13 @@ def field_to_envelope(grid, dim, phase_unwrap_nd=False):
         is_hilbert=True,
         phase_unwrap_nd=phase_unwrap_nd,
     )
-    grid.field *= np.exp(1j * omg0_h * grid.axes[-1])
+    field *= np.exp(1j * omg0_h * grid.axes[-1])
+    grid.set_temporal_field(field)
 
     return grid, omg0_h
 
 
-def hilbert_transform(grid):
+def hilbert_transform(field):
     """Make a hilbert transform of the grid field.
 
     Currently the arrays need to be flipped along t (both the input field and
@@ -599,7 +611,7 @@ def hilbert_transform(grid):
     grid : Grid
         The lasy grid whose field should be transformed.
     """
-    return hilbert(grid.field[:, :, ::-1])[:, :, ::-1]
+    return hilbert(field[:, :, ::-1])[:, :, ::-1]
 
 
 def get_grid_cell_volume(grid, dim):
@@ -675,8 +687,7 @@ def create_grid(array, axes, dim):
         assert np.all(grid.axes[0] == axes["x"])
         assert np.all(grid.axes[1] == axes["y"])
         assert np.allclose(grid.axes[2], axes["t"], rtol=1.0e-14)
-        assert grid.field.shape == array.shape
-        grid.field = array
+        grid.set_temporal_field(array)
     else:  # dim == "rt":
         lo = (axes["r"][0], axes["t"][0])
         hi = (axes["r"][-1], axes["t"][-1])
@@ -684,8 +695,7 @@ def create_grid(array, axes, dim):
         grid = Grid(dim, lo, hi, npoints, n_azimuthal_modes=1)
         assert np.all(grid.axes[0] == axes["r"])
         assert np.allclose(grid.axes[1], axes["t"], rtol=1.0e-14)
-        assert grid.field.shape == array[np.newaxis].shape
-        grid.field = array[np.newaxis]
+        grid.set_temporal_field(array)
     return grid
 
 
@@ -732,6 +742,8 @@ def export_to_z(dim, grid, omega0, z_axis=None, z0=0.0, t0=0.0, backend="NP"):
 
     FieldAxprp = ScalarFieldEnvelope(omega0 / c, t_axis)
 
+    field = grid.get_temporal_field()
+
     if dim == "rt":
         # Construct the propagator
         prop = []
@@ -747,20 +759,20 @@ def export_to_z(dim, grid, omega0, z_axis=None, z0=0.0, t0=0.0, backend="NP"):
             )
 
         field_z = np.zeros(
-            (grid.field.shape[0], grid.field.shape[1], z_axis.size),
-            dtype=grid.field.dtype,
+            (field.shape[0], field.shape[1], z_axis.size),
+            dtype=field.dtype,
         )
 
         # Convert the spectral image to the spatial field representation
         for i_m in range(grid.azimuthal_modes.size):
-            FieldAxprp.import_field(np.transpose(grid.field[i_m]).copy())
+            FieldAxprp.import_field(np.transpose(field[i_m]).copy())
 
             field_z[i_m] = prop[i_m].t2z(FieldAxprp.Field_ft, z_axis, z0=z0, t0=t0).T
 
             field_z[i_m] *= np.exp(-1j * (z_axis / c + t0) * omega0)
     else:
         # Construct the propagator
-        Nx, Ny, Nt = grid.field.shape
+        Nx, Ny, Nt = field.shape
         Lx = grid.hi[0] - grid.lo[0]
         Ly = grid.hi[1] - grid.lo[1]
         prop = PropagatorFFT2(
@@ -771,7 +783,7 @@ def export_to_z(dim, grid, omega0, z_axis=None, z0=0.0, t0=0.0, backend="NP"):
             verbose=False,
         )
         # Convert the spectral image to the spatial field representation
-        FieldAxprp.import_field(np.moveaxis(grid.field, -1, 0).copy())
+        FieldAxprp.import_field(np.moveaxis(field, -1, 0).copy())
         field_z = prop.t2z(FieldAxprp.Field_ft, z_axis, z0=z0, t0=t0)
         field_z = np.moveaxis(field_z, 0, -1)
         field_z *= np.exp(-1j * (z_axis / c + t0) * omega0)
@@ -840,14 +852,16 @@ def import_from_z(dim, grid, omega0, field_z, z_axis, z0=0.0, t0=0.0, backend="N
             )
 
         # Convert the spectral image to the spatial field representation
+        field = np.zeros(grid.shape, dtype=np.complex128)
         for i_m in range(grid.azimuthal_modes.size):
             transform_data = np.transpose(field_fft[i_m]).copy()
             transform_data *= np.exp(-1j * z_axis[0] * (k_z[:, None] - omega0 / c))
-            grid.field[i_m] = prop[i_m].z2t(transform_data, t_axis, z0=z0, t0=t0).T
-            grid.field[i_m] *= np.exp(1j * (z0 / c + t_axis) * omega0)
+            field[i_m] = prop[i_m].z2t(transform_data, t_axis, z0=z0, t0=t0).T
+            field[i_m] *= np.exp(1j * (z0 / c + t_axis) * omega0)
+        grid.set_temporal_field(field)
     else:
         # Construct the propagator
-        Nx, Ny, Nt = grid.field.shape
+        Nx, Ny, _ = grid.npoints
         Lx = grid.hi[0] - grid.lo[0]
         Ly = grid.hi[1] - grid.lo[1]
         prop = PropagatorFFT2(
@@ -860,5 +874,6 @@ def import_from_z(dim, grid, omega0, field_z, z_axis, z0=0.0, t0=0.0, backend="N
         # Convert the spectral image to the spatial field representation
         transform_data = np.moveaxis(field_fft, -1, 0).copy()
         transform_data *= np.exp(-1j * z_axis[0] * (k_z[:, None, None] - omega0 / c))
-        grid.field = np.moveaxis(prop.z2t(transform_data, t_axis, z0=z0, t0=t0), 0, -1)
-        grid.field *= np.exp(1j * (z0 / c + t_axis) * omega0)
+        field = np.moveaxis(prop.z2t(transform_data, t_axis, z0=z0, t0=t0), 0, -1)
+        field *= np.exp(1j * (z0 / c + t_axis) * omega0)
+        grid.set_temporal_field(field)
diff --git a/lasy/utils/openpmd_output.py b/lasy/utils/openpmd_output.py
index 485b7ba9..b0b72814 100644
--- a/lasy/utils/openpmd_output.py
+++ b/lasy/utils/openpmd_output.py
@@ -53,7 +53,7 @@ def write_to_openpmd_file(
         Whether the envelope is converted to normalized vector potential
         before writing to file.
     """
-    array = grid.field
+    array = grid.get_temporal_field()
 
     # Create file
     series = io.Series("{}_%05T.{}".format(file_prefix, file_format), io.Access.create)
diff --git a/lasy/utils/phase_retrieval.py b/lasy/utils/phase_retrieval.py
index 1aa79331..4d9f915d 100644
--- a/lasy/utils/phase_retrieval.py
+++ b/lasy/utils/phase_retrieval.py
@@ -51,9 +51,9 @@ def gerchberg_saxton_algo(
     """
     laser1 = copy.deepcopy(laserPos1)
     laser2 = copy.deepcopy(laserPos2)
-    amp1 = np.abs(laser1.grid.field)
+    amp1 = np.abs(laser1.grid.get_temporal_field())
     amp1_summed = np.sum(amp1)
-    amp2 = np.abs(laser2.grid.field)
+    amp2 = np.abs(laser2.grid.get_temporal_field())
     phase1 = np.zeros_like(amp1)
 
     if condition == "max_iterations":
@@ -65,15 +65,16 @@ def gerchberg_saxton_algo(
 
     i = 0
     while breakout(cond):
-        laser1.grid.field = amp1 * np.exp(1j * phase1)
+        laser1.grid.set_temporal_field(amp1 * np.exp(1j * phase1))
         laser1.propagate(dz, show_progress=False)
 
-        phase2 = np.angle(laser1.grid.field)
-        laser2.grid.field = amp2 * np.exp(1j * phase2)
+        phase2 = np.angle(laser1.grid.get_temporal_field())
+        laser2.grid.set_temporal_field(amp2 * np.exp(1j * phase2))
         laser2.propagate(-dz, show_progress=False)
 
-        phase1 = np.angle(laser2.grid.field)
-        amp1_calc = np.abs(laser2.grid.field)
+        field2 = laser2.grid.get_temporal_field()
+        phase1 = np.angle(field2)
+        amp1_calc = np.abs(field2)
         amp_error_summed = np.sum(np.abs(amp1_calc) - amp1)
         if debug:
             i += 1
diff --git a/tests/test_gaussian_propagator.py b/tests/test_gaussian_propagator.py
index 81b2b980..e669eff6 100644
--- a/tests/test_gaussian_propagator.py
+++ b/tests/test_gaussian_propagator.py
@@ -23,12 +23,13 @@ def gaussian():
 
 def get_w0(laser):
     # Calculate the laser waist
+    field = laser.grid.get_temporal_field()
     if laser.dim == "xyt":
-        Nx, Ny, Nt = laser.grid.field.shape
-        A2 = (np.abs(laser.grid.field[Nx // 2 - 1, :, :]) ** 2).sum(-1)
+        Nx, Ny, Nt = field.shape
+        A2 = (np.abs(field[Nx // 2 - 1, :, :]) ** 2).sum(-1)
         ax = laser.grid.axes[1]
     else:
-        A2 = (np.abs(laser.grid.field[0, :, :]) ** 2).sum(-1)
+        A2 = (np.abs(field[0, :, :]) ** 2).sum(-1)
         ax = laser.grid.axes[0]
         if ax[0] > 0:
             A2 = np.r_[A2[::-1], A2]
diff --git a/tests/test_gsalgo.py b/tests/test_gsalgo.py
index 7907b6fa..bd913800 100644
--- a/tests/test_gsalgo.py
+++ b/tests/test_gsalgo.py
@@ -49,7 +49,8 @@ def test_3D_case(gaussian):
 
     # 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)
+    field = laser.grid.get_temporal_field()
+    laser.grid.set_temporal_field(np.abs(field) * np.exp(1j * phase3D))
 
     # PROPAGATE THE FIELD FIELD FOWARDS AND BACKWARDS BY 1 MM
     propDist = 2e-3
diff --git a/tests/test_parabolic_mirror.py b/tests/test_parabolic_mirror.py
index 123ace07..1c60a4b1 100644
--- a/tests/test_parabolic_mirror.py
+++ b/tests/test_parabolic_mirror.py
@@ -26,12 +26,13 @@
 
 def get_w0(laser):
     # Calculate the laser waist
+    field = laser.grid.get_temporal_field()
     if laser.dim == "xyt":
-        Nx, Ny, Nt = laser.grid.field.shape
-        A2 = (np.abs(laser.grid.field[Nx // 2 - 1, :, :]) ** 2).sum(-1)
+        Nx = field.shape[0]
+        A2 = (np.abs(field[Nx // 2 - 1, :, :]) ** 2).sum(-1)
         ax = laser.grid.axes[1]
     else:
-        A2 = (np.abs(laser.grid.field[0, :, :]) ** 2).sum(-1)
+        A2 = (np.abs(field[0, :, :]) ** 2).sum(-1)
         ax = laser.grid.axes[0]
         if ax[0] > 0:
             A2 = np.r_[A2[::-1], A2]
diff --git a/tests/test_speckles.py b/tests/test_speckles.py
index 22eb3cd5..f7e7c5bc 100644
--- a/tests/test_speckles.py
+++ b/tests/test_speckles.py
@@ -56,7 +56,7 @@ def test_intensity_distribution(temporal_smoothing_type):
 
     laser = Laser(dimensions, lo, hi, num_points, profile)
 
-    F = laser.grid.field
+    F = laser.grid.get_temporal_field()
 
     # get spatial statistics
     # <real env> = 0 = <imag env> = <er * ei>
@@ -126,7 +126,7 @@ def test_spatial_correlation(temporal_smoothing_type):
     num_points = (200, 200, 300)
 
     laser = Laser(dimensions, lo, hi, num_points, profile)
-    F = laser.grid.field
+    F = laser.grid.get_temporal_field()
 
     # compare speckle profile / autocorrelation
     # compute autocorrelation using Wiener-Khinchin Theorem
@@ -209,7 +209,7 @@ def test_sinc_zeros(temporal_smoothing_type):
     num_points = (300, 300, 10)
 
     laser = Laser(dimensions, lo, hi, num_points, profile)
-    F = laser.grid.field
+    F = laser.grid.get_temporal_field()
 
     assert abs(F[0, :, :]).max() / abs(F).max() < 1.0e-8
     assert abs(F[-1, :, :]).max() / abs(F).max() < 1.0e-8
@@ -270,6 +270,6 @@ def test_FM_SSD_periodicity():
     num_points = (160, 200, 400)  # Number of points in each dimension
 
     laser = Laser(dimensions, lo, hi, num_points, laser_profile)
-    F = laser.grid.field
+    F = laser.grid.get_temporal_field()
     period_error = abs(F[:, :, 0] - F[:, :, -1]).max() / abs(F).max()
     assert period_error < 1.0e-8
diff --git a/tests/test_t2z2t.py b/tests/test_t2z2t.py
index d0bacbf2..4ad2c7f1 100644
--- a/tests/test_t2z2t.py
+++ b/tests/test_t2z2t.py
@@ -54,10 +54,11 @@ def check_correctness(laser_t_in, laser_t_out, laser_z_analytic, z_axis):
         laser_t_out.dim, laser_t_out.grid, laser_t_out.profile.omega0, laser_z, z_axis
     )
 
-    ind0 = laser_t_in.grid.field.shape[0] // 2 - 1
+    field = laser_t_in.grid.get_temporal_field()
+    ind0 = field.shape[0] // 2 - 1
 
-    laser_t_in_2d = laser_t_in.grid.field[ind0]
-    laser_t_out_2d = laser_t_out.grid.field[ind0]
+    laser_t_in_2d = field[ind0]
+    laser_t_out_2d = field[ind0]
     laser_z_2d = laser_z[ind0]
 
     assert np.allclose(laser_t_in_2d, laser_t_out_2d, atol=2e-7, rtol=0)