Enable initialisation and estimation of optional spatio-temperal coup…

…ling parameters to gaussian profile (#307)

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

docs/source/tutorials/axiparabola.ipynb

@@ -40,7 +40,7 @@
    "outputs": [],
    "source": [
     "from lasy.laser import Laser\n",
-    "from lasy.profiles.gaussian_profile import CombinedLongitudinalTransverseProfile\n",
+    "from lasy.profiles.combined_profile import CombinedLongitudinalTransverseProfile\n",
     "from lasy.profiles.longitudinal import GaussianLongitudinalProfile\n",
     "from lasy.profiles.transverse import SuperGaussianTransverseProfile\n",
     "\n",

lasy/profiles/gaussian_profile.py

@@ -1,9 +1,9 @@
-from . import CombinedLongitudinalTransverseProfile
-from .longitudinal import GaussianLongitudinalProfile
-from .transverse import GaussianTransverseProfile
+import numpy as np
 
+from .profile import Profile
 
-class GaussianProfile(CombinedLongitudinalTransverseProfile):
+
+class GaussianProfile(Profile):
     r"""
     Class for the analytic profile of a Gaussian laser pulse.
 
@@ -21,6 +21,7 @@ class GaussianProfile(CombinedLongitudinalTransverseProfile):
     :math:`\boldsymbol{x}_\perp` is the transverse coordinate (orthogonal
     to the propagation direction). The other parameters in this formula
     are defined below.
+    This profile also supports some chirp parameters that are omitted in the expression above for clarity.
 
     Parameters
     ----------
@@ -61,6 +62,20 @@ class GaussianProfile(CombinedLongitudinalTransverseProfile):
     z_foc : float (in meter), optional
         Position of the focal plane. (The laser pulse is initialized at `z=0`.)
 
+    phi2 : float (in second^2), optional (default: '0')
+        The group-delay dispersion defined as :math:`\phi^{(2)} = \frac{dt_0}{d\omega}`. Here :math:`t_0` is the temporal position of this frequency component.
+
+    beta : float (in second), optional (default: '0')
+        The angular dispersion defined as :math:`\beta = \frac{d\theta_0}{d\omega}`. Here :math:`\theta_0` is the propagation angle of this frequency component.
+
+    zeta : float (in meter * second), optional (default: '0')
+        A spatial chirp defined as :math:`\zeta = \frac{dx_0}{d\omega}`. Here :math:`x_0` is the transverse beam center position of this frequency component.
+        The definitions of beta, phi2, and zeta are taken from [S. Akturk et al., Optics Express 12, 4399 (2004)].
+
+    stc_theta : float (in radian), optional (default: '0')
+        Transverse direction along which there are chirps and spatio-temporal couplings.
+        A value of 0 corresponds to the x-axis.
+
     Examples
     --------
     >>> import matplotlib.pyplot as plt
@@ -104,12 +119,93 @@ class GaussianProfile(CombinedLongitudinalTransverseProfile):
     """
 
     def __init__(
-        self, wavelength, pol, laser_energy, w0, tau, t_peak, cep_phase=0, z_foc=0
+        self,
+        wavelength,
+        pol,
+        laser_energy,
+        w0,
+        tau,
+        t_peak,
+        cep_phase=0,
+        z_foc=0,
+        phi2=0,
+        beta=0,
+        zeta=0,
+        stc_theta=0,
     ):
-        super().__init__(
-            wavelength,
-            pol,
-            laser_energy,
-            GaussianLongitudinalProfile(wavelength, tau, t_peak, cep_phase),
-            GaussianTransverseProfile(w0, z_foc, wavelength),
+        super().__init__(wavelength, pol)
+        self.laser_energy = laser_energy
+        self.w0 = w0
+        self.tau = tau
+        self.t_peak = t_peak
+        self.cep_phase = cep_phase
+        self.z_foc = z_foc
+        self.z_foc_over_zr = z_foc * wavelength / (np.pi * w0**2)
+        self.phi2 = phi2
+        self.beta = beta
+        self.zeta = zeta
+        self.stc_theta = stc_theta
+
+    def evaluate(self, x, y, t):
+        """
+        Return the longitudinal envelope.
+
+        Parameters
+        ----------
+        t : ndarrays of floats
+            Define longitudinal points on which to evaluate the envelope
+
+        x,y : ndarrays of floats
+            Define transverse points on which to evaluate the envelope
+
+        Returns
+        -------
+        envelope : ndarray of complex numbers
+            Contains the value of the longitudinal envelope at the
+            specified points. This array has the same shape as the array t.
+        """
+        inv_tau2 = self.tau ** (-2)
+        inv_complex_waist_2 = 1.0 / (
+            self.w0**2 * (1.0 + 2.0j * self.z_foc / (self.k0 * self.w0**2))
         )
+        stretch_factor = (
+            1
+            + 4.0
+            * (-self.zeta + self.beta * self.z_foc)
+            * inv_tau2
+            * (-self.zeta + self.beta * self.z_foc)
+            * inv_complex_waist_2
+            + 2.0j * (self.phi2 - self.beta**2 * self.k0 * self.z_foc) * inv_tau2
+        )
+        stc_exponent = (
+            1.0
+            / stretch_factor
+            * inv_tau2
+            * (
+                t
+                - self.t_peak
+                - self.beta
+                * self.k0
+                * (x * np.cos(self.stc_theta) + y * np.sin(self.stc_theta))
+                - 2.0j
+                * (x * np.cos(self.stc_theta) + y * np.sin(self.stc_theta))
+                * (-self.zeta - self.beta * self.z_foc)
+                * inv_complex_waist_2
+            )
+            ** 2
+        )
+        # Term for wavefront curvature + Gouy phase
+        diffract_factor = 1.0 - 1j * self.z_foc_over_zr
+        # Calculate the argument of the complex exponential
+        exp_argument = -(x**2 + y**2) / (self.w0**2 * diffract_factor)
+        # Get the profile
+        envelope = (
+            np.exp(
+                -stc_exponent
+                + exp_argument
+                + 1.0j * (self.cep_phase + self.omega0 * self.t_peak)
+            )
+            / diffract_factor
+        )
+
+        return envelope

lasy/profiles/longitudinal/gaussian_profile.py

@@ -31,7 +31,7 @@ class GaussianLongitudinalProfile(LongitudinalProfile):
     cep_phase : float (in radian), optional
         The Carrier Enveloppe Phase (CEP), i.e. :math:`\phi_{cep}`
         in the above formula (i.e. the phase of the laser
-        oscillation, at the time where the laser envelope is maximum)
+        oscillation, at the time where the laser envelope is maximum).
     """
 
     def __init__(self, wavelength, tau, t_peak, cep_phase=0):
@@ -47,7 +47,7 @@ def evaluate(self, t):
         Parameters
         ----------
         t : ndarrays of floats
-            Define points on which to evaluate the envelope
+            Define longitudinal points on which to evaluate the envelope
 
         Returns
         -------
@@ -59,5 +59,4 @@ def evaluate(self, t):
             -((t - self.t_peak) ** 2) / self.tau**2
             + 1.0j * (self.cep_phase + self.omega0 * self.t_peak)
         )
-
         return envelope

lasy/profiles/profile.py

@@ -38,6 +38,7 @@ def __init__(self, wavelength, pol):
         self.pol = np.array([pol[0] / norm_pol, pol[1] / norm_pol])
         self.lambda0 = wavelength
         self.omega0 = 2 * np.pi * c / self.lambda0
+        self.k0 = 2.0 * np.pi / wavelength
 
     def evaluate(self, x, y, t):
         """