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,126 @@
         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.
     """
 
-    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
 
-    def evaluate(self, t):
+        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, x, y):
         """
         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)
+        assert (
+            x is not None and y is not None
+        ), "transverse points should be specified if spatio-temperal coupling exits"
+        inv_complex_waist_2 = (
+            1.0
+            / (self.w0**2 * (1.0 + 2.0j * self.z_foc_over_zr / (self.k0 * self.w0**2)))
+            if self.beta
+            else 0
+        )
+        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 * self.t_peak)
         )
-
         return envelope

lasy/profiles/longitudinal/longitudinal_profile.py

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