recalculation of grid correction factors

Author: dbochkov-flexcompute

CHANGELOG.md

@@ -33,6 +33,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added a Gaussian inverse design filter option with autograd gradients and complete padding mode coverage.
 - Added support for argument passing to DRC file when running checks with `DRCRunner.run(..., drc_args={key: value})` in klayout plugin.
 - Added support for `nonlinear_spec` in `CustomMedium` and `CustomDispersiveMedium`.
+- Added `interp_spec` in `ModeSpec` to allow downsampling and interpolation of waveguide modes in frequency.
 
 ### Breaking Changes
 - Edge singularity correction at PEC and lossy metal edges defaults to `True`.

docs/api/mode.rst

@@ -7,4 +7,6 @@ Mode Specifications
    :toctree: _autosummary/
    :template: module.rst
 
-   tidy3d.ModeSpec
+   tidy3d.ModeSpec
+   tidy3d.ModeSortSpec
+   tidy3d.ModeInterpSpec

tidy3d/components/data/monitor_data.py

@@ -18,7 +18,7 @@
 from tidy3d.components.base_sim.data.monitor_data import AbstractMonitorData
 from tidy3d.components.grid.grid import Coords, Grid
 from tidy3d.components.medium import Medium, MediumType
-from tidy3d.components.mode_spec import ModeInterpSpec, ModeSortSpec
+from tidy3d.components.mode_spec import ModeInterpSpec, ModeSortSpec, ModeSpec
 from tidy3d.components.monitor import (
     AuxFieldTimeMonitor,
     DiffractionMonitor,
@@ -45,6 +45,7 @@
     ArrayFloat1D,
     ArrayFloat2D,
     Coordinate,
+    Direction,
     EMField,
     EpsSpecType,
     FreqArray,
@@ -79,6 +80,7 @@
     MixedModeDataArray,
     ModeAmpsDataArray,
     ModeDispersionDataArray,
+    ModeIndexDataArray,
     ScalarFieldDataArray,
     ScalarFieldTimeDataArray,
     TimeDataArray,
@@ -2458,6 +2460,36 @@ class ModeSolverData(ModeData):
         None, title="Amplitudes", description="Unused for ModeSolverData."
     )
 
+    grid_distances_primal: Union[tuple[float], tuple[float, float]] = pd.Field(
+        (0.0,),
+        title="Distances to the Primal Grid",
+        description="Relative distances to the primal grid locations along the normal direction in "
+        "the original simulation grid. Needed to recalculate grid corrections after "
+        "interpolating in frequency.",
+    )
+
+    grid_distances_dual: Union[tuple[float], tuple[float, float]] = pd.Field(
+        (0.0,),
+        title="Distances to the Dual Grid",
+        description="Relative distances to the dual grid locations along the normal direction in "
+        "the original simulation grid. Needed to recalculate grid corrections after "
+        "interpolating in frequency.",
+    )
+
+    @pd.validator("eps_spec", always=True)
+    @skip_if_fields_missing(["monitor"])
+    def eps_spec_match_mode_spec(cls, val, values):
+        """Raise validation error if frequencies in eps_spec does not match frequency list"""
+        if val:
+            mnt = values["monitor"]
+            if (mnt.reduce_data and len(val) != mnt.mode_spec.interp_spec.num_points) or (
+                not mnt.reduce_data and len(val) != len(mnt.freqs)
+            ):
+                raise ValidationError(
+                    "eps_spec must be provided at the same frequencies as mode solver data."
+                )
+        return val
+
     def normalize(self, source_spectrum_fn: Callable[[float], complex]) -> ModeSolverData:
         """Return copy of self after normalization is applied using source spectrum function."""
         return self.copy()
@@ -2468,6 +2500,63 @@ def _normalize_modes(self):
         for field in self.field_components.values():
             field /= scaling
 
+    @staticmethod
+    def _grid_correction_factors(
+        primal_distances: tuple[float, ...],
+        dual_distances: tuple[float, ...],
+        mode_spec: ModeSpec,
+        n_complex: ModeIndexDataArray,
+        direction: Direction,
+        normal_dim: str,
+    ) -> tuple[FreqModeDataArray, FreqModeDataArray]:
+        """Calculate the grid correction factors for the primal and dual grid.
+
+        Parameters
+        ----------
+        primal_distances : tuple[float, ...]
+            Relative distances to the primal grid locations along the normal direction in the original simulation grid.
+        dual_distances : tuple[float, ...]
+            Relative distances to the dual grid locations along the normal direction in the original simulation grid.
+        mode_spec : ModeSpec
+            Mode specification.
+        n_complex : ModeIndexDataArray
+            Effective indices of the modes.
+        direction : Direction
+            Direction of the propagation.
+        normal_dim : str
+            Name of the normal dimension.
+
+        Returns
+        -------
+        tuple[FreqModeDataArray, FreqModeDataArray]
+            Grid correction factors for the primal and dual grid.
+        """
+
+        distances_primal = xr.DataArray(primal_distances, coords={normal_dim: primal_distances})
+        distances_dual = xr.DataArray(dual_distances, coords={normal_dim: dual_distances})
+
+        # Propagation phase at the primal and dual locations. The k-vector is along the propagation
+        # direction, so angle_theta has to be taken into account. The distance along the propagation
+        # direction is the distance along the normal direction over cosine(theta).
+        cos_theta = np.cos(mode_spec.angle_theta)
+        k_vec = cos_theta * 2 * np.pi * n_complex * n_complex.f / C_0
+        if direction == "-":
+            k_vec *= -1
+        phase_primal = np.exp(1j * k_vec * distances_primal)
+        phase_dual = np.exp(1j * k_vec * distances_dual)
+
+        # Fields are modified by a linear interpolation to the exact monitor position
+        if distances_primal.size > 1:
+            phase_primal = phase_primal.interp(**{normal_dim: 0})
+        else:
+            phase_primal = phase_primal.squeeze(dim=normal_dim)
+        if distances_dual.size > 1:
+            phase_dual = phase_dual.interp(**{normal_dim: 0})
+        else:
+            phase_dual = phase_dual.squeeze(dim=normal_dim)
+
+        return FreqModeDataArray(phase_primal), FreqModeDataArray(phase_dual)
+
     @staticmethod
     def _validate_cheb_nodes(freqs: np.ndarray) -> None:
         """Validate that frequencies are approximately at Chebyshev nodes.
@@ -2504,7 +2593,8 @@ def interp_in_freq(
         self,
         freqs: FreqArray,
         method: Literal["linear", "cubic", "cheb"] = "linear",
-        renormalize: Optional[bool] = True,
+        renormalize: bool = True,
+        recalculate_grid_correction: bool = True,
     ) -> ModeSolverData:
         """Interpolate mode data to new frequency points.
 
@@ -2524,10 +2614,10 @@ def interp_in_freq(
             frequencies), ``"cheb"`` for Chebyshev polynomial interpolation using barycentric
             formula (requires 3+ source frequencies at Chebyshev nodes).
             For complex-valued data, real and imaginary parts are interpolated independently.
-        renormalize : Optional[bool] = True
+        renormalize : bool = True
             Whether to renormalize the mode profiles to unity power after interpolation.
-        recalculate_grid_correction : Optional[bool] = True
-            Whether to recalculate the grid correction after interpolation or use interpolated
+        recalculate_grid_correction : bool = True
+            Whether to recalculate the grid correction factors after interpolation or use interpolated
             grid corrections.
 
         Returns
@@ -2563,7 +2653,11 @@ def interp_in_freq(
         """
         # Validate input
         freqs = np.array(freqs)
+
         source_freqs = np.array(self.monitor.freqs)
+        if self.monitor.reduce_data:
+            # it is validated that if reduce_data is True, then interp_spec is not None
+            source_freqs = self.monitor.mode_spec.interp_spec.sampling_points(source_freqs)
 
         # Validate method-specific requirements
         if method == "cubic" and len(source_freqs) < 4:
@@ -2607,14 +2701,30 @@ def interp_in_freq(
         if self.eps_spec is not None:
             update_dict["eps_spec"] = list(
                 self._interp_dataarray_in_freq(
-                    FreqDataArray(self.eps_spec, coords={"f": self.monitor.freqs}),
+                    FreqDataArray(self.eps_spec, coords={"f": source_freqs}),
                     freqs,
                     "nearest",
                 ).data
             )
 
-        # Update monitor with new frequencies
-        update_dict["monitor"] = self.monitor.updated_copy(freqs=list(freqs))
+        # Update monitor with new frequencies, remove interp_spec and set reduce_data to False
+        update_dict["monitor"] = self.monitor.updated_copy(
+            freqs=list(freqs),
+            mode_spec=self.monitor.mode_spec.updated_copy(interp_spec=None),
+            reduce_data=False,
+        )
+
+        if recalculate_grid_correction:
+            update_dict["grid_primal_correction"], update_dict["grid_dual_correction"] = (
+                self._grid_correction_factors(
+                    list(self.grid_distances_primal),
+                    list(self.grid_distances_dual),
+                    self.monitor.mode_spec,
+                    update_dict["n_complex"],
+                    self.monitor.direction,
+                    "xyz"[self.monitor._normal_axis],
+                )
+            )
 
         updated_data = self.updated_copy(**update_dict)
         if renormalize:
@@ -2627,12 +2737,13 @@ def interpolated_copy(self) -> ModeSolverData:
         """Return a copy of the data with interpolated fields."""
         if self.monitor.mode_spec.interp_spec is None or not self.monitor.reduce_data:
             return self
-        return self.interp_in_freq(
+        interpolated_data = self.interp_in_freq(
             freqs=self.monitor.freqs,
             method=self.monitor.mode_spec.interp_spec.method,
             renormalize=True,
-            monitor=self.monitor.updated_copy(reduce_data=False),
+            recalculate_grid_correction=True,
         )
+        return interpolated_data.updated_copy(monitor=self.monitor.updated_copy(reduce_data=False))
 
     @property
     def time_reversed_copy(self) -> FieldData:

tidy3d/components/mode/mode_solver.py

@@ -537,14 +537,17 @@ def _get_data_with_interp(self) -> ModeSolverData:
         # Get data at reduced frequencies
         data_reduced = mode_solver_reduced.data_raw
 
+        # restore original mode_spec
+        data_reduced = data_reduced.updated_copy(
+            monitor=data_reduced.monitor.updated_copy(
+                freqs=self.freqs, mode_spec=self.mode_spec, reduce_data=True
+            )
+        )
+
         if self.reduce_data:
             return data_reduced
 
-        return data_reduced.interp_in_freq(
-            freqs=self.freqs,
-            method=self.mode_spec.interp_spec.method,
-            renormalize=True,
-        )
+        return data_reduced.interpolated_copy
 
     @cached_property
     def grid_snapped(self) -> Grid:
@@ -670,7 +673,7 @@ def rotated_mode_solver_data(self) -> ModeSolverData:
         for _ in self.freqs:
             eps_spec.append("tensorial_complex")
         # finite grid corrections
-        grid_factors = solver._grid_correction(
+        grid_factors, relative_grid_distances = solver._grid_correction(
             simulation=solver.simulation,
             plane=solver.plane,
             mode_spec=solver.mode_spec,
@@ -689,6 +692,8 @@ def rotated_mode_solver_data(self) -> ModeSolverData:
             grid_primal_correction=grid_factors[0],
             grid_dual_correction=grid_factors[1],
             eps_spec=eps_spec,
+            grid_distances_primal=relative_grid_distances[0],
+            grid_distances_dual=relative_grid_distances[1],
             **rotated_mode_fields,
         )
 
@@ -1257,7 +1262,7 @@ def _data_on_yee_grid(self) -> ModeSolverData:
             data_dict[field_name] = scalar_field_data
 
         # finite grid corrections
-        grid_factors = solver._grid_correction(
+        grid_factors, relative_grid_distances = solver._grid_correction(
             simulation=solver.simulation,
             plane=solver.plane,
             mode_spec=solver.mode_spec,
@@ -1275,6 +1280,8 @@ def _data_on_yee_grid(self) -> ModeSolverData:
             grid_expanded=grid_expanded,
             grid_primal_correction=grid_factors[0],
             grid_dual_correction=grid_factors[1],
+            grid_distances_primal=relative_grid_distances[0],
+            grid_distances_dual=relative_grid_distances[1],
             eps_spec=eps_spec,
             **data_dict,
         )
@@ -1329,7 +1336,7 @@ def _data_on_yee_grid_relative(self, basis: ModeSolverData) -> ModeSolverData:
             data_dict[field_name] = scalar_field_data
 
         # finite grid corrections
-        grid_factors = self._grid_correction(
+        grid_factors, relative_grid_distances = self._grid_correction(
             simulation=self.simulation,
             plane=self.plane,
             mode_spec=self.mode_spec,
@@ -1347,6 +1354,8 @@ def _data_on_yee_grid_relative(self, basis: ModeSolverData) -> ModeSolverData:
             grid_expanded=grid_expanded,
             grid_primal_correction=grid_factors[0],
             grid_dual_correction=grid_factors[1],
+            grid_distances_primal=relative_grid_distances[0],
+            grid_distances_dual=relative_grid_distances[1],
             eps_spec=eps_spec,
             **data_dict,
         )
@@ -1892,7 +1901,9 @@ def _grid_correction(
         mode_spec: ModeSpec,
         n_complex: ModeIndexDataArray,
         direction: Direction,
-    ) -> [FreqModeDataArray, FreqModeDataArray]:
+    ) -> tuple[
+        tuple[FreqModeDataArray, FreqModeDataArray], tuple[tuple[float, ...], tuple[float, ...]]
+    ]:
         """
         Compute grid correction factors for the mode fields.
 
@@ -1935,27 +1946,43 @@ def _grid_correction(
         normal_dual = grid.centers.to_list[normal_axis]
         normal_dual = xr.DataArray(normal_dual, coords={normal_dim: normal_dual})
 
-        # Propagation phase at the primal and dual locations. The k-vector is along the propagation
-        # direction, so angle_theta has to be taken into account. The distance along the propagation
-        # direction is the distance along the normal direction over cosine(theta).
-        cos_theta = np.cos(mode_spec.angle_theta)
-        k_vec = cos_theta * 2 * np.pi * n_complex * n_complex.f / C_0
-        if direction == "-":
-            k_vec *= -1
-        phase_primal = np.exp(1j * k_vec * (normal_primal - normal_pos))
-        phase_dual = np.exp(1j * k_vec * (normal_dual - normal_pos))
-
-        # Fields are modified by a linear interpolation to the exact monitor position
-        if normal_primal.size > 1:
-            phase_primal = phase_primal.interp(**{normal_dim: normal_pos})
-        else:
-            phase_primal = phase_primal.squeeze(dim=normal_dim)
-        if normal_dual.size > 1:
-            phase_dual = phase_dual.interp(**{normal_dim: normal_pos})
-        else:
-            phase_dual = phase_dual.squeeze(dim=normal_dim)
+        def find_closest_distances_to_grid_points(
+            normal_pos: float, grid_coords: ArrayFloat1D
+        ) -> tuple[float, float]:
+            """Find the closest points to the normal position in the grid coordinates."""
+
+            if grid_coords.size == 1:
+                return [float(grid_coords.data[0] - normal_pos)]
+
+            distances = grid_coords.data - normal_pos
+            # First, find the signed distance to the closest grid point
+            closest_distance_ind = np.argmin(np.abs(distances))
+            closest_distance = distances[closest_distance_ind]
+
+            # Then, if the closest distance is positive, take the previous point, otherwise take the next point
+            if closest_distance > 0:
+                first_dist = distances[closest_distance_ind - 1]
+                second_dist = distances[closest_distance_ind]
+            else:
+                first_dist = distances[closest_distance_ind]
+                second_dist = distances[closest_distance_ind + 1]
+
+            # Return the two closest points
+            return [first_dist, second_dist]
+
+        primal_closest_distances = find_closest_distances_to_grid_points(normal_pos, normal_primal)
+        dual_closest_distances = find_closest_distances_to_grid_points(normal_pos, normal_dual)
+
+        grid_correction_factors = ModeSolverData._grid_correction_factors(
+            primal_closest_distances,
+            dual_closest_distances,
+            mode_spec,
+            n_complex,
+            direction,
+            normal_dim,
+        )
 
-        return FreqModeDataArray(phase_primal), FreqModeDataArray(phase_dual)
+        return grid_correction_factors, (primal_closest_distances, dual_closest_distances)
 
     @property
     def _is_tensorial(self) -> bool: