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 Coords, Grid, ModeIndexDataArray, ScalarFieldDataArray, ScalarModeFieldDataArray
 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
@@ -38,6 +38,69 @@
 SOLVER_ID = "Solver-ID"
 
 
+def make_fill_fraction_mode_data():
+    freq = np.array([2e14])
+    mode_spec = td.ModeSpec(num_modes=2)
+    monitor = td.ModeSolverMonitor(
+        size=(3.0, 0.0, 3.0),
+        center=(0.0, 0.0, 0.0),
+        freqs=freq,
+        mode_spec=mode_spec,
+        name="fill_fraction",
+    )
+
+    grid = Grid(
+        boundaries=Coords(
+            x=np.array([-1.5, -0.5, 0.5, 1.5]),
+            y=np.array([-0.5, 0.5]),
+            z=np.array([-1.5, -0.5, 0.5, 1.5]),
+        )
+    )
+
+    coords = {
+        "x": np.array([-1.0, 0.0, 1.0]),
+        "y": np.array([0.0]),
+        "z": np.array([-1.0, 0.0, 1.0]),
+        "f": freq,
+        "mode_index": np.arange(2),
+    }
+    shape = (3, 1, 3, 1, 2)
+
+    ex_data = np.zeros(shape, dtype=complex)
+    ex_data[1, 0, 1, 0, 0] = 2.0
+    for ix in (0, 2):
+        for iz in (0, 2):
+            ex_data[ix, 0, iz, 0, 1] = 1.0
+
+    zero_data = np.zeros(shape, dtype=complex)
+
+    fields = {
+        "Ex": ScalarModeFieldDataArray(ex_data, coords=coords),
+        "Ey": ScalarModeFieldDataArray(np.copy(zero_data), coords=coords),
+        "Ez": ScalarModeFieldDataArray(np.copy(zero_data), coords=coords),
+        "Hx": ScalarModeFieldDataArray(np.copy(zero_data), coords=coords),
+        "Hy": ScalarModeFieldDataArray(np.copy(zero_data), coords=coords),
+        "Hz": ScalarModeFieldDataArray(np.copy(zero_data), coords=coords),
+    }
+
+    n_complex = ModeIndexDataArray(
+        np.array([[1.6 + 0.0j, 1.3 + 0.0j]]),
+        coords={"f": freq, "mode_index": np.arange(2)},
+    )
+
+    data = ModeSolverData(
+        monitor=monitor,
+        symmetry=(0, 0, 0),
+        symmetry_center=(0.0, 0.0, 0.0),
+        grid_expanded=grid,
+        n_complex=n_complex,
+        **fields,
+    )
+
+    bounding_box = td.Box(center=(0.0, 0.0, 0.0), size=(1.0, 2.0, 1.0))
+    return data, bounding_box
+
+
 @pytest.fixture
 def mock_remote_api(monkeypatch):
     def void(*args, **kwargs):
@@ -1343,21 +1406,41 @@ def test_modes_filter_sort():
     for key in get_args(MODE_DATA_KEYS):
         print(key)
         # Test ascending
-        sort_spec = td.ModeSortSpec(sort_key=key, sort_order="ascending", track_freq=None)
+        sort_kwargs = {
+            "sort_key": key,
+            "sort_order": "ascending",
+            "track_freq": None,
+        }
+        if key == "fill_fraction_box":
+            sort_kwargs["bounding_box"] = td.Box(center=PLANE.center, size=(5.0, 4.0, 5.0))
+        sort_spec = td.ModeSortSpec(**sort_kwargs)
         modes = modes.sort_modes(sort_spec)
         metric = getattr(modes, key)
         assert np.all(metric.diff(dim="mode_index") >= 0)
 
         # Test descending
-        sort_spec = td.ModeSortSpec(sort_key=key, sort_order="descending", track_freq=None)
+        sort_kwargs = {
+            "sort_key": key,
+            "sort_order": "descending",
+            "track_freq": None,
+        }
+        if key == "fill_fraction_box":
+            sort_kwargs["bounding_box"] = td.Box(center=PLANE.center, size=(5.0, 4.0, 5.0))
+        sort_spec = td.ModeSortSpec(**sort_kwargs)
         modes = modes.sort_modes(sort_spec)
         metric = getattr(modes, key)
         assert np.all(metric.diff(dim="mode_index") <= 0)
 
         # Test descending with a large reference value should be the same as ascending
-        sort_spec = td.ModeSortSpec(
-            sort_key=key, sort_order="descending", sort_reference=100, track_freq=None
-        )
+        sort_kwargs = {
+            "sort_key": key,
+            "sort_order": "descending",
+            "sort_reference": 100,
+            "track_freq": None,
+        }
+        if key == "fill_fraction_box":
+            sort_kwargs["bounding_box"] = td.Box(center=PLANE.center, size=(5.0, 4.0, 5.0))
+        sort_spec = td.ModeSortSpec(**sort_kwargs)
         modes = modes.sort_modes(sort_spec)
         metric = getattr(modes, key)
         assert np.all(metric.diff(dim="mode_index") >= 0)
@@ -1503,3 +1586,110 @@ 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)
+
+
+def test_mode_sort_spec_fill_fraction_box_filter_drops_modes():
+    data, bounding_box = make_fill_fraction_mode_data()
+
+    sort_spec = td.ModeSortSpec(
+        filter_key="fill_fraction_box",
+        filter_reference=0.5,
+        filter_order="over",
+        drop_modes=True,
+        bounding_box=bounding_box,
+    )
+
+    filtered = data.sort_modes(sort_spec)
+
+    assert filtered.n_eff.sizes["mode_index"] == 1
+    assert filtered.monitor.mode_spec.num_modes == 1
+
+    fills = data.fill_fraction(bounding_box)
+    assert np.isclose(fills.isel(mode_index=0, f=0).item(), 1.0)
+    assert np.isclose(fills.isel(mode_index=1, f=0).item(), 0.0)
+
+
+def test_mode_sort_spec_fill_fraction_box_requires_bounding_box():
+    with pytest.raises(pydantic.ValidationError):
+        td.ModeSortSpec(filter_key="fill_fraction_box")
+
+
+def test_mode_data_fill_fraction_box_requires_intersection():
+    data, _ = make_fill_fraction_mode_data()
+    with pytest.raises(ValidationError):
+        data.fill_fraction(td.Box(center=(0.0, 2.0, 0.0), size=(1.0, 1.0, 1.0)))