Update the ModeSolver by adding perturbation-based group index and group velocity dispersion calculatoion method

tidy3d/components/data/dataset.py

@@ -390,13 +390,26 @@ class ModeSolverDataset(ElectromagneticFieldDataset):
         description="Index associated with group velocity of the mode.",
     )
 
+    n_group_analytic: GroupIndexDataArray = pd.Field(
+        None,
+        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,
     )
 
+    dispersion_analytic: GroupIndexDataArray = 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 their associated data."""
@@ -431,6 +444,17 @@ def n_group(self) -> GroupIndexDataArray:
             )
         return self.n_group_raw
 
+    @property
+    def n_group_new(self) -> GroupIndexDataArray:
+        """Group index."""
+        if self.n_group_analytic is None:
+            log.warning(
+                "The group index was not computed. To calculate group index, pass "
+                "'calculate_group_index = True' in the 'ModeSpec'.",
+                log_once=True,
+            )
+        return self.n_group_analytic
+
     @property
     def dispersion(self) -> ModeDispersionDataArray:
         r"""Dispersion parameter.

tidy3d/components/mode.py

@@ -145,6 +145,14 @@ class ModeSpec(Tidy3dBaseModel):
         f"default of {GROUP_INDEX_STEP} is used.",
     )
 
+    calculate_group_index: bool = pd.Field(
+        False,
+        title="Perturbation-based calculation of the group index and the GVD",
+        description="Control the computation of the group index and the group velocity dispersion"
+        "alongside the effective index. If set to 'True', the perturbation theory based algorithm"
+        " is used for calculation.",
+    )
+
     @pd.validator("bend_axis", always=True)
     @skip_if_fields_missing(["bend_radius"])
     def bend_axis_given(cls, val, values):
@@ -207,3 +215,43 @@ def check_precision(cls, values):
                 )
 
         return values
+
+class SingleFreqModesData(Tidy3dBaseModel):
+    """A data class to store the modes data for a single frequency"""
+
+    E_vectors: np.ndarray = pd.Field(
+        None, title="E field", description="Electric field of the eigenmodes, shape (3, N, num_modes)."
+    )
+
+    H_vectors: np.ndarray = pd.Field(
+        None, title="H field", description="Magnetic field of the eigenmodes, shape (3, N, num_modes)."
+    )
+
+    E_fields: np.ndarray = pd.Field(
+        None, title="E field", description="Electric field of the eigenmodes, shape (3, Nx, Ny, 1, num_modes)."
+    )
+
+    H_fields: np.ndarray = pd.Field(
+        None, title="H field", description="Magnetic field of the eigenmodes, shape (3, Nx, Ny, 1, num_modes)."
+    )
+
+    n_eff: np.ndarray = pd.Field(
+        None, title="Mode refractive index", description="Real part of the effective index, shape (num_modes, )."
+    )
+
+    k_eff: np.ndarray = pd.Field(
+        None, title="Mode absorption index", description="Imaginary part of the effective index, shape (num_modes, )."
+    )
+
+    n_group: np.ndarray = pd.Field(
+        None, title="Mode group index", description="Real part of the effective group index, shape (num_modes, )."
+    )
+
+    GVD: np.ndarray = pd.Field(
+        None, title="Group velocity dispersion", description="Group velocity dispersion data, shape (num_modes, )."
+    )
+
+    eps_spec : Literal["diagonal", "tensorial_real", "tensorial_complex"] = pd.Field(
+        None,
+        title="Permittivity characterization on the mode solver's plane",
+    )

tidy3d/plugins/mode/mode_solver.py

@@ -19,6 +19,8 @@
 from ...components.data.data_array import (
     FreqModeDataArray,
     ModeIndexDataArray,
+    GroupIndexDataArray,
+    ModeDispersionDataArray,
     ScalarModeFieldDataArray,
 )
 from ...components.data.monitor_data import ModeSolverData
@@ -370,7 +372,7 @@ def _data_on_yee_grid(self) -> ModeSolverData:
         )
 
         # Compute and store the modes at all frequencies
-        n_complex, fields, eps_spec = solver._solve_all_freqs(
+        n_complex, fields, n_group, n_GVD, eps_spec = solver._solve_all_freqs(
             coords=_solver_coords, symmetry=solver.solver_symmetry
         )
 
@@ -384,6 +386,29 @@ def _data_on_yee_grid(self) -> ModeSolverData:
         )
         data_dict = {"n_complex": index_data}
 
+        if n_group is not None:
+            if n_group[0] is not None:
+                n_group_data = GroupIndexDataArray(
+                    np.stack(n_group, axis=0),
+                    coords=dict(
+                        f=list(solver.freqs),
+                        mode_index=np.arange(solver.mode_spec.num_modes),
+                    ),
+                )
+
+                data_dict["n_group_analytic"] = n_group_data
+
+        if n_GVD is not None:
+            if n_GVD[0] is not None:
+                n_GVD_data = ModeDispersionDataArray(
+                    np.stack(n_GVD, axis=0),
+                    coords=dict(
+                        f=list(solver.freqs),
+                        mode_index=np.arange(solver.mode_spec.num_modes),
+                    ),
+                )
+                data_dict["dispersion_analytic"] = n_GVD_data
+
         # Construct the field data on Yee grid
         for field_name in ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz"):
             xyz_coords = solver.grid_snapped[field_name].to_list
@@ -670,15 +695,19 @@ def _solve_all_freqs(
 
         fields = []
         n_complex = []
+        n_group = []
+        n_GVD = []
         eps_spec = []
         for freq in self.freqs:
-            n_freq, fields_freq, eps_spec_freq = self._solve_single_freq(
+            n_freq, fields_freq, n_group_freq, GVD_freq, eps_spec_freq = self._solve_single_freq(
                 freq=freq, coords=coords, symmetry=symmetry
             )
             fields.append(fields_freq)
             n_complex.append(n_freq)
+            n_group.append(n_group_freq)
+            n_GVD.append(GVD_freq)
             eps_spec.append(eps_spec_freq)
-        return n_complex, fields, eps_spec
+        return n_complex, fields, n_group, n_GVD, eps_spec
 
     def _solve_all_freqs_relative(
         self,
@@ -731,7 +760,8 @@ def _solve_single_freq(
         if not LOCAL_SOLVER_IMPORTED:
             raise ImportError(IMPORT_ERROR_MSG)
 
-        solver_fields, n_complex, eps_spec = compute_modes(
+        # solver_fields, n_complex, eps_spec = compute_modes(
+        modes_data = compute_modes(
             eps_cross=self._solver_eps(freq),
             coords=coords,
             freq=freq,
@@ -740,8 +770,11 @@ def _solve_single_freq(
             direction=self.direction,
         )
 
+        solver_fields = fields = np.stack((modes_data.E_fields, modes_data.H_fields), axis=0)
         fields = self._postprocess_solver_fields(solver_fields)
-        return n_complex, fields, eps_spec
+
+        n_complex = modes_data.n_eff + 1j * modes_data.k_eff
+        return n_complex, fields, modes_data.n_group, modes_data.GVD, modes_data.eps_spec
 
     def _rotate_field_coords_inverse(self, field: FIELD) -> FIELD:
         """Move the propagation axis to the z axis in the array."""
@@ -782,7 +815,7 @@ def _solve_single_freq_relative(
 
         solver_basis_fields = self._postprocess_solver_fields_inverse(basis_fields)
 
-        solver_fields, n_complex, eps_spec = compute_modes(
+        modes_data = compute_modes(
             eps_cross=self._solver_eps(freq),
             coords=coords,
             freq=freq,
@@ -792,8 +825,12 @@ def _solve_single_freq_relative(
             solver_basis_fields=solver_basis_fields,
         )
 
+        solver_fields = fields = np.stack((modes_data.E_fields, modes_data.H_fields), axis=0)
         fields = self._postprocess_solver_fields(solver_fields)
-        return n_complex, fields, eps_spec
+
+        n_complex = modes_data.n_eff + 1j * modes_data.k_eff
+
+        return n_complex, fields, modes_data.eps_spec
 
     def _rotate_field_coords(self, field: FIELD) -> FIELD:
         """Move the propagation axis=z to the proper order in the array."""