Speed up layer refinement by caching and skip exterior structures

Author: weiliangjin2021

tidy3d/components/grid/grid_spec.py

@@ -1226,8 +1226,8 @@ def center_axis(self) -> float:
 
     @cached_property
     def _is_inplane_unbounded(self) -> bool:
-        """Whether the layer is unbounded in inplane dimensions."""
-        return np.isinf(self.size[(self.axis + 1) % 3]) and np.isinf(self.size[(self.axis + 2) % 3])
+        """Whether the layer is unbounded in any of the inplane dimensions."""
+        return np.isinf(self.size[(self.axis + 1) % 3]) or np.isinf(self.size[(self.axis + 2) % 3])
 
     def _unpop_axis(self, ax_coord: float, plane_coord: Any) -> CoordinateOptional:
         """Combine coordinate along axis with identical coordinates on the plane tangential to the axis.
@@ -1318,7 +1318,14 @@ def _corners(self, structure_list: List[Structure]) -> List[CoordinateOptional]:
         if self.corner_finder is None:
             return []
 
-        inplane_points = self.corner_finder.corners(self.axis, self.center_axis, structure_list)
+        # filter structures outside the layer
+        structures_intersect = structure_list
+        if not self._is_inplane_unbounded:
+            structures_intersect = [s for s in structure_list if self.intersects(s.geometry)]
+        inplane_points = self.corner_finder.corners(
+            self.axis, self.center_axis, structures_intersect
+        )
+
         # filter corners outside the inplane bounds
         if not self._is_inplane_unbounded:
             inplane_points = [point for point in inplane_points if self._inplane_inside(point)]
@@ -1625,22 +1632,30 @@ def internal_snapping_points(self, structures: List[Structure]) -> List[Coordina
             snapping_points += layer_spec.generate_snapping_points(list(structures))
         return snapping_points
 
-    def all_snapping_points(self, structures: List[Structure]) -> List[CoordinateOptional]:
+    def all_snapping_points(
+        self,
+        structures: List[Structure],
+        internal_snapping_points: List[CoordinateOptional] = None,
+    ) -> List[CoordinateOptional]:
         """Internal and external snapping points. External snapping points take higher priority.
         So far, internal snapping points are generated by `layer_refinement_specs`.
 
         Parameters
         ----------
         structures : List[Structure]
             List of physical structures.
+        internal_snapping_points : List[CoordinateOptional]
+            If `None`, recomputes internal snapping points.
 
         Returns
         -------
         List[CoordinateOptional]
             List of snapping points coordinates.
         """
 
-        return self.internal_snapping_points(structures) + list(self.snapping_points)
+        if internal_snapping_points is None:
+            return self.internal_snapping_points(structures) + list(self.snapping_points)
+        return internal_snapping_points + list(self.snapping_points)
 
     @property
     def external_override_structures(self) -> List[StructureType]:
@@ -1677,7 +1692,11 @@ def internal_override_structures(
         return override_structures
 
     def all_override_structures(
-        self, structures: List[Structure], wavelength: pd.PositiveFloat, sim_size: Tuple[float, 3]
+        self,
+        structures: List[Structure],
+        wavelength: pd.PositiveFloat,
+        sim_size: Tuple[float, 3],
+        internal_override_structures: List[MeshOverrideStructure] = None,
     ) -> List[StructureType]:
         """Internal and external mesh override structures. External override structures take higher priority.
         So far, internal override structures all come from `layer_refinement_specs`.
@@ -1688,17 +1707,22 @@ def all_override_structures(
             List of structures, with the simulation structure being the first item.
         wavelength : pd.PositiveFloat
             Wavelength to use for minimal step size in vaccum.
+        internal_override_structures : List[MeshOverrideStructure]
+            If `None`, recomputes internal override structures.
 
         Returns
         -------
         List[StructureType]
             List of override structures.
         """
 
-        return (
-            self.internal_override_structures(structures, wavelength, sim_size)
-            + self.external_override_structures
-        )
+        if internal_override_structures is None:
+            return (
+                self.internal_override_structures(structures, wavelength, sim_size)
+                + self.external_override_structures
+            )
+
+        return internal_override_structures + self.external_override_structures
 
     def _min_vacuum_dl_in_autogrid(self, wavelength: float, sim_size: Tuple[float, 3]) -> float:
         """Compute grid step size in vacuum for Autogrd. If AutoGrid is applied along more than 1 dimension,
@@ -1757,6 +1781,8 @@ def make_grid(
         periodic: Tuple[bool, bool, bool],
         sources: List[SourceType],
         num_pml_layers: List[Tuple[pd.NonNegativeInt, pd.NonNegativeInt]],
+        internal_override_structures: List[MeshOverrideStructure] = None,
+        internal_snapping_points: List[CoordinateOptional] = None,
     ) -> Grid:
         """Make the entire simulation grid based on some simulation parameters.
 
@@ -1772,6 +1798,10 @@ def make_grid(
             List of sources.
         num_pml_layers : List[Tuple[float, float]]
             List containing the number of absorber layers in - and + boundaries.
+        internal_override_structures : List[MeshOverrideStructure]
+            If `None`, recomputes internal override structures.
+        internal_snapping_points : List[CoordinateOptional]
+            If `None`, recomputes internal snapping points.
 
         Returns
         -------
@@ -1830,7 +1860,10 @@ def make_grid(
 
         sim_size = list(structures[0].geometry.size)
         all_structures = list(structures) + self.all_override_structures(
-            list(structures), wavelength, sim_size
+            list(structures),
+            wavelength,
+            sim_size,
+            internal_override_structures,
         )
 
         # apply internal `dl_min` if any AutoGrid has unset `dl_min`
@@ -1856,7 +1889,7 @@ def make_grid(
                 periodic=periodic[idim],
                 wavelength=wavelength,
                 num_pml_layers=num_pml_layers[idim],
-                snapping_points=self.all_snapping_points(structures),
+                snapping_points=self.all_snapping_points(structures, internal_snapping_points),
             )
 
         coords = Coords(**coords_dict)

tidy3d/components/simulation.py

@@ -98,7 +98,17 @@
 from .source.utils import SourceType
 from .structure import MeshOverrideStructure, Structure
 from .subpixel_spec import SubpixelSpec
-from .types import TYPE_TAG_STR, Ax, Axis, FreqBound, InterpMethod, Literal, Symmetry, annotate_type
+from .types import (
+    TYPE_TAG_STR,
+    Ax,
+    Axis,
+    CoordinateOptional,
+    FreqBound,
+    InterpMethod,
+    Literal,
+    Symmetry,
+    annotate_type,
+)
 from .validators import (
     assert_objects_in_sim_bounds,
     validate_mode_objects_symmetry,
@@ -722,6 +732,31 @@ def pml_thicknesses(self) -> List[Tuple[float, float]]:
 
         return pml_thicknesses
 
+    @cached_property
+    def internal_override_structures(self) -> List[MeshOverrideStructure]:
+        """Internal mesh override structures. So far, internal override structures all come from `layer_refinement_specs`.
+
+        Returns
+        -------
+        List[MeshOverrideSructure]
+            List of override structures.
+        """
+        wavelength = self.grid_spec.get_wavelength(self.sources)
+        return self.grid_spec.internal_override_structures(
+            self.scene.all_structures, wavelength, self.geometry.size
+        )
+
+    @cached_property
+    def internal_snapping_points(self) -> List[CoordinateOptional]:
+        """Internal snapping points. So far, internal snapping points are generated by `layer_refinement_specs`.
+
+        Returns
+        -------
+        List[CoordinateOptional]
+            List of snapping points coordinates.
+        """
+        return self.grid_spec.internal_snapping_points(self.scene.all_structures)
+
     @equal_aspect
     @add_ax_if_none
     def plot_lumped_elements(
@@ -842,12 +877,8 @@ def plot_grid(
         plot_params[0] = plot_params[0].include_kwargs(edgecolor=kwargs["colors_internal"])
 
         if self.grid_spec.auto_grid_used:
-            wavelength = self.grid_spec.get_wavelength(self.sources)
-            internal_override_structures = self.grid_spec.internal_override_structures(
-                self.scene.all_structures, wavelength, self.geometry.size
-            )
             all_override_structures = [
-                internal_override_structures,
+                self.internal_override_structures,
                 self.grid_spec.external_override_structures,
             ]
             for structures, plot_param in zip(all_override_structures, plot_params):
@@ -867,11 +898,8 @@ def plot_grid(
                     ax.add_patch(rect)
 
         # Plot snapping points
-        internal_snapping_points = self.grid_spec.internal_snapping_points(
-            self.scene.all_structures
-        )
         for points, plot_param in zip(
-            [internal_snapping_points, self.grid_spec.snapping_points], plot_params
+            [self.internal_snapping_points, self.grid_spec.snapping_points], plot_params
         ):
             for point in points:
                 _, (x_point, y_point) = Geometry.pop_axis(point, axis=axis)
@@ -1056,20 +1084,27 @@ def grid(self) -> Grid:
 
         # Add a simulation Box as the first structure
         structures = [Structure(geometry=self.geometry, medium=self.medium)]
-        structures += self.structures
+        structures += self.static_structures
 
         grid = self.grid_spec.make_grid(
             structures=structures,
             symmetry=self.symmetry,
             periodic=self._periodic,
             sources=self.sources,
             num_pml_layers=self.num_pml_layers,
+            internal_snapping_points=self.internal_snapping_points,
+            internal_override_structures=self.internal_override_structures,
         )
 
         # This would AutoGrid the in-plane directions of the 2D materials
         # return self._grid_corrections_2dmaterials(grid)
         return grid
 
+    @cached_property
+    def static_structures(self) -> list[Structure]:
+        """Structures in simulation with all autograd tracers removed."""
+        return [structure.to_static() for structure in self.structures]
+
     @cached_property
     def num_cells(self) -> int:
         """Number of cells in the simulation.
@@ -3774,11 +3809,6 @@ def _validate_time_monitors_num_steps(self) -> None:
                     "at which the monitor stores data."
                 )
 
-    @cached_property
-    def static_structures(self) -> list[Structure]:
-        """Structures in simulation with all autograd tracers removed."""
-        return [structure.to_static() for structure in self.structures]
-
     @cached_property
     def monitors_data_size(self) -> Dict[str, float]:
         """Dictionary mapping monitor names to their estimated storage size in bytes."""
@@ -4550,6 +4580,8 @@ def _grid_corrections_2dmaterials(self, grid: Grid) -> Grid:
             symmetry=self.symmetry,
             sources=self.sources,
             num_pml_layers=self.num_pml_layers,
+            internal_snapping_points=self.internal_snapping_points,
+            internal_override_structures=self.internal_override_structures,
         )
 
         # Handle 2D materials if ``AutoGrid`` is used for in-plane directions
@@ -4587,33 +4619,6 @@ def _grid_corrections_2dmaterials(self, grid: Grid) -> Grid:
 
         return Grid(boundaries=Coords(**dict(zip("xyz", coords_all))))
 
-    @cached_property
-    def grid(self) -> Grid:
-        """FDTD grid spatial locations and information.
-
-        Returns
-        -------
-        :class:`.Grid`
-            :class:`.Grid` storing the spatial locations relevant to the simulation.
-        """
-
-        # Add a simulation Box as the first structure
-        structures = [Structure(geometry=self.geometry, medium=self.medium)]
-
-        structures += self.static_structures
-
-        grid = self.grid_spec.make_grid(
-            structures=structures,
-            symmetry=self.symmetry,
-            periodic=self._periodic,
-            sources=self.sources,
-            num_pml_layers=self.num_pml_layers,
-        )
-
-        # This would AutoGrid the in-plane directions of the 2D materials
-        # return self._grid_corrections_2dmaterials(grid)
-        return grid
-
     @cached_property
     def num_cells(self) -> int:
         """Number of cells in the simulation grid.