Calculate mode dispersion when group index is calculated (flexcompute#1294)

Signed-off-by: Lucas Heitzmann Gabrielli <lucas@flexcompute.com>

Author: lucas-flexcompute

CHANGELOG.md

@@ -5,6 +5,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- `ModeData.dispersion` and `ModeSolverData.dispersion` are calculated together with the group index.
+
 ## [2.5.0] - 2023-12-13
 
 ### Added

tests/test_plugins/test_mode_solver.py

@@ -556,7 +556,7 @@ def test_mode_solver_2D():
 @pytest.mark.parametrize("local", [True, False])
 @responses.activate
 def test_group_index(mock_remote_api, log_capture, local):
-    """Test group index calculation"""
+    """Test group index and dispersion calculation"""
 
     simulation = td.Simulation(
         size=(5, 5, 1),
@@ -595,10 +595,15 @@ def test_group_index(mock_remote_api, log_capture, local):
     if local:
         assert modes.n_group is None
         assert len(log_capture) == 1
-        assert log_capture[0][0] == 30
-        assert "ModeSpec" in log_capture[0][1]
+        assert modes.dispersion is None
+        assert len(log_capture) == 2
+        for log_msg in log_capture:
+            assert log_msg[0] == 30
+            assert "ModeSpec" in log_msg[1]
         _ = modes.n_group
-        assert len(log_capture) == 1
+        assert len(log_capture) == 2
+        _ = modes.dispersion
+        assert len(log_capture) == 2
 
     # Group index calculated
     ms = ModeSolver(
@@ -613,6 +618,10 @@ def test_group_index(mock_remote_api, log_capture, local):
         assert (modes.n_group.sel(mode_index=0).values < 4.2).all()
         assert (modes.n_group.sel(mode_index=1).values > 3.7).all()
         assert (modes.n_group.sel(mode_index=1).values < 4.0).all()
+        assert (modes.dispersion.sel(mode_index=0).values > 1400).all()
+        assert (modes.dispersion.sel(mode_index=0).values < 1500).all()
+        assert (modes.dispersion.sel(mode_index=1).values > -16500).all()
+        assert (modes.dispersion.sel(mode_index=1).values < -15000).all()
 
 
 def test_pml_params():

tidy3d/components/data/data_array.py

@@ -7,7 +7,14 @@
 import dask
 import h5py
 
-from ...constants import HERTZ, SECOND, MICROMETER, RADIAN
+from ...constants import (
+    HERTZ,
+    MICROMETER,
+    PICOSECOND_PER_NANOMETER_PER_KILOMETER,
+    RADIAN,
+    SECOND,
+    WATT,
+)
 from ...exceptions import DataError, FileError
 from ..types import Bound, Axis
 
@@ -440,7 +447,7 @@ class FluxDataArray(DataArray):
 
     __slots__ = ()
     _dims = ("f",)
-    _data_attrs = {"units": "W", "long_name": "flux"}
+    _data_attrs = {"units": WATT, "long_name": "flux"}
 
 
 class FluxTimeDataArray(DataArray):
@@ -455,7 +462,7 @@ class FluxTimeDataArray(DataArray):
 
     __slots__ = ()
     _dims = ("t",)
-    _data_attrs = {"units": "W", "long_name": "flux"}
+    _data_attrs = {"units": WATT, "long_name": "flux"}
 
 
 class ModeAmpsDataArray(DataArray):
@@ -491,6 +498,41 @@ class ModeIndexDataArray(DataArray):
     _data_attrs = {"long_name": "Propagation index"}
 
 
+class GroupIndexDataArray(DataArray):
+    """Group index of a mode.
+
+    Example
+    -------
+    >>> f = [2e14, 3e14]
+    >>> mode_index = np.arange(4)
+    >>> coords = dict(f=f, mode_index=mode_index)
+    >>> data = GroupIndexDataArray((1+1j) * np.random.random((2,4)), coords=coords)
+    """
+
+    __slots__ = ()
+    _dims = ("f", "mode_index")
+    _data_attrs = {"long_name": "Group index"}
+
+
+class ModeDispersionDataArray(DataArray):
+    """Dispersion parameter of a mode.
+
+    Example
+    -------
+    >>> f = [2e14, 3e14]
+    >>> mode_index = np.arange(4)
+    >>> coords = dict(f=f, mode_index=mode_index)
+    >>> data = ModeDispersionDataArray((1+1j) * np.random.random((2,4)), coords=coords)
+    """
+
+    __slots__ = ()
+    _dims = ("f", "mode_index")
+    _data_attrs = {
+        "long_name": "Dispersion parameter",
+        "units": PICOSECOND_PER_NANOMETER_PER_KILOMETER,
+    }
+
+
 class FieldProjectionAngleDataArray(DataArray):
     """Far fields in frequency domain as a function of angles theta and phi.
 
@@ -654,6 +696,8 @@ class IndexedDataArray(DataArray):
     FluxTimeDataArray,
     ModeAmpsDataArray,
     ModeIndexDataArray,
+    GroupIndexDataArray,
+    ModeDispersionDataArray,
     FieldProjectionAngleDataArray,
     FieldProjectionCartesianDataArray,
     FieldProjectionKSpaceDataArray,

tidy3d/components/data/dataset.py

@@ -13,7 +13,7 @@
 
 from .data_array import DataArray
 from .data_array import ScalarFieldDataArray, ScalarFieldTimeDataArray, ScalarModeFieldDataArray
-from .data_array import ModeIndexDataArray
+from .data_array import ModeIndexDataArray, GroupIndexDataArray, ModeDispersionDataArray
 from .data_array import TriangleMeshDataArray
 from .data_array import TimeDataArray
 from .data_array import PointDataArray, IndexedDataArray, CellDataArray, SpatialDataArray
@@ -23,6 +23,7 @@
 from ..types import Axis, Bound, ArrayLike, Ax, Coordinate, Literal
 from ..types import vtk, requires_vtk
 from ...exceptions import DataError, ValidationError, Tidy3dNotImplementedError
+from ...constants import PICOSECOND_PER_NANOMETER_PER_KILOMETER
 from ...log import log
 
 
@@ -354,31 +355,38 @@ class ModeSolverDataset(ElectromagneticFieldDataset):
         description="Complex-valued effective propagation constants associated with the mode.",
     )
 
-    n_group_raw: ModeIndexDataArray = pd.Field(
+    n_group_raw: GroupIndexDataArray = pd.Field(
         None,
-        alias="n_group",
+        alias="n_group",  # This is for backwards compatibility only when loading old data
         title="Group Index",
         description="Index associated with group velocity of the mode.",
     )
 
+    dispersion_raw: ModeDispersionDataArray = pd.Field(
+        None,
+        title="Dispersion",
+        description="Dispersion parameter for the mode.",
+        units=PICOSECOND_PER_NANOMETER_PER_KILOMETER,
+    )
+
     @property
     def field_components(self) -> Dict[str, DataArray]:
         """Maps the field components to thier associated data."""
 
         return {field: getattr(self, field) for field in ["Ex", "Ey", "Ez", "Hx", "Hy", "Hz"]}
 
     @property
-    def n_eff(self):
+    def n_eff(self) -> ModeIndexDataArray:
         """Real part of the propagation index."""
         return self.n_complex.real
 
     @property
-    def k_eff(self):
+    def k_eff(self) -> ModeIndexDataArray:
         """Imaginary part of the propagation index."""
         return self.n_complex.imag
 
     @property
-    def n_group(self):
+    def n_group(self) -> GroupIndexDataArray:
         """Group index."""
         if self.n_group_raw is None:
             log.warning(
@@ -388,6 +396,22 @@ def n_group(self):
             )
         return self.n_group_raw
 
+    @property
+    def dispersion(self) -> ModeDispersionDataArray:
+        r"""Dispersion parameter.
+
+        .. math::
+
+           D = -\frac{\lambda}{c_0} \frac{{\rm d}^2 n_{\text{eff}}}{{\rm d}\lambda^2}
+        """
+        if self.dispersion_raw is None:
+            log.warning(
+                "The dispersion was not computed. To calculate dispersion, pass "
+                "'group_index_step = True' in the 'ModeSpec'.",
+                log_once=True,
+            )
+        return self.dispersion_raw
+
     def plot_field(self, *args, **kwargs):
         """Warn user to use the :class:`.ModeSolver` ``plot_field`` function now."""
         raise DeprecationWarning(

tidy3d/components/data/monitor_data.py

@@ -11,7 +11,8 @@
 from pandas import DataFrame
 
 from .data_array import FluxTimeDataArray, FluxDataArray
-from .data_array import MixedModeDataArray, ModeIndexDataArray, ModeAmpsDataArray
+from .data_array import MixedModeDataArray, ModeAmpsDataArray
+from .data_array import ModeIndexDataArray, GroupIndexDataArray, ModeDispersionDataArray
 from .data_array import FieldProjectionAngleDataArray, FieldProjectionCartesianDataArray
 from .data_array import FieldProjectionKSpaceDataArray
 from .data_array import DataArray, DiffractionDataArray
@@ -32,7 +33,7 @@
 from ..source import SourceTimeType, CustomFieldSource
 from ..medium import Medium, MediumType
 from ...exceptions import SetupError, DataError, Tidy3dNotImplementedError, ValidationError
-from ...constants import ETA_0, C_0, MICROMETER
+from ...constants import ETA_0, C_0, MICROMETER, PICOSECOND_PER_NANOMETER_PER_KILOMETER
 from ...log import log
 
 from ..base_sim.data.monitor_data import AbstractMonitorData
@@ -1160,24 +1161,44 @@ def _group_index_post_process(self, frequency_step: float) -> ModeSolverData:
         back = slice(0, num_freqs, 3)
         center = slice(1, num_freqs, 3)
         fwd = slice(2, num_freqs, 3)
+        freqs = freqs[center]
 
         # calculate group index
         n_center = self.n_eff.isel(f=center).values
         n_backward = self.n_eff.isel(f=back).values
         n_forward = self.n_eff.isel(f=fwd).values
 
-        n_group_data = n_center + (n_forward - n_backward) / (2 * frequency_step)
-        n_group = ModeIndexDataArray(
+        inv_step = 1 / frequency_step
+        # n_g = n + f * df/dn
+        # dn/df = (n+ - n-) / (2 f df)
+        n_group_data = n_center + (n_forward - n_backward) * inv_step * 0.5
+        # D = -2 * pi * c / lda^2 * d(v_g^-1)/dw = -(f / c)^2 * (2 * dn/df + f * d2n/df2)
+        # d2n/df2 = (n+ - 2n + n-) / (f df)^2
+        # The '1e18' factor converts from s/um^2 to ps/(nm km)
+        dispersion_data = (
+            (n_forward * (inv_step + 1) + n_backward * (inv_step - 1) - n_center * inv_step * 2)
+            * freqs.reshape((-1, 1))
+            * (-1e18 * inv_step / C_0**2)
+        )
+
+        mode_index = list(self.n_complex.coords["mode_index"].values)
+        f = list(freqs)
+        n_group = GroupIndexDataArray(
             n_group_data,
-            coords={
-                "f": list(freqs[center]),
-                "mode_index": list(self.n_complex.coords["mode_index"].values),
-            },
-            attrs={"long name": "Group index"},
+            coords={"f": f, "mode_index": mode_index},
+        )
+
+        dispersion = ModeDispersionDataArray(
+            dispersion_data,
+            coords={"f": f, "mode_index": mode_index},
         )
 
         # remove data corresponding to frequencies used only for group index calculation
-        update_dict = {"n_complex": self.n_complex.isel(f=center), "n_group_raw": n_group}
+        update_dict = {
+            "n_complex": self.n_complex.isel(f=center),
+            "n_group_raw": n_group,
+            "dispersion_raw": dispersion,
+        }
 
         for key, field in self.field_components.items():
             update_dict[key] = field.isel(f=center)
@@ -1188,7 +1209,7 @@ def _group_index_post_process(self, frequency_step: float) -> ModeSolverData:
         if self.eps_spec:
             update_dict["eps_spec"] = self.eps_spec[center]
 
-        update_dict["monitor"] = self.monitor.updated_copy(freqs=freqs[center])
+        update_dict["monitor"] = self.monitor.updated_copy(freqs=freqs)
 
         return self.copy(update=update_dict)
 
@@ -1392,25 +1413,32 @@ class ModeData(MonitorData):
         description="Complex-valued effective propagation constants associated with the mode.",
     )
 
-    n_group_raw: ModeIndexDataArray = pd.Field(
+    n_group_raw: GroupIndexDataArray = pd.Field(
         None,
-        alias="n_group",
+        alias="n_group",  # This is for backwards compatibility only when loading old data
         title="Group Index",
         description="Index associated with group velocity of the mode.",
     )
 
+    dispersion_raw: ModeDispersionDataArray = pd.Field(
+        None,
+        title="Dispersion",
+        description="Dispersion parameter for the mode.",
+        units=PICOSECOND_PER_NANOMETER_PER_KILOMETER,
+    )
+
     @property
-    def n_eff(self):
+    def n_eff(self) -> ModeIndexDataArray:
         """Real part of the propagation index."""
         return self.n_complex.real
 
     @property
-    def k_eff(self):
+    def k_eff(self) -> ModeIndexDataArray:
         """Imaginary part of the propagation index."""
         return self.n_complex.imag
 
     @property
-    def n_group(self):
+    def n_group(self) -> GroupIndexDataArray:
         """Group index."""
         if self.n_group_raw is None:
             log.warning(
@@ -1420,6 +1448,22 @@ def n_group(self):
             )
         return self.n_group_raw
 
+    @property
+    def dispersion(self) -> ModeDispersionDataArray:
+        r"""Dispersion parameter.
+
+        .. math::
+
+           D = -\frac{\lambda}{c_0} \frac{{\rm d}^2 n_{\text{eff}}}{{\rm d}\lambda^2}
+        """
+        if self.dispersion_raw is None:
+            log.warning(
+                "The dispersion was not computed. To calculate dispersion, pass "
+                "'group_index_step = True' in the 'ModeSpec'.",
+                log_once=True,
+            )
+        return self.dispersion_raw
+
     def normalize(self, source_spectrum_fn) -> ModeData:
         """Return copy of self after normalization is applied using source spectrum function."""
         source_freq_amps = source_spectrum_fn(self.amps.f)[None, :, None]

tidy3d/constants.py

@@ -49,6 +49,7 @@
 PERCMCUBE = "1/cm^3"
 WATT = "W"
 VOLT = "V"
+PICOSECOND_PER_NANOMETER_PER_KILOMETER = "ps/(nm km)"
 
 THERMAL_CONDUCTIVITY = "W/(um*K)"
 SPECIFIC_HEAT_CAPACITY = "J/(kg*K)"