Repository: LASY-org/lasy

Issue:
Overloading Addition and Multiplication Operators of TransverseProfile
Currently in the definition of the laser `profile` the addition and multiplication operators have been overloaded thanks to some additions by Sören. However, we should also consider adding this functionality to the `transverse_profile` as for the decomposition/reconstruction of the laser profile into Hermite Gaussians, only the transverse profile is required and currently performing these calculations is quite slow and cumbersome with the full `profile` object.

Discussion:

Solution patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 0de5f27..96c868c 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -101,7 +101,7 @@ repos:
 
 # Python Formatting
 - repo: https://github.com/psf/black
-  rev: 23.11.0  # Keep in sync with blacken-docs
+  rev: 23.12.1  # Keep in sync with blacken-docs
   hooks:
   - id: black
 - repo: https://github.com/asottile/blacken-docs
diff --git a/docs/source/api/profiles/longitudinal/cosine.rst b/docs/source/api/profiles/longitudinal/cosine.rst
new file mode 100644
index 0000000..aaadf52
--- /dev/null
+++ b/docs/source/api/profiles/longitudinal/cosine.rst
@@ -0,0 +1,10 @@
+Cosine Longitudinal Profile
+===========================
+
+Used to define a truncated cosine longitudinal laser profile (i.e., a profile described by the function :math:`\cos{\left( {\frac{\pi}{2} \frac{t - t_{peak}}{t_{fwhm}} } \right)}`,
+in the interval :math:`[t_{peak} - t_{fwhm}/2, t_{peak} + t_{fwhm}/2]`).
+
+------------
+
+.. autoclass:: lasy.profiles.longitudinal.cosine_profile.CosineLongitudinalProfile
+    :members:
diff --git a/docs/source/api/profiles/longitudinal/index.rst b/docs/source/api/profiles/longitudinal/index.rst
index dea0f7f..1a4a3cb 100644
--- a/docs/source/api/profiles/longitudinal/index.rst
+++ b/docs/source/api/profiles/longitudinal/index.rst
@@ -5,5 +5,6 @@ Longitudinal Profiles
    :maxdepth: 4
    :hidden:
 
+   cosine
    gaussian
    longitudinal_profile_from_data
diff --git a/lasy/laser.py b/lasy/laser.py
index 653bde3..2cd8049 100644
--- a/lasy/laser.py
+++ b/lasy/laser.py
@@ -233,11 +233,11 @@ class Laser:
                     verbose=False,
                 )
             # Propagate the spectral image
-            transform_data = np.transpose(field_fft).copy()
+            transform_data = np.moveaxis(field_fft, -1, 0).copy()
             self.prop.step(
                 transform_data, distance, overwrite=True, show_progress=show_progress
             )
-            field_fft[:, :, :] = np.transpose(transform_data).copy()
+            field_fft[:, :, :] = np.moveaxis(transform_data, 0, -1).copy()
 
         # Choose the time translation assuming propagation at v=c
         translate_time = distance / c
diff --git a/lasy/profiles/longitudinal/__init__.py b/lasy/profiles/longitudinal/__init__.py
index 22f405a..6b0dc4b 100644
--- a/lasy/profiles/longitudinal/__init__.py
+++ b/lasy/profiles/longitudinal/__init__.py
@@ -1,4 +1,9 @@
+from .cosine_profile import CosineLongitudinalProfile
 from .gaussian_profile import GaussianLongitudinalProfile
 from .longitudinal_profile_from_data import LongitudinalProfileFromData
 
-__all__ = ["GaussianLongitudinalProfile", "LongitudinalProfileFromData"]
+__all__ = [
+    "CosineLongitudinalProfile",
+    "GaussianLongitudinalProfile",
+    "LongitudinalProfileFromData",
+]
diff --git a/lasy/profiles/longitudinal/cosine_profile.py b/lasy/profiles/longitudinal/cosine_profile.py
new file mode 100644
index 0000000..92c927d
--- /dev/null
+++ b/lasy/profiles/longitudinal/cosine_profile.py
@@ -0,0 +1,71 @@
+import numpy as np
+
+from .longitudinal_profile import LongitudinalProfile
+
+
+class CosineLongitudinalProfile(LongitudinalProfile):
+    r"""
+    Derived 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)
+    corresponds to:
+
+    .. math::
+
+        \mathcal{L}(t) = \cos\left({ \frac{\pi}{2} \frac{t-t_{peak}}{\tau_{fwhm}} }\right)
+            \theta\left({ \frac{t-t_{peak}}{\tau_{fwhm}} + 1 }\right)
+            \theta\left({ 1 - \frac{t-t_{peak}}{\tau_{fwhm}}} \right)
+            \exp\left({ + i (\phi_{cep} + \omega_0 t_{peak} ) }\right)
+
+
+    Parameters
+    ----------
+    wavelength : float (in meter)
+        The main laser wavelength :math:`\lambda_0` of the laser.
+
+    tau_fwhm : float (in second)
+        The Full-Width-Half-Maximum duration of the intensity distribution of the pulse,
+        i.e. :math:`\tau_{fwhm}` in the above formula.
+
+    t_peak : float (in second)
+        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
+        The Carrier Enveloppe Phase (CEP)
+        (i.e. the phase of the laser oscillation, at the time where the
+        laser envelope is maximum,  :math:`\phi_{cep}` in the above formula).
+    """
+
+    def __init__(self, wavelength, tau_fwhm, t_peak, cep_phase=0):
+        super().__init__(wavelength)
+        self.tau_fwhm = tau_fwhm
+        self.t_peak = t_peak
+        self.cep_phase = cep_phase
+
+    def evaluate(self, t):
+        """
+        Return the longitudinal envelope.
+
+        Parameters
+        ----------
+        t: ndarrays of floats
+            Define 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.
+        """
+        tn = (t - self.t_peak) / self.tau_fwhm
+
+        envelope = (
+            np.cos(0.5 * np.pi * tn)
+            * (tn > -1)
+            * (tn < 1)
+            * np.exp(+1.0j * (self.cep_phase + self.omega0 * self.t_peak))
+        )
+
+        return envelope
diff --git a/lasy/profiles/longitudinal/longitudinal_profile_from_data.py b/lasy/profiles/longitudinal/longitudinal_profile_from_data.py
index 673a5e1..e8a699e 100644
--- a/lasy/profiles/longitudinal/longitudinal_profile_from_data.py
+++ b/lasy/profiles/longitudinal/longitudinal_profile_from_data.py
@@ -63,7 +63,10 @@ class LongitudinalProfileFromData(LongitudinalProfile):
             # First find central frequency
             wavelength = data["axis"]
             spectral_intensity = data["intensity"]
-            spectral_phase = data["phase"]
+            if data.get("phase") is None:
+                spectral_phase = np.zeros_like(wavelength)
+            else:
+                spectral_phase = data["phase"]
             dt = data["dt"]
             cwl = np.sum(spectral_intensity * wavelength) / np.sum(spectral_intensity)
             cfreq = c / cwl
@@ -106,7 +109,10 @@ class LongitudinalProfileFromData(LongitudinalProfile):
         elif data["datatype"] == "temporal":
             time = data["axis"]
             temporal_intensity = data["intensity"]
-            temporal_phase = data["phase"]
+            if data.get("phase") is None:
+                temporal_phase = np.zeros_like(time)
+            else:
+                temporal_phase = data["phase"]
             cwl = data["wavelength"]
 
         else:
diff --git a/lasy/profiles/transverse/__init__.py b/lasy/profiles/transverse/__init__.py
index 1a6109a..781f10f 100644
--- a/lasy/profiles/transverse/__init__.py
+++ b/lasy/profiles/transverse/__init__.py
@@ -4,6 +4,11 @@ from .laguerre_gaussian_profile import LaguerreGaussianTransverseProfile
 from .super_gaussian_profile import SuperGaussianTransverseProfile
 from .jinc_profile import JincTransverseProfile
 from .transverse_profile_from_data import TransverseProfileFromData
+from .transverse_profile import (
+    TransverseProfile,
+    SummedTransverseProfile,
+    ScaledTransverseProfile,
+)
 
 __all__ = [
     "GaussianTransverseProfile",
@@ -12,4 +17,7 @@ __all__ = [
     "SuperGaussianTransverseProfile",
     "JincTransverseProfile",
     "TransverseProfileFromData",
+    "TransverseProfile",
+    "SummedTransverseProfile",
+    "ScaledTransverseProfile",
 ]
diff --git a/lasy/profiles/transverse/transverse_profile.py b/lasy/profiles/transverse/transverse_profile.py
index a30fdc9..4f4c332 100644
--- a/lasy/profiles/transverse/transverse_profile.py
+++ b/lasy/profiles/transverse/transverse_profile.py
@@ -34,6 +34,18 @@ class TransverseProfile(object):
         # (This should be replaced by any class that inherits from this one.)
         return np.zeros(x.shape, dtype="complex128")
 
+    def __add__(self, other):
+        """Return the sum of two transverse profiles."""
+        return SummedTransverseProfile(self, other)
+
+    def __mul__(self, factor):
+        """Return the scaled transverse profile."""
+        return ScaledTransverseProfile(self, factor)
+
+    def __rmul__(self, factor):
+        """Return the scaled transverse profile."""
+        return ScaledTransverseProfile(self, factor)
+
     def evaluate(self, x, y):
         """
         Return the transverse envelope modified by any spatial offsets.
@@ -72,3 +84,62 @@ class TransverseProfile(object):
         self.y_offset = y_offset
 
         return self
+
+
+class SummedTransverseProfile(TransverseProfile):
+    """
+    Base class for transverse profiles that are the sum of several other transverse profiles.
+
+    Transverse Profile class that represents the sum of multiple transverse profiles.
+
+    Parameters
+    ----------
+    transverse_profiles: list of TransverseProfile objects
+        List of transverse profiles to be summed.
+    """
+
+    def __init__(self, *transverse_profiles):
+        """Initialize the summed profile."""
+        TransverseProfile.__init__(self)
+        # Check that all transverse_profiles are TransverseProfile objects
+        assert all(
+            [isinstance(tp, TransverseProfile) for tp in transverse_profiles]
+        ), "All summands must be Profile objects."
+        self.transverse_profiles = transverse_profiles
+
+    def evaluate(self, x, y):
+        """Return the envelope field of the summed profile."""
+        # Sum the fields of each profile
+        return sum([tp.evaluate(x, y) for tp in self.transverse_profiles])
+
+
+class ScaledTransverseProfile(TransverseProfile):
+    """
+    Base class for transverse profiles that are scaled by a factor.
+
+    Transverse Profile class that represents scaled transverse profiles.
+
+    Parameters
+    ----------
+    transverse_profile: TrasnverseProfile object
+        Trasnverse profile to be scaled.
+    factor: int or float
+        Factor by which to scale the profile.
+    """
+
+    def __init__(self, transverse_profile, factor):
+        """Initialize the summed profile."""
+        TransverseProfile.__init__(self)
+        # Check that the factor is a number
+        assert isinstance(factor, (int, float, complex)), "The factor must be a number."
+        # Check that the profile is a Profile object
+        assert isinstance(
+            transverse_profile, TransverseProfile
+        ), "The profile must be a TransverseProfile object."
+        self.transverse_profile = transverse_profile
+        self.factor = factor
+
+    def evaluate(self, x, y):
+        """Return the envelope field of the scaled transverse profile."""
+        # Sum the fields of each profile
+        return self.transverse_profile.evaluate(x, y) * self.factor
diff --git a/lasy/utils/laser_utils.py b/lasy/utils/laser_utils.py
index 0bc7c66..36ba206 100644
--- a/lasy/utils/laser_utils.py
+++ b/lasy/utils/laser_utils.py
@@ -764,8 +764,9 @@ def export_to_z(dim, grid, omega0, z_axis=None, z0=0.0, t0=0.0, backend="NP"):
             verbose=False,
         )
         # Convert the spectral image to the spatial field representation
-        FieldAxprp.import_field(np.transpose(grid.field).copy())
-        field_z = prop.t2z(FieldAxprp.Field_ft, z_axis, z0=z0, t0=t0).T
+        FieldAxprp.import_field(np.moveaxis(grid.field, -1, 0).copy())
+        field_z = prop.t2z(FieldAxprp.Field_ft, z_axis, z0=z0, t0=t0)
+        field_z = np.moveaxis(field_z, 0, -1)
         field_z *= np.exp(-1j * (z_axis / c + t0) * omega0)
 
     return field_z
@@ -850,7 +851,7 @@ def import_from_z(dim, grid, omega0, field_z, z_axis, z0=0.0, t0=0.0, backend="N
             verbose=False,
         )
         # Convert the spectral image to the spatial field representation
-        transform_data = np.transpose(field_fft).copy()
+        transform_data = np.moveaxis(field_fft, -1, 0).copy()
         transform_data *= np.exp(-1j * z_axis[0] * (k_z[:, None, None] - omega0 / c))
-        grid.field = prop.z2t(transform_data, t_axis, z0=z0, t0=t0).T
+        grid.field = np.moveaxis(prop.z2t(transform_data, t_axis, z0=z0, t0=t0), 0, -1)
         grid.field *= np.exp(1j * (z0 / c + t_axis) * omega0)
diff --git a/requirements.txt b/requirements.txt
index a53579d..6325cda 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,4 @@
-axiprop==0.5.1
+axiprop
 numpy
 openpmd-api
 openpmd-viewer


Test patch:
diff --git a/tests/test_gaussian_propagator.py b/tests/test_gaussian_propagator.py
index b84237a..01ef32e 100644
--- a/tests/test_gaussian_propagator.py
+++ b/tests/test_gaussian_propagator.py
@@ -66,7 +66,7 @@ def test_3D_case(gaussian):
     dim = "xyt"
     lo = (-25e-6, -25e-6, -60e-15)
     hi = (+25e-6, +25e-6, +60e-15)
-    npoints = (100, 100, 100)
+    npoints = (100, 102, 100)
 
     laser = Laser(dim, lo, hi, npoints, gaussian)
     check_gaussian_propagation(laser)
diff --git a/tests/test_laser_profiles.py b/tests/test_laser_profiles.py
index 2073743..4367183 100644
--- a/tests/test_laser_profiles.py
+++ b/tests/test_laser_profiles.py
@@ -3,11 +3,15 @@
 import pytest
 import numpy as np
 from scipy.special import gamma as gamma
+from scipy.constants import c
 
 from lasy.laser import Laser
 from lasy.profiles.profile import Profile, SummedProfile, ScaledProfile
 from lasy.profiles import GaussianProfile, FromArrayProfile
-from lasy.profiles.longitudinal import GaussianLongitudinalProfile
+from lasy.profiles.longitudinal import (
+    GaussianLongitudinalProfile,
+    CosineLongitudinalProfile,
+)
 from lasy.profiles.transverse import (
     GaussianTransverseProfile,
     LaguerreGaussianTransverseProfile,
@@ -15,6 +19,9 @@ from lasy.profiles.transverse import (
     HermiteGaussianTransverseProfile,
     JincTransverseProfile,
     TransverseProfileFromData,
+    TransverseProfile,
+    SummedTransverseProfile,
+    ScaledTransverseProfile,
 )
 from lasy.utils.exp_data_utils import find_center_of_mass
 
@@ -32,6 +39,19 @@ class MockProfile(Profile):
         return np.ones_like(x, dtype="complex128") * self.value
 
 
+class MockTransverseProfile(TransverseProfile):
+    """
+    A mock TransverseProfile class that always returns a constant value.
+    """
+
+    def __init__(self, value):
+        super().__init__()
+        self.value = value
+
+    def evaluate(self, x, y):
+        return np.ones_like(x, dtype="complex128") * self.value
+
+
 @pytest.fixture(scope="function")
 def gaussian():
     # Cases with Gaussian laser
@@ -132,6 +152,63 @@ def test_transverse_profiles_3d():
     assert (x0_test - x0) / x0 < 0.1
 
 
+def test_longitudinal_profiles():
+    npoints = 10000
+
+    wavelength = 800e-9
+    tau_fwhm = 30.0e-15
+    t_peak = 1.0 * tau_fwhm
+    cep_phase = 0.5 * np.pi
+    omega_0 = 2.0 * np.pi * c / wavelength
+
+    t = np.linspace(t_peak - 4 * tau_fwhm, t_peak + 4 * tau_fwhm, npoints)
+
+    # GaussianLongitudinalProfile
+    print("GaussianLongitudinalProfile")
+    tau = tau_fwhm / np.sqrt(2 * np.log(2))
+    profile_gaussian = GaussianLongitudinalProfile(wavelength, tau, t_peak, cep_phase)
+    field_gaussian = profile_gaussian.evaluate(t)
+
+    std_gauss = np.sqrt(np.average((t - t_peak) ** 2, weights=np.abs(field_gaussian)))
+    std_gauss_th = tau / np.sqrt(2.0)
+    print("std_th = ", std_gauss_th)
+    print("std = ", std_gauss)
+    assert np.abs(std_gauss - std_gauss_th) / std_gauss_th < 0.01
+
+    t_peak_gaussian = t[np.argmax(np.abs(field_gaussian))]
+    print("t_peak_th = ", t_peak)
+    print("t_peak = ", t_peak_gaussian)
+    assert np.abs(t_peak_gaussian - t_peak) / t_peak < 0.01
+
+    ff_gaussian = field_gaussian * np.exp(-1.0j * omega_0 * t)
+    cep_phase_gaussian = np.angle(ff_gaussian[np.argmax(np.abs(field_gaussian))])
+    print("cep_phase_th = ", cep_phase)
+    print("cep_phase = ", cep_phase_gaussian)
+    assert np.abs(cep_phase_gaussian - cep_phase) / cep_phase < 0.02
+
+    # CosineLongitudinalProfile
+    print("CosineLongitudinalProfile")
+    profile_cos = CosineLongitudinalProfile(wavelength, tau_fwhm, t_peak, cep_phase)
+    field_cos = profile_cos.evaluate(t)
+
+    std_cos = np.sqrt(np.average((t - t_peak) ** 2, weights=np.abs(field_cos)))
+    std_cos_th = tau_fwhm * np.sqrt(1 - 8 / np.pi**2)
+    print("std_th = ", std_cos_th)
+    print("std = ", std_cos)
+    assert np.abs(std_cos - std_cos_th) / std_cos_th < 0.01
+
+    t_peak_cos = t[np.argmax(np.abs(field_cos))]
+    print("t_peak_th = ", t_peak)
+    print("t_peak = ", t_peak_cos)
+    assert np.abs(t_peak_cos - t_peak) / t_peak < 0.01
+
+    ff_cos = field_cos * np.exp(-1.0j * omega_0 * t)
+    cep_phase_cos = np.angle(ff_cos[np.argmax(np.abs(field_cos))])
+    print("cep_phase_th = ", cep_phase)
+    print("cep_phase = ", cep_phase_cos)
+    assert np.abs(cep_phase_cos - cep_phase) / cep_phase < 0.02
+
+
 def test_profile_gaussian_3d_cartesian(gaussian):
     # - 3D Cartesian case
     dim = "xyt"
@@ -255,3 +332,46 @@ def test_scale_error_if_not_scalar():
         profile_1 * profile_1
     with pytest.raises(AssertionError):
         profile_1 * [1.0, 2.0]
+
+
+def test_add_transverse_profiles():
+    # Add the two profiles together
+    trans_profile_1 = MockTransverseProfile(1.0)
+    trans_profile_2 = MockTransverseProfile(2.0)
+    summed_trans_profile = trans_profile_1 + trans_profile_2
+    # Check that the result is a SummedTransverseProfile object
+    assert isinstance(summed_trans_profile, SummedTransverseProfile)
+    # Check that the profiles are stored correctly
+    assert summed_trans_profile.transverse_profiles[0] == trans_profile_1
+    assert summed_trans_profile.transverse_profiles[1] == trans_profile_2
+    # Check that the evaluate method works
+    assert np.allclose(summed_trans_profile.evaluate(0, 0), 3.0)
+
+
+def test_add_transverse_error_if_not_all_transverse_profiles():
+    trans_profile_1 = MockTransverseProfile(1.0)
+    with pytest.raises(AssertionError):
+        trans_profile_1 + 1.0
+
+
+def test_scale_transverse_profiles():
+    # Add the two profiles together
+    trans_profile_1 = MockTransverseProfile(1.0)
+    scaled_trans_profile = 2.0 * trans_profile_1
+    scaled_trans_profile_right = trans_profile_1 * 2.0
+    # Check that the result is a ScaledProfile object
+    assert isinstance(scaled_trans_profile, ScaledTransverseProfile)
+    assert isinstance(scaled_trans_profile_right, ScaledTransverseProfile)
+    # Check that the profiles are stored correctly
+    assert scaled_trans_profile.transverse_profile == trans_profile_1
+    # Check that the evaluate method works
+    assert np.allclose(scaled_trans_profile.evaluate(0, 0), 2.0)
+    assert np.allclose(scaled_trans_profile.evaluate(0, 0), 2.0)
+
+
+def test_scale_trans_error_if_not_scalar():
+    trans_profile_1 = MockTransverseProfile(1.0)
+    with pytest.raises(AssertionError):
+        trans_profile_1 * trans_profile_1
+    with pytest.raises(AssertionError):
+        trans_profile_1 * [1.0, 2.0]
diff --git a/tests/test_t2z2t.py b/tests/test_t2z2t.py
index 6fc8713..6c31156 100644
--- a/tests/test_t2z2t.py
+++ b/tests/test_t2z2t.py
@@ -94,7 +94,7 @@ def test_3D_case(gaussian):
     tau = gaussian.long_profile.tau
     lo = (-5 * w0, -5 * w0, -3.5 * tau)
     hi = (5 * w0, 5 * w0, 3.5 * tau)
-    npoints = (160, 160, 65)
+    npoints = (160, 164, 65)
 
     laser_t_in = Laser(dim, lo, hi, npoints, gaussian)
     laser_t_out = Laser(dim, lo, hi, npoints, gaussian)