Add interpolation to EME

Author: caseyflex

tidy3d/components/eme/data/sim_data.py

@@ -197,11 +197,23 @@ def smatrix_in_basis(
             modes1 = port_modes1
         if not modes2_provided:
             modes2 = port_modes2
-        f1 = list(modes1.field_components.values())[0].f.values
-        f2 = list(modes2.field_components.values())[0].f.values
+        f1 = list(modes1.monitor.freqs)
+        f2 = list(modes2.monitor.freqs)
 
         f = np.array(sorted(set(f1).intersection(f2).intersection(self.simulation.freqs)))
 
+        interp_spec1 = modes1.monitor.mode_spec.interp_spec
+        interp_spec2 = modes2.monitor.mode_spec.interp_spec
+
+        interp_overlaps = False
+        if interp_spec1 is not None and interp_spec2 is not None and interp_spec1 == interp_spec2:
+            interp_overlaps = True
+        else:
+            if interp_spec1 is not None:
+                modes1 = modes1.interpolated_copy
+            if interp_spec2 is not None:
+                modes2 = modes2.interpolated_copy
+
         modes_in_1 = "mode_index" in list(modes1.field_components.values())[0].coords
         modes_in_2 = "mode_index" in list(modes2.field_components.values())[0].coords
 
@@ -259,6 +271,10 @@ def smatrix_in_basis(
                 overlaps1 = modes1.outer_dot(port_modes1, conjugate=False)
                 if not modes_in_1:
                     overlaps1 = overlaps1.expand_dims(dim={"mode_index_0": mode_index_1}, axis=1)
+                if interp_overlaps:
+                    overlaps1 = modes1._interp_dataarray_in_freq(
+                        overlaps1, freqs=f, method=interp_spec1.method
+                    )
                 O1 = overlaps1.sel(f=f, mode_index_1=keep_mode_inds1)
 
                 O1out = O1.rename(mode_index_0="mode_index_out", mode_index_1="mode_index_out_old")
@@ -288,6 +304,9 @@ def smatrix_in_basis(
                 overlaps2 = modes2.outer_dot(port_modes2, conjugate=False)
                 if not modes_in_2:
                     overlaps2 = overlaps2.expand_dims(dim={"mode_index_0": mode_index_2}, axis=1)
+                overlaps2 = modes2._interp_dataarray_in_freq(
+                    overlaps2, freqs=f, method=interp_spec2.method
+                )
                 O2 = overlaps2.sel(f=f, mode_index_1=keep_mode_inds2)
 
                 O2out = O2.rename(mode_index_0="mode_index_out", mode_index_1="mode_index_out_old")

tidy3d/components/eme/grid.py

@@ -11,9 +11,9 @@
 from tidy3d.components.base import Tidy3dBaseModel, skip_if_fields_missing
 from tidy3d.components.geometry.base import Box
 from tidy3d.components.grid.grid import Coords1D
-from tidy3d.components.mode_spec import ModeSpec
+from tidy3d.components.mode_spec import ModeInterpSpec, ModeSpec
 from tidy3d.components.structure import Structure
-from tidy3d.components.types import ArrayFloat1D, Axis, Coordinate, Size, TrackFreq
+from tidy3d.components.types import ArrayFloat1D, Axis, Coordinate, Size
 from tidy3d.constants import RADIAN, fp_eps, inf
 from tidy3d.exceptions import SetupError, ValidationError
 
@@ -26,13 +26,14 @@
 class EMEModeSpec(ModeSpec):
     """Mode spec for EME cells. Overrides some of the defaults and allowed values."""
 
-    track_freq: Union[TrackFreq, None] = pd.Field(
-        None,
-        title="Mode Tracking Frequency",
-        description="Parameter that turns on/off mode tracking based on their similarity. "
-        "Can take values ``'lowest'``, ``'central'``, or ``'highest'``, which correspond to "
-        "mode tracking based on the lowest, central, or highest frequency. "
-        "If ``None`` no mode tracking is performed, which is the default for best performance.",
+    interp_spec: Optional[ModeInterpSpec] = pd.Field(
+        ModeInterpSpec(method="cheb", num_points=3),
+        title="Mode frequency interpolation specification",
+        description="Specification for computing modes at a reduced set of frequencies and "
+        "interpolating to obtain results at all requested frequencies. This can significantly "
+        "reduce computational cost for broadband simulations where modes vary smoothly with "
+        "frequency. Requires frequency tracking to be enabled (``sort_spec.track_freq`` must "
+        "not be ``None``) to ensure consistent mode ordering across frequencies.",
     )
 
     angle_theta: Literal[0.0] = pd.Field(

tidy3d/components/eme/simulation.py

@@ -561,6 +561,7 @@ def mode_solver_monitors(self) -> list[ModeSolverMonitor]:
                 freqs=freqs,
                 mode_spec=mode_specs[i],
                 colocate=False,
+                reduce_data=mode_specs[i].interp_spec is not None,
             )
             monitors.append(monitor)
         return monitors
@@ -1009,6 +1010,18 @@ def _monitor_freqs(self, monitor: Monitor) -> list[pd.NonNegativeFloat]:
             return list(self.freqs)
         return list(monitor.freqs)
 
+    def _monitor_mode_freqs(self, monitor: EMEModeSolverMonitor) -> list[pd.NonNegativeFloat]:
+        """Monitor frequencies."""
+        freqs = set()
+        cell_inds = self._monitor_eme_cell_indices(monitor=monitor)
+        for cell_ind in cell_inds:
+            interp_spec = self.eme_grid.mode_specs[cell_ind].interp_spec
+            if interp_spec is None:
+                freqs |= set(self.freqs)
+            else:
+                freqs |= set(interp_spec.sampling_points(self.freqs))
+        return list(freqs)
+
     def _monitor_num_freqs(self, monitor: Monitor) -> int:
         """Total number of freqs included in monitor."""
         return len(self._monitor_freqs(monitor=monitor))