New resampling (#269)

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

lasy/laser.py

@@ -202,7 +202,9 @@ def apply_optics(self, optical_element):
             )
         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, grid=None, backend="NP", show_progress=True
+    ):
         """
         Propagate the laser pulse by the distance specified.
 
@@ -216,8 +218,13 @@ def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True
             will be attenuated (to assert proper Hankel transform).
             Only used for ``'rt'``.
 
+        grid : Grid object (optional)
+            Resample the field onto a new grid of different radial size and/or different number
+            of radial grid points. Only works 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
         """
@@ -237,32 +244,80 @@ 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)
 
+        # Retrieve the spectral field from the current grid
+        spectral_field = self.grid.get_spectral_field()
+
         if self.dim == "rt":
-            # Construct the propagator (check if exists)
-            if not hasattr(self, "prop"):
-                spatial_axes = (self.grid.axes[0],)
-                self.prop = []
+            # Resampling onto new grid
+            if grid is not None:
+                # Overwrite time information from current grid
+                grid.lo[time_axis_indx] = self.grid.lo[time_axis_indx]
+                grid.hi[time_axis_indx] = self.grid.hi[time_axis_indx]
+                grid.axes[time_axis_indx] = self.grid.axes[time_axis_indx]
+                # Get radial axis of current and resampled grid
+                spatial_axes = (self.grid.axes[0],)  # current radial axis
+                spatial_axes_n = (grid.axes[0],)  # resampled axis
+                # Overwrite grid with new grid (for the resampled field)
+                self.grid = grid
+                # Creating an empty array to store the resampled spectral field.
+                # This will be needed in the propagation step below.
+                spectral_field_n = np.zeros(
+                    (len(self.grid.azimuthal_modes), grid.npoints[0], grid.npoints[1]),
+                    dtype="complex128",
+                )
+                self.prop = []  # Delete existing propagator
+                # Construct the propagator and pass resampled axis
                 for m in self.grid.azimuthal_modes:
                     self.prop.append(
                         PropagatorResampling(
                             *spatial_axes,
                             self.omega_1d / c,
+                            *spatial_axes_n,
                             mode=m,
                             backend=backend,
                             verbose=False,
                         )
                     )
-            # Propagate the spectral image
-            spectral_field = self.grid.get_spectral_field()
+            # No resampling (propagating on existing grid)
+            else:
+                if not hasattr(self, "prop"):
+                    spatial_axes = (self.grid.axes[0],)
+                    self.prop = []
+                    # Construct the propagator
+                    for m in self.grid.azimuthal_modes:
+                        self.prop.append(
+                            PropagatorResampling(
+                                *spatial_axes,
+                                self.omega_1d / c,
+                                mode=m,
+                                backend=backend,
+                                verbose=False,
+                            )
+                        )
+            # Propagate the spectral image in time
             for i_m in range(self.grid.azimuthal_modes.size):
                 transform_data = np.transpose(spectral_field[i_m]).copy()
-                self.prop[i_m].step(
+                # Propagation step. The spectral field is explicitely
+                # overwritten with the updated field.
+                transform_data = self.prop[i_m].step(
                     transform_data,
                     distance,
-                    overwrite=True,
+                    overwrite=False,
                     show_progress=show_progress,
                 )
-                spectral_field[i_m, :, :] = np.transpose(transform_data).copy()
+                if grid is not None:
+                    # Store the updated and resampled field.
+                    spectral_field_n[i_m, :, :] = np.transpose(transform_data).copy()
+                else:
+                    # Store the updated field.
+                    spectral_field[i_m, :, :] = np.transpose(transform_data).copy()
+
+            if grid is not None:
+                # Define the resampled field as new spectral field
+                spectral_field = spectral_field_n
+                # Delete Propagator if resampling and propagation was done
+                del self.prop
+
         else:
             # Construct the propagator (check if exists)
             if not hasattr(self, "prop"):
@@ -277,7 +332,6 @@ def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True
                     verbose=False,
                 )
             # Propagate the spectral image
-            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

requirements.txt

@@ -1,4 +1,4 @@
-axiprop
+axiprop>=0.5.3
 numpy
 openpmd-api
 openpmd-viewer

tests/test_resampling.py

@@ -0,0 +1,133 @@
+# -*- coding: utf-8 -*-
+"""Checking the implementation of the resampling propagator.
+
+Initializing a Gaussian pulse in the near field, and
+propagating it through a parabolic mirror, and then to
+the focal position (radial axis is resampled to accomodate the new size of the beam);
+we then check that the waist has the expected value in the far field (i.e. in the focal plane)
+"""
+
+import numpy as np
+
+from lasy.laser import Grid, Laser
+from lasy.optical_elements import ParabolicMirror
+from lasy.profiles.combined_profile import CombinedLongitudinalTransverseProfile
+from lasy.profiles.gaussian_profile import GaussianProfile
+from lasy.profiles.longitudinal.gaussian_profile import GaussianLongitudinalProfile
+from lasy.profiles.transverse.laguerre_gaussian_profile import (
+    LaguerreGaussianTransverseProfile,
+)
+
+# Laser parameters
+wavelength = 0.8e-6
+w0 = 5.0e-3  # m, initialized in near field
+pol = (1, 0)
+laser_energy = 1.0  # J
+t_peak = 0.0e-15  # s
+tau = 30.0e-15  # s
+
+# Define the initial grid for the laser
+dim = "rt"
+lo = (0e-3, -90e-15)
+hi = (15e-3, +90e-15)
+npoints = (500, 100)
+
+
+def get_w0(laser, m):
+    # Calculate the laser waist
+    field = laser.grid.get_temporal_field()
+    A2 = (np.abs(field[m, :, :]) ** 2).sum(-1)
+    ax = laser.grid.axes[0]
+    if ax[0] > 0:
+        A2 = np.r_[A2[::-1], A2]
+        ax = np.r_[-ax[::-1], ax]
+    else:
+        A2 = np.r_[A2[::-1][:-1], A2]
+        ax = np.r_[-ax[::-1][:-1], ax]
+
+    sigma = 2 * np.sqrt(np.average(ax**2, weights=A2))
+
+    return sigma
+
+
+def check_resampling(laser, new_grid, m=0):
+    # Focus down the laser and propagate
+    f0 = 2.0  # focal distance in m
+    laser.apply_optics(ParabolicMirror(f=f0))
+    laser.propagate(f0, nr_boundary=128, grid=new_grid)  # resample the radial grid
+
+    # Check that the value is the expected one in the near field
+    w0_num = get_w0(laser, m)
+    assert m in [0, 1]
+    w0_theor = wavelength * f0 / (np.pi * w0)
+    if m == 1:
+        w0_theor *= np.sqrt(3)
+    err = 2 * np.abs(w0_theor - w0_num) / (w0_theor + w0_num)
+    assert err < 1e-3
+
+
+def test_resampling_gaussian():
+    # Initialize the laser (gaussian profile)
+    gaussian_profile = GaussianProfile(wavelength, pol, laser_energy, w0, tau, t_peak)
+    laser = Laser(dim, lo, hi, npoints, gaussian_profile)
+
+    # Define the new grid for the laser
+    new_r_max = 300e-6
+    npoints_new = (50, 100)
+    new_grid = Grid(
+        dim,
+        lo,
+        (new_r_max, hi[1]),
+        npoints_new,
+        n_azimuthal_modes=laser.grid.n_azimuthal_modes,
+    )
+    # Check resampling propagator
+    check_resampling(laser, new_grid)
+
+
+def test_resampling_multimode_gaussian():
+    # Initialize the laser (gaussian profile)
+    gaussian_profile = GaussianProfile(wavelength, pol, laser_energy, w0, tau, t_peak)
+    laser = Laser(
+        dim, lo, hi, npoints, gaussian_profile, n_azimuthal_modes=3, n_theta_evals=20
+    )
+
+    # Define the new grid for the laser
+    new_r_max = 300e-6
+    npoints_new = (50, 100)
+    new_grid = Grid(
+        dim,
+        lo,
+        (new_r_max, hi[1]),
+        npoints_new,
+        n_azimuthal_modes=laser.grid.n_azimuthal_modes,
+    )
+    # Check resampling propagator
+    check_resampling(laser, new_grid)
+
+
+def test_resampling_laguerre():
+    # Initialize the laser (LaguerreGaussian)
+    p = 0  # Radial order of Generalized Laguerre polynomial
+    m = 1  # Phase Rotation
+
+    LongitProfile = GaussianLongitudinalProfile(wavelength, tau, t_peak)
+    TransvProfile = LaguerreGaussianTransverseProfile(w0, p, m)
+    pulseProfile = CombinedLongitudinalTransverseProfile(
+        wavelength, pol, laser_energy, LongitProfile, TransvProfile
+    )
+
+    laser = Laser(dim, lo, hi, npoints, pulseProfile, n_azimuthal_modes=2)
+
+    # Define the new grid for the laser
+    new_r_max = 300e-6
+    npoints_new = (50, 100)
+    new_grid = Grid(
+        dim,
+        lo,
+        (new_r_max, hi[1]),
+        npoints_new,
+        n_azimuthal_modes=laser.grid.n_azimuthal_modes,
+    )
+    # Check resampling propagator
+    check_resampling(laser, new_grid, m)