From 359a322b58ad9c11e879ebac71627803cba3a573 Mon Sep 17 00:00:00 2001
From: Igor Andriyash <igor.andriyash@gmail.com>
Date: Fri, 14 Apr 2023 15:51:48 +0200
Subject: [PATCH] Add time axis refinements to get_full_field method (#114)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 lasy/laser.py             | 46 -----------------------
 lasy/utils/laser_utils.py | 77 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 46 deletions(-)

diff --git a/lasy/laser.py b/lasy/laser.py
index 26d2a1fd..b3155318 100644
--- a/lasy/laser.py
+++ b/lasy/laser.py
@@ -266,49 +266,3 @@ def write_to_file(self, file_prefix="laser", file_format="h5"):
             self.profile.lambda0,
             self.profile.pol,
         )
-
-    def get_full_field(self, theta=0, slice=0, slice_axis="x"):
-        """
-        Reconstruct the laser pulse with carrier frequency on the default grid
-
-        Parameters
-        ----------
-        theta : float (rad) (optional)
-            Azimuthal angle
-        slice : float (optional)
-            Normalised position of the slice from -0.5 to 0.5
-
-        Returns
-        -------
-            Et : ndarray (V/m)
-                The reconstructed field, with shape (Nr, Nt_new) (for `rt`)
-                or (Nx, Nt_new) (for `xyt`)
-            extent : ndarray (Xmin, Xmax, Tmin, Tmax)
-                Physical extent of the reconstructed field
-        """
-        omega0 = self.profile.omega0
-        field = self.field.field.copy()
-        time_axis = self.box.axes[-1][None, :]
-
-        if self.dim == "rt":
-            azimuthal_phase = np.exp(-1j * self.box.azimuthal_modes * theta)
-            field *= azimuthal_phase[:, None, None]
-            field = field.sum(0)
-        elif slice_axis == "x":
-            Nx_middle = field.shape[0] // 2 - 1
-            Nx_slice = int((1 + slice) * Nx_middle)
-            field = field[Nx_slice, :]
-        elif slice_axis == "y":
-            Ny_middle = field.shape[1] // 2 - 1
-            Ny_slice = int((1 + slice) * Ny_middle)
-            field = field[:, Ny_slice, :]
-        else:
-            return None
-
-        field *= np.exp(-1j * omega0 * time_axis)
-        field = np.real(field)
-        ext = np.array(
-            [self.box.lo[-1], self.box.hi[-1], self.box.lo[0], self.box.hi[0]]
-        )
-
-        return field, ext
diff --git a/lasy/utils/laser_utils.py b/lasy/utils/laser_utils.py
index 74473899..2635c802 100644
--- a/lasy/utils/laser_utils.py
+++ b/lasy/utils/laser_utils.py
@@ -1,5 +1,6 @@
 import numpy as np
 from scipy.constants import c, epsilon_0
+from scipy.interpolate import interp1d
 
 
 def compute_laser_energy(dim, grid):
@@ -121,3 +122,79 @@ def normalize_peak_intensity(peak_intensity, grid):
     input_peak_intensity = intensity.max()
 
     grid.field *= np.sqrt(peak_intensity / input_peak_intensity)
+
+
+def get_full_field(laser, theta=0, slice=0, slice_axis="x", Nt=None):
+    """
+    Reconstruct the laser pulse with carrier frequency on the default grid
+
+    Parameters
+    ----------
+    theta : float (rad) (optional)
+        Azimuthal angle
+    slice : float (optional)
+        Normalised position of the slice from -0.5 to 0.5
+    Nt: int (optional)
+        Number of time points on which field should be sampled. If is None,
+        the orignal time grid is used, otherwise field is interpolated on a
+        new grid.
+
+    Returns
+    -------
+        Et : ndarray (V/m)
+            The reconstructed field, with shape (Nr, Nt) (for `rt`)
+            or (Nx, Nt) (for `xyt`)
+        extent : ndarray (Tmin, Tmax, Xmin, Xmax)
+            Physical extent of the reconstructed field
+    """
+    omega0 = laser.profile.omega0
+    field = laser.field.field.copy()
+    time_axis = laser.box.axes[-1]
+
+    if laser.dim == "rt":
+        azimuthal_phase = np.exp(-1j * laser.box.azimuthal_modes * theta)
+        field_upper = field * azimuthal_phase[:, None, None]
+        field_upper = field_upper.sum(0)
+        azimuthal_phase = np.exp(1j * laser.box.azimuthal_modes * theta)
+        field_lower = field * azimuthal_phase[:, None, None]
+        field_lower = field_lower.sum(0)
+        field = np.vstack((field_lower[::-1][:-1], field_upper))
+    elif slice_axis == "x":
+        Nx_middle = field.shape[0] // 2 - 1
+        Nx_slice = int((1 + slice) * Nx_middle)
+        field = field[Nx_slice, :]
+    elif slice_axis == "y":
+        Ny_middle = field.shape[1] // 2 - 1
+        Ny_slice = int((1 + slice) * Ny_middle)
+        field = field[:, Ny_slice, :]
+    else:
+        return None
+
+    if Nt is not None:
+        Nr = field.shape[0]
+        time_axis_new = np.linspace(laser.box.lo[-1], laser.box.hi[-1], Nt)
+        field_new = np.zeros((Nr, Nt), dtype=field.dtype)
+
+        for ir in range(Nr):
+            interp_fu_abs = interp1d(time_axis, np.abs(field[ir]))
+            slice_abs = interp_fu_abs(time_axis_new)
+            interp_fu_angl = interp1d(time_axis, np.unwrap(np.angle(field[ir])))
+            slice_angl = interp_fu_angl(time_axis_new)
+            field_new[ir] = slice_abs * np.exp(1j * slice_angl)
+
+        time_axis = time_axis_new
+        field = field_new
+
+    field *= np.exp(-1j * omega0 * time_axis[None, :])
+    field = np.real(field)
+
+    if laser.dim == "rt":
+        ext = np.array(
+            [laser.box.lo[-1], laser.box.hi[-1], -laser.box.hi[0], laser.box.hi[0]]
+        )
+    else:
+        ext = np.array(
+            [laser.box.lo[-1], laser.box.hi[-1], laser.box.lo[0], laser.box.hi[0]]
+        )
+
+    return field, ext