Add thin optical elements in lasy

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.

docs/source/api/optical_elements/index.rst

@@ -0,0 +1,7 @@
+Optical elements
+================
+
+.. toctree::
+   :maxdepth: 1
+
+   parabolic_mirror

docs/source/api/optical_elements/parabolic_mirror.rst

@@ -0,0 +1,5 @@
+Parabolic mirror
+================
+
+.. autoclass:: lasy.optical_elements.ParabolicMirror
+    :members:

lasy/laser.py

@@ -160,6 +160,51 @@ def normalize(self, value, kind="energy"):
         else:
             raise ValueError(f'kind "{kind}" not recognized')
 
+    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_1d = 2 * np.pi * np.fft.fftfreq(Nt, dt) + omega0
+
+        # Apply optical element
+        if self.dim == "rt":
+            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, 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:
+            x, y, omega = np.meshgrid(
+                self.grid.axes[0], self.grid.axes[1], omega_1d, indexing="ij"
+            )
+            field_fft *= optical_element.amplitude_multiplier(x, y, omega)
+
+        # Transform field from frequency to temporal domain
+        self.grid.field[:, :, :] = np.fft.fft(
+            field_fft, axis=time_axis_indx, norm="backward"
+        )
+
     def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True):
         """
         Propagate the laser pulse by the distance specified.
@@ -173,6 +218,7 @@ def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True
             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).
         show_progress : bool (optional)
@@ -212,6 +258,7 @@ def propagate(self, distance, nr_boundary=None, backend="NP", show_progress=True
         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])
 

lasy/optical_elements/__init__.py

@@ -0,0 +1,3 @@
+from .parabolic_mirror import ParabolicMirror
+
+__all__ = ["ParabolicMirror"]

lasy/optical_elements/optical_element.py

@@ -0,0 +1,41 @@
+import numpy as np
+
+
+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):
+        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::
+
+            \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.ones_like(x, dtype="complex128")

lasy/optical_elements/parabolic_mirror.py

@@ -0,0 +1,46 @@
+from .optical_element import OpticalElement
+import numpy as np
+from scipy.constants import c
+
+
+class ParabolicMirror(OpticalElement):
+    r"""
+    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}/2cf)
+
+    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 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))

tests/test_parabolic_mirror.py

@@ -0,0 +1,80 @@
+# -*- 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 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 = 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
+    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.0  # focal distance in m
+    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)
+    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
+    dim = "xyt"
+    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_profile)
+    check_parabolic_mirror(laser)
+
+
+def test_RT_case():
+    # - Cylindrical case
+    # The laser is initialized in the near field
+    dim = "rt"
+    lo = (0e-6, -60e-15)
+    hi = (15e-3, +60e-15)
+    npoints = (750, 100)
+
+    laser = Laser(dim, lo, hi, npoints, gaussian_profile)
+    check_parabolic_mirror(laser)