Repository: LASY-org/lasy

Issue:
Flattened Gaussian Transverse before and after focusing
The flattened Gaussian transverse profile #349 adds a much needed profile to `lasy`, however, this profile only defines after focusing (ie. near the focus).

For a variety of use cases, it would make sense to have access to both the near field (ie before focusing with an OAP) and far field (i.e near focus) profiles. 

I think we should add the near field profile too, either via an option in the profile or via the creation of a second profile.

Discussion:

Solution patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index b025cdf..29780a9 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -56,7 +56,7 @@ repos:
 # Python: Ruff linter & formatter
 #   https://docs.astral.sh/ruff/
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.9.6
+  rev: v0.11.0
   hooks:
     # Run the linter
     - id: ruff
diff --git a/README.md b/README.md
index d204113..6c50ec8 100644
--- a/README.md
+++ b/README.md
@@ -6,7 +6,25 @@
 
 More specifically, `lasy` offers many ways to define complex laser pulses (e.g. from commonly-known analytical formulas, from experimental measurements, etc.) and offers pre-processing functionalities (e.g. propagation, re-normalization, geometry conversion). The laser field is then exported in a standardized file, that can be read by external simulation codes.
 
-For an extended tutorial on `lasy` please check out [our tutorial](https://lasydoc.readthedocs.io/en/latest/tutorials/index.html).
+## Installation
+
+For the standard release of the code simply run:
+```
+python3 -m pip install lasy
+```
+
+If you would prefer the most recent version of the code with the latest functionalities, then run:
+```
+python3 -m pip install git+https://github.com/LASY-org/lasy.git
+```
+
+## Tutorials
+An interactive Google Colab tutorial showing some of the main functionalities of lasy can be found [here](https://colab.research.google.com/drive/1nPwgIUea6Jhzc9CSPDZXCcviWCmebjF1?usp=sharing).
+
+This document is updated on a best-effort basis to keep track with developments in the lasy code base.
+
+Additionally, a set of static (automatically tested) examples can be found [here](https://lasydoc.readthedocs.io/en/latest/tutorials/index.html).
+
 
 ## Documentation
 
@@ -53,19 +71,15 @@ Then, for each contribution:
 - Functions in `utils/laser_utils.py` only depend on standard types (Python & Numpy) and on the `Grid` class. That way, they are relatively stand-alone and can be used on different data structures. A simple Grid factory is provided for that purpose.
 - A PR should be open for any contribution: the description helps to explain the code and open dicussion.
 
-## Install
 
-```
-python3 -m pip install lasy
-```
+
+## Testing
 
 For tests, you need to have a few extra packages, such as `pytest` and `openpmd-viewer` installed:
 ```bash
 python3 -m pip install -r tests/requirements.txt
 ```
 
-## Test
-
 After successful installation, you can run the unit tests:
 ```bash
 # Run all tests
diff --git a/docs/source/api/profiles/transverse/flattened_gaussian.rst b/docs/source/api/profiles/transverse/flattened_gaussian.rst
new file mode 100644
index 0000000..ad37f0f
--- /dev/null
+++ b/docs/source/api/profiles/transverse/flattened_gaussian.rst
@@ -0,0 +1,10 @@
+Flattened Gaussian Transverse Profile
+=====================================
+
+Implements a flattened transverse Gaussian profile.
+This profile can either be defined close to a focus, in a far-field regime or in a collimated beam in the near field.
+
+------------
+
+.. autoclass:: lasy.profiles.transverse.FlattenedGaussianTransverseProfile
+    :members:
diff --git a/docs/source/api/profiles/transverse/index.rst b/docs/source/api/profiles/transverse/index.rst
index 4576744..a327b18 100644
--- a/docs/source/api/profiles/transverse/index.rst
+++ b/docs/source/api/profiles/transverse/index.rst
@@ -14,3 +14,4 @@ Transverse Laser Profiles
    super_gaussian_profile
    jinc_profile
    transverse_profile_from_data
+   flattened_gaussian
diff --git a/docs/source/api/utils/index.rst b/docs/source/api/utils/index.rst
index 817885d..7862662 100644
--- a/docs/source/api/utils/index.rst
+++ b/docs/source/api/utils/index.rst
@@ -7,3 +7,4 @@ Utils
 
    grid
    gerchberg_saxton
+   laser_utils
diff --git a/docs/source/api/utils/laser_utils.rst b/docs/source/api/utils/laser_utils.rst
new file mode 100644
index 0000000..a7a2e34
--- /dev/null
+++ b/docs/source/api/utils/laser_utils.rst
@@ -0,0 +1,6 @@
+Laser Utils
+===========
+
+.. automodule:: lasy.utils.laser_utils
+   :members:
+   :undoc-members:
diff --git a/docs/source/tutorials/index.rst b/docs/source/tutorials/index.rst
index 49d7db7..40fc03b 100644
--- a/docs/source/tutorials/index.rst
+++ b/docs/source/tutorials/index.rst
@@ -1,6 +1,15 @@
 Tutorials
 =========
 
+An interactive Google Colab tutorial showing some of the main functionalities of lasy can be found `here`_.
+
+.. _here: https://colab.research.google.com/drive/1nPwgIUea6Jhzc9CSPDZXCcviWCmebjF1?usp=sharing
+
+This document is updated on a best-effort basis to keep track with developments in the lasy code base.
+
+Additionally, a set of static (automatically tested) examples can be found below:
+
+
 .. toctree::
    :maxdepth: 1
 
diff --git a/lasy/optical_elements/__init__.py b/lasy/optical_elements/__init__.py
index 4f464f9..4841404 100644
--- a/lasy/optical_elements/__init__.py
+++ b/lasy/optical_elements/__init__.py
@@ -1,5 +1,6 @@
 from .axicon import Axicon
 from .axiparabola import Axiparabola
+from .intensity_mask import IntensityMask
 from .parabolic_mirror import ParabolicMirror
 from .polynomial_spectral_phase import PolynomialSpectralPhase
 from .zernike_aberrations import ZernikeAberrations
@@ -9,5 +10,6 @@ __all__ = [
     "PolynomialSpectralPhase",
     "Axiparabola",
     "Axicon",
+    "IntensityMask",
     "ZernikeAberrations",
 ]
diff --git a/lasy/optical_elements/intensity_mask.py b/lasy/optical_elements/intensity_mask.py
new file mode 100644
index 0000000..ccba133
--- /dev/null
+++ b/lasy/optical_elements/intensity_mask.py
@@ -0,0 +1,54 @@
+from lasy.optical_elements.optical_element import OpticalElement
+
+
+class IntensityMask(OpticalElement):
+    """
+    Class for a radially symmetric intensity mask.
+
+    This creates an optical element which acts to mask out the intensity of a laser pulse.
+    The mask is radially symmetric and can either mask intensity beyond a user defined radius, thus acting as
+    an aperture. Alternatively, the optica can mask intensity within a user defined radius, acting in this case as
+    an optic with a hole.
+
+    Parameters
+    ----------
+    R : float (in meter)
+        The radius of the mask
+    center: tuple (floats)
+        Center of the mask. Default is (0,0)
+    mask_type: string
+        Should be 'aperture' (default, allows light inside) or 'hole' (allows light outside).
+
+    """
+
+    def __init__(self, R, center=(0, 0), mask_type="aperture"):
+        assert mask_type in ["aperture", "hole"], (
+            "mask_type must be 'aperture' or 'hole'"
+        )
+        self.R = R
+        self.center = center
+        self.mask_type = mask_type
+
+    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 array omega.
+        """
+        r_squared = (x - self.center[0]) ** 2 + (y - self.center[1]) ** 2
+        mask = r_squared <= self.R**2  # True inside, False outside
+
+        if self.mask_type == "aperture":
+            return mask.astype(float)  # 1 inside, 0 outside
+        else:  # "hole"
+            return (~mask).astype(float)  # 0 inside, 1 outside
diff --git a/lasy/optical_elements/polynomial_spectral_phase.py b/lasy/optical_elements/polynomial_spectral_phase.py
index ea92488..5df44b5 100644
--- a/lasy/optical_elements/polynomial_spectral_phase.py
+++ b/lasy/optical_elements/polynomial_spectral_phase.py
@@ -17,12 +17,16 @@ class PolynomialSpectralPhase(OpticalElement):
 
     .. math::
 
-        \phi(\omega) = \frac{\text{GDD}}{2!} (\omega - \omega_0)^2 + \frac{\text{TOD}}{3!} (\omega - \omega_0)^3 + \frac{\text{FOD}}{4!} (\omega - \omega_0)^4
+        \phi(\omega) = \text{delay} (\omega - \omega_0) + \frac{\text{GDD}}{2!} (\omega - \omega_0)^2 + \frac{\text{TOD}}{3!} (\omega - \omega_0)^3 + \frac{\text{FOD}}{4!} (\omega - \omega_0)^4
 
     The other parameters in this formula are defined below.
 
     Parameters
     ----------
+    omega0 : float (in rad/s)
+        Central angular frequency about which the polynomial is expanded
+    delay : float (in s), optional
+        Group delay (by default: ``delay=0``). Positive value delays the pulse, i.e. it arrives at a later time
     gdd : float (in s^2), optional
         Group Delay Dispersion (by default: ``gdd=0``). ``gdd > 0`` corresponds to a positive
         chirp, i.e. the low-frequency part of the spectrum arriving earlier than the
@@ -33,12 +37,11 @@ class PolynomialSpectralPhase(OpticalElement):
         arriving after the main pulse.
     fod : float (in s^4), optional
         Fourth-order Dispersion (by default: ``fod=0``).
-    omega0 : float (in rad/s)
-        Central angular frequency about which the polynomial is expanded
     """
 
-    def __init__(self, omega0, gdd=0, tod=0, fod=0):
+    def __init__(self, omega0, delay=0, gdd=0, tod=0, fod=0):
         self.omega0 = omega0
+        self.delay = delay
         self.gdd = gdd
         self.tod = tod
         self.fod = fod
@@ -61,7 +64,8 @@ class PolynomialSpectralPhase(OpticalElement):
             This array has the same shape as the array omega.
         """
         spectral_phase = (
-            self.gdd / 2 * (omega - self.omega0) ** 2
+            self.delay * (omega - self.omega0)
+            + self.gdd / 2 * (omega - self.omega0) ** 2
             + self.tod / 6 * (omega - self.omega0) ** 3
             + self.fod / 24 * (omega - self.omega0) ** 4
         )
diff --git a/lasy/profiles/from_openpmd_profile.py b/lasy/profiles/from_openpmd_profile.py
index 5786566..0f7a9f3 100644
--- a/lasy/profiles/from_openpmd_profile.py
+++ b/lasy/profiles/from_openpmd_profile.py
@@ -109,7 +109,7 @@ class FromOpenPMDProfile(FromArrayProfile):
         # extract the envelope with a Hilbert transform
         if not is_envelope:
             grid = create_grid(F, axes, dim, is_envelope=is_envelope)
-            grid, omg0 = field_to_envelope(grid, dim, phase_unwrap_nd)
+            omg0 = field_to_envelope(grid, dim, phase_unwrap_nd)
             array = grid.get_temporal_field()[0]
         else:
             s = io.Series(path + "/" + prefix + "_%T.h5", io.Access.read_only)
diff --git a/lasy/profiles/longitudinal/cosine_profile.py b/lasy/profiles/longitudinal/cosine_profile.py
index 108ddff..86dfb29 100644
--- a/lasy/profiles/longitudinal/cosine_profile.py
+++ b/lasy/profiles/longitudinal/cosine_profile.py
@@ -8,7 +8,7 @@ class CosineLongitudinalProfile(LongitudinalProfile):
     Class for the analytic longitudinal truncated cosine profile of a laser pulse.
 
     More precisely, the longitudinal envelope
-    (to be used in the :class:CombinedLongitudinalTransverseProfile class)
+    (to be used in the :class:`.CombinedLongitudinalTransverseProfile` class)
     corresponds to:
 
     .. math::
diff --git a/lasy/profiles/transverse/__init__.py b/lasy/profiles/transverse/__init__.py
index af9daa7..e0d8019 100644
--- a/lasy/profiles/transverse/__init__.py
+++ b/lasy/profiles/transverse/__init__.py
@@ -1,3 +1,4 @@
+from .flattened_gaussian_profile import FlattenedGaussianTransverseProfile
 from .gaussian_profile import GaussianTransverseProfile
 from .hermite_gaussian_profile import HermiteGaussianTransverseProfile
 from .jinc_profile import JincTransverseProfile
@@ -20,4 +21,5 @@ __all__ = [
     "TransverseProfile",
     "SummedTransverseProfile",
     "ScaledTransverseProfile",
+    "FlattenedGaussianTransverseProfile",
 ]
diff --git a/lasy/profiles/transverse/flattened_gaussian_profile.py b/lasy/profiles/transverse/flattened_gaussian_profile.py
index 6d1b6ff..ec078ca 100644
--- a/lasy/profiles/transverse/flattened_gaussian_profile.py
+++ b/lasy/profiles/transverse/flattened_gaussian_profile.py
@@ -1,3 +1,5 @@
+import math
+
 import numpy as np
 from scipy.special import binom
 
@@ -9,7 +11,7 @@ class FlattenedGaussianTransverseProfile(TransverseProfile):
     Class for the analytic profile of a Flattened-Gaussian laser pulse.
 
     Define a complex transverse profile with a flattened Gaussian intensity
-    distribution **far from focus** that transform into a distribution
+    distribution **far from focus** that transforms into a distribution
     with rings **in the focal plane**. (See `Santarsiero et al., J.
     Modern Optics, 1997 <http://doi.org/10.1080/09500349708232927>`_)
 
@@ -17,84 +19,104 @@ class FlattenedGaussianTransverseProfile(TransverseProfile):
     flatness of the transverse profile **far from focus**,
     and increases the number of rings **in the focal plane**.
 
-    The implementation of this class is directly copied from that in `FBPIC`
+    The implementation of this class is based on that from `FBPIC`
     <https://github.com/fbpic/fbpic/blob/dev/fbpic/lpa_utils/laser/transverse_laser_profiles.py>.
 
-    **In the focal plane** (:math:`z=z_f`), the profile translates to a
+    **In the focal plane** (:math:`z=0`), or in the far field, the profile translates to a
     laser with a transverse electric field:
 
     .. math::
 
-        E(x,y,z=zf) \propto
-        \exp\left(-\frac{r^2}{(N+1)w0^2}\right)
-        \sum_{n=0}^N c'_n L^0_n\left(\frac{2\,r^2}{(N+1)w0^2}\right)
+        E(x,y,z=0) \propto
+        \exp\left(-\frac{r^2}{(N+1)w_0^2}\right)
+        \sum_{n=0}^N c'_n L^0_n\left(\frac{2\,r^2}{(N+1)w_0^2}\right)
+
+
+    with Laguerre polynomials :math:`L^0_n` and
+
+    .. math::
 
-        \mathrm{with} Laguerre polynomials :math:`L^0_n` and
-        \qquad c'_n = \sum_{m=n}^{N}\frac{1}{2^m}\binom{m}{n}
+        \qquad c'_n=\sum_{m=n}^{N}\frac{1}{2^m}\binom{m}{n}
 
-    - For :math:`N=0`, this is a Gaussian profile: :math:`E\propto\exp\left(-\frac{r^2}{w0^2}\right)`.
+    - For :math:`N=0`, this is a Gaussian profile: :math:`E\propto\exp\left(-\frac{r^2}{w_0^2}\right)`.
 
-    - For :math:`N\rightarrow\infty`, this is a Jinc profile: :math:`E\propto \frac{J_1(r/w0)}{r/w0}`.
+    - For :math:`N\rightarrow\infty`, this is a Jinc profile: :math:`E\propto\frac{J_1(r/w_0)}{r/w_0}`.
 
-    The equivalent expression **far from focus** is
+    The equivalent expression for **the collimated beam** in the near field which produces this focus is
+    given by:
 
     .. math::
 
-        E(x,y,z=\infty) \propto
-        \exp\left(-\frac{(N+1)r^2}{w(z)^2}\right)
-        \sum_{n=0}^N \frac{1}{n!}\left(\frac{(N+1)\,r^2}{w(z)^2}\right)^n
+        E(x,y) \propto
+        \exp\left(-\frac{(N+1)r^2}{w^2}\right)
+        \sum_{n=0}^N \frac{1}{n!}\left(\frac{(N+1)\,r^2}{w^2}\right)^n
+
+    with the relationship between the spot sizes of the beams in the far field and in the near field given by:
+
+    .. math::
 
-        \mathrm{with} \qquad w(z) = \frac{\lambda_0}{\pi w0}|z-z_{foc}|
+        w = \frac{\lambda_0}{\pi w_0}|z_{\mathrm{foc}}|
 
-    - For :math:`N=0`, this is a Gaussian profile: :math:`E\propto\exp\left(-\frac{r^2}{w_(z)^2}\right)`.
+    where :math:`z_{\mathrm{foc}}` is the distance between the far field and near field planes.
 
-    - For :math:`N\rightarrow\infty`, this is a flat profile: :math:`E\propto \Theta(w(z)-r)`.
+    - Note that a beam defined using the near field definition would be
+      equivalent to a beam defined with the corresponding parameters far from focus in
+      the far field, but without the parabolic phase arising from being
+      defined far from the focus.
+
+    - For :math:`N=0`, the near field profile is a Gaussian profile: :math:`E\propto\exp\left(-\frac{r^2}{w^2}\right)`.
+    - For :math:`N\rightarrow\infty`, the near field profile is a flat profile: :math:`E\propto\Theta(w-r)`.
 
     Parameters
     ----------
-    w0 : float (in meter)
-        The waist of the laser pulse,
-        i.e. :math:`w_{0}` in the above formula.
-    N: int
+    field_type : string
+        Options: 'nearfield', when the beam is defined far from focus (e.g., right before the focusing optics), or 'farfield', when the beam is in the vicinity of the focus.
+
+    w : float (in meter)
+        The waist of the laser pulse. If ``field_type == 'farfield'`` then this
+        variable corresponds to :math:`w_{0}` in the above far field formula.
+        If ``field_type == 'nearfield'`` then this variable corresponds to
+        :math:`w` in the above near field formula.
+
+    N : int
         Determines the "flatness" of the transverse profile, far from
         focus (see the above formula).
         Default: ``N=6`` ; somewhat close to an 8th order supergaussian.
+
     wavelength : float (in meter)
         The main laser wavelength :math:`\lambda_0` of the laser.
-    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)``
+    z_foc : float (in meter), optional
+        Only required if defining the pulse in the far field. Gives the position
+        of the focal plane. (The laser pulse is initialized at ``z=0``.)
 
-    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, wavelength, z_foc=0):
+    def __init__(self, field_type, w, N, wavelength, z_foc=0):
         super().__init__()
         # Ensure that N is an integer
-        self.N = int(round(N))
-        # Calculate effective waist of the Laguerre-Gauss modes, at focus
-        self.w_foc = w0 * (self.N + 1) ** 0.5
-        # Calculate Rayleigh Length
-        self.zr = np.pi * self.w_foc**2 / wavelength
-        # Evaluation distance w.r.t focal position
-        self.z_eval = z_foc
-        # Calculate the coefficients for the Laguerre-Gaussian modes
-        self.cn = np.empty(self.N + 1)
-        for n in range(self.N + 1):
-            m_values = np.arange(n, self.N + 1)
-            self.cn[n] = np.sum((1.0 / 2) ** m_values * binom(m_values, n)) / (
-                self.N + 1
-            )
+        assert isinstance(N, int) and N >= 0
+        assert field_type in ["nearfield", "farfield"]
+        self.field_type = field_type
+        self.N = N
+
+        if field_type == "farfield":
+            w0 = w
+            # Calculate effective waist of the Laguerre-Gauss modes, at focus
+            self.w_foc = w0 * (self.N + 1) ** 0.5
+            # Calculate Rayleigh Length
+            self.zr = np.pi * self.w_foc**2 / wavelength
+            # Evaluation distance w.r.t focal position
+            self.z_eval = z_foc
+            # Calculate the coefficients for the Laguerre-Gaussian modes
+            self.cn = np.empty(self.N + 1)
+            for n in range(self.N + 1):
+                m_values = np.arange(n, self.N + 1)
+                self.cn[n] = np.sum((1.0 / 2) ** m_values * binom(m_values, n)) / (
+                    self.N + 1
+                )
+        else:
+            self.w = w
 
     def _evaluate(self, x, y):
         """
@@ -112,34 +134,49 @@ class FlattenedGaussianTransverseProfile(TransverseProfile):
             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 + Gouy phase
-        diffract_factor = 1.0 - 1j * self.z_eval / self.zr
-        w = self.w_foc * np.abs(diffract_factor)
-        psi = np.angle(diffract_factor)
-        # Argument for the Laguerre polynomials
-        scaled_radius_squared = 2 * (x**2 + y**2) / w**2
-
-        # Sum recursively over the Laguerre polynomials
-        laguerre_sum = np.zeros_like(x, dtype=np.complex128)
-        for n in range(0, self.N + 1):
-            # Recursive calculation of the Laguerre polynomial
-            # - `L` represents $L_n$
-            # - `L1` represents $L_{n-1}$
-            # - `L2` represents $L_{n-2}$
-            if n == 0:
-                L = 1.0
-            elif n == 1:
-                L1 = L
-                L = 1.0 - scaled_radius_squared
-            else:
-                L2 = L1
-                L1 = L
-                L = (((2 * n - 1) - scaled_radius_squared) * L1 - (n - 1) * L2) / n
-            # Add to the sum, including the term for the additional Gouy phase
-            laguerre_sum += self.cn[n] * np.exp(-(2j * n) * psi) * L
-
-        # Final envelope: multiply by n-independent propagation factors
-        exp_argument = -(x**2 + y**2) / (self.w_foc**2 * diffract_factor)
-        envelope = laguerre_sum * np.exp(exp_argument) / diffract_factor
-
-        return envelope
+        if self.field_type == "farfield":
+            # Term for wavefront curvature + Gouy phase
+            diffract_factor = 1.0 - 1j * self.z_eval / self.zr
+            w = self.w_foc * np.abs(diffract_factor)
+            psi = np.angle(diffract_factor)
+            # Argument for the Laguerre polynomials
+            scaled_radius_squared = 2 * (x**2 + y**2) / w**2
+
+            # Sum recursively over the Laguerre polynomials
+            laguerre_sum = np.zeros_like(x, dtype=np.complex128)
+            for n in range(0, self.N + 1):
+                # Recursive calculation of the Laguerre polynomial
+                # - `L` represents $L_n$
+                # - `L1` represents $L_{n-1}$
+                # - `L2` represents $L_{n-2}$
+                if n == 0:
+                    L = 1.0
+                elif n == 1:
+                    L1 = L
+                    L = 1.0 - scaled_radius_squared
+                else:
+                    L2 = L1
+                    L1 = L
+                    L = (((2 * n - 1) - scaled_radius_squared) * L1 - (n - 1) * L2) / n
+                # Add to the sum, including the term for the additional Gouy phase
+                laguerre_sum += self.cn[n] * np.exp(-(2j * n) * psi) * L
+
+            # Final envelope: multiply by n-independent propagation factors
+            exp_argument = -(x**2 + y**2) / (self.w_foc**2 * diffract_factor)
+            envelope = laguerre_sum * np.exp(exp_argument) / diffract_factor
+
+            return envelope
+
+        else:
+            N = self.N
+            w = self.w
+
+            sumseries = 0
+            for n in range(N + 1):
+                sumseries += (
+                    1 / math.factorial(n) * ((N + 1) * (x**2 + y**2) / w**2) ** n
+                )
+
+            envelope = np.exp(-(N + 1) * (x**2 + y**2) / w**2) * sumseries
+
+            return envelope
diff --git a/lasy/utils/laser_utils.py b/lasy/utils/laser_utils.py
index 87eea72..b9f930c 100644
--- a/lasy/utils/laser_utils.py
+++ b/lasy/utils/laser_utils.py
@@ -19,10 +19,10 @@ def compute_laser_energy(dim, grid):
     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.
+        - ``'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
@@ -65,10 +65,10 @@ def normalize_energy(dim, energy, grid):
     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.
+        - ``'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.
 
     energy : scalar (J)
         Energy of the laser pulse after normalization.
@@ -257,7 +257,8 @@ def get_spectrum(grid, dim, range=None, bins=20, omega0=None, method="sum"):
     Initially, the spectrum is calculated as the Fourier transform of the
     electric field :math:`E(t)`.
 
-    ..math::
+    .. math::
+
         \int E(t) e^{-i \omega t} dt
 
     neglecting the negative frequencies. If ``method=="raw"``, no further
@@ -267,7 +268,8 @@ def get_spectrum(grid, dim, range=None, bins=20, omega0=None, method="sum"):
 
     For the other methods, the spectral energy density is calculated as
 
-    ..math::
+    .. math::
+
         \frac{\epsilon_0 c}{2\pi} |\int E(t) e^{-i \omega t} dt| ^ 2
 
     If ``method=="on_axis"``, a 1D real array with on-axis value of the
@@ -276,7 +278,8 @@ def get_spectrum(grid, dim, range=None, bins=20, omega0=None, method="sum"):
     Otherwise, if ``method=="sum"`` (default), the transverse integral of the
     spectral energy density is calculated:
 
-    ..math::
+    .. math::
+
         \frac{\epsilon_0 c}{2\pi} \int |\int E(t) e^{-i \omega t} dt| ^ 2 dx dy
 
     The units of this array are :math:`\mathrm{J / (rad Hz)}`
@@ -288,13 +291,13 @@ def get_spectrum(grid, dim, range=None, bins=20, omega0=None, method="sum"):
         spectrum is computed, and the associated metadata. The last axis must
         be the longitudinal dimension.
 
-    dim : string (optional)
+    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.
+        - ``'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.
 
     range : list of float (optional)
         List of two values indicating the minimum and maximum frequency of the
@@ -399,14 +402,13 @@ def get_frequency(
         be the longitudinal dimension.
         Can be the full electric field or the envelope.
 
-    dim : string (optional)
-        Dimensionality of the array.
-        Options are:
+    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.
+        - ``'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.
 
     is_hilbert : boolean (optional)
         If True, the field argument is assumed to be a Hilbert transform, and
@@ -455,8 +457,8 @@ def get_frequency(
         if dim == "xyt" and phase_unwrap_nd:
             print("WARNING: using 3D phase unwrapping, this can be expensive")
 
-        h = field if is_hilbert else hilbert_transform(grid)
-        h = np.squeeze(field)
+        h = field if is_hilbert else hilbert_transform(field)
+
         if phase_unwrap_nd:
             try:
                 from skimage.restoration import unwrap_phase
@@ -495,8 +497,14 @@ def get_duration(grid, dim):
     ----------
     grid : Grid
         The grid with the envelope to analyze.
-    dim : str
-        Dimensionality of the grid.
+
+    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.
 
     Returns
     -------
@@ -577,15 +585,21 @@ def vector_potential_to_field(grid, omega0, direct=True):
 
 
 def field_to_envelope(grid, dim, phase_unwrap_nd=False):
-    """Get the complex envelope of a field by applying a Hilbert transform.
+    """Convert a field to its complex envelope representation by applying a Hilbert transform.
 
     Parameters
     ----------
     grid : Grid
-        The field from which to extract the envelope.
+        The Grid object on which the field is replaced with an envelope.
+        This object is modified by the function.
 
-    dim : str
-        Dimensions of the field. Possible values are `'xyt'` or `'rt'`.
+    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.
 
     phase_unwrap_nd : boolean (optional)
         If True, the phase unwrapping is n-dimensional (2- or 3-D depending on dim).
@@ -595,32 +609,31 @@ def field_to_envelope(grid, dim, phase_unwrap_nd=False):
 
     Returns
     -------
-    tuple
-        A tuple with the envelope array and the central wavelength.
+    scalar
+        The central angular frequency.
     """
     assert not grid.is_envelope
 
-    field = grid.get_temporal_field()
-
-    # hilbert transform needs inverted time axis.
-    field = hilbert_transform(field)
-
     # Get central wavelength from array
-    omg_h, omg0_h = get_frequency(
+    _, omg0_h = get_frequency(
         grid,
         dim=dim,
-        is_hilbert=True,
+        is_hilbert=False,
         phase_unwrap_nd=phase_unwrap_nd,
     )
+    # Hilbert transform
+    field = hilbert_transform(grid.get_temporal_field())
+    # Remove carrier frequency
     field *= np.exp(1j * omg0_h * grid.axes[-1])
+    # Store envelope
     grid.set_is_envelope(True)
     grid.set_temporal_field(field)
 
-    return grid, omg0_h
+    return omg0_h
 
 
 def hilbert_transform(field):
-    """Make a hilbert transform of the grid field.
+    """Make a hilbert transform of the field.
 
     Currently the arrays need to be flipped along t (both the input field and
     its transform) to get the imaginary part (and thus the phase) with the
@@ -628,8 +641,8 @@ def hilbert_transform(field):
 
     Parameters
     ----------
-    grid : Grid
-        The lasy grid whose field should be transformed.
+    field : 3d numpy array
+        The field whose field should be transformed.
     """
     return hilbert(field[:, :, ::-1])[:, :, ::-1]
 
@@ -641,8 +654,14 @@ def get_grid_cell_volume(grid, dim):
     ----------
     grid : Grid
         The grid form which to compute the cell volume
-    dim : str
-        Dimensionality of the grid.
+
+    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.
 
     Returns
     -------
@@ -688,10 +707,17 @@ def create_grid(array, axes, dim, is_envelope=True):
     ----------
     array : ndarray
         The input field array.
+
     axes : dict
         Dictionary with the information of the array axes.
-    dim : {'xyt, 'rt'}
-        The dimensionality of the array.
+
+    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.
 
     Returns
     -------
@@ -731,10 +757,11 @@ def export_to_z(dim, grid, omega0, z_axis=None, z0=0.0, t0=0.0, backend="NP"):
     ----------
     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.
+
+        - ``'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 a ndarrays (V/m) with
@@ -823,10 +850,11 @@ def import_from_z(dim, grid, omega0, field_z, z_axis, z0=0.0, t0=0.0, backend="N
     ----------
     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.
+
+        - ``'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
@@ -911,16 +939,17 @@ def get_w0(grid, dim):
     ----------
     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.
+
+        - ``'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.
 
-    Return
-    ----------
+    Returns
+    -------
     sigma : Standard deviation of a**2 in m
     """
     field = grid.get_temporal_field()
@@ -951,10 +980,11 @@ def get_phi2(dim, grid):
     ----------
     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.
+
+        - ``'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.
@@ -984,16 +1014,17 @@ def get_zeta(dim, grid, k0):
     ----------
     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.
+
+        - ``'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.
 
-     Return
-    ----------
+    Returns
+    -------
     zeta_x, zeta_y : Spatial chirp in :math:`\zeta=\frac{dx_0}{d\omega}` (meter * second)
     nu_x, nu_y: Spatial chirp in :math:`\nu=\frac{d\omega_0}{dx}` (meter^-1 * second^-1)
     """
@@ -1038,17 +1069,18 @@ def get_beta(dim, grid, k0):
     ----------
     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.
+
+        - ``'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.
 
-    Return
-    ----------
+    Returns
+    -------
     beta_x, beta_y : Angular dispersion in :math:` \beta = \frac{d\theta_0}{d\omega}` (second)
     """
     assert dim == "xyt", "No angular chirp for axis-symmetric dimension."
@@ -1079,18 +1111,19 @@ def get_pft(dim, grid):
     ----------
     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.
+
+        - ``'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.
 
-    Return
-    ----------
+    Returns
+    -------
     pft_x, pft_y : Pulse front tilt in :math:`p=\frac{dt_0}{dx}` (second * meter^-1).
     """
     assert dim == "xyt", "No pulse front tilt for cylindrical symmetry."
@@ -1113,16 +1146,17 @@ def get_propation_angle(dim, grid, k0):
     ----------
     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.
+
+        - ``'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.
 
-    Return
-    ----------
+    Returns
+    -------
     angle_x, angle_y : Propagating angle in :math:`p = \frac{k_x}{k_z}` or :math:`p = \frac{k_y}{k_z}` (in radians).
     """
     assert dim == "xyt", "Propagation is always on-axis for axis-symmetric dimension."


Test patch:
diff --git a/tests/test_laser_profiles.py b/tests/test_laser_profiles.py
index d328f47..04ca528 100644
--- a/tests/test_laser_profiles.py
+++ b/tests/test_laser_profiles.py
@@ -1,11 +1,19 @@
 # -*- coding: utf-8 -*-
 
+import copy
+
 import numpy as np
 import pytest
 from scipy.constants import c
 
 from lasy.laser import Laser
-from lasy.profiles import FromArrayProfile, GaussianProfile, SpeckleProfile
+from lasy.optical_elements.parabolic_mirror import ParabolicMirror
+from lasy.profiles import (
+    CombinedLongitudinalTransverseProfile,
+    FromArrayProfile,
+    GaussianProfile,
+    SpeckleProfile,
+)
 from lasy.profiles.longitudinal import (
     CosineLongitudinalProfile,
     GaussianLongitudinalProfile,
@@ -14,6 +22,7 @@ from lasy.profiles.longitudinal import (
 )
 from lasy.profiles.profile import Profile, ScaledProfile, SummedProfile
 from lasy.profiles.transverse import (
+    FlattenedGaussianTransverseProfile,
     GaussianTransverseProfile,
     HermiteGaussianTransverseProfile,
     JincTransverseProfile,
@@ -25,6 +34,8 @@ from lasy.profiles.transverse import (
     TransverseProfileFromData,
 )
 from lasy.utils.exp_data_utils import find_center_of_mass
+from lasy.utils.grid import Grid
+from lasy.utils.laser_utils import compute_laser_energy, get_w0
 
 
 class MockProfile(Profile):
@@ -533,3 +544,65 @@ def test_scale_trans_error_if_not_scalar():
         trans_profile_1 * trans_profile_1
     with pytest.raises(AssertionError):
         trans_profile_1 * [1.0, 2.0]
+
+
+def test_flattened_gaussian_profile():
+    w = 20e-3
+    N = 25
+    wl = 800e-9
+    tau = 30e-15
+    pol = (1, 0)
+    energy = 1.0
+    focal_length = 1.0
+
+    w0 = focal_length * wl / np.pi / w
+
+    nf = FlattenedGaussianTransverseProfile(
+        field_type="nearfield", w=w, N=N, wavelength=wl
+    )
+    ff = FlattenedGaussianTransverseProfile(
+        field_type="farfield", w=w0, N=N, wavelength=wl
+    )
+
+    long = GaussianLongitudinalProfile(wl, tau, 0)
+
+    nf_prof = CombinedLongitudinalTransverseProfile(wl, pol, energy, long, nf)
+    ff_prof = CombinedLongitudinalTransverseProfile(wl, pol, energy, long, ff)
+
+    dim = "rt"
+    lo = (0, -100e-15)
+    hi_ff = (500e-6, 100e-15)
+    hi_nf = (40e-3, 100e-15)
+    npoints = (5000, 200)
+
+    las_nf = Laser(dim, lo, hi_nf, npoints, nf_prof)
+    las_ff = Laser(dim, lo, hi_ff, npoints, ff_prof)
+
+    las_nf_cp = copy.deepcopy(las_nf)
+
+    OAP = ParabolicMirror(f=focal_length)
+    las_nf_cp.apply_optics(OAP)
+    las_nf_cp.propagate(
+        focal_length, grid=Grid(dim, lo, hi_ff, npoints, n_azimuthal_modes=1)
+    )
+
+    radlineout_nf_cp = (
+        np.abs(las_nf_cp.grid.get_temporal_field()[0, :, int(npoints[1] / 2)]) ** 2
+    )
+    radlineout_ff = (
+        np.abs(las_ff.grid.get_temporal_field()[0, :, int(npoints[1] / 2)]) ** 2
+    )
+
+    err = np.sum(
+        np.abs(np.abs(radlineout_nf_cp) ** 2 - np.abs(radlineout_ff) ** 2)
+    ) / np.sum(np.abs(radlineout_ff) ** 2)
+
+    assert err < 1e-2
+
+    energy_ff = compute_laser_energy(dim, las_ff.grid)
+    energy_nf_cp = compute_laser_energy(dim, las_nf_cp.grid)
+    assert np.abs(energy_ff - energy_nf_cp) / energy_ff < 1e-4
+
+    w0_ff = get_w0(las_ff.grid, dim)
+    w0_nf_cp = get_w0(las_nf_cp.grid, dim)
+    assert np.abs(w0_nf_cp - w0_ff) / w0_ff < 1e-2
diff --git a/tests/test_polynomial_spectral_phase.py b/tests/test_polynomial_spectral_phase.py
index fc28258..6173248 100644
--- a/tests/test_polynomial_spectral_phase.py
+++ b/tests/test_polynomial_spectral_phase.py
@@ -4,7 +4,7 @@
 Test checks the implementation of the polynomial spectral phase
 by initializing a Gaussian pulse (with flat spectral phase),
 adding spectral phase to it, and checking the corresponding
-temporal shape of the laser pulse again analytical formulas.
+temporal shape of the laser pulse against analytical formulas.
 """
 
 import numpy as np
@@ -31,6 +31,42 @@ hi = (+12e-6, +12e-6, +100e-15)
 npoints = (100, 100, 200)
 
 
+def test_delay():
+    """Add and compare group delay for laser.
+
+    Add group delay to the laser and compare the on-axis field with the
+    analytical formula for a Gaussian pulse with group delay.
+    """
+    delay = 50e-15
+    dazzler = PolynomialSpectralPhase(omega0, delay=delay)
+
+    # Initialize the laser
+    dim = "xyt"
+    lo = (-12e-6, -12e-6, -100e-15)
+    hi = (+12e-6, +12e-6, +100e-15)
+    npoints = (100, 100, 400)
+    laser = Laser(dim, lo, hi, npoints, gaussian_profile)
+
+    # Get field before and after dazzler
+    E_before = laser.grid.get_temporal_field()
+    laser.apply_optics(dazzler)
+    E_after = laser.grid.get_temporal_field()
+
+    # Extract peak field before dazzler
+    E0 = abs(E_before[50, 50]).max()
+
+    # Compute the analtical expression in real space for a Gaussian
+    t = np.linspace(laser.grid.lo[-1], laser.grid.hi[-1], laser.grid.npoints[-1])
+
+    E_analytical = E0 * np.exp(-1.0 * ((t - delay) / tau) ** 2)
+
+    # Compare the on-axis field with the analytical formula
+    tol = 1.2e-3
+    assert np.all(
+        abs(E_after[50, 50, :] - E_analytical) / abs(E_analytical).max() < tol
+    )
+
+
 def test_gdd():
     """Add and compare GDD laser.