functionality to select simulation subsection

tests/test_components/test_custom.py

@@ -351,7 +351,7 @@ def test_n_cfl():
     assert med.n_cfl >= 2
 
 
-def verify_custom_medium_methods(mat):
+def verify_custom_medium_methods(mat, reduced_fields=[]):
     """Verify that the methods in custom medium is producing expected results."""
     freq = 1.0
     assert isinstance(mat, AbstractCustomMedium)
@@ -362,6 +362,41 @@ def verify_custom_medium_methods(mat):
     for i in range(3):
         assert np.allclose(eps_grid[i].shape, [len(f) for f in coord_interp.to_list])
 
+    # check reducing data
+    subsection = td.Box(size=(0.3, 0.4, 0.35), center=(0.4, 0.4, 0.4))
+
+    mat_reduced = mat.sel_inside(subsection.bounds)
+
+    for field in reduced_fields:
+        original = getattr(mat, field)
+        reduced = getattr(mat_reduced, field)
+
+        if original is None:
+            assert reduced is None
+            continue
+
+        # data fields in medium classes could be SpatialArrays or 2d tuples of spatial arrays
+        # lets convert everything into 2d tuples of spatial arrays for uniform handling
+        if isinstance(original, td.SpatialDataArray):
+            original = [
+                [
+                    original,
+                ],
+            ]
+            reduced = [
+                [
+                    reduced,
+                ],
+            ]
+
+        for or_set, re_set in zip(original, reduced):
+            assert len(or_set) == len(re_set)
+
+            for ind in range(len(or_set)):
+                diff = (or_set[ind] - re_set[ind]).abs
+                assert diff.does_cover(subsection.bounds)
+                assert np.allclose(diff, 0)
+
     # construct sim
     struct = td.Structure(
         geometry=td.Box(size=(0.5, 0.5, 0.5)),
@@ -375,6 +410,8 @@ def verify_custom_medium_methods(mat):
         structures=(struct,),
     )
     _ = sim.grid
+    sim_reduced = sim.subsection(subsection, remove_outside_custom_mediums=False)
+    sim_reduced = sim.subsection(subsection, remove_outside_custom_mediums=True)
 
     # bkg
     sim = td.Simulation(
@@ -384,6 +421,8 @@ def verify_custom_medium_methods(mat):
         medium=mat,
     )
     _ = sim.grid
+    sim_reduced = sim.subsection(subsection, remove_outside_custom_mediums=False)
+    sim_reduced = sim.subsection(subsection, remove_outside_custom_mediums=True)
 
 
 def test_anisotropic_custom_medium():
@@ -438,7 +477,7 @@ def test_custom_isotropic_medium():
     with pytest.raises(pydantic.ValidationError):
         mat = CustomMedium(permittivity=permittivity, conductivity=sigmatmp)
     mat = CustomMedium(permittivity=permittivity, conductivity=sigmatmp, allow_gain=True)
-    verify_custom_medium_methods(mat)
+    verify_custom_medium_methods(mat, ["permittivity", "conductivity"])
 
     # inconsistent coords
     with pytest.raises(pydantic.ValidationError):
@@ -448,15 +487,15 @@ def test_custom_isotropic_medium():
         mat = CustomMedium(permittivity=permittivity, conductivity=sigmatmp)
 
     mat = CustomMedium(permittivity=permittivity, conductivity=conductivity)
-    verify_custom_medium_methods(mat)
+    verify_custom_medium_methods(mat, ["permittivity", "conductivity"])
 
     mat = CustomMedium(permittivity=permittivity)
-    verify_custom_medium_methods(mat)
+    verify_custom_medium_methods(mat, ["permittivity", "conductivity"])
 
 
-def verify_custom_dispersive_medium_methods(mat):
+def verify_custom_dispersive_medium_methods(mat, reduced_fields=[]):
     """Verify that the methods in custom dispersive medium is producing expected results."""
-    verify_custom_medium_methods(mat)
+    verify_custom_medium_methods(mat, reduced_fields)
     freq = 1.0
     for i in range(3):
         assert mat.eps_dataarray_freq(freq)[i].shape == (Nx, Ny, Nz)
@@ -515,7 +554,7 @@ def test_custom_pole_residue():
 
     eps_inf = td.SpatialDataArray(np.random.random((Nx, Ny, Nz)) + 1, coords=dict(x=X, y=Y, z=Z))
     mat = CustomPoleResidue(eps_inf=eps_inf, poles=((a, c),))
-    verify_custom_dispersive_medium_methods(mat)
+    verify_custom_dispersive_medium_methods(mat, ["eps_inf", "poles"])
     assert mat.n_cfl > 1
 
     # to custom non-dispersive medium
@@ -529,12 +568,12 @@ def test_custom_pole_residue():
         mat_medium = mat.to_medium()
     mat = CustomPoleResidue(eps_inf=eps_inf, poles=((a, c - 0.1),), allow_gain=True)
     mat_medium = mat.to_medium()
-    verify_custom_medium_methods(mat_medium)
+    verify_custom_medium_methods(mat_medium, ["permittivity", "conductivity"])
     assert mat_medium.n_cfl > 1
 
     # custom medium to pole residue
     mat = CustomPoleResidue.from_medium(mat_medium)
-    verify_custom_dispersive_medium_methods(mat)
+    verify_custom_dispersive_medium_methods(mat, ["eps_inf", "poles"])
     assert mat.n_cfl > 1
 
 
@@ -578,14 +617,14 @@ def test_custom_sellmeier():
         mat = CustomSellmeier(coeffs=((b1, c2), (btmp, c2)))
 
     mat = CustomSellmeier(coeffs=((b1, c1), (b2, c2)))
-    verify_custom_dispersive_medium_methods(mat)
+    verify_custom_dispersive_medium_methods(mat, ["coeffs"])
     assert mat.n_cfl == 1
 
     # from dispersion
     n = td.SpatialDataArray(2 + np.random.random((Nx, Ny, Nz)), coords=dict(x=X, y=Y, z=Z))
     dn_dwvl = td.SpatialDataArray(-np.random.random((Nx, Ny, Nz)), coords=dict(x=X, y=Y, z=Z))
     mat = CustomSellmeier.from_dispersion(n=n, dn_dwvl=dn_dwvl, freq=2, interp_method="linear")
-    verify_custom_dispersive_medium_methods(mat)
+    verify_custom_dispersive_medium_methods(mat, ["coeffs"])
     assert mat.n_cfl == 1
 
 
@@ -638,13 +677,13 @@ def test_custom_lorentz():
     mat = CustomLorentz(
         eps_inf=eps_inf, coeffs=((de1, f1, delta1), (detmp, f2, delta2)), allow_gain=True
     )
-    verify_custom_dispersive_medium_methods(mat)
+    verify_custom_dispersive_medium_methods(mat, ["eps_inf", "coeffs"])
     assert mat.n_cfl > 1
 
     mat = CustomLorentz(
         eps_inf=eps_inf, coeffs=((de1, f1, delta1), (de2, f2, delta2)), subpixel=True
     )
-    verify_custom_dispersive_medium_methods(mat)
+    verify_custom_dispersive_medium_methods(mat, ["eps_inf", "coeffs"])
     assert mat.n_cfl > 1
     assert mat.pole_residue.subpixel
 
@@ -681,7 +720,7 @@ def test_custom_drude():
         mat = CustomDrude(eps_inf=eps_inf, coeffs=((f1, delta1), (ftmp, delta2)))
 
     mat = CustomDrude(eps_inf=eps_inf, coeffs=((f1, delta1), (f2, delta2)))
-    verify_custom_dispersive_medium_methods(mat)
+    verify_custom_dispersive_medium_methods(mat, ["eps_inf", "coeffs"])
     assert mat.n_cfl > 1
 
 
@@ -728,11 +767,11 @@ def test_custom_debye():
         )
         mat = CustomDebye(eps_inf=eps_inf, coeffs=((eps1, tau1), (epstmp, tau2)))
     mat = CustomDebye(eps_inf=eps_inf, coeffs=((eps1, tau1), (epstmp, tau2)), allow_gain=True)
-    verify_custom_dispersive_medium_methods(mat)
+    verify_custom_dispersive_medium_methods(mat, ["eps_inf", "coeffs"])
     assert mat.n_cfl > 1
 
     mat = CustomDebye(eps_inf=eps_inf, coeffs=((eps1, tau1), (eps2, tau2)))
-    verify_custom_dispersive_medium_methods(mat)
+    verify_custom_dispersive_medium_methods(mat, ["eps_inf", "coeffs"])
     assert mat.n_cfl > 1
 
 

tests/test_components/test_simulation.py

@@ -2248,3 +2248,76 @@ def test_to_gds(tmp_path):
     assert np.allclose(areas[(2, 1)], 0.5)
     assert np.allclose(areas[(1, 0)], 0.25 * np.pi * 1.4**2, atol=1e-2)
     assert np.allclose(areas[(0, 0)], 0.25 * np.pi * 1.4**2, atol=1e-2)
+
+
+def test_sim_subsection():
+    region = td.Box(size=(0.3, 0.5, 0.7), center=(0.1, 0.05, 0.02))
+
+    sim_red = SIM_FULL.subsection(region=region)
+    assert sim_red.structures != SIM_FULL.structures
+    sim_red = SIM_FULL.subsection(
+        region=region,
+        symmetry=(1, 0, -1),
+        monitors=[
+            mnt
+            for mnt in SIM_FULL.monitors
+            if not isinstance(mnt, (td.ModeMonitor, td.ModeSolverMonitor))
+        ],
+    )
+    assert sim_red.symmetry == (1, 0, -1)
+    sim_red = SIM_FULL.subsection(
+        region=region, boundary_spec=td.BoundarySpec.all_sides(td.Periodic())
+    )
+    sim_red = SIM_FULL.subsection(region=region, sources=[], grid_spec=td.GridSpec.uniform(dl=20))
+    assert len(sim_red.sources) == 0
+    sim_red = SIM_FULL.subsection(region=region, monitors=[])
+    assert len(sim_red.monitors) == 0
+    sim_red = SIM_FULL.subsection(region=region, remove_outside_structures=False)
+    assert sim_red.structures == SIM_FULL.structures
+    sim_red = SIM_FULL.subsection(region=region, remove_outside_custom_mediums=True)
+
+    fine_custom_medium = td.CustomMedium(
+        permittivity=td.SpatialDataArray(
+            1 + np.random.random((11, 12, 13)),
+            coords=dict(
+                x=np.linspace(-0.51, 0.52, 11),
+                y=np.linspace(-1.02, 1.04, 12),
+                z=np.linspace(-1.51, 1.51, 13),
+            ),
+        )
+    )
+
+    sim = SIM_FULL.updated_copy(
+        structures=[
+            td.Structure(
+                geometry=td.Box(size=(1, 2, 3)),
+                medium=fine_custom_medium,
+            )
+        ],
+        medium=fine_custom_medium,
+    )
+    sim_red = sim.subsection(region=region, remove_outside_custom_mediums=True)
+
+    # check automatic symmetry expansion
+    sim_sym = SIM_FULL.updated_copy(
+        symmetry=(-1, 0, 1),
+        sources=[src for src in SIM_FULL.sources if not isinstance(src, td.TFSF)],
+    )
+    sim_red = sim_sym.subsection(region=region)
+    assert np.allclose(sim_red.center, (0, 0.05, 0.0))
+
+    # check grid is preserved when requested
+    sim_red = SIM_FULL.subsection(
+        region=region, grid_spec="identical", boundary_spec=td.BoundarySpec.all_sides(td.Periodic())
+    )
+    grids_1d = SIM_FULL.grid.boundaries
+    grids_1d_red = sim_red.grid.boundaries
+    tol = 1e-8
+    for full_grid, red_grid in zip(
+        [grids_1d.x, grids_1d.y, grids_1d.z], [grids_1d_red.x, grids_1d_red.y, grids_1d_red.z]
+    ):
+        # find index into full grid at which reduced grid is starting
+        start = red_grid[0]
+        ind = np.argmax(np.logical_and(full_grid >= start - tol, full_grid <= start + tol))
+        # compare
+        assert np.allclose(red_grid, full_grid[ind : ind + len(red_grid)])