Enable initialisation and estimation of optional spatio-temperal coupling parameters to gaussian profile

lasy/profiles/combined_profile.py

@@ -68,5 +68,7 @@ def evaluate(self, x, y, t):
             Contains the value of the envelope at the specified points
             This array has the same shape as the arrays x, y, t
         """
-        envelope = self.trans_profile.evaluate(x, y) * self.long_profile.evaluate(t)
+        envelope = self.trans_profile.evaluate(x, y) * self.long_profile.evaluate(
+            t, x, y
+        )
         return envelope

lasy/profiles/gaussian_profile.py

@@ -61,6 +61,34 @@ class GaussianProfile(CombinedLongitudinalTransverseProfile):
     z_foc : float (in meter), optional
         Position of the focal plane. (The laser pulse is initialized at `z=0`.)
 
+    beta : float (in second), optional
+        The angular dispersion parameterized by
+
+        .. math::
+
+            \beta = \frac{d\theta_0}{d\omega}
+        Here :math:`\theta_0` is the propagation angle of this component.
+
+    phi2 : float (in second^2), optional (default '0')
+        The group-delay dispertion parameterized by
+
+        .. math::
+
+            \phi^{(2)} = \frac{dt}{d\omega}
+
+    zeta : float (in meter * second) optional (default '0')
+        The spatio-chirp parameterized by
+
+        .. math::
+
+         \zeta = \frac{x_0}{d\omega}
+        Here :math:`x_0` is the beam center position.
+
+    stc_theta :  float (in rad) optional (default '0')
+        Transeverse direction along which spatio-temperal field couples.
+        0 is along x axis.
+
+
     Examples
     --------
     >>> import matplotlib.pyplot as plt
@@ -104,12 +132,35 @@ 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),
+            GaussianLongitudinalProfile(
+                wavelength,
+                tau,
+                t_peak,
+                cep_phase,
+                beta=beta,
+                phi2=phi2,
+                zeta=zeta,
+                stc_theta=stc_theta,
+                w0=w0,
+                z_foc=z_foc,
+            ),
             GaussianTransverseProfile(w0, z_foc, wavelength),
         )

lasy/profiles/longitudinal/gaussian_profile.py

@@ -28,36 +28,122 @@ class GaussianLongitudinalProfile(LongitudinalProfile):
         The time at which the laser envelope reaches its maximum amplitude,
         i.e. :math:`t_{peak}` in the above formula.
 
-    cep_phase : float (in radian), optional
+    cep_phase : float (in radian), optional(default '0')
         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).
+
+    beta : float (in second), optional
+        The angular dispersion parameterized by
+        .. math::
+            \beta = \frac{d\theta_0}{d\omega}
+        Here :math:`\theta_0` is the propagation angle of this component.
+
+    phi2 : float (in second^2), optional (default '0')
+        The group-delay dispertion parameterized by
+        .. math::
+            \phi^{(2)} = \frac{dt}{d\omega}
+
+    zeta : float (in meter * second) optional (defalut '0')
+        The spatio-chirp parameterized by
+        .. math::
+         \zeta = \frac{x_0}{d\omega}
+        Here :math:`x_0` is the beam center position.
+
+    stc_theta :  float (in rad) optional (default '0')
+        Transeverse direction along which spatio-temperal field couples.
+        0 is along x axis.
+
+    z_foc : float (in meter), necessary if beta is not 0
+        Position of the focal plane. (The laser pulse is initialized at
+        ``z=0``.)
+
+     w0 : float (in meter), necessary if beta is not 0
+        The waist of the laser pulse.
+
+    All those above STC units and definiations are taken from `S. Akturk
+    et al., Optics Express 12, 4399 (2004) <https://doi.org/10.1364/OPEX.12.004399>`__.
     """
 
-    def __init__(self, wavelength, tau, t_peak, cep_phase=0):
+    def __init__(
+        self,
+        wavelength,
+        tau,
+        t_peak,
+        cep_phase=0,
+        beta=0,
+        phi2=0,
+        zeta=0,
+        stc_theta=0,
+        w0=0,
+        z_foc=0,
+    ):
         super().__init__(wavelength)
         self.tau = tau
         self.t_peak = t_peak
         self.cep_phase = cep_phase
+        self.beta = beta
+        self.phi2 = phi2
+        self.zeta = zeta
+        self.w0 = w0
+        self.stc_theta = stc_theta
+        if z_foc == 0:
+            self.z_foc_over_zr = 0
+        else:
+            assert (
+                wavelength is not None
+            ), "You need to pass the wavelength, when `z_foc` is non-zero."
+            self.z_foc_over_zr = z_foc * wavelength / (np.pi * w0**2)
 
-    def evaluate(self, t):
+    def evaluate(self, t, x=0, y=0):
         """
         Return the longitudinal envelope.
 
         Parameters
         ----------
         t : ndarrays of floats
-            Define points on which to evaluate the envelope
+            Define longitudinal points on which to evaluate the envelope
+
+        x,y : ndarrays of floats, necessray if spatio-temperal coulping exists
+            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_over_zr / (self.k0 * self.w0**2))
+        )
+        stretch_factor = (
+            1
+            + 4.0
+            * (self.zeta + self.beta * self.z_foc_over_zr * inv_tau2)
+            * (self.zeta + self.beta * self.z_foc_over_zr * inv_complex_waist_2)
+            + 2.0j
+            * (self.phi2 - self.beta**2 * self.k0 * self.z_foc_over_zr)
+            * 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_over_zr)
+                * inv_complex_waist_2
+            )
+            ** 2
+        )
         envelope = np.exp(
-            -((t - self.t_peak) ** 2) / self.tau**2
-            + 1.0j * (self.cep_phase + self.omega0 * self.t_peak)
+            -stc_exponent + 1.0j * (self.cep_phase + self.omega0 * (t - self.t_peak))
         )
-
         return envelope

lasy/profiles/longitudinal/longitudinal_profile.py

@@ -13,8 +13,9 @@ class LongitudinalProfile(object):
     def __init__(self, wavelength):
         self.lambda0 = wavelength
         self.omega0 = 2 * pi * c / self.lambda0
+        self.k0 = 2.0 * pi / wavelength
 
-    def evaluate(self, t):
+    def evaluate(self, t, x=0, y=0):
         """
         Return the longitudinal envelope.
 

lasy/utils/laser_utils.py

@@ -873,3 +873,125 @@
         field = np.moveaxis(prop.z2t(transform_data, t_axis, z0=z0, t0=t0), 0, -1)
         field *= np.exp(1j * (z0 / c + t_axis) * omega0)
         grid.set_temporal_field(field)
+
+
+def get_STC(dim, grid, tau, w0, k0):
+    """
+    Calculate the spatio-temperal coupling factors of the laser.
+
+    Parameters
+    ----------
+    laser:
+    dim : string
+        Dimensionality of the array. Options are:
+        - 'xyt': The laser pulse is represented on a 3D grid:
+                 Cartesian (x,y) transversely, and temporal (t) longitudinally.
+        - 'rt' : The laser pulse is represented on a 2D grid:
+                 Cylindrical (r) transversely, and temporal (t) longitudinally.
+
+    grid : a Grid object.
+        It contains an ndarray (V/m) with
+        the value of the envelope field and the associated metadata
+        that defines the points at which the laser is defined.
+
+    tau : scalar
+        Duration of the laser pulse in s.
+
+    w0 : scalar
+        Waist of laser in m.
+
+    k0 : scalar
+        Wavenumber of the field
+
+    Return
+    ----------
+    STC_fac : dict of floats
+        A dictionary of floats corresponding to the STC factors. The keys are:
+            Phi2: Group-delayed dispersion in :math:`\Phi^{(2)}=d(\omega_0)/dt`
+            phi2: Group-delayed dispersion in :math:`\phi^{(2)}=dt_0/d(\omega)`
+            nu: Spatio-chirp in :math:`\nu=d(\omega_0)/dx`
+            zeta: Spatio-chirp in :math:`\zeta=dx_0/d(\omega_0)`
+            stc_theta_zeta: The direction of the linear spatial chirp on xoy plane\
+            in rad (0 is along x)
+            beta: Angular dispersion in :math:` \beta = d\theta_0/d\omega`
+            pft: Pulse front tilt in :math:` p=dt/dx`
+            stc_theta_beta: The direction of the linear angular chirp on xoy plane\
+            in rad (0 is along x)
+    All those above units and definiations are taken from `S. Akturk et al., Optics\
+    Express 12, 4399 (2004) <https://doi.org/10.1364/OPEX.12.004399>`__.
+    """
+    # Initialise the returned dictionary
+    STC_fac = {
+        "Phi2": 0,
+        "phi2": 0,
+        "nu": 0,
+        "zeta": 0,
+        "stc_theta_zeta": 0,
+        "beta": 0,
+        "pft": 0,
+        "stc_theta_beta": 0,
+    }
+    env = grid.get_temporal_field()
+    env_abs = np.abs(env)
+    phi_envelop = np.unwrap(np.array(np.arctan2(env.imag, env.real)), axis=2)
+    pphi_pz = (np.diff(phi_envelop, axis=2)) / (grid.dx[-1])
+    # Calculate goup-delayed dispersion
+    pphi_pz2 = (np.diff(pphi_pz, axis=2)) / (grid.dx[-1])
+    STC_fac["Phi2"] = np.sum(pphi_pz2 * env_abs[:, :, : env_abs.shape[2] - 2]) / np.sum(
+        env_abs[:, :, : env_abs.shape[2] - 2]
+    )
+    STC_fac["phi2"] = np.max(np.roots([4 * Phi2, -4, tau**4 * Phi2]))
+    # Calculate spatio- and angular dispersion
+    if dim == "rt":
+        pphi_pzpr = (np.diff(pphi_pz, axis=1)) / grid.dx[0]
+        STC_fac["nu"] = np.sum(
+            pphi_pzpr * env_abs[:, : env_abs.shape[1] - 1, : env_abs.shape[2] - 1]
+        ) / np.sum(env_abs[:, : env_abs.shape[1] - 1, : env_abs.shape[2] - 1])
+        # Transfer the unit from nu to zeta
+        STC_fac["zeta"] = np.min(np.roots([4 * nu, -4, nu * w0**2 * tau**2]))
+        # No angular dispersion in 2D and the direction of spatio-chirp is certain
+        return STC_fac
+    if dim == "xyt":
+        pphi_pzpy = (np.diff(pphi_pz, axis=1)) / grid.dx[1]
+        pphi_pzpx = (np.diff(pphi_pz, axis=0)) / grid.dx[0]
+        # Calculate the STC angle in XOY for spatio coupling
+        theta = np.arctan2(
+            pphi_pzpy[: env_abs.shape[0] - 1, : env_abs.shape[1] - 1, :],
+            pphi_pzpx[: env_abs.shape[0] - 1, : env_abs.shape[1] - 1, :],
+        )
+        STC_fac["stc_theta_zeta"] = np.sum(
+            theta
+            * env_abs[
+                : env_abs.shape[0] - 1, : env_abs.shape[1] - 1, : env_abs.shape[2] - 1
+            ]
+        ) / np.sum(
+            env_abs[
+                : env_abs.shape[0] - 1, : env_abs.shape[1] - 1, : env_abs.shape[2] - 1
+            ]
+        )
+        pphi_pzpr = (
+            pphi_pzpy[: env_abs.shape[0] - 1, : env_abs.shape[1] - 1, :] ** 2
+            + pphi_pzpx[:: env_abs.shape[0] - 1, : env_abs.shape[1] - 1, :] ** 2
+        ) ** 0.5
+        STC_fac["nu"] = np.sum(
+            pphi_pzpr
+            * env_abs[
+                : env_abs.shape[0] - 1, : env_abs.shape[1] - 1, : env_abs.shape[2] - 1
+            ]
+        ) / np.sum(
+            env_abs[
+                : env_abs.shape[0] - 1, : env_abs.shape[1] - 1, : env_abs.shape[2] - 1
+            ]
+        )
+        STC_fac["zeta"] = np.min(np.roots([4 * nu, -4, nu * w0**2 * tau**2]))
+        # calculate angular dispersion and pulse front tilt
+        z_centroids = np.sum(grid.axes[2] * env_abs, axis=2) / np.sum(env_abs, axis=2)
+        weight = np.mean(env_abs**2, axis=2)
+        derivative_x = np.gradient(z_centroids, axis=0) / grid.dx[0]
+        derivative_y = np.gradient(z_centroids, axis=1) / grid.dx[1]
+        pft_x = np.sum(derivative_x * weight) / np.sum(weight)
+        pft_y = np.sum(derivative_y * weight) / np.sum(weight)
+        STC_fac["pft"] = np.sqrt((pft_x**2 + pft_y**2))
+        STC_fac["stc_theta_beta"] = np.arctan2(pft_y, pft_x)
+        STC_fac["beta"] = (np.sqrt((pft_x**2 + pft_y**2)) - temp_chirp * nu) / k0
+        return STC_fac