Update definition of Flattened Gaussian to include both NF and FF

lasy/profiles/transverse/flattened_gaussian_profile.py

@@ -1,3 +1,5 @@
+import math
+
 import numpy as np
 from scipy.special import binom
 
@@ -17,10 +19,10 @@ 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=z_f`), or in the far field, the profile translates to a
     laser with a transverse electric field:
 
     .. math::
@@ -36,7 +38,8 @@ class FlattenedGaussianTransverseProfile(TransverseProfile):
 
     - For :math:`N\rightarrow\infty`, this is a Jinc profile: :math:`E\propto \frac{J_1(r/w0)}{r/w0}`.
 
-    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::
 
@@ -46,23 +49,35 @@ class FlattenedGaussianTransverseProfile(TransverseProfile):
 
         \mathrm{with} \qquad w(z) = \frac{\lambda_0}{\pi w0}|z-z_{foc}|
 
+    - Note that a beam defined using the near field definition would be equivalent to a beam defined with
+    the corresponding parameters in the far field, but without the parabolic phase arising from being defined
+    far from the focus.
+
     - For :math:`N=0`, this is a Gaussian profile: :math:`E\propto\exp\left(-\frac{r^2}{w_(z)^2}\right)`.
 
     - For :math:`N\rightarrow\infty`, this is a flat profile: :math:`E\propto \Theta(w(z)-r)`.
 
     Parameters
     ----------
-    w0 : float (in meter)
-        The waist of the laser pulse,
-        i.e. :math:`w_{0}` in the above formula.
+    field_type : string
+        Options: 'nearfield', when the beam is defined far from focus and
+        has been collimated, or 'farfield', when the beam is in the vincinity
+        of or has been directly propagated from the focus. In this case there
+        can be a large defocus in the spatial phase.
+    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(z)` 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
+        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``.)
 
     Warnings
@@ -78,23 +93,30 @@ class FlattenedGaussianTransverseProfile(TransverseProfile):
     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 field_type in ["nearfield", "farfield"]
+        self.field_type = field_type
+
+        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,50 @@ 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 + 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
+            if N > 0:
+                for n in range(N):
+                    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

tests/test_laser_profiles.py

@@ -1,11 +1,18 @@
 # -*- 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.profiles import (
+    CombinedLongitudinalTransverseProfile,
+    FromArrayProfile,
+    GaussianProfile,
+    SpeckleProfile,
+)
 from lasy.profiles.longitudinal import (
     CosineLongitudinalProfile,
     GaussianLongitudinalProfile,
@@ -14,6 +21,7 @@
 )
 from lasy.profiles.profile import Profile, ScaledProfile, SummedProfile
 from lasy.profiles.transverse import (
+    FlattenedGaussianTransverseProfile,
     GaussianTransverseProfile,
     HermiteGaussianTransverseProfile,
     JincTransverseProfile,
@@ -533,3 +541,61 @@
         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 = (1000e-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 = (
+        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(
+            las_nf_cp.grid.get_temporal_field()[0, :, :]
+            - las_ff.grid.get_temporal_field()[0, :, :]
+        )
+        ** 2
+    )
+
+    assert err < 1