Repository: LASY-org/lasy

Issue:
Color map in laser.show()
Laser.show() works by default, but when I try to change colormap I get this error

```py
TypeError                                 Traceback (most recent call last)
Cell In[88], line 1
----> 1 laser.show(show_intensity=True, cmap =ffwd_cmap)

File ~/.conda/envs/icsenv/lib/python3.9/site-packages/lasy/laser.py:438, in Laser.show(self, show_intensity, **kw)
    425 def show(self, show_intensity=False, **kw):
    426     """
    427     Show a 2D image of the laser amplitude or intensity.
    428 
   (...)
    436     **kw : additional arguments to be passed to matplotlib's imshow command
    437     """
--> 438     show_laser(self.grid, self.dim, show_intensity, **kw)

File ~/.conda/envs/icsenv/lib/python3.9/site-packages/lasy/utils/plotting.py:96, in show_laser(grid, dim, show_intensity, **kw)
     94 divider = make_axes_locatable(ax)
     95 cax = divider.append_axes("right", size="5%", pad=0.05)
---> 96 im = ax.imshow(
     97     F_plot, extent=extent, cmap="Reds", aspect="auto", origin="lower", **kw
     98 )
     99 cb = fig.colorbar(im, cax=cax)
    100 cb.set_label(cbar_label)

TypeError: matplotlib.axes._axes.Axes.imshow() got multiple values for keyword argument 'cmap'
```

reproducer using the Gaussian tutorial
```py
from lasy.laser import Laser
from lasy.profiles.gaussian_profile import GaussianProfile

wavelength = 800e-9  # Laser wavelength in meters
polarization = (1, 0)  # Linearly polarized in the x direction
energy = 1.5  # Energy of the laser pulse in joules
spot_size = 25e-6  # Waist of the laser pulse in meters
pulse_duration = 30e-15  # Pulse duration of the laser in seconds
t_peak = 0.0  # Location of the peak of the laser pulse in time

laser_profile = GaussianProfile(
    wavelength, polarization, energy, spot_size, pulse_duration, t_peak
)

dimensions = "rt"  # Use cylindrical geometry
lo = (0, -2.5 * pulse_duration)  # Lower bounds of the simulation box
hi = (5 * spot_size, 2.5 * pulse_duration)  # Upper bounds of the simulation box
num_points = (300, 500)  # Number of points in each dimension

laser = Laser(dimensions, lo, hi, num_points, laser_profile)

laser.show(cmap='plasma')

```

Hints from issue comments:
@MaxThevenet: Could be related to https://github.com/LASY-org/lasy/pull/335.@MaxThevenet: Fixed in #401

Solution patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index ebb72ad..5441d98 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.11.5
+  rev: v0.11.6
   hooks:
     # Run the linter
     - id: ruff
diff --git a/lasy/__init__.py b/lasy/__init__.py
index 43c4ab0..22049ab 100644
--- a/lasy/__init__.py
+++ b/lasy/__init__.py
@@ -1 +1 @@
-__version__ = "0.6.1"
+__version__ = "0.6.2"
diff --git a/lasy/profiles/from_openpmd_profile.py b/lasy/profiles/from_openpmd_profile.py
index 96c7946..7fb2631 100644
--- a/lasy/profiles/from_openpmd_profile.py
+++ b/lasy/profiles/from_openpmd_profile.py
@@ -82,12 +82,22 @@ class FromOpenPMDProfile(FromArrayProfile):
             #           to Ex & Ey at LASY mode decomposition.
             array_list = convert_modes(array_list, geometry, is_envelope, verbose)
             dim = "xyt" if geometry == "cartesian" else "rt"
-            env_array_list = []
-            for array in array_list:
-                grid = create_grid(array, axes, dim, is_envelope=False)
-                omg0 = field_to_envelope(grid, dim)
-                array = grid.get_temporal_field()
-                env_array_list.append(array)
+            # Detect whether Ex (index 0) or Ey (index 1) is strongest (major)
+            if np.max(np.abs(array_list[0])) >= np.max(np.abs(array_list[1])):
+                imajor = 0
+            else:
+                imajor = 1
+            iminor = 1 - imajor
+            # Convert major field to envelope, and measure frequency
+            grid_major = create_grid(array_list[imajor], axes, dim, is_envelope=False)
+            omg0 = field_to_envelope(grid_major, dim)
+            # Convert other field to envelope, for the same frequency
+            grid_minor = create_grid(array_list[iminor], axes, dim, is_envelope=False)
+            field_to_envelope(grid_minor, dim, omg0)
+            # Measure polarization state from Ex and Ey
+            env_array_list = [None, None]
+            env_array_list[imajor] = grid_major.get_temporal_field()
+            env_array_list[iminor] = grid_minor.get_temporal_field()
             array, pol = isolate_polarization(env_array_list, dim)
         wavelength = 2 * np.pi * c / omg0
 
diff --git a/lasy/utils/laser_utils.py b/lasy/utils/laser_utils.py
index 1d4e59e..b691275 100644
--- a/lasy/utils/laser_utils.py
+++ b/lasy/utils/laser_utils.py
@@ -699,7 +699,7 @@ def vector_potential_to_field(grid, omega0, direct=True):
         return 1j * m_e * omega * c * field / e
 
 
-def field_to_envelope(grid, dim, phase_unwrap_nd=False):
+def field_to_envelope(grid, dim, omega0=None, phase_unwrap_nd=False):
     """Convert a field to its complex envelope representation by applying a Hilbert transform.
 
     Parameters
@@ -716,6 +716,10 @@ def field_to_envelope(grid, dim, phase_unwrap_nd=False):
         - ``'rt'`` : The laser pulse is represented on a 2D grid:
                     Cylindrical (r) transversely, and temporal (t) longitudinally.
 
+    omega0 : scalar
+        Central frequency at which the envelope will be defined.
+        If None, the central frequency measured from the field is used.
+
     phase_unwrap_nd : boolean (optional)
         If True, the phase unwrapping is n-dimensional (2- or 3-D depending on dim).
         If False, the phase unwrapping is done in t, treating each transverse cell
@@ -725,26 +729,30 @@ def field_to_envelope(grid, dim, phase_unwrap_nd=False):
     Returns
     -------
     scalar
-        The central angular frequency.
+        If input parameter omega0 was None, return the central angular frequency.
+        Otherwise, return None.
     """
     assert not grid.is_envelope
 
     # Get central wavelength from array
-    _, omg0_h = get_frequency(
-        grid,
-        dim=dim,
-        is_hilbert=False,
-        phase_unwrap_nd=phase_unwrap_nd,
-    )
+    if omega0 is None:
+        _, omg0 = get_frequency(
+            grid,
+            dim=dim,
+            is_hilbert=False,
+            phase_unwrap_nd=phase_unwrap_nd,
+        )
+    else:
+        omg0 = omega0
     # Hilbert transform
     field = hilbert_transform(grid.get_temporal_field())
     # Remove carrier frequency
-    field *= np.exp(1j * omg0_h * grid.axes[-1])
+    field *= np.exp(1j * omg0 * grid.axes[-1])
     # Store envelope
     grid.set_is_envelope(True)
     grid.set_temporal_field(field)
 
-    return omg0_h
+    return omg0 if omega0 is None else None
 
 
 def hilbert_transform(field):
@@ -1385,13 +1393,16 @@ def get_spectral_phase(grid, dim, omega0, method="sum", ordering="zero_center"):
     return phase, omega
 
 
-def get_gdd(grid, dim, omega0, omega_gdd=None, method="sum"):
+def get_dispersion(grid, dim, omega0, order, omega_eval=None, method="sum"):
     r"""
-    Calculate the group delay dispersion (GDD) of the laser.
+    Calculate the n-th order dispersion polynomial of the laser.
 
     .. math::
-        GDD = \frac{\partial^2 \phi(\omega)}{\partial \omega^2}
+        \Phi^{(n)} = \frac{\partial^n \phi(\omega)}{\partial \omega^n}
 
+    where n is the order to which the disperison is calculated (in s^n/rad).
+
+    E.g. order=1, calculates the group delay (GD), order=2 calculates group delay dispersion (GDD) and order=3 calculates the third-order dispersion (TOD).
 
     Parameters
     ----------
@@ -1409,32 +1420,38 @@ def get_gdd(grid, dim, omega0, omega_gdd=None, method="sum"):
     omega0 : float
         Angular frequency at which the laser envelope is defined.
 
-    omega_gdd : float, optional
-        Central angular frequency at which the GDD is calculated, if `None`, `omega0` is used.
+    order : integer
+        Dispersion polynomial order that should be calculated.
+
+    omega_eval : float, optional
+        Central angular frequency at which the dispersion polynomial is calculated, if `None`, `omega0` is used.
 
     method : string, optional
-        Method of retrieving the phase that is used for calculating the GDD. Options are:
+        Method of retrieving the phase that is used for calculating the dispersion. Options are:
 
         - ``'sum'``: Calculates the spectral phase of the spatially summed field (default).
         - ``'on-axis'``: Calculates the on-axis spectral phase.
 
     Returns
     -------
-    gdd: ndarray of floats (1D)
-        Group delay dispersion over the entire spectral range (in s^2)
+    disp: ndarray of floats (1D)
+        n-th order dispersion over the entire spectral range (in s^n/rad)
 
-    gdd0: float
-        Group delay dispersion at the center frequency (in s^2)
+    disp0: float
+        n-th order dispersion at the center frequency (in s^n/rad)
 
     """
     # calculate the spectral phase of the laser pulse
     phase, omega = get_spectral_phase(grid, dim, method=method, omega0=omega0)
 
-    # calculate the second derivative wrt. angular frequency
-    gd = np.gradient(phase, omega, axis=-1)
-    gdd = np.gradient(gd, omega, axis=-1)
+    # calculate the n-th order derivative wrt. angular frequency
+    disp = np.gradient(phase, omega, axis=-1)
+    for _ in range(order - 1):
+        disp = np.gradient(disp, omega, axis=-1)
+
+    # get the dispersion at the specified frequency of the envelope's frequency
+    omega_eval = omega_eval if omega_eval is not None else omega0
+
+    disp0 = interp1d(omega, disp, bounds_error=True)(omega_eval)
 
-    # get the GDD at the specified frequency or the envelope's frequency
-    omega_eval = omega_gdd if omega_gdd is not None else omega0
-    gdd0 = np.interp(omega_eval, omega, gdd)
-    return gdd, gdd0
+    return disp, disp0
diff --git a/lasy/utils/plotting.py b/lasy/utils/plotting.py
index bb85e8c..c89a7b9 100644
--- a/lasy/utils/plotting.py
+++ b/lasy/utils/plotting.py
@@ -30,6 +30,11 @@ def show_laser(grid, dim, show_intensity, **kw):
 
     **kw : additional arguments to be passed to matplotlib's imshow command
     """
+    if "cmap" in kw.keys():
+        pass
+    else:
+        kw["cmap"] = "Reds"  # Set default colormap
+
     if show_intensity:
         F = epsilon_0 * c / 2 * np.abs(grid.get_temporal_field()) ** 2 / 1e4
         cbar_label = r"I (W/cm$^2$)"
@@ -93,9 +98,7 @@ def show_laser(grid, dim, show_intensity, **kw):
     fig, ax = plt.subplots()
     divider = make_axes_locatable(ax)
     cax = divider.append_axes("right", size="5%", pad=0.05)
-    im = ax.imshow(
-        F_plot, extent=extent, cmap="Reds", aspect="auto", origin="lower", **kw
-    )
+    im = ax.imshow(F_plot, extent=extent, aspect="auto", origin="lower", **kw)
     cb = fig.colorbar(im, cax=cax)
     cb.set_label(cbar_label)
     ax.set_xlabel(r"t " + temporal_unit)


Test patch:
diff --git a/tests/test_STC.py b/tests/test_STC.py
index f152123..f5487b4 100644
--- a/tests/test_STC.py
+++ b/tests/test_STC.py
@@ -14,7 +14,7 @@ from lasy.profiles.combined_profile import CombinedLongitudinalTransverseProfile
 from lasy.profiles.gaussian_profile import GaussianProfile
 from lasy.profiles.longitudinal import GaussianLongitudinalProfile
 from lasy.profiles.transverse import GaussianTransverseProfile
-from lasy.utils.laser_utils import get_beta, get_gdd, get_zeta
+from lasy.utils.laser_utils import get_beta, get_dispersion, get_zeta
 
 wavelength = 0.6e-6  # m
 pol = (1, 0)
@@ -99,8 +99,8 @@ err_imag = np.average(
     / np.array(env_combined.imag)
 )
 
-gdd_3d, gdd0_3d = get_gdd(
-    grid=laser_3d.grid, dim=laser_3d.dim, omega0=laser_3d.profile.omega0
+gdd_3d, gdd0_3d = get_dispersion(
+    grid=laser_3d.grid, dim=laser_3d.dim, omega0=laser_3d.profile.omega0, order=2
 )
 [zeta_x, zeta_y], [nu_x, nu_y] = get_zeta(
     laser_3d.dim, laser_3d.grid, 2.0 * np.pi / 0.6e-6
diff --git a/tests/test_laser_utils.py b/tests/test_laser_utils.py
index 5446622..648429c 100644
--- a/tests/test_laser_utils.py
+++ b/tests/test_laser_utils.py
@@ -2,8 +2,14 @@ import numpy as np
 from scipy.constants import c, epsilon_0
 
 from lasy.laser import Laser
+from lasy.optical_elements.polynomial_spectral_phase import PolynomialSpectralPhase
 from lasy.profiles.gaussian_profile import GaussianProfile
-from lasy.utils.laser_utils import compute_laser_energy, get_duration, get_spectrum
+from lasy.utils.laser_utils import (
+    compute_laser_energy,
+    get_dispersion,
+    get_duration,
+    get_spectrum,
+)
 
 
 def get_gaussian_profile():
@@ -48,6 +54,29 @@ def test_laser_analysis_utils():
         tau_rms = get_duration(laser.grid, dim)
         np.testing.assert_approx_equal(2 * tau_rms, laser.profile.tau, significant=3)
 
+        # Check that the spectral phase terms are calculated correctly.
+        gd = 10e-15
+        gdd = 50e-30
+        tod = 100e-45
+        laser_chirped = get_gaussian_laser(dim)
+        dazzler = PolynomialSpectralPhase(
+            omega0=laser_chirped.profile.omega0, delay=gd, gdd=gdd, tod=tod
+        )
+        laser_chirped.apply_optics(dazzler)
+        _, gd_evaluated = get_dispersion(
+            laser_chirped.grid, dim, omega0=laser_chirped.profile.omega0, order=1
+        )
+        _, gdd_evaluated = get_dispersion(
+            laser_chirped.grid, dim, omega0=laser_chirped.profile.omega0, order=2
+        )
+        _, tod_evaluated = get_dispersion(
+            laser_chirped.grid, dim, omega0=laser_chirped.profile.omega0, order=3
+        )
+
+        assert np.isclose(gd, gd_evaluated, atol=laser_chirped.grid.dx[-1])
+        assert np.isclose(gdd, gdd_evaluated, atol=0)
+        assert np.isclose(tod, tod_evaluated, atol=0)
+
 
 def test_laser_normalization_utils():
     """Test the different laser normalization utilities in both geometries."""