From a641a0c3b264c060073c4437b4cf3b716303f578 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Thu, 26 Oct 2023 11:06:55 -0700
Subject: [PATCH 01/26] Add `show` method to visualize the laser

---
 docs/source/tutorials/gaussian_laser.ipynb | 30 ++++++++++++++++++++-
 lasy/laser.py                              | 31 ++++++++++++++++++++++
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/docs/source/tutorials/gaussian_laser.ipynb b/docs/source/tutorials/gaussian_laser.ipynb
index 850bfa26..e9b99f20 100644
--- a/docs/source/tutorials/gaussian_laser.ipynb
+++ b/docs/source/tutorials/gaussian_laser.ipynb
@@ -78,6 +78,24 @@
     "laser = Laser(dimensions,lo,hi,num_points,laser_profile)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "8ac8f985-4a09-4b9a-a9ff-11ce2ef94caa",
+   "metadata": {},
+   "source": [
+    "The laser pulse can be visualized with the `show` method."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bb8dc634-b3c4-4448-834b-fe9acb6c9645",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "laser.show()"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "292c6e63-0480-4fb9-b1ad-6b679a3472d0",
@@ -97,6 +115,16 @@
     "laser.propagate(-z_R)                               # Propagate the pulse upstream of the focal plane"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3158e51d-e331-438d-90c1-5f8c63bf9841",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "laser.show()"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "8df816dc-9b47-4cb0-8716-600352fafb72",
@@ -143,7 +171,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.5"
+   "version": "3.11.6"
   }
  },
  "nbformat": 4,
diff --git a/lasy/laser.py b/lasy/laser.py
index 7b7503ea..e093f879 100644
--- a/lasy/laser.py
+++ b/lasy/laser.py
@@ -276,3 +276,34 @@ def write_to_file(
             self.profile.pol,
             save_as_vector_potential,
         )
+
+
+    def show( self, **kw ):
+        """
+        Show a 2D image of the laser amplitude
+
+        Parameters:
+        ----------
+        **kw: additional arguments to be passed to matplotlib's imshow command
+        """
+        if self.dim == "rt":
+            # Show field above and below axis, with proper sign
+            E = [
+                np.concatenate( ( (-1)**m * self.grid.field[0,::-1], self.grid.field[0] ) ) \
+                for m in self.grid.azimuthal_modes
+            ]
+            E = sum(E)
+            extent = [self.grid.lo[-1], self.grid.hi[-1], -self.grid.hi[0], self.grid.hi[0]]
+
+        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,:]
+            extent=[self.grid.lo[-1], self.grid.hi[-1], self.grid.lo[0], self.grid.hi[0]]
+
+        import matplotlib.pyplot as plt
+        plt.imshow( abs(E), extent=extent, aspect='auto', origin='lower', **kw )
+        cb = plt.colorbar()
+        cb.set_label('$|E_{envelope}|$ (V/m)')
+        plt.xlabel('t (s)')
+        plt.ylabel('x (m)')
\ No newline at end of file

From 8c611351739df2a909af0561085f3691a3a3c1bf Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 26 Oct 2023 18:10:56 +0000
Subject: [PATCH 02/26] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 lasy/laser.py | 34 +++++++++++++++++++++++-----------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/lasy/laser.py b/lasy/laser.py
index e093f879..2713153e 100644
--- a/lasy/laser.py
+++ b/lasy/laser.py
@@ -277,8 +277,7 @@ def write_to_file(
             save_as_vector_potential,
         )
 
-
-    def show( self, **kw ):
+    def show(self, **kw):
         """
         Show a 2D image of the laser amplitude
 
@@ -289,21 +288,34 @@ def show( self, **kw ):
         if self.dim == "rt":
             # Show field above and below axis, with proper sign
             E = [
-                np.concatenate( ( (-1)**m * self.grid.field[0,::-1], self.grid.field[0] ) ) \
+                np.concatenate(
+                    ((-1) ** m * self.grid.field[0, ::-1], self.grid.field[0])
+                )
                 for m in self.grid.azimuthal_modes
             ]
             E = sum(E)
-            extent = [self.grid.lo[-1], self.grid.hi[-1], -self.grid.hi[0], self.grid.hi[0]]
+            extent = [
+                self.grid.lo[-1],
+                self.grid.hi[-1],
+                -self.grid.hi[0],
+                self.grid.hi[0],
+            ]
 
         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,:]
-            extent=[self.grid.lo[-1], self.grid.hi[-1], self.grid.lo[0], self.grid.hi[0]]
+            i_slice = int(self.grid.field.shape[1] // 2)
+            E = self.grid.field[:, i_slice, :]
+            extent = [
+                self.grid.lo[-1],
+                self.grid.hi[-1],
+                self.grid.lo[0],
+                self.grid.hi[0],
+            ]
 
         import matplotlib.pyplot as plt
-        plt.imshow( abs(E), extent=extent, aspect='auto', origin='lower', **kw )
+
+        plt.imshow(abs(E), extent=extent, aspect="auto", origin="lower", **kw)
         cb = plt.colorbar()
-        cb.set_label('$|E_{envelope}|$ (V/m)')
-        plt.xlabel('t (s)')
-        plt.ylabel('x (m)')
\ No newline at end of file
+        cb.set_label("$|E_{envelope}|$ (V/m)")
+        plt.xlabel("t (s)")
+        plt.ylabel("x (m)")

From 8c4ba85d579aee55c59122c5ddbaff63d59f9713 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Thu, 26 Oct 2023 11:12:32 -0700
Subject: [PATCH 03/26] Fix doc

---
 lasy/laser.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/lasy/laser.py b/lasy/laser.py
index 2713153e..0e7d9520 100644
--- a/lasy/laser.py
+++ b/lasy/laser.py
@@ -279,9 +279,9 @@ def write_to_file(
 
     def show(self, **kw):
         """
-        Show a 2D image of the laser amplitude
+        Show a 2D image of the laser amplitude.
 
-        Parameters:
+        Parameters
         ----------
         **kw: additional arguments to be passed to matplotlib's imshow command
         """

From 773381f7693c2ad5e1c0d64868d23103327ecd62 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Thu, 26 Oct 2023 11:24:04 -0700
Subject: [PATCH 04/26] Apply suggestions from code review

---
 lasy/laser.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lasy/laser.py b/lasy/laser.py
index 0e7d9520..8a2e70c4 100644
--- a/lasy/laser.py
+++ b/lasy/laser.py
@@ -286,7 +286,7 @@ def show(self, **kw):
         **kw: additional arguments to be passed to matplotlib's imshow command
         """
         if self.dim == "rt":
-            # Show field above and below axis, with proper sign
+            # Show field in the plane y=0, above and below axis, with proper sign for each mode
             E = [
                 np.concatenate(
                     ((-1) ** m * self.grid.field[0, ::-1], self.grid.field[0])

From 2f0752c313628d9f385b706520a35e240ab6f651 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Fri, 20 Oct 2023 14:42:36 -0700
Subject: [PATCH 05/26] Base class for optical elements

---
 lasy/optical_elements/__init__.py        |  0
 lasy/optical_elements/optical_element.py | 38 ++++++++++++++++++++++++
 2 files changed, 38 insertions(+)
 create mode 100644 lasy/optical_elements/__init__.py
 create mode 100644 lasy/optical_elements/optical_element.py

diff --git a/lasy/optical_elements/__init__.py b/lasy/optical_elements/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/lasy/optical_elements/optical_element.py b/lasy/optical_elements/optical_element.py
new file mode 100644
index 00000000..64b78145
--- /dev/null
+++ b/lasy/optical_elements/optical_element.py
@@ -0,0 +1,38 @@
+
+class OpticalElement(object):
+    """
+    Base class to model thin optical elements.
+
+    Any optical element should inherit from this class, and define its own
+    `amplitude_multiplier` method, using the same signature as the method below.
+    """
+
+    def __init__(self):
+        pass
+
+    def amplitude_multiplier(self, x, y, omega):
+        """
+        Return the complex number :math:`T` with which to multiply the
+        complex amplitude of the laser just before this thin element,
+        in order to obtain the complex amplitude output laser just
+        after this thin element:
+
+        .. math::
+
+            \tilde{\mathcal{E}}_{out}(x, y, \omega) = T(x, y, \omega)\tilde{\mathcal{E}}_{in}(x, y, \omega)
+
+        Parameters
+        ----------
+        x, y, omega: ndarrays of floats
+            Define points on which to evaluate the multiplier.
+            These arrays need to all have the same shape.
+
+        Returns
+        -------
+        multiplier: ndarray of complex numbers
+            Contains the value of the multiplier at the specified points
+            This array has the same shape as the arrays x, y, omega
+        """
+        # The base class only defines dummy multiplier
+        # (This should be replaced by any class that inherits from this one.)
+        return np.zeros_like(x, dtype="complex128")

From b9b22f3cc87951558a3b9ecf49f5ae0e96a1f8d6 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Fri, 20 Oct 2023 15:25:52 -0700
Subject: [PATCH 06/26] Add documentation

---
 docs/source/api/index.rst |  1 +
 lasy/laser.py             | 27 ++++++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst
index d2d1099c..708bbc5a 100644
--- a/docs/source/api/index.rst
+++ b/docs/source/api/index.rst
@@ -9,6 +9,7 @@ The sections below describe the main objects that are accessible in the ``lasy``
 
    laser
    profiles/index
+   optical_elements/index
    utils/index
 
 If you are looking for a specific class or function, see the :ref:`genindex` or use the search bar of this website.
diff --git a/lasy/laser.py b/lasy/laser.py
index 8a2e70c4..9e5ddd21 100644
--- a/lasy/laser.py
+++ b/lasy/laser.py
@@ -139,7 +139,8 @@ def normalize(self, value, kind="energy"):
         else:
             raise ValueError(f'kind "{kind}" not recognized')
 
-    def propagate(self, distance, nr_boundary=None, backend="NP"):
+    def propagate(self, distance, initial_optical_element=None,
+                  nr_boundary=None, backend="NP"):
         """
         Propagate the laser pulse by the distance specified.
 
@@ -148,10 +149,16 @@ def propagate(self, distance, nr_boundary=None, backend="NP"):
         distance : scalar
             Distance by which the laser pulse should be propagated
 
+        initial_optical_element: an :class:`.OpticalElement` object (optional)
+            Represents a thin optical element, through which the laser
+            propagates, before propagating for `distance` in free space.
+            If this is `None`, no optical element is used.
+
         nr_boundary : integer (optional)
             Number of cells at the end of radial axis, where the field
             will be attenuated (to assert proper Hankel transform).
             Only used for ``'rt'``.
+
         backend : string (optional)
             Backend used by axiprop (see axiprop documentation).
         """
@@ -189,10 +196,22 @@ def propagate(self, distance, nr_boundary=None, backend="NP"):
         omega0 = self.profile.omega0
         Nt = self.grid.field.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":
+
+           # Apply optical element
+            if initial_optical_element is not None:
+                r, w = np.meshgrid(self.grid.r, omega, indexing="ij")
+                # The line below assumes that amplitude_multiplier
+                # is cylindrically-symmetric, hence we pass
+                # `r` as `x` and 0 as `y`
+                multiplier = initial_optical_element.amplitude_multiplier(r, 0, w)
+                for m in self.grid.azimuthal_modes:
+                    field_fft[i_m, :, :] *= multiplier
+
             # Construct the propagator (check if exists)
             if not hasattr(self, "prop"):
                 spatial_axes = (self.grid.axes[0],)
@@ -213,6 +232,12 @@ def propagate(self, distance, nr_boundary=None, backend="NP"):
                 self.prop[i_m].step(transform_data, distance, overwrite=True)
                 field_fft[i_m, :, :] = np.transpose(transform_data).copy()
         else:
+
+           # Apply optical element
+            if initial_optical_element is not None:
+                x, y, w = np.meshgrid(self.grid.x, self.grid.y, omega, indexing="ij")
+                field_fft *= initial_optical_element.amplitude_multiplier(x, y, w)
+
             # Construct the propagator (check if exists)
             if not hasattr(self, "prop"):
                 Nx, Ny, Nt = self.grid.field.shape

From bbb95b2b549b0a7c99c9eb312f7b9f0bfcb366c0 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Fri, 20 Oct 2023 15:58:30 -0700
Subject: [PATCH 07/26] Add parabolic mirror

---
 lasy/optical_elements/parabolic_mirror.py | 29 +++++++++++++++++++++++
 1 file changed, 29 insertions(+)
 create mode 100644 lasy/optical_elements/parabolic_mirror.py

diff --git a/lasy/optical_elements/parabolic_mirror.py b/lasy/optical_elements/parabolic_mirror.py
new file mode 100644
index 00000000..ebbfe03f
--- /dev/null
+++ b/lasy/optical_elements/parabolic_mirror.py
@@ -0,0 +1,29 @@
+from . import OpticalElement
+import numpy as np
+
+class ParabolicMirror(OpticalElement):
+    r"""
+    Derived class for a parabolic mirror.
+
+    More precisely, the amplitude multiplier corresponds to:
+
+    .. math::
+
+        T(\boldsymbol{x}_\perp,\omega) = \exp(i\omega \sqrt{x^2+y^2}/2f)
+
+    where
+    :math:`\boldsymbol{x}_\perp` is the transverse coordinate (orthogonal
+    to the propagation direction). The other parameters in this formula
+    are defined below.
+
+    Parameters
+    ----------
+    f : float (in meter)
+        The focal length of the parabolic mirror.
+    """
+
+    def __init__(self, f):
+        self.f = f
+
+    def amplitude_multiplier(self, x, y, omega):
+        return np.exp(1j*omega*(x**2 + y**2)/(2*self.f))
\ No newline at end of file

From 63ad71e2c082ed63535dfa8b8c77740c02ebe8a6 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Tue, 24 Oct 2023 15:05:42 -0700
Subject: [PATCH 08/26] Add automated test for parabolic mirror

---
 lasy/laser.py                             |  7 +-
 lasy/optical_elements/__init__.py         |  3 +
 lasy/optical_elements/parabolic_mirror.py |  8 ++-
 tests/test_parabolic_mirror.py            | 85 +++++++++++++++++++++++
 4 files changed, 97 insertions(+), 6 deletions(-)
 create mode 100644 tests/test_parabolic_mirror.py

diff --git a/lasy/laser.py b/lasy/laser.py
index 9e5ddd21..1900404d 100644
--- a/lasy/laser.py
+++ b/lasy/laser.py
@@ -204,12 +204,12 @@ def propagate(self, distance, initial_optical_element=None,
 
            # Apply optical element
             if initial_optical_element is not None:
-                r, w = np.meshgrid(self.grid.r, omega, indexing="ij")
+                r, w = np.meshgrid( self.grid.axes[0], omega, indexing="ij")
                 # The line below assumes that amplitude_multiplier
                 # is cylindrically-symmetric, hence we pass
                 # `r` as `x` and 0 as `y`
                 multiplier = initial_optical_element.amplitude_multiplier(r, 0, w)
-                for m in self.grid.azimuthal_modes:
+                for i_m in range(self.grid.azimuthal_modes.size):
                     field_fft[i_m, :, :] *= multiplier
 
             # Construct the propagator (check if exists)
@@ -235,7 +235,8 @@ def propagate(self, distance, initial_optical_element=None,
 
            # Apply optical element
             if initial_optical_element is not None:
-                x, y, w = np.meshgrid(self.grid.x, self.grid.y, omega, indexing="ij")
+                x, y, w = np.meshgrid(self.grid.axes[0], self.grid.axes[1],
+                                      omega, indexing="ij")
                 field_fft *= initial_optical_element.amplitude_multiplier(x, y, w)
 
             # Construct the propagator (check if exists)
diff --git a/lasy/optical_elements/__init__.py b/lasy/optical_elements/__init__.py
index e69de29b..ebc3c4c2 100644
--- a/lasy/optical_elements/__init__.py
+++ b/lasy/optical_elements/__init__.py
@@ -0,0 +1,3 @@
+from .parabolic_mirror import ParabolicMirror
+
+__all__ = ['ParabolicMirror']
\ No newline at end of file
diff --git a/lasy/optical_elements/parabolic_mirror.py b/lasy/optical_elements/parabolic_mirror.py
index ebbfe03f..058d9615 100644
--- a/lasy/optical_elements/parabolic_mirror.py
+++ b/lasy/optical_elements/parabolic_mirror.py
@@ -1,5 +1,6 @@
-from . import OpticalElement
+from .optical_element import OpticalElement
 import numpy as np
+from scipy.constants import c
 
 class ParabolicMirror(OpticalElement):
     r"""
@@ -9,7 +10,7 @@ class ParabolicMirror(OpticalElement):
 
     .. math::
 
-        T(\boldsymbol{x}_\perp,\omega) = \exp(i\omega \sqrt{x^2+y^2}/2f)
+        T(\boldsymbol{x}_\perp,\omega) = \exp(i\omega \sqrt{x^2+y^2}/2cf)
 
     where
     :math:`\boldsymbol{x}_\perp` is the transverse coordinate (orthogonal
@@ -26,4 +27,5 @@ def __init__(self, f):
         self.f = f
 
     def amplitude_multiplier(self, x, y, omega):
-        return np.exp(1j*omega*(x**2 + y**2)/(2*self.f))
\ No newline at end of file
+        # TODO: add reference
+        return np.exp( -1j*omega*(x**2 + y**2)/(2*c*self.f) )
\ No newline at end of file
diff --git a/tests/test_parabolic_mirror.py b/tests/test_parabolic_mirror.py
new file mode 100644
index 00000000..a87a3ce8
--- /dev/null
+++ b/tests/test_parabolic_mirror.py
@@ -0,0 +1,85 @@
+# -*- coding: utf-8 -*-
+"""
+This test checks the implementation of the parabolic mirror
+by initializing a Gaussian pulse in the near field, and
+propagating it through a parabolic mirror, and then to
+the focal position ; we then check that the waist as the
+expected value in the far field (i.e. in the focal plane)
+"""
+
+import pytest
+
+import numpy as np
+from lasy.laser import Laser
+from lasy.profiles.gaussian_profile import GaussianProfile
+from lasy.optical_elements import ParabolicMirror
+
+wavelength = 0.8e-6
+w0 = 10.0e-3  # m, initialized in near field
+
+@pytest.fixture(scope="function")
+def gaussian():
+    # Cases with Gaussian laser
+    pol = (1, 0)
+    laser_energy = 1.0  # J
+    t_peak = 0.0e-15  # s
+    tau = 30.0e-15  # s
+    profile = GaussianProfile(wavelength, pol, laser_energy, w0, tau, t_peak)
+
+    return profile
+
+
+def get_w0(laser):
+    # Calculate the laser waist
+    if laser.dim == "xyt":
+        Nx, Ny, Nt = laser.grid.field.shape
+        A2 = (np.abs(laser.grid.field[Nx // 2 - 1, :, :]) ** 2).sum(-1)
+        ax = laser.grid.axes[1]
+    else:
+        A2 = (np.abs(laser.grid.field[0, :, :]) ** 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_parabolic_mirror(laser):
+    # Propagate laser after parabolic mirror + vacuum
+    f0 = 8. # focal distance in m
+    laser.propagate( f0, initial_optical_element=ParabolicMirror( f=f0 ) )
+    # Check that the value is the expected one in the near field
+    w0_num = get_w0(laser)
+    w0_theor = wavelength * f0 / (np.pi*w0)
+    err = 2 * np.abs(w0_theor - w0_num) / (w0_theor + w0_num)
+    assert err < 1e-3
+
+"""
+def test_3D_case(gaussian):
+    # - 3D case
+    # The laser is initialized in the near field
+    dim = "xyt"
+    lo = (-25e-3, -25e-3, -60e-15)
+    hi = (+25e-3, +25e-3, +60e-15)
+    npoints = (3000, 3000, 100)
+
+    laser = Laser(dim, lo, hi, npoints, gaussian)
+    check_parabolic_mirror(laser)
+"""
+
+def test_RT_case(gaussian):
+    # - Cylindrical case
+    # The laser is initialized in the near field
+    dim = "rt"
+    lo = (0e-6, -60e-15)
+    hi = (25e-3, +60e-15)
+    npoints = (1500, 100)
+
+    laser = Laser(dim, lo, hi, npoints, gaussian)
+    check_parabolic_mirror(laser)

From 751c6a634d947e211a05c7939662839b39c8a1df Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Tue, 24 Oct 2023 16:33:15 -0700
Subject: [PATCH 09/26] Add 3D automated test

---
 tests/test_parabolic_mirror.py | 40 ++++++++++++++--------------------
 1 file changed, 16 insertions(+), 24 deletions(-)

diff --git a/tests/test_parabolic_mirror.py b/tests/test_parabolic_mirror.py
index a87a3ce8..b3c8ec70 100644
--- a/tests/test_parabolic_mirror.py
+++ b/tests/test_parabolic_mirror.py
@@ -15,19 +15,14 @@
 from lasy.optical_elements import ParabolicMirror
 
 wavelength = 0.8e-6
-w0 = 10.0e-3  # m, initialized in near field
-
-@pytest.fixture(scope="function")
-def gaussian():
-    # Cases with Gaussian laser
-    pol = (1, 0)
-    laser_energy = 1.0  # J
-    t_peak = 0.0e-15  # s
-    tau = 30.0e-15  # s
-    profile = GaussianProfile(wavelength, pol, laser_energy, w0, tau, t_peak)
-
-    return profile
+w0 = 5.0e-3  # m, initialized in near field
 
+# The laser is initialized in the near field
+pol = (1, 0)
+laser_energy = 1.0  # J
+t_peak = 0.0e-15  # s
+tau = 30.0e-15  # s
+gaussian_profile = GaussianProfile(wavelength, pol, laser_energy, w0, tau, t_peak)
 
 def get_w0(laser):
     # Calculate the laser waist
@@ -49,7 +44,6 @@ def get_w0(laser):
 
     return sigma
 
-
 def check_parabolic_mirror(laser):
     # Propagate laser after parabolic mirror + vacuum
     f0 = 8. # focal distance in m
@@ -60,26 +54,24 @@ def check_parabolic_mirror(laser):
     err = 2 * np.abs(w0_theor - w0_num) / (w0_theor + w0_num)
     assert err < 1e-3
 
-"""
-def test_3D_case(gaussian):
+def test_3D_case():
     # - 3D case
     # The laser is initialized in the near field
     dim = "xyt"
-    lo = (-25e-3, -25e-3, -60e-15)
-    hi = (+25e-3, +25e-3, +60e-15)
-    npoints = (3000, 3000, 100)
+    lo = (-12e-3, -12e-3, -60e-15)
+    hi = (+12e-3, +12e-3, +60e-15)
+    npoints = (500, 500, 100)
 
-    laser = Laser(dim, lo, hi, npoints, gaussian)
+    laser = Laser(dim, lo, hi, npoints, gaussian_profile)
     check_parabolic_mirror(laser)
-"""
 
-def test_RT_case(gaussian):
+def test_RT_case():
     # - Cylindrical case
     # The laser is initialized in the near field
     dim = "rt"
     lo = (0e-6, -60e-15)
-    hi = (25e-3, +60e-15)
-    npoints = (1500, 100)
+    hi = (15e-3, +60e-15)
+    npoints = (750, 100)
 
-    laser = Laser(dim, lo, hi, npoints, gaussian)
+    laser = Laser(dim, lo, hi, npoints, gaussian_profile)
     check_parabolic_mirror(laser)

From 4458be0840a9dd64473d7c793a34231cdb38d7d6 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Tue, 24 Oct 2023 16:34:42 -0700
Subject: [PATCH 10/26] Update comments

---
 lasy/optical_elements/parabolic_mirror.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/lasy/optical_elements/parabolic_mirror.py b/lasy/optical_elements/parabolic_mirror.py
index 058d9615..5600fdf8 100644
--- a/lasy/optical_elements/parabolic_mirror.py
+++ b/lasy/optical_elements/parabolic_mirror.py
@@ -10,7 +10,7 @@ class ParabolicMirror(OpticalElement):
 
     .. math::
 
-        T(\boldsymbol{x}_\perp,\omega) = \exp(i\omega \sqrt{x^2+y^2}/2cf)
+        T(\boldsymbol{x}_\perp,\omega) = \exp(-i\omega \sqrt{x^2+y^2}/2cf)
 
     where
     :math:`\boldsymbol{x}_\perp` is the transverse coordinate (orthogonal
@@ -27,5 +27,4 @@ def __init__(self, f):
         self.f = f
 
     def amplitude_multiplier(self, x, y, omega):
-        # TODO: add reference
         return np.exp( -1j*omega*(x**2 + y**2)/(2*c*self.f) )
\ No newline at end of file

From d6bd90a0453c6e0d69412641368050ba1afe0aa4 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Wed, 25 Oct 2023 14:49:33 -0700
Subject: [PATCH 11/26] Add axiparabola

---
 lasy/optical_elements/axiparabola.py | 56 ++++++++++++++++++++++++++++
 1 file changed, 56 insertions(+)
 create mode 100644 lasy/optical_elements/axiparabola.py

diff --git a/lasy/optical_elements/axiparabola.py b/lasy/optical_elements/axiparabola.py
new file mode 100644
index 00000000..b07d7d16
--- /dev/null
+++ b/lasy/optical_elements/axiparabola.py
@@ -0,0 +1,56 @@
+from .optical_element import OpticalElement
+import numpy as np
+from scipy.constants import c
+
+class AxiParabolaWithDelay(OpticalElement):
+    r"""
+    Class that represents the combination of an axiparabola, with
+    an additional optical element that provides a radially-dependent
+    delay (e.g. an optical echelon) to tune the group velocity.
+
+    The rays that impinge the axiparabola at different radii are focused
+    to different positions on the axis (resulting in an extended "focal
+    range"). An additional radially-dependent delay is usually applied,
+    in order to tune the effective group velocity on axis.
+
+    For more details, see K. Oubrerie et al, "Axiparabola: a new tool
+    for high-intensity optics", J. Opt. 24 045503 (2022)
+
+    Parameters
+    ----------
+    f0: float (in meter)
+        The focal distance, i.e. the distance, from the axiparabola,
+        where the focal range starts.
+
+    delta: float (in meter)
+        The length of the focal range.
+
+    R: float (in meter)
+        The radius of the axiparabola. Rays coming from r=0 focus
+        at z=f0 ; rays coming from r=R focus at z=f0+delta
+    """
+
+    def __init__(self, f0, delta, R):
+        self.f0 = f0
+        self.delta = delta
+        self.R = R
+
+        # Function that defines the z position where rays that impinge at r focus.
+        # Assuming uniform intensity on the axiparabola, and in order to get
+        # a z-independent intensity over the focal range, we need
+        # (see Eq. 6 in Oubrerie et al.)
+        z_foc = lambda r: self.f0 + self.delta * (r/self.R)**2
+
+        # Solve Eq. 2 in Oubrerie et al. to find the sag function
+
+    def amplitude_multiplier(self, x, y, omega):
+
+
+        # Interpolation
+
+        # Calculate phase shift
+        T = np.exp( -2j*(omega/c)*sag )
+        # Remove intensity beyond R
+        T[ x**2 + y**2 > self.R**2 ] = 0
+
+        return T
\ No newline at end of file

From 4b5bb3e8f2a5a620c0071ea91ede4789f51c8747 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Wed, 25 Oct 2023 16:45:48 -0700
Subject: [PATCH 12/26] Fix minor bugs

---
 lasy/optical_elements/axiparabola.py     | 3 ++-
 lasy/optical_elements/optical_element.py | 1 +
 tests/test_parabolic_mirror.py           | 2 --
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/lasy/optical_elements/axiparabola.py b/lasy/optical_elements/axiparabola.py
index b07d7d16..9b2e30fa 100644
--- a/lasy/optical_elements/axiparabola.py
+++ b/lasy/optical_elements/axiparabola.py
@@ -4,7 +4,7 @@
 
 class AxiParabolaWithDelay(OpticalElement):
     r"""
-    Class that represents the combination of an axiparabola, with
+    Class that represents the combination of an axiparabola with
     an additional optical element that provides a radially-dependent
     delay (e.g. an optical echelon) to tune the group velocity.
 
@@ -47,6 +47,7 @@ def amplitude_multiplier(self, x, y, omega):
 
 
         # Interpolation
+        sag = np.zeros_like(x)
 
         # Calculate phase shift
         T = np.exp( -2j*(omega/c)*sag )
diff --git a/lasy/optical_elements/optical_element.py b/lasy/optical_elements/optical_element.py
index 64b78145..b40406f9 100644
--- a/lasy/optical_elements/optical_element.py
+++ b/lasy/optical_elements/optical_element.py
@@ -1,3 +1,4 @@
+import numpy as np
 
 class OpticalElement(object):
     """
diff --git a/tests/test_parabolic_mirror.py b/tests/test_parabolic_mirror.py
index b3c8ec70..3d6659e3 100644
--- a/tests/test_parabolic_mirror.py
+++ b/tests/test_parabolic_mirror.py
@@ -7,8 +7,6 @@
 expected value in the far field (i.e. in the focal plane)
 """
 
-import pytest
-
 import numpy as np
 from lasy.laser import Laser
 from lasy.profiles.gaussian_profile import GaussianProfile

From 85d1fde44d96f824eeee9804d4c6f08eab79e066 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 25 Oct 2023 23:48:36 +0000
Subject: [PATCH 13/26] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 lasy/laser.py                             | 18 +++++++++---------
 lasy/optical_elements/__init__.py         |  2 +-
 lasy/optical_elements/axiparabola.py      | 11 +++++------
 lasy/optical_elements/optical_element.py  |  1 +
 lasy/optical_elements/parabolic_mirror.py |  3 ++-
 tests/test_parabolic_mirror.py            | 10 +++++++---
 6 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/lasy/laser.py b/lasy/laser.py
index 1900404d..ccf11bfa 100644
--- a/lasy/laser.py
+++ b/lasy/laser.py
@@ -139,8 +139,9 @@ def normalize(self, value, kind="energy"):
         else:
             raise ValueError(f'kind "{kind}" not recognized')
 
-    def propagate(self, distance, initial_optical_element=None,
-                  nr_boundary=None, backend="NP"):
+    def propagate(
+        self, distance, initial_optical_element=None, nr_boundary=None, backend="NP"
+    ):
         """
         Propagate the laser pulse by the distance specified.
 
@@ -201,10 +202,9 @@ def propagate(self, distance, initial_optical_element=None,
         omega_shape = (1, 1, self.grid.field.shape[time_axis_indx])
 
         if self.dim == "rt":
-
-           # Apply optical element
+            # Apply optical element
             if initial_optical_element is not None:
-                r, w = np.meshgrid( self.grid.axes[0], omega, indexing="ij")
+                r, w = np.meshgrid(self.grid.axes[0], omega, indexing="ij")
                 # The line below assumes that amplitude_multiplier
                 # is cylindrically-symmetric, hence we pass
                 # `r` as `x` and 0 as `y`
@@ -232,11 +232,11 @@ def propagate(self, distance, initial_optical_element=None,
                 self.prop[i_m].step(transform_data, distance, overwrite=True)
                 field_fft[i_m, :, :] = np.transpose(transform_data).copy()
         else:
-
-           # Apply optical element
+            # Apply optical element
             if initial_optical_element is not None:
-                x, y, w = np.meshgrid(self.grid.axes[0], self.grid.axes[1],
-                                      omega, indexing="ij")
+                x, y, w = np.meshgrid(
+                    self.grid.axes[0], self.grid.axes[1], omega, indexing="ij"
+                )
                 field_fft *= initial_optical_element.amplitude_multiplier(x, y, w)
 
             # Construct the propagator (check if exists)
diff --git a/lasy/optical_elements/__init__.py b/lasy/optical_elements/__init__.py
index ebc3c4c2..af8ba752 100644
--- a/lasy/optical_elements/__init__.py
+++ b/lasy/optical_elements/__init__.py
@@ -1,3 +1,3 @@
 from .parabolic_mirror import ParabolicMirror
 
-__all__ = ['ParabolicMirror']
\ No newline at end of file
+__all__ = ["ParabolicMirror"]
diff --git a/lasy/optical_elements/axiparabola.py b/lasy/optical_elements/axiparabola.py
index 9b2e30fa..06269bb4 100644
--- a/lasy/optical_elements/axiparabola.py
+++ b/lasy/optical_elements/axiparabola.py
@@ -2,6 +2,7 @@
 import numpy as np
 from scipy.constants import c
 
+
 class AxiParabolaWithDelay(OpticalElement):
     r"""
     Class that represents the combination of an axiparabola with
@@ -39,19 +40,17 @@ def __init__(self, f0, delta, R):
         # Assuming uniform intensity on the axiparabola, and in order to get
         # a z-independent intensity over the focal range, we need
         # (see Eq. 6 in Oubrerie et al.)
-        z_foc = lambda r: self.f0 + self.delta * (r/self.R)**2
+        z_foc = lambda r: self.f0 + self.delta * (r / self.R) ** 2
 
         # Solve Eq. 2 in Oubrerie et al. to find the sag function
 
     def amplitude_multiplier(self, x, y, omega):
-
-
         # Interpolation
         sag = np.zeros_like(x)
 
         # Calculate phase shift
-        T = np.exp( -2j*(omega/c)*sag )
+        T = np.exp(-2j * (omega / c) * sag)
         # Remove intensity beyond R
-        T[ x**2 + y**2 > self.R**2 ] = 0
+        T[x**2 + y**2 > self.R**2] = 0
 
-        return T
\ No newline at end of file
+        return T
diff --git a/lasy/optical_elements/optical_element.py b/lasy/optical_elements/optical_element.py
index b40406f9..0142f3bd 100644
--- a/lasy/optical_elements/optical_element.py
+++ b/lasy/optical_elements/optical_element.py
@@ -1,5 +1,6 @@
 import numpy as np
 
+
 class OpticalElement(object):
     """
     Base class to model thin optical elements.
diff --git a/lasy/optical_elements/parabolic_mirror.py b/lasy/optical_elements/parabolic_mirror.py
index 5600fdf8..b2b62a02 100644
--- a/lasy/optical_elements/parabolic_mirror.py
+++ b/lasy/optical_elements/parabolic_mirror.py
@@ -2,6 +2,7 @@
 import numpy as np
 from scipy.constants import c
 
+
 class ParabolicMirror(OpticalElement):
     r"""
     Derived class for a parabolic mirror.
@@ -27,4 +28,4 @@ def __init__(self, f):
         self.f = f
 
     def amplitude_multiplier(self, x, y, omega):
-        return np.exp( -1j*omega*(x**2 + y**2)/(2*c*self.f) )
\ No newline at end of file
+        return np.exp(-1j * omega * (x**2 + y**2) / (2 * c * self.f))
diff --git a/tests/test_parabolic_mirror.py b/tests/test_parabolic_mirror.py
index 3d6659e3..ae3fa062 100644
--- a/tests/test_parabolic_mirror.py
+++ b/tests/test_parabolic_mirror.py
@@ -22,6 +22,7 @@
 tau = 30.0e-15  # s
 gaussian_profile = GaussianProfile(wavelength, pol, laser_energy, w0, tau, t_peak)
 
+
 def get_w0(laser):
     # Calculate the laser waist
     if laser.dim == "xyt":
@@ -42,16 +43,18 @@ def get_w0(laser):
 
     return sigma
 
+
 def check_parabolic_mirror(laser):
     # Propagate laser after parabolic mirror + vacuum
-    f0 = 8. # focal distance in m
-    laser.propagate( f0, initial_optical_element=ParabolicMirror( f=f0 ) )
+    f0 = 8.0  # focal distance in m
+    laser.propagate(f0, initial_optical_element=ParabolicMirror(f=f0))
     # Check that the value is the expected one in the near field
     w0_num = get_w0(laser)
-    w0_theor = wavelength * f0 / (np.pi*w0)
+    w0_theor = wavelength * f0 / (np.pi * w0)
     err = 2 * np.abs(w0_theor - w0_num) / (w0_theor + w0_num)
     assert err < 1e-3
 
+
 def test_3D_case():
     # - 3D case
     # The laser is initialized in the near field
@@ -63,6 +66,7 @@ def test_3D_case():
     laser = Laser(dim, lo, hi, npoints, gaussian_profile)
     check_parabolic_mirror(laser)
 
+
 def test_RT_case():
     # - Cylindrical case
     # The laser is initialized in the near field

From 3cf7182f08ac0a91e68cbddf9ec9a585bf075bf8 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Thu, 26 Oct 2023 13:35:29 -0700
Subject: [PATCH 14/26] Implement simplified axiparabola

---
 docs/source/tutorials/gaussian_laser.ipynb |  2 +-
 docs/source/tutorials/index.rst            |  1 +
 lasy/optical_elements/__init__.py          |  3 ++-
 lasy/optical_elements/axiparabola.py       | 21 ++++++++-------------
 4 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/docs/source/tutorials/gaussian_laser.ipynb b/docs/source/tutorials/gaussian_laser.ipynb
index e9b99f20..bc705ccd 100644
--- a/docs/source/tutorials/gaussian_laser.ipynb
+++ b/docs/source/tutorials/gaussian_laser.ipynb
@@ -5,7 +5,7 @@
    "id": "996ae31d-192b-4988-bd6b-649c833ae967",
    "metadata": {},
    "source": [
-    "# Gaussian laser pulse"
+    "# Creating and saving a Gaussian laser pulse"
    ]
   },
   {
diff --git a/docs/source/tutorials/index.rst b/docs/source/tutorials/index.rst
index ed7c4913..d03510c5 100644
--- a/docs/source/tutorials/index.rst
+++ b/docs/source/tutorials/index.rst
@@ -5,3 +5,4 @@ Tutorials
    :maxdepth: 1
 
    gaussian_laser.ipynb
+   axiparabola.ipynb
diff --git a/lasy/optical_elements/__init__.py b/lasy/optical_elements/__init__.py
index af8ba752..ac9530c9 100644
--- a/lasy/optical_elements/__init__.py
+++ b/lasy/optical_elements/__init__.py
@@ -1,3 +1,4 @@
 from .parabolic_mirror import ParabolicMirror
+from .axiparabola import AxiParabola
 
-__all__ = ["ParabolicMirror"]
+__all__ = ["ParabolicMirror", "AxiParabola"]
diff --git a/lasy/optical_elements/axiparabola.py b/lasy/optical_elements/axiparabola.py
index 06269bb4..215f0011 100644
--- a/lasy/optical_elements/axiparabola.py
+++ b/lasy/optical_elements/axiparabola.py
@@ -3,7 +3,7 @@
 from scipy.constants import c
 
 
-class AxiParabolaWithDelay(OpticalElement):
+class AxiParabola(OpticalElement):
     r"""
     Class that represents the combination of an axiparabola with
     an additional optical element that provides a radially-dependent
@@ -14,8 +14,8 @@ class AxiParabolaWithDelay(OpticalElement):
     range"). An additional radially-dependent delay is usually applied,
     in order to tune the effective group velocity on axis.
 
-    For more details, see K. Oubrerie et al, "Axiparabola: a new tool
-    for high-intensity optics", J. Opt. 24 045503 (2022)
+    For more details, see S. Smartsev et al, "Axiparabola: a long-focal-depth,
+    high-resolution mirror for broadband high-intensity lasers", Optics Letters 44, 14 (2019)
 
     Parameters
     ----------
@@ -36,17 +36,12 @@ def __init__(self, f0, delta, R):
         self.delta = delta
         self.R = R
 
-        # Function that defines the z position where rays that impinge at r focus.
-        # Assuming uniform intensity on the axiparabola, and in order to get
-        # a z-independent intensity over the focal range, we need
-        # (see Eq. 6 in Oubrerie et al.)
-        z_foc = lambda r: self.f0 + self.delta * (r / self.R) ** 2
-
-        # Solve Eq. 2 in Oubrerie et al. to find the sag function
-
     def amplitude_multiplier(self, x, y, omega):
-        # Interpolation
-        sag = np.zeros_like(x)
+        # Implement Eq. 4 in Smatsev et al.
+        r2 = x**2 + y**2
+        sag = (1./(4*self.f0)) * r2 \
+            - (self.delta / (8*self.f0**2*self.R**2) ) * r2**2 \
+            + self.delta * (self.R**2 + 8*self.f0*self.delta) / (96*self.f0**4*self.R**4) * r2**3
 
         # Calculate phase shift
         T = np.exp(-2j * (omega / c) * sag)

From 3d57f18cf0af15396957d2d53dcc79d4d5865bfe Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 26 Oct 2023 20:37:07 +0000
Subject: [PATCH 15/26] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 lasy/optical_elements/axiparabola.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/lasy/optical_elements/axiparabola.py b/lasy/optical_elements/axiparabola.py
index 215f0011..95d3945c 100644
--- a/lasy/optical_elements/axiparabola.py
+++ b/lasy/optical_elements/axiparabola.py
@@ -39,9 +39,14 @@ def __init__(self, f0, delta, R):
     def amplitude_multiplier(self, x, y, omega):
         # Implement Eq. 4 in Smatsev et al.
         r2 = x**2 + y**2
-        sag = (1./(4*self.f0)) * r2 \
-            - (self.delta / (8*self.f0**2*self.R**2) ) * r2**2 \
-            + self.delta * (self.R**2 + 8*self.f0*self.delta) / (96*self.f0**4*self.R**4) * r2**3
+        sag = (
+            (1.0 / (4 * self.f0)) * r2
+            - (self.delta / (8 * self.f0**2 * self.R**2)) * r2**2
+            + self.delta
+            * (self.R**2 + 8 * self.f0 * self.delta)
+            / (96 * self.f0**4 * self.R**4)
+            * r2**3
+        )
 
         # Calculate phase shift
         T = np.exp(-2j * (omega / c) * sag)

From 452dc5fa961644ab8afd3a47b5da3da2d0d64927 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Thu, 26 Oct 2023 13:40:57 -0700
Subject: [PATCH 16/26] Fix documentations

---
 lasy/optical_elements/axiparabola.py      | 15 +++++++++++++++
 lasy/optical_elements/optical_element.py  | 11 ++++++-----
 lasy/optical_elements/parabolic_mirror.py | 15 +++++++++++++++
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/lasy/optical_elements/axiparabola.py b/lasy/optical_elements/axiparabola.py
index 95d3945c..31f79e20 100644
--- a/lasy/optical_elements/axiparabola.py
+++ b/lasy/optical_elements/axiparabola.py
@@ -37,6 +37,21 @@ def __init__(self, f0, delta, R):
         self.R = R
 
     def amplitude_multiplier(self, x, y, omega):
+        """
+        Return the amplitude multiplier.
+
+        Parameters
+        ----------
+        x, y, omega: ndarrays of floats
+            Define points on which to evaluate the multiplier.
+            These arrays need to all have the same shape.
+
+        Returns
+        -------
+        multiplier: ndarray of complex numbers
+            Contains the value of the multiplier at the specified points
+            This array has the same shape as the arrays x, y, omega
+        """
         # Implement Eq. 4 in Smatsev et al.
         r2 = x**2 + y**2
         sag = (
diff --git a/lasy/optical_elements/optical_element.py b/lasy/optical_elements/optical_element.py
index 0142f3bd..51de8e08 100644
--- a/lasy/optical_elements/optical_element.py
+++ b/lasy/optical_elements/optical_element.py
@@ -13,11 +13,12 @@ def __init__(self):
         pass
 
     def amplitude_multiplier(self, x, y, omega):
-        """
-        Return the complex number :math:`T` with which to multiply the
-        complex amplitude of the laser just before this thin element,
-        in order to obtain the complex amplitude output laser just
-        after this thin element:
+        r"""
+        Return the amplitude multiplier :math:`T`.
+
+        This number multiplies the complex amplitude of the laser
+        just before this thin element, in order to obtain the complex
+        amplitude output laser just after this thin element:
 
         .. math::
 
diff --git a/lasy/optical_elements/parabolic_mirror.py b/lasy/optical_elements/parabolic_mirror.py
index b2b62a02..f58a884a 100644
--- a/lasy/optical_elements/parabolic_mirror.py
+++ b/lasy/optical_elements/parabolic_mirror.py
@@ -28,4 +28,19 @@ def __init__(self, f):
         self.f = f
 
     def amplitude_multiplier(self, x, y, omega):
+        """
+        Return the amplitude multiplier.
+
+        Parameters
+        ----------
+        x, y, omega: ndarrays of floats
+            Define points on which to evaluate the multiplier.
+            These arrays need to all have the same shape.
+
+        Returns
+        -------
+        multiplier: ndarray of complex numbers
+            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))

From c1a199604a27b686444219ac58f0ac3aeaeca8e5 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Fri, 27 Oct 2023 06:50:54 -0700
Subject: [PATCH 17/26] Add tutorial with axiparabola

---
 docs/source/tutorials/axiparabola.ipynb | 240 ++++++++++++++++++++++++
 1 file changed, 240 insertions(+)
 create mode 100644 docs/source/tutorials/axiparabola.ipynb

diff --git a/docs/source/tutorials/axiparabola.ipynb b/docs/source/tutorials/axiparabola.ipynb
new file mode 100644
index 00000000..72b33c5e
--- /dev/null
+++ b/docs/source/tutorials/axiparabola.ipynb
@@ -0,0 +1,240 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "996ae31d-192b-4988-bd6b-649c833ae967",
+   "metadata": {},
+   "source": [
+    "# Initializing a flying-focus laser from an axiparabola"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f25efcc9-1ed0-4f9a-8c62-da36f56ba02f",
+   "metadata": {},
+   "source": [
+    "In this example, we generate a \"flying-focus\" laser from an axiparabola. This is done by sending a super-Gaussian laser (in the near-field) onto an axiparabola and propagating it to the far field."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7f076564-8330-4af0-ad92-708e1825596f",
+   "metadata": {},
+   "source": [
+    "## Generate a super-Gaussian laser"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eb5e1364-1b92-42de-9518-dedf62be4544",
+   "metadata": {},
+   "source": [
+    "Define the physical profile, as combination of a longitudinal and transverse profile."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ed0a409c-3c02-4fc4-b41b-ef41d0fb5a3a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from lasy.laser import Laser\n",
+    "from lasy.profiles.gaussian_profile import CombinedLongitudinalTransverseProfile\n",
+    "from lasy.profiles.longitudinal import GaussianLongitudinalProfile\n",
+    "from lasy.profiles.transverse import SuperGaussianTransverseProfile\n",
+    "\n",
+    "wavelength     = 800e-9  # Laser wavelength in meters\n",
+    "polarization   = (1,0)   # Linearly polarized in the x direction\n",
+    "energy         = 1.5     # Energy of the laser pulse in joules\n",
+    "spot_size      = 1e-3    # Spot size in the near-field: millimeter-scale\n",
+    "pulse_duration = 30e-15  # Pulse duration of the laser in seconds\n",
+    "t_peak         = 0.0     # Location of the peak of the laser pulse in time\n",
+    "\n",
+    "laser_profile = CombinedLongitudinalTransverseProfile(\n",
+    "    wavelength, polarization, energy,\n",
+    "    GaussianLongitudinalProfile(wavelength, pulse_duration, t_peak),\n",
+    "    SuperGaussianTransverseProfile(spot_size, n_order=16)\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c303d3b6-2422-445b-a332-bc3f48388612",
+   "metadata": {},
+   "source": [
+    "Define the grid on which this profile is evaluated. \n",
+    "\n",
+    "**The grid needs to be wide enough to contain the millimeter-scale spot size, but also fine enough to resolve the micron-scale laser wavelength.**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "84cb6fd2-2f44-4cc9-a5c1-6dd1e0d72c67",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dimensions     = 'rt'                              # Use cylindrical geometry\n",
+    "lo             = (0,-2.5*pulse_duration)           # Lower bounds of the simulation box\n",
+    "hi             = (1.1*spot_size,2.5*pulse_duration)  # Upper bounds of the simulation box\n",
+    "num_points     = (3000, 30)                         # Number of points in each dimension\n",
+    "\n",
+    "laser = Laser(dimensions,lo,hi,num_points,laser_profile)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bb8dc634-b3c4-4448-834b-fe9acb6c9645",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "laser.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1f588477-90c0-45c3-a47c-2b57733b462f",
+   "metadata": {},
+   "source": [
+    "## Propagate the laser through the axiparabola, and to the far field.\n",
+    "\n",
+    "First, define the parameters of the axiparabola."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e3eae91f-edb7-44ac-9b92-5c1a1328a4b3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from lasy.optical_elements import AxiParabola\n",
+    "f0 = 3e-2     # Focal distance\n",
+    "delta = 1.5e-2  # Focal range\n",
+    "R = spot_size # Radius\n",
+    "axiparabola = AxiParabola( f0, delta, R )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "292c6e63-0480-4fb9-b1ad-6b679a3472d0",
+   "metadata": {},
+   "source": [
+    "Propagate the laser through the axiparabola, and for a distance z=f0 (beginning of the focal range)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3158e51d-e331-438d-90c1-5f8c63bf9841",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "laser.propagate( f0, initial_optical_element=axiparabola )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fa581cef-fc3d-414b-9f19-df8545d3b62f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "laser.show()\n",
+    "plt.ylim(-0.25e-3, 0.25e-3)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a5582bfe-ce2d-4aec-a757-03a48d4f16f2",
+   "metadata": {},
+   "source": [
+    "At this point, the laser can be saved to file, and used e.g. as input to a PIC simulation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9b8c068a-715b-400f-b6ed-7711a60c0f1f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "laser.write_to_file('flying_focus', 'h5')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "192cabcd-b96b-4bff-80b0-2c59d5a780cb",
+   "metadata": {},
+   "source": [
+    "## Check that the electric field on axis remains high over many Rayleigh ranges\n",
+    "\n",
+    "An axiparabola can maintain a high laser field over a long distance (larger than the Rayleigh length).\n",
+    "Here, we can check that the laser field remains high over several Rayleigh length."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ef0d6f2e-3aac-47a0-baef-39b42d4c5749",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import math\n",
+    "ZR = math.pi*wavelength*f0**2/spot_size**2\n",
+    "print('Rayleigh length: ', ZR)\n",
+    "assert delta > 5*ZR"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e2113f14-3ec7-4392-8d91-96f6e28dcc71",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "laser.propagate(2*ZR)\n",
+    "laser.show()\n",
+    "plt.ylim(-0.25e-3, 0.25e-3)\n",
+    "plt.title('Laser field after 2 Rayleigh range')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3cb4bd3b-ad8d-471f-a2c0-1257c1b3bd39",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "laser.propagate(2*ZR)\n",
+    "laser.show()\n",
+    "plt.ylim(-0.25e-3, 0.25e-3)\n",
+    "plt.title('Laser field after 4 Rayleigh range')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

From 50ba7465815d88766c2daaf362bdfe24d04318ef Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 27 Oct 2023 21:53:48 +0000
Subject: [PATCH 18/26] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 lasy/laser.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/lasy/laser.py b/lasy/laser.py
index 4f64bf44..86472a8b 100644
--- a/lasy/laser.py
+++ b/lasy/laser.py
@@ -140,7 +140,12 @@ def normalize(self, value, kind="energy"):
             raise ValueError(f'kind "{kind}" not recognized')
 
     def propagate(
-        self, distance, initial_optical_element=None, nr_boundary=None, backend="NP", show_progress=True
+        self,
+        distance,
+        initial_optical_element=None,
+        nr_boundary=None,
+        backend="NP",
+        show_progress=True,
     ):
         """
         Propagate the laser pulse by the distance specified.

From 343e05e6f61557ffbd37bb6f522eb2c84ce71bd5 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Mon, 15 Jul 2024 11:56:46 -0700
Subject: [PATCH 19/26] Remove axiparabola

---
 docs/source/tutorials/axiparabola.ipynb    | 240 ---------------------
 docs/source/tutorials/gaussian_laser.ipynb |   2 +-
 docs/source/tutorials/index.rst            |   1 -
 lasy/optical_elements/__init__.py          |   3 +-
 lasy/optical_elements/axiparabola.py       |  71 ------
 lasy/optical_elements/parabolic_mirror.py  |   2 +-
 6 files changed, 3 insertions(+), 316 deletions(-)
 delete mode 100644 docs/source/tutorials/axiparabola.ipynb
 delete mode 100644 lasy/optical_elements/axiparabola.py

diff --git a/docs/source/tutorials/axiparabola.ipynb b/docs/source/tutorials/axiparabola.ipynb
deleted file mode 100644
index 72b33c5e..00000000
--- a/docs/source/tutorials/axiparabola.ipynb
+++ /dev/null
@@ -1,240 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "id": "996ae31d-192b-4988-bd6b-649c833ae967",
-   "metadata": {},
-   "source": [
-    "# Initializing a flying-focus laser from an axiparabola"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "f25efcc9-1ed0-4f9a-8c62-da36f56ba02f",
-   "metadata": {},
-   "source": [
-    "In this example, we generate a \"flying-focus\" laser from an axiparabola. This is done by sending a super-Gaussian laser (in the near-field) onto an axiparabola and propagating it to the far field."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "7f076564-8330-4af0-ad92-708e1825596f",
-   "metadata": {},
-   "source": [
-    "## Generate a super-Gaussian laser"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "eb5e1364-1b92-42de-9518-dedf62be4544",
-   "metadata": {},
-   "source": [
-    "Define the physical profile, as combination of a longitudinal and transverse profile."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "ed0a409c-3c02-4fc4-b41b-ef41d0fb5a3a",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from lasy.laser import Laser\n",
-    "from lasy.profiles.gaussian_profile import CombinedLongitudinalTransverseProfile\n",
-    "from lasy.profiles.longitudinal import GaussianLongitudinalProfile\n",
-    "from lasy.profiles.transverse import SuperGaussianTransverseProfile\n",
-    "\n",
-    "wavelength     = 800e-9  # Laser wavelength in meters\n",
-    "polarization   = (1,0)   # Linearly polarized in the x direction\n",
-    "energy         = 1.5     # Energy of the laser pulse in joules\n",
-    "spot_size      = 1e-3    # Spot size in the near-field: millimeter-scale\n",
-    "pulse_duration = 30e-15  # Pulse duration of the laser in seconds\n",
-    "t_peak         = 0.0     # Location of the peak of the laser pulse in time\n",
-    "\n",
-    "laser_profile = CombinedLongitudinalTransverseProfile(\n",
-    "    wavelength, polarization, energy,\n",
-    "    GaussianLongitudinalProfile(wavelength, pulse_duration, t_peak),\n",
-    "    SuperGaussianTransverseProfile(spot_size, n_order=16)\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "c303d3b6-2422-445b-a332-bc3f48388612",
-   "metadata": {},
-   "source": [
-    "Define the grid on which this profile is evaluated. \n",
-    "\n",
-    "**The grid needs to be wide enough to contain the millimeter-scale spot size, but also fine enough to resolve the micron-scale laser wavelength.**"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "84cb6fd2-2f44-4cc9-a5c1-6dd1e0d72c67",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "dimensions     = 'rt'                              # Use cylindrical geometry\n",
-    "lo             = (0,-2.5*pulse_duration)           # Lower bounds of the simulation box\n",
-    "hi             = (1.1*spot_size,2.5*pulse_duration)  # Upper bounds of the simulation box\n",
-    "num_points     = (3000, 30)                         # Number of points in each dimension\n",
-    "\n",
-    "laser = Laser(dimensions,lo,hi,num_points,laser_profile)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "bb8dc634-b3c4-4448-834b-fe9acb6c9645",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "laser.show()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "1f588477-90c0-45c3-a47c-2b57733b462f",
-   "metadata": {},
-   "source": [
-    "## Propagate the laser through the axiparabola, and to the far field.\n",
-    "\n",
-    "First, define the parameters of the axiparabola."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "e3eae91f-edb7-44ac-9b92-5c1a1328a4b3",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from lasy.optical_elements import AxiParabola\n",
-    "f0 = 3e-2     # Focal distance\n",
-    "delta = 1.5e-2  # Focal range\n",
-    "R = spot_size # Radius\n",
-    "axiparabola = AxiParabola( f0, delta, R )"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "292c6e63-0480-4fb9-b1ad-6b679a3472d0",
-   "metadata": {},
-   "source": [
-    "Propagate the laser through the axiparabola, and for a distance z=f0 (beginning of the focal range)."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "3158e51d-e331-438d-90c1-5f8c63bf9841",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "laser.propagate( f0, initial_optical_element=axiparabola )"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "fa581cef-fc3d-414b-9f19-df8545d3b62f",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import matplotlib.pyplot as plt\n",
-    "laser.show()\n",
-    "plt.ylim(-0.25e-3, 0.25e-3)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "a5582bfe-ce2d-4aec-a757-03a48d4f16f2",
-   "metadata": {},
-   "source": [
-    "At this point, the laser can be saved to file, and used e.g. as input to a PIC simulation."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "9b8c068a-715b-400f-b6ed-7711a60c0f1f",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "laser.write_to_file('flying_focus', 'h5')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "192cabcd-b96b-4bff-80b0-2c59d5a780cb",
-   "metadata": {},
-   "source": [
-    "## Check that the electric field on axis remains high over many Rayleigh ranges\n",
-    "\n",
-    "An axiparabola can maintain a high laser field over a long distance (larger than the Rayleigh length).\n",
-    "Here, we can check that the laser field remains high over several Rayleigh length."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "ef0d6f2e-3aac-47a0-baef-39b42d4c5749",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import math\n",
-    "ZR = math.pi*wavelength*f0**2/spot_size**2\n",
-    "print('Rayleigh length: ', ZR)\n",
-    "assert delta > 5*ZR"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "e2113f14-3ec7-4392-8d91-96f6e28dcc71",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "laser.propagate(2*ZR)\n",
-    "laser.show()\n",
-    "plt.ylim(-0.25e-3, 0.25e-3)\n",
-    "plt.title('Laser field after 2 Rayleigh range')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "3cb4bd3b-ad8d-471f-a2c0-1257c1b3bd39",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "laser.propagate(2*ZR)\n",
-    "laser.show()\n",
-    "plt.ylim(-0.25e-3, 0.25e-3)\n",
-    "plt.title('Laser field after 4 Rayleigh range')"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.6"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}
diff --git a/docs/source/tutorials/gaussian_laser.ipynb b/docs/source/tutorials/gaussian_laser.ipynb
index bc705ccd..e9b99f20 100644
--- a/docs/source/tutorials/gaussian_laser.ipynb
+++ b/docs/source/tutorials/gaussian_laser.ipynb
@@ -5,7 +5,7 @@
    "id": "996ae31d-192b-4988-bd6b-649c833ae967",
    "metadata": {},
    "source": [
-    "# Creating and saving a Gaussian laser pulse"
+    "# Gaussian laser pulse"
    ]
   },
   {
diff --git a/docs/source/tutorials/index.rst b/docs/source/tutorials/index.rst
index d03510c5..ed7c4913 100644
--- a/docs/source/tutorials/index.rst
+++ b/docs/source/tutorials/index.rst
@@ -5,4 +5,3 @@ Tutorials
    :maxdepth: 1
 
    gaussian_laser.ipynb
-   axiparabola.ipynb
diff --git a/lasy/optical_elements/__init__.py b/lasy/optical_elements/__init__.py
index ac9530c9..af8ba752 100644
--- a/lasy/optical_elements/__init__.py
+++ b/lasy/optical_elements/__init__.py
@@ -1,4 +1,3 @@
 from .parabolic_mirror import ParabolicMirror
-from .axiparabola import AxiParabola
 
-__all__ = ["ParabolicMirror", "AxiParabola"]
+__all__ = ["ParabolicMirror"]
diff --git a/lasy/optical_elements/axiparabola.py b/lasy/optical_elements/axiparabola.py
deleted file mode 100644
index 31f79e20..00000000
--- a/lasy/optical_elements/axiparabola.py
+++ /dev/null
@@ -1,71 +0,0 @@
-from .optical_element import OpticalElement
-import numpy as np
-from scipy.constants import c
-
-
-class AxiParabola(OpticalElement):
-    r"""
-    Class that represents the combination of an axiparabola with
-    an additional optical element that provides a radially-dependent
-    delay (e.g. an optical echelon) to tune the group velocity.
-
-    The rays that impinge the axiparabola at different radii are focused
-    to different positions on the axis (resulting in an extended "focal
-    range"). An additional radially-dependent delay is usually applied,
-    in order to tune the effective group velocity on axis.
-
-    For more details, see S. Smartsev et al, "Axiparabola: a long-focal-depth,
-    high-resolution mirror for broadband high-intensity lasers", Optics Letters 44, 14 (2019)
-
-    Parameters
-    ----------
-    f0: float (in meter)
-        The focal distance, i.e. the distance, from the axiparabola,
-        where the focal range starts.
-
-    delta: float (in meter)
-        The length of the focal range.
-
-    R: float (in meter)
-        The radius of the axiparabola. Rays coming from r=0 focus
-        at z=f0 ; rays coming from r=R focus at z=f0+delta
-    """
-
-    def __init__(self, f0, delta, R):
-        self.f0 = f0
-        self.delta = delta
-        self.R = R
-
-    def amplitude_multiplier(self, x, y, omega):
-        """
-        Return the amplitude multiplier.
-
-        Parameters
-        ----------
-        x, y, omega: ndarrays of floats
-            Define points on which to evaluate the multiplier.
-            These arrays need to all have the same shape.
-
-        Returns
-        -------
-        multiplier: ndarray of complex numbers
-            Contains the value of the multiplier at the specified points
-            This array has the same shape as the arrays x, y, omega
-        """
-        # Implement Eq. 4 in Smatsev et al.
-        r2 = x**2 + y**2
-        sag = (
-            (1.0 / (4 * self.f0)) * r2
-            - (self.delta / (8 * self.f0**2 * self.R**2)) * r2**2
-            + self.delta
-            * (self.R**2 + 8 * self.f0 * self.delta)
-            / (96 * self.f0**4 * self.R**4)
-            * r2**3
-        )
-
-        # Calculate phase shift
-        T = np.exp(-2j * (omega / c) * sag)
-        # Remove intensity beyond R
-        T[x**2 + y**2 > self.R**2] = 0
-
-        return T
diff --git a/lasy/optical_elements/parabolic_mirror.py b/lasy/optical_elements/parabolic_mirror.py
index f58a884a..2bd21f00 100644
--- a/lasy/optical_elements/parabolic_mirror.py
+++ b/lasy/optical_elements/parabolic_mirror.py
@@ -5,7 +5,7 @@
 
 class ParabolicMirror(OpticalElement):
     r"""
-    Derived class for a parabolic mirror.
+    Class for a parabolic mirror.
 
     More precisely, the amplitude multiplier corresponds to:
 

From c63a1222407f6ab2093553d7bc9211ef48a4338f Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Mon, 15 Jul 2024 12:16:48 -0700
Subject: [PATCH 20/26] Add new API

---
 lasy/laser.py                  | 67 +++++++++++++++++++++-------------
 tests/test_parabolic_mirror.py |  3 +-
 2 files changed, 44 insertions(+), 26 deletions(-)

diff --git a/lasy/laser.py b/lasy/laser.py
index 4a3e0285..94962d0c 100644
--- a/lasy/laser.py
+++ b/lasy/laser.py
@@ -160,14 +160,48 @@ def normalize(self, value, kind="energy"):
         else:
             raise ValueError(f'kind "{kind}" not recognized')
 
-    def propagate(
-        self,
-        distance,
-        initial_optical_element=None,
-        nr_boundary=None,
-        backend="NP",
-        show_progress=True,
-    ):
+    def apply_optics( self, optical_element ):
+        """
+        Propagate the laser pulse through a thin optical element.
+
+        Parameter
+        ---------
+        optical_element: an :class:`.OpticalElement` object (optional)
+            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]
+        omega = 2 * np.pi * np.fft.fftfreq(Nt, dt) + omega0
+
+        # Apply optical element
+        if self.dim == "rt":
+            r, w = np.meshgrid(self.grid.axes[0], omega, indexing="ij")
+            # The line below assumes that amplitude_multiplier
+            # is cylindrically-symmetric, hence we pass
+            # `r` as `x` and 0 as `y`
+            multiplier = optical_element.amplitude_multiplier(r, 0, w)
+            for i_m in range(self.grid.azimuthal_modes.size):
+                field_fft[i_m, :, :] *= multiplier
+        else:
+            x, y, w = np.meshgrid(
+                self.grid.axes[0], self.grid.axes[1], omega, indexing="ij"
+            )
+            field_fft *= optical_element.amplitude_multiplier(x, y, w)
+
+        # Transform field from frequency to temporal domain
+        self.grid.field[:, :, :] = np.fft.fft(
+            field_fft, axis=time_axis_indx, norm="backward"
+        )
+
+
+    def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True):
         """
         Propagate the laser pulse by the distance specified.
 
@@ -230,16 +264,6 @@ def propagate(
         omega_shape = (1, 1, self.grid.field.shape[time_axis_indx])
 
         if self.dim == "rt":
-            # Apply optical element
-            if initial_optical_element is not None:
-                r, w = np.meshgrid(self.grid.axes[0], omega, indexing="ij")
-                # The line below assumes that amplitude_multiplier
-                # is cylindrically-symmetric, hence we pass
-                # `r` as `x` and 0 as `y`
-                multiplier = initial_optical_element.amplitude_multiplier(r, 0, w)
-                for i_m in range(self.grid.azimuthal_modes.size):
-                    field_fft[i_m, :, :] *= multiplier
-
             # Construct the propagator (check if exists)
             if not hasattr(self, "prop"):
                 spatial_axes = (self.grid.axes[0],)
@@ -265,13 +289,6 @@ def propagate(
                 )
                 field_fft[i_m, :, :] = np.transpose(transform_data).copy()
         else:
-            # Apply optical element
-            if initial_optical_element is not None:
-                x, y, w = np.meshgrid(
-                    self.grid.axes[0], self.grid.axes[1], omega, indexing="ij"
-                )
-                field_fft *= initial_optical_element.amplitude_multiplier(x, y, w)
-
             # Construct the propagator (check if exists)
             if not hasattr(self, "prop"):
                 Nx, Ny, Nt = self.grid.field.shape
diff --git a/tests/test_parabolic_mirror.py b/tests/test_parabolic_mirror.py
index ae3fa062..bfee3568 100644
--- a/tests/test_parabolic_mirror.py
+++ b/tests/test_parabolic_mirror.py
@@ -47,7 +47,8 @@ def get_w0(laser):
 def check_parabolic_mirror(laser):
     # Propagate laser after parabolic mirror + vacuum
     f0 = 8.0  # focal distance in m
-    laser.propagate(f0, initial_optical_element=ParabolicMirror(f=f0))
+    laser.apply_optics( ParabolicMirror(f=f0) )
+    laser.propagate(f0)
     # Check that the value is the expected one in the near field
     w0_num = get_w0(laser)
     w0_theor = wavelength * f0 / (np.pi * w0)

From 12d6ee8ce25981060a7c84878843d9769a208a9e Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Mon, 15 Jul 2024 12:26:43 -0700
Subject: [PATCH 21/26] Add missing documentation

---
 docs/source/api/optical_elements/index.rst            | 7 +++++++
 docs/source/api/optical_elements/parabolic_mirror.rst | 5 +++++
 2 files changed, 12 insertions(+)
 create mode 100644 docs/source/api/optical_elements/index.rst
 create mode 100644 docs/source/api/optical_elements/parabolic_mirror.rst

diff --git a/docs/source/api/optical_elements/index.rst b/docs/source/api/optical_elements/index.rst
new file mode 100644
index 00000000..8be108e0
--- /dev/null
+++ b/docs/source/api/optical_elements/index.rst
@@ -0,0 +1,7 @@
+Optical elements
+================
+
+.. toctree::
+   :maxdepth: 1
+
+   parabolic_mirror
diff --git a/docs/source/api/optical_elements/parabolic_mirror.rst b/docs/source/api/optical_elements/parabolic_mirror.rst
new file mode 100644
index 00000000..e20685d9
--- /dev/null
+++ b/docs/source/api/optical_elements/parabolic_mirror.rst
@@ -0,0 +1,5 @@
+Parabolic mirror
+================
+
+.. autoclass:: lasy.optical_elements.ParabolicMirror
+    :members:

From 7c082c8f4f93ba07723b67999d28ff2b873aba73 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 15 Jul 2024 19:26:04 +0000
Subject: [PATCH 22/26] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 lasy/laser.py                  | 3 +--
 tests/test_parabolic_mirror.py | 2 +-
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/lasy/laser.py b/lasy/laser.py
index 94962d0c..eb0bbc3c 100644
--- a/lasy/laser.py
+++ b/lasy/laser.py
@@ -160,7 +160,7 @@ def normalize(self, value, kind="energy"):
         else:
             raise ValueError(f'kind "{kind}" not recognized')
 
-    def apply_optics( self, optical_element ):
+    def apply_optics(self, optical_element):
         """
         Propagate the laser pulse through a thin optical element.
 
@@ -200,7 +200,6 @@ def apply_optics( self, optical_element ):
             field_fft, axis=time_axis_indx, norm="backward"
         )
 
-
     def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True):
         """
         Propagate the laser pulse by the distance specified.
diff --git a/tests/test_parabolic_mirror.py b/tests/test_parabolic_mirror.py
index bfee3568..93dab003 100644
--- a/tests/test_parabolic_mirror.py
+++ b/tests/test_parabolic_mirror.py
@@ -47,7 +47,7 @@ def get_w0(laser):
 def check_parabolic_mirror(laser):
     # Propagate laser after parabolic mirror + vacuum
     f0 = 8.0  # focal distance in m
-    laser.apply_optics( ParabolicMirror(f=f0) )
+    laser.apply_optics(ParabolicMirror(f=f0))
     laser.propagate(f0)
     # Check that the value is the expected one in the near field
     w0_num = get_w0(laser)

From 0c09ed8312ede0f7f2b56756eb431c42b670c74c Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Tue, 16 Jul 2024 01:33:20 -0700
Subject: [PATCH 23/26] Update lasy/optical_elements/optical_element.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Maxence Thévenet <maxence.thevenet@desy.de>
---
 lasy/optical_elements/optical_element.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lasy/optical_elements/optical_element.py b/lasy/optical_elements/optical_element.py
index 51de8e08..1a1251b7 100644
--- a/lasy/optical_elements/optical_element.py
+++ b/lasy/optical_elements/optical_element.py
@@ -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.zeros_like(x, dtype="complex128")
+        return np.ones_like(x, dtype="complex128")

From 411747b14855f02d69c3aa8eb0343a70d405d4a4 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Tue, 16 Jul 2024 01:35:34 -0700
Subject: [PATCH 24/26] Apply suggestions from code review
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Maxence Thévenet <maxence.thevenet@desy.de>
---
 lasy/laser.py | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/lasy/laser.py b/lasy/laser.py
index eb0bbc3c..4d937160 100644
--- a/lasy/laser.py
+++ b/lasy/laser.py
@@ -184,7 +184,7 @@ def apply_optics(self, optical_element):
         if self.dim == "rt":
             r, w = np.meshgrid(self.grid.axes[0], omega, indexing="ij")
             # The line below assumes that amplitude_multiplier
-            # is cylindrically-symmetric, hence we pass
+            # is cylindrically symmetric, hence we pass
             # `r` as `x` and 0 as `y`
             multiplier = optical_element.amplitude_multiplier(r, 0, w)
             for i_m in range(self.grid.azimuthal_modes.size):
@@ -209,11 +209,6 @@ def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True
         distance : scalar
             Distance by which the laser pulse should be propagated
 
-        initial_optical_element: an :class:`.OpticalElement` object (optional)
-            Represents a thin optical element, through which the laser
-            propagates, before propagating for `distance` in free space.
-            If this is `None`, no optical element is used.
-
         nr_boundary : integer (optional)
             Number of cells at the end of radial axis, where the field
             will be attenuated (to assert proper Hankel transform).

From 0d76c51fb65b01dd007345598268f896981f040c Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Tue, 16 Jul 2024 01:40:10 -0700
Subject: [PATCH 25/26] Avoid confusion between omega and w

---
 lasy/laser.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/lasy/laser.py b/lasy/laser.py
index 4d937160..23704ff4 100644
--- a/lasy/laser.py
+++ b/lasy/laser.py
@@ -178,22 +178,22 @@ def apply_optics(self, optical_element):
         dt = self.grid.dx[time_axis_indx]
         omega0 = self.profile.omega0
         Nt = self.grid.field.shape[time_axis_indx]
-        omega = 2 * np.pi * np.fft.fftfreq(Nt, dt) + omega0
+        omega_1d = 2 * np.pi * np.fft.fftfreq(Nt, dt) + omega0
 
         # Apply optical element
         if self.dim == "rt":
-            r, w = np.meshgrid(self.grid.axes[0], omega, indexing="ij")
+            r, omega = np.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`
-            multiplier = optical_element.amplitude_multiplier(r, 0, w)
+            multiplier = optical_element.amplitude_multiplier(r, 0, omega)
             for i_m in range(self.grid.azimuthal_modes.size):
                 field_fft[i_m, :, :] *= multiplier
         else:
-            x, y, w = np.meshgrid(
-                self.grid.axes[0], self.grid.axes[1], omega, indexing="ij"
+            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, w)
+            field_fft *= optical_element.amplitude_multiplier(x, y, omega)
 
         # Transform field from frequency to temporal domain
         self.grid.field[:, :, :] = np.fft.fft(

From 66178cf87e0d49dbba1f96a530ece07b51709563 Mon Sep 17 00:00:00 2001
From: Remi Lehe <remi.lehe@normalesup.org>
Date: Tue, 16 Jul 2024 01:44:07 -0700
Subject: [PATCH 26/26] Add comment explaining why we can multiply each
 azimuthal component

---
 lasy/laser.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/lasy/laser.py b/lasy/laser.py
index 23704ff4..bdb35e76 100644
--- a/lasy/laser.py
+++ b/lasy/laser.py
@@ -187,6 +187,11 @@ def apply_optics(self, optical_element):
             # is cylindrically symmetric, hence we pass
             # `r` as `x` and 0 as `y`
             multiplier = optical_element.amplitude_multiplier(r, 0, omega)
+            # The azimuthal modes are the components of the Fourier transform
+            # along theta (FT_theta). Because the multiplier is assumed to be
+            # cylindrically symmetric (i.e. theta-independent):
+            # 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
         else: