Fix bug in subsection with nonlinear materials source detection

Author: caseyflex

CHANGELOG.md

@@ -20,6 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 - Ensure `path` argument in `run()` function is respected when running under autograd or the adjoint plugin.
+- Bug in `Simulation.subsection` (used in the mode solver) when nonlinear materials rely on information about sources outside of the region.
 
 
 ## [2.7.3] - 2024-09-12

tests/test_components/test_medium.py

@@ -616,6 +616,15 @@ def test_nonlinear_medium(log_capture):
     assert n0 == nonlinear_spec.models[0]._get_n0(n0=None, medium=medium, freqs=[freq0])
     assert freq0 == nonlinear_spec.models[0]._get_freq0(freq0=None, freqs=[freq0])
 
+    # subsection with nonlinear materials needs to hardcode source info
+    sim2 = sim.updated_copy(center=(-4, -4, -4), path="sources/0")
+    sim2 = sim2.updated_copy(
+        models=[td.TwoPhotonAbsorption(beta=1)], path="structures/0/medium/nonlinear_spec"
+    )
+    sim2 = sim2.subsection(region=td.Box(center=(0, 0, 0), size=(1, 1, 0)))
+    assert sim2.structures[0].medium.nonlinear_spec.models[0].n0 == n0
+    assert sim2.structures[0].medium.nonlinear_spec.models[0].freq0 == freq0
+
     # can't detect n0 with different source freqs
     source_time2 = source_time.updated_copy(freq0=2 * freq0)
     source2 = source.updated_copy(source_time=source_time2)

tests/test_components/test_simulation.py

@@ -2452,7 +2452,10 @@ def test_sim_subsection(unstructured, nz):
     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
+    assert len(sim_red.structures) == len(SIM_FULL.structures)
+    for strc_red, strc in zip(sim_red.structures, SIM_FULL.structures):
+        if strc.medium.nonlinear_spec is None:
+            assert strc == strc_red
     sim_red = SIM_FULL.subsection(region=region, remove_outside_custom_mediums=True)
 
     perm = td.SpatialDataArray(

tidy3d/components/medium.py

@@ -159,6 +159,12 @@ def _validate_medium_freqs(self, medium: AbstractMedium, freqs: List[pd.Positive
         """Any additional validation that depends on the central frequencies of the sources."""
         pass
 
+    def _hardcode_medium_freqs(
+        self, medium: AbstractMedium, freqs: List[pd.PositiveFloat]
+    ) -> NonlinearSpec:
+        """Update the nonlinear model to hardcode information on medium and freqs."""
+        return self
+
     def _get_freq0(self, freq0, freqs: List[pd.PositiveFloat]) -> float:
         """Get a single value for freq0."""
 
@@ -426,6 +432,14 @@ def _validate_medium_freqs(self, medium: AbstractMedium, freqs: List[pd.Positive
                     "gain medium are unstable, and are likely to diverge."
                 )
 
+    def _hardcode_medium_freqs(
+        self, medium: AbstractMedium, freqs: List[pd.PositiveFloat]
+    ) -> TwoPhotonAbsorption:
+        """Update the nonlinear model to hardcode information on medium and freqs."""
+        n0 = self._get_n0(n0=self.n0, medium=medium, freqs=freqs)
+        freq0 = self._get_freq0(freq0=self.freq0, freqs=freqs)
+        return self.updated_copy(n0=n0, freq0=freq0)
+
     def _validate_medium(self, medium: AbstractMedium):
         """Check that the model is compatible with the medium."""
         # if n0 is specified, we can go ahead and validate passivity
@@ -515,6 +529,13 @@ def _validate_medium_freqs(self, medium: AbstractMedium, freqs: List[pd.Positive
                     "gain medium are unstable, and are likely to diverge."
                 )
 
+    def _hardcode_medium_freqs(
+        self, medium: AbstractMedium, freqs: List[pd.PositiveFloat]
+    ) -> KerrNonlinearity:
+        """Update the nonlinear model to hardcode information on medium and freqs."""
+        n0 = self._get_n0(n0=self.n0, medium=medium, freqs=freqs)
+        return self.updated_copy(n0=n0)
+
     def _validate_medium(self, medium: AbstractMedium):
         """Check that the model is compatible with the medium."""
         # if n0 is specified, we can go ahead and validate passivity
@@ -584,6 +605,16 @@ def _validate_num_iters(cls, val, values):
             )
         return val
 
+    def _hardcode_medium_freqs(
+        self, medium: AbstractMedium, freqs: List[pd.PositiveFloat]
+    ) -> NonlinearSpec:
+        """Update the nonlinear spec to hardcode information on medium and freqs."""
+        new_models = []
+        for model in self.models:
+            new_model = model._hardcode_medium_freqs(medium=medium, freqs=freqs)
+            new_models.append(new_model)
+        return self.updated_copy(models=new_models)
+
 
 class AbstractMedium(ABC, Tidy3dBaseModel):
     """A medium within which electromagnetic waves propagate."""

tidy3d/components/simulation.py

@@ -1464,11 +1464,26 @@ def subsection(
         if remove_outside_structures:
             new_structures = [strc for strc in self.structures if new_box.intersects(strc.geometry)]
         else:
-            new_structures = self.structures
+            new_structures = list(self.structures)
 
         if sources is None:
             sources = [src for src in self.sources if new_box.intersects(src)]
 
+        # some nonlinear materials depend on the central frequency
+        # we update them with hardcoded freq0
+        freqs = np.array([source.source_time.freq0 for source in self.sources])
+        for i, structure in enumerate(new_structures):
+            medium = structure.medium
+            nonlinear_spec = medium.nonlinear_spec
+            if nonlinear_spec is not None:
+                new_nonlinear_spec = nonlinear_spec._hardcode_medium_freqs(
+                    medium=medium, freqs=freqs
+                )
+                new_structure = structure.updated_copy(
+                    nonlinear_spec=new_nonlinear_spec, path="medium"
+                )
+                new_structures[i] = new_structure
+
         if monitors is None:
             monitors = [mnt for mnt in self.monitors if new_box.intersects(mnt)]
 

tidy3d/plugins/mode/mode_solver.py

@@ -341,7 +341,20 @@ def data_raw(self) -> ModeSolverData:
 
     def _data_on_yee_grid(self) -> ModeSolverData:
         """Solve for all modes, and construct data with fields on the Yee grid."""
-        solver = self.reduced_simulation_copy
+
+        # we try to do reduced simulation copy for efficiency
+        # it should never fail -- if it does, this is likely due to an oversight
+        # in the Simulation.subsection method. but falling back to non-reduced
+        # simulation prevents unneeded errors in this case
+        try:
+            solver = self.reduced_simulation_copy
+        except Exception as e:
+            solver = self
+            log.warning(
+                "Mode solver reduced_simulation_copy failed. "
+                "Falling back to non-reduced simulation, which may be slower. "
+                f"Exception: {str(e)}"
+            )
 
         _, _solver_coords = solver.plane.pop_axis(
             solver._solver_grid.boundaries.to_list, axis=solver.normal_axis