Add option to drop filtered modes

Author: momchil-flex

tests/test_plugins/test_mode_solver.py

@@ -10,12 +10,12 @@
 
 import tidy3d as td
 import tidy3d.plugins.mode.web as msweb
-from tidy3d import ScalarFieldDataArray
+from tidy3d import ModeIndexDataArray, ScalarFieldDataArray
 from tidy3d.components.data.monitor_data import ModeSolverData
 from tidy3d.components.mode.derivatives import create_sfactor_b, create_sfactor_f
 from tidy3d.components.mode.solver import compute_modes
 from tidy3d.components.mode_spec import MODE_DATA_KEYS
-from tidy3d.exceptions import DataError, SetupError
+from tidy3d.exceptions import DataError, SetupError, ValidationError
 from tidy3d.plugins.mode import ModeSolver
 from tidy3d.plugins.mode.mode_solver import MODE_MONITOR_NAME
 from tidy3d.web.core.environment import Env
@@ -1503,3 +1503,78 @@ def test_sort_spec_track_freq():
     assert np.allclose(modes_lowest.Ex.abs, modes_lowest_retracked.Ex.abs)
     assert np.all(modes_lowest.n_eff == modes_lowest_retracked.n_eff)
     assert np.all(modes_lowest.n_group == modes_lowest_retracked.n_group)
+
+
+def test_mode_sort_spec_drop_modes_reduces_modes():
+    freqs = np.array([2e14, 4e14])
+    mode_spec = td.ModeSpec(num_modes=3)
+    monitor = td.ModeSolverMonitor(
+        size=(1.0, 0.0, 1.0),
+        center=(0.0, 0.0, 0.0),
+        freqs=freqs,
+        mode_spec=mode_spec,
+        name="drop_modes",
+    )
+    n_complex = ModeIndexDataArray(
+        np.array(
+            [
+                [1.6 + 0.6j, 1.5 + 0.2j, 1.1 + 0.5j],
+                [1.7 + 0.4j, 1.4 + 0.3j, 1.0 + 0.1j],
+            ]
+        ),
+        coords={"f": freqs, "mode_index": np.arange(3)},
+    )
+    data = ModeSolverData(monitor=monitor, n_complex=n_complex)
+
+    sort_spec = td.ModeSortSpec(
+        filter_key="n_eff",
+        filter_reference=1.3,
+        filter_order="over",
+        sort_key="k_eff",
+        sort_order="ascending",
+        drop_modes=True,
+    )
+
+    sorted_data = data.sort_modes(sort_spec)
+
+    assert sorted_data.n_eff.sizes["mode_index"] == 2
+    assert np.allclose(sorted_data.n_eff.isel(f=0).values, [1.5, 1.6])
+    assert np.allclose(sorted_data.n_eff.isel(f=1).values, [1.4, 1.7])
+    assert sorted_data.monitor.mode_spec.num_modes == 2
+    assert sorted_data.monitor.mode_spec.sort_spec.drop_modes is True
+
+
+def test_mode_sort_spec_drop_modes_all_filtered():
+    freqs = np.array([2e14, 4e14])
+    mode_spec = td.ModeSpec(num_modes=3)
+    monitor = td.ModeSolverMonitor(
+        size=(1.0, 0.0, 1.0),
+        center=(0.0, 0.0, 0.0),
+        freqs=freqs,
+        mode_spec=mode_spec,
+        name="drop_all",
+    )
+    n_complex = ModeIndexDataArray(
+        np.array(
+            [
+                [1.1 + 0.1j, 1.05 + 0.05j, 1.0 + 0.01j],
+                [1.1 + 0.1j, 1.05 + 0.05j, 1.0 + 0.01j],
+            ]
+        ),
+        coords={"f": freqs, "mode_index": np.arange(3)},
+    )
+    data = ModeSolverData(monitor=monitor, n_complex=n_complex)
+
+    sort_spec = td.ModeSortSpec(
+        filter_key="n_eff",
+        filter_reference=2.0,
+        drop_modes=True,
+    )
+
+    with pytest.raises(ValidationError):
+        _ = data.sort_modes(sort_spec)
+
+
+def test_mode_sort_spec_drop_modes_requires_filter():
+    with pytest.raises(pydantic.ValidationError):
+        td.ModeSortSpec(drop_modes=True)

tidy3d/components/data/monitor_data.py

@@ -2322,6 +2322,72 @@ def _apply_mode_reorder(self, sort_inds_2d):
 
         return self.updated_copy(**modify_data)
 
+    def _apply_mode_subset(self, subset_inds_2d: np.ndarray) -> ModeSolverData:
+        """Return copy of self containing only the selected modes.
+
+        Parameters
+        ----------
+        subset_inds_2d : np.ndarray
+            Array of shape ``(num_freqs, num_modes_keep)`` containing the indices of the original
+            modes to retain at each frequency.
+
+        Returns
+        -------
+        :class:`.ModeSolverData`
+            Copy of self with only the retained modes.
+        """
+
+        subset_inds_2d = np.asarray(subset_inds_2d, dtype=int)
+        if subset_inds_2d.ndim != 2:
+            raise DataError(
+                "subset_inds_2d must be a 2D array of shape (num_freqs, num_modes_keep)."
+            )
+
+        num_freqs, num_keep = subset_inds_2d.shape
+        if num_keep == 0:
+            raise DataError("Cannot create a mode subset with zero modes.")
+
+        num_modes_full = self.n_eff["mode_index"].size
+
+        modify_data = {}
+        new_mode_index_coord = np.arange(num_keep)
+
+        for key, data in self.data_arrs.items():
+            if "mode_index" not in data.dims or "f" not in data.dims:
+                continue
+
+            dims_orig = tuple(data.dims)
+            coords_out = {
+                k: (v.values if hasattr(v, "values") else np.asarray(v))
+                for k, v in data.coords.items()
+            }
+
+            f_axis = data.get_axis_num("f")
+            m_axis = data.get_axis_num("mode_index")
+            src_order = (
+                [f_axis] + [ax for ax in range(data.ndim) if ax not in (f_axis, m_axis)] + [m_axis]
+            )
+
+            arr = np.moveaxis(data.data, src_order, range(data.ndim))
+            nf, nm = arr.shape[0], arr.shape[-1]
+            if nf != num_freqs or nm != num_modes_full:
+                raise DataError(
+                    "subset_inds_2d shape does not match array shape in _apply_mode_subset."
+                )
+
+            arr2 = arr.reshape(nf, -1, nm)
+            inds = subset_inds_2d[:, None, :]
+            arr2_subset = np.take_along_axis(arr2, inds, axis=2)
+            arr_subset = arr2_subset.reshape(arr.shape[:-1] + (num_keep,))
+            arr_subset = np.moveaxis(arr_subset, range(data.ndim), src_order)
+
+            coords_out["mode_index"] = new_mode_index_coord
+            coords_out["f"] = data.coords["f"].values
+
+            modify_data[key] = DataArray(arr_subset, coords=coords_out, dims=dims_orig)
+
+        return self.updated_copy(**modify_data)
+
     def sort_modes(
         self, sort_spec: Optional[ModeSortSpec] = None, track_freq: Optional[TrackFreq] = None
     ) -> ModeSolverData:
@@ -2356,8 +2422,9 @@ def sort_modes(
         num_freqs = self.n_eff["f"].size
         num_modes = self.n_eff["mode_index"].size
         all_inds = np.arange(num_modes)
-        identity = np.arange(num_modes)
-        sort_inds_2d = np.tile(identity, (num_freqs, 1))
+        drop_modes = getattr(sort_spec, "drop_modes", False)
+        if drop_modes and sort_spec.filter_key is None:
+            raise ValidationError("ModeSortSpec.drop_modes requires 'filter_key' to be set.")
 
         # Helper to compute ordered indices within a subset
         def _order_indices(indices, vals_all):
@@ -2376,12 +2443,13 @@ def _order_indices(indices, vals_all):
             filter_metric = getattr(self, sort_spec.filter_key)
         if sort_spec.sort_key is not None:
             sort_metric = getattr(self, sort_spec.sort_key)
+        identity = np.arange(num_modes)
+        sort_inds_2d = np.tile(identity, (num_freqs, 1))
 
         for ifreq in range(num_freqs):
             # Build groups according to filter if requested
             if filter_metric is not None:
                 vals_filt = filter_metric.isel(f=ifreq).values
-                # Boolean mask for modes in the first group
                 if sort_spec.filter_order == "over":
                     mask_first = vals_filt >= sort_spec.filter_reference
                 else:
@@ -2406,11 +2474,10 @@ def _order_indices(indices, vals_all):
 
             sort_inds_2d[ifreq, : len(sort_inds)] = sort_inds
 
-        # If all rows are identity, skip
         if np.all(sort_inds_2d == np.tile(identity, (num_freqs, 1))):
             data_sorted = self
         else:
-            data_sorted = self._apply_mode_reorder(sort_inds_2d)  # this creates a copy
+            data_sorted = self._apply_mode_reorder(sort_inds_2d)
             data_sorted = data_sorted.updated_copy(
                 path="monitor/mode_spec", sort_spec=sort_spec, deep=False, validate=False
             )
@@ -2422,6 +2489,38 @@ def _order_indices(indices, vals_all):
         if track_freq and num_freqs > 1:
             data_sorted = data_sorted.overlap_sort(track_freq)
 
+        if drop_modes:
+            # Re-evaluate the filter after sorting/tracking so modes are dropped consistently.
+            filter_metric_sorted = getattr(data_sorted, sort_spec.filter_key)
+            masks_after = []
+            for ifreq in range(num_freqs):
+                vals = filter_metric_sorted.isel(f=ifreq).values
+                if sort_spec.filter_order == "over":
+                    mask = vals >= sort_spec.filter_reference
+                else:
+                    mask = vals <= sort_spec.filter_reference
+                masks_after.append(mask)
+
+            keep_mask = np.all(np.stack(masks_after, axis=0), axis=0)
+            if not np.any(keep_mask):
+                raise ValidationError(
+                    "Filtering removes all modes; relax the filter threshold or disable drop_modes."
+                )
+
+            num_modes_sorted = filter_metric_sorted.sizes["mode_index"]
+            if keep_mask.sum() < num_modes_sorted:
+                keep_inds = np.where(keep_mask)[0]
+                subset_inds_2d = np.tile(keep_inds, (num_freqs, 1))
+                data_subset = data_sorted._apply_mode_subset(subset_inds_2d)
+                mspec = data_subset.monitor.mode_spec
+                mspec_updated = mspec.updated_copy(num_modes=keep_inds.size, validate=False)
+                monitor_updated = data_subset.monitor.updated_copy(
+                    mode_spec=mspec_updated, validate=False
+                )
+                data_sorted = data_subset.updated_copy(
+                    monitor=monitor_updated, deep=False, validate=False
+                )
+
         return data_sorted
 
 

tidy3d/components/mode_spec.py

@@ -35,7 +35,8 @@ class ModeSortSpec(Tidy3dBaseModel):
     applied to ``filter_key``: modes "over" or "under" ``filter_reference`` are placed first,
     with the remaining modes placed next. Second, an optional sorting step orders modes within
     each group according to ``sort_key``, optionally with respect to ``sort_reference`` and in
-    the specified ``sort_order``.
+    the specified ``sort_order``. If ``drop_modes`` is set to ``True``, the modes that do not meet
+    the filter criterion are removed instead of being appended as a second group.
     """
 
     # Filtering stage
@@ -54,6 +55,15 @@ class ModeSortSpec(Tidy3dBaseModel):
         title="Filtering order",
         description="Select whether the first group contains values over or under the reference.",
     )
+    drop_modes: bool = pd.Field(
+        False,
+        title="Drop filtered modes",
+        description=(
+            "If ``True``, modes that do not satisfy the filter criterion are removed entirely "
+            "instead of being appended after the filtered group. Only modes passing the filter "
+            "at every tracked frequency are kept."
+        ),
+    )
 
     # Sorting stage
     sort_key: Optional[MODE_DATA_KEYS] = pd.Field(
@@ -85,6 +95,12 @@ class ModeSortSpec(Tidy3dBaseModel):
         "while at other frequencies it can change depending on the mode tracking.",
     )
 
+    @pd.validator("drop_modes", always=True)
+    def _drop_requires_filter(cls, val, values):
+        if val and values.get("filter_key") is None:
+            raise ValidationError("ModeSortSpec.drop_modes requires 'filter_key' to be set.")
+        return val
+
 
 class FrequencySamplingSpec(Tidy3dBaseModel, ABC):
     """Abstract base class for frequency sampling specifications."""