Add option to initialize other profiles out of focus (#143)

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

lasy/profiles/transverse/hermite_gaussian_profile.py

@@ -1,5 +1,5 @@
 import numpy as np
-from scipy.special.orthogonal import hermite
+from scipy.special import hermite
 from math import factorial
 
 from .transverse_profile import TransverseProfile
@@ -32,13 +32,38 @@ class HermiteGaussianTransverseProfile(TransverseProfile):
         The order of hermite polynomial in the x direction
     n_y : int (dimensionless)
         The order of hermite polynomial in the y direction
+    wavelength : float (in meter), optional
+        The main laser wavelength :math:`\lambda_0` of the laser.
+        (Only needed if ``z_foc`` is different than 0.)
+    z_foc : float (in meter), optional
+        Position of the focal plane. (The laser pulse is initialized at
+        ``z=0``.)
+
+    Warnings
+    --------
+    In order to initialize the pulse out of focus, you can either:
+
+    - Use a non-zero ``z_foc``
+    - Use ``z_foc=0`` (i.e. initialize the pulse at focus) and then call
+      ``laser.propagate(-z_foc)``
+
+    Both methods are in principle equivalent, but note that the first
+    method uses the paraxial approximation, while the second method does
+    not make this approximation.
     """
 
-    def __init__(self, w0, n_x, n_y):
+    def __init__(self, w0, n_x, n_y, wavelength=None, z_foc=0):
         super().__init__()
         self.w0 = w0
         self.n_x = n_x
         self.n_y = n_y
+        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, x, y):
         """
@@ -56,13 +81,22 @@ def _evaluate(self, x, y):
             Contains the value of the envelope at the specified points
             This array has the same shape as the arrays x, y
         """
+        # Term for wavefront curvature, waist and Gouy phase
+        diffract_factor = 1.0 - 1j * self.z_foc_over_zr
+        w = self.w0 * abs(diffract_factor)
+        psi = np.angle(diffract_factor)
+
         envelope = (
             np.sqrt(2 / np.pi)
             * np.sqrt(1 / (2 ** (self.n_x) * factorial(self.n_x) * self.w0))
             * np.sqrt(1 / (2 ** (self.n_y) * factorial(self.n_y) * self.w0))
-            * hermite(self.n_x)(np.sqrt(2) * x / self.w0)
-            * hermite(self.n_y)(np.sqrt(2) * y / self.w0)
-            * np.exp(-(x**2 + y**2) / self.w0**2)
+            * hermite(self.n_x)(np.sqrt(2) * x / w)
+            * hermite(self.n_y)(np.sqrt(2) * y / w)
+            * np.exp(
+                -(x**2 + y**2) / (self.w0**2 * diffract_factor)
+                - 1.0j * (self.n_x + self.n_y) * psi
+            )
+            # Additional Gouy phase
+            * (1.0 / diffract_factor)
         )
-
         return envelope

lasy/profiles/transverse/laguerre_gaussian_profile.py

@@ -9,7 +9,7 @@ class LaguerreGaussianTransverseProfile(TransverseProfile):
     High-order Gaussian laser pulse expressed in the Laguerre-Gaussian formalism.
 
     Derived class for an analytic profile.
-    More precisely, the transverse envelope (to be used in the
+    More precisely, at focus (`z_foc=0`), the transverse envelope (to be used in the
     :class:CombinedLongitudinalTransverseLaser class) corresponds to:
 
     .. math::
@@ -31,13 +31,37 @@ class LaguerreGaussianTransverseProfile(TransverseProfile):
         The radial order of Generalized Laguerre polynomial
     m : int (dimensionless)
         Defines the phase rotation, i.e. :math:`m` in the above formula.
+    wavelength : float (in meter), optional
+        The main laser wavelength :math:`\lambda_0` of the laser.
+        (Only needed if ``z_foc`` is different than 0.)
+    z_foc : float (in meter), optional
+        Position of the focal plane. (The laser pulse is initialized at `z=0`.)
+
+    Warnings
+    --------
+    In order to initialize the pulse out of focus, you can either:
+
+    - Use a non-zero ``z_foc``
+    - Use ``z_foc=0`` (i.e. initialize the pulse at focus) and then call
+      ``laser.propagate(-z_foc)``
+
+    Both methods are in principle equivalent, but note that the first
+    method uses the paraxial approximation, while the second method does
+    not make this approximation.
     """
 
-    def __init__(self, w0, p, m):
+    def __init__(self, w0, p, m, wavelength=None, z_foc=0):
         super().__init__()
         self.w0 = w0
         self.p = p
         self.m = m
+        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, x, y):
         """
@@ -55,17 +79,24 @@ def _evaluate(self, x, y):
             Contains the value of the envelope at the specified points
             This array has the same shape as the arrays x, y
         """
+        # Term for wavefront curvature, waist and Gouy phase
+        diffract_factor = 1.0 - 1j * self.z_foc_over_zr
+        w = self.w0 * abs(diffract_factor)
+        psi = np.angle(diffract_factor)
         # complex_position corresponds to r e^{+/-i\theta}
         if self.m > 0:
             complex_position = x - 1j * y
         else:
             complex_position = x + 1j * y
         radius = abs(complex_position)
-        scaled_rad_squared = (radius**2) / self.w0**2
         envelope = (
             complex_position ** abs(self.m)
-            * genlaguerre(self.p, abs(self.m))(2 * scaled_rad_squared)
-            * np.exp(-scaled_rad_squared)
+            * genlaguerre(self.p, abs(self.m))(2 * radius**2 / w**2)
+            * np.exp(
+                -(radius**2) / (self.w0**2 * diffract_factor)
+                - 1.0j * (2 * self.p + self.m) * psi
+            )  # Additional Gouy phase
+            * (1.0 / diffract_factor)
         )
 
         return envelope