Enable snapping_points in grid generation

CHANGELOG.md

@@ -34,6 +34,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `RectangularWaveguide` supports layered cladding above and below core.
 - `SubpixelSpec` accepted by `Simulation.subpixel` to select subpixel averaging methods separately for dielectric, metal, and PEC materials. Specifically, added support for conformal mesh methods near PEC structures that can be specified through the field `pec` in the `SubpixelSpec` class. Note: previously, `subpixel=False` was implementing staircasing for every material except PEC. Now, `subpixel=False` implements direct staircasing for all materials. For PEC, the behavior of `subpixel=False` in Tidy3D < 2.7 is now achieved through `subpixel=SubpixelSpec(pec=HeuristicPECStaircasing())`, while `subpixel=True` in Tidy3D < 2.7 is now achieved through `subpixel=SubpixelSpec(pec=Staircasing())`. The default is `subpixel=SubpixelSpec(pec=PECConformal())` for more accurate PEC modelling.
 - Lossless `Green2008` variant for crystalline silicon added to material library.
+- `GridSpec` supports `snapping_points` that enforce grid boundaries to pass through them.
 
 ### Changed
 - IMPORTANT NOTE: differentiable fields in the `adjoint` plugin (`JaxBox.size`, `JaxBox.center`, `JaxPolySlab.vertices`) no longer store the derivative information after the object is initialized. For example, if using JaxPolySlab.vertices directly in an objective function, the vertices will have no effect on the gradient. Instead, this information is now stored in a field of the same name with `_jax` suffix, eg. `JaxPolySlab.vertices_jax`. For some applications, such as evaluating penalty values, please change to `radius_penalty.evaluate(polyslab.vertices_jax)` or use the vertices as generated by your parameterization functions (`make_vertices(params)`).

tests/test_components/test_grid_spec.py

@@ -29,9 +29,73 @@ def test_make_coords():
         periodic=(True, False, False),
         wavelength=1.0,
         num_pml_layers=(10, 4),
+        snapping_points=(),
     )
 
 
+def test_make_coords_with_snapping_points():
+    """Test the behavior of snapping points"""
+    gs = make_grid_spec()
+    make_coords_args = dict(
+        structures=[
+            td.Structure(geometry=td.Box(size=(2, 2, 1)), medium=td.Medium()),
+            td.Structure(geometry=td.Box(size=(1, 1, 1)), medium=td.Medium(permittivity=4)),
+        ],
+        symmetry=(0, 0, 0),
+        periodic=(False, False, False),
+        wavelength=1.0,
+        num_pml_layers=(0, 0),
+        axis=0,
+    )
+
+    # 1) no snapping points, 0.85 is not on any grid boundary
+    coord_original = gs.grid_x.make_coords(
+        snapping_points=(),
+        **make_coords_args,
+    )
+    assert not np.any(np.isclose(coord_original, 0.85))
+
+    # 2) with snapping points at 0.85, grid should pass through 0.85
+    coord = gs.grid_x.make_coords(
+        snapping_points=((0.85, 0, 0),),
+        **make_coords_args,
+    )
+    assert np.any(np.isclose(coord, 0.85))
+
+    #  snapping still takes effect if the point is completely outside along other axes
+    coord = gs.grid_x.make_coords(
+        snapping_points=((0.85, 10, 0),),
+        **make_coords_args,
+    )
+    assert np.any(np.isclose(coord, 0.85))
+
+    coord = gs.grid_x.make_coords(
+        snapping_points=((0.85, 0, -10),),
+        **make_coords_args,
+    )
+    assert np.any(np.isclose(coord, 0.85))
+
+    # 3) snapping takes no effect if it's too close to interval boundaries
+    coord = gs.grid_x.make_coords(
+        snapping_points=((0.98, 0, 0),),
+        **make_coords_args,
+    )
+    assert np.allclose(coord_original, coord)
+
+    # and no snapping if it's compeletely outside the simulation domain
+    coord = gs.grid_x.make_coords(
+        snapping_points=((10, 0, 0),),
+        **make_coords_args,
+    )
+    assert np.allclose(coord_original, coord)
+
+    coord = gs.grid_x.make_coords(
+        snapping_points=((-10, 0, 0),),
+        **make_coords_args,
+    )
+    assert np.allclose(coord_original, coord)
+
+
 def test_make_coords_2d():
     gs = make_grid_spec()
     _ = gs.grid_x.make_coords(
@@ -44,6 +108,7 @@ def test_make_coords_2d():
         periodic=(True, True, False),
         wavelength=1.0,
         num_pml_layers=(10, 4),
+        snapping_points=(),
     )
 
 

tidy3d/components/grid/grid_spec.py

@@ -15,7 +15,7 @@
 from ..geometry.base import Box
 from ..source import SourceType
 from ..structure import Structure, StructureType
-from ..types import TYPE_TAG_STR, Axis, Symmetry, annotate_type
+from ..types import TYPE_TAG_STR, Axis, Coordinate, Symmetry, annotate_type
 from .grid import Coords, Coords1D, Grid
 from .mesher import GradedMesher, MesherType
 
@@ -31,6 +31,7 @@ def make_coords(
         periodic: bool,
         wavelength: pd.PositiveFloat,
         num_pml_layers: Tuple[pd.NonNegativeInt, pd.NonNegativeInt],
+        snapping_points: Tuple[Coordinate, ...],
     ) -> Coords1D:
         """Generate 1D coords to be used as grid boundaries, based on simulation parameters.
         Symmetry, and PML layers will be treated here.
@@ -48,6 +49,8 @@ def make_coords(
             Free-space wavelength.
         num_pml_layers : Tuple[int, int]
             number of layers in the absorber + and - direction along one dimension.
+        snapping_points : Tuple[Coordinate, ...]
+            A set of points that enforce grid boundaries to pass through them.
 
         Returns
         -------
@@ -66,6 +69,7 @@ def make_coords(
             wavelength=wavelength,
             symmetry=symmetry,
             is_periodic=is_periodic,
+            snapping_points=snapping_points,
         )
 
         # incorporate symmetries
@@ -366,6 +370,7 @@ def _make_coords_initial(
         wavelength: float,
         symmetry: Symmetry,
         is_periodic: bool,
+        snapping_points: Tuple[Coordinate, ...],
     ) -> Coords1D:
         """Customized 1D coords to be used as grid boundaries.
 
@@ -382,6 +387,8 @@ def _make_coords_initial(
             normal to each of the three axes.
         is_periodic : bool
             Apply periodic boundary condition or not.
+        snapping_points : Tuple[Coordinate, ...]
+            A set of points that enforce grid boundaries to pass through them.
 
         Returns
         -------
@@ -411,6 +418,11 @@ def _make_coords_initial(
             self.min_steps_per_wvl,
             self.dl_min,
         )
+        # insert snapping_points
+        interval_coords, max_dl_list = self.mesher.insert_snapping_points(
+            axis, interval_coords, max_dl_list, snapping_points
+        )
+
         # Put just a single pixel if 2D-like simulation
         if interval_coords.size == 1:
             dl = wavelength / self.min_steps_per_wvl
@@ -514,6 +526,15 @@ class GridSpec(Tidy3dBaseModel):
         "uses :class:`.AutoGrid`.",
     )
 
+    snapping_points: Tuple[Coordinate, ...] = pd.Field(
+        (),
+        title="Grid specification snapping_points",
+        description="A set of points that enforce grid boundaries to pass through them. "
+        "However, some points might be skipped if they are too close. "
+        "Note: it only takes effect when at least one of the three dimensions "
+        "uses :class:`.AutoGrid`.",
+    )
+
     @property
     def auto_grid_used(self) -> bool:
         """True if any of the three dimensions uses :class:`.AutoGrid`."""
@@ -630,6 +651,7 @@ def make_grid(
                 periodic=periodic[idim],
                 wavelength=wavelength,
                 num_pml_layers=num_pml_layers[idim],
+                snapping_points=self.snapping_points,
             )
 
         coords = Coords(**coords_dict)
@@ -642,6 +664,7 @@ def auto(
         min_steps_per_wvl: pd.PositiveFloat = 10.0,
         max_scale: pd.PositiveFloat = 1.4,
         override_structures: List[StructureType] = (),
+        snapping_points: Tuple[Coordinate, ...] = (),
         dl_min: pd.NonNegativeFloat = 0.0,
         mesher: MesherType = GradedMesher(),
     ) -> GridSpec:
@@ -661,6 +684,8 @@ def auto(
             A list of structures that is added on top of the simulation structures in
             the process of generating the grid. This can be used to refine the grid or make it
             coarser depending than the expected need for higher/lower resolution regions.
+        snapping_points : Tuple[Coordinate, ...]
+            A set of points that enforce grid boundaries to pass through them.
         dl_min: pd.NonNegativeFloat
             Lower bound of grid size.
         mesher : MesherType = GradedMesher()
@@ -684,6 +709,7 @@ def auto(
             grid_y=grid_1d,
             grid_z=grid_1d,
             override_structures=override_structures,
+            snapping_points=snapping_points,
         )
 
     @classmethod

tidy3d/components/grid/mesher.py

@@ -19,7 +19,7 @@
 from ..base import Tidy3dBaseModel
 from ..medium import AnisotropicMedium, Medium2D, PECMedium
 from ..structure import MeshOverrideStructure, Structure, StructureType
-from ..types import ArrayFloat1D, Axis, Bound
+from ..types import ArrayFloat1D, Axis, Bound, Coordinate
 
 _ROOTS_TOL = 1e-10
 
@@ -43,6 +43,16 @@ def parse_structures(
     ) -> Tuple[ArrayFloat1D, ArrayFloat1D]:
         """Calculate the positions of all bounding box interfaces along a given axis."""
 
+    @abstractmethod
+    def insert_snapping_points(
+        self,
+        axis: Axis,
+        interval_coords: ArrayFloat1D,
+        max_dl_list: ArrayFloat1D,
+        snapping_points: List[Coordinate],
+    ) -> Tuple[ArrayFloat1D, ArrayFloat1D]:
+        """Insert snapping_points to the intervals."""
+
     @abstractmethod
     def make_grid_multiple_intervals(
         self,
@@ -63,6 +73,63 @@ class GradedMesher(Mesher):
     """Implements automatic nonuniform meshing with a set minimum steps per wavelength and
     a graded mesh expanding from higher- to lower-resolution regions."""
 
+    def insert_snapping_points(
+        self,
+        axis: Axis,
+        interval_coords: ArrayFloat1D,
+        max_dl_list: ArrayFloat1D,
+        snapping_points: List[Coordinate],
+    ) -> Tuple[ArrayFloat1D, ArrayFloat1D]:
+        """Insert snapping_points to the intervals.
+
+        Parameters
+        ----------
+        axis : Axis
+            Axis index along which to operate.
+        interval_coords : ArrayFloat1D
+            Coordinate of interval boundaries.
+        max_dl_list : ArrayFloat1D
+            Maximal allowed step size of each interval generated from `parse_structures`.
+        snapping_points : List[Coordinate]
+            A set of points that enforce grid boundaries to pass through them.
+
+        Returns
+        -------
+        interval_coords : Array
+            An array of coordinates, where the first element is the simulation min boundary, the
+            last element is the simulation max boundary, and the intermediate coordinates are all
+            locations with a fixpoint or a structure has a bounding box edge along the specified axis.
+            The boundaries are filtered such that no interval is smaller than the smallest
+            of the ``max_steps``.
+        max_steps : array_like
+            An array of size ``interval_coords.size - 1`` giving the maximum grid step required in
+            each ``interval_coords[i]:interval_coords[i+1]`` interval, depending on the materials
+            in that interval, the supplied wavelength, and the minimum required step per wavelength.
+
+        """
+        # Don't do anything for a 2D-like simulation, or no fix points
+        if interval_coords.size == 1 or len(snapping_points) < 1:
+            return interval_coords, max_dl_list
+
+        min_step = np.amin(max_dl_list) * 0.5
+        for point in snapping_points:
+            new_coord = point[axis]
+            # Skip if the point is outside the domain
+            if new_coord >= interval_coords[-1] or new_coord <= interval_coords[0]:
+                continue
+            # search insertion location
+            ind = np.searchsorted(interval_coords, new_coord, side="left")
+            # Skip snapping_points if the distance to the existing interval boundarires are
+            # smaller than `min_step` defined above.
+            if abs(new_coord - interval_coords[ind]) < min_step:
+                continue
+            if abs(new_coord - interval_coords[ind - 1]) < min_step:
+                continue
+            # insertion
+            interval_coords = np.insert(interval_coords, ind, new_coord)
+            max_dl_list = np.insert(max_dl_list, ind - 1, max_dl_list[ind - 1])
+        return interval_coords, max_dl_list
+
     def parse_structures(
         self,
         axis: Axis,

tidy3d/components/simulation.py

@@ -774,6 +774,12 @@ def plot_grid(
             )
             ax.add_patch(rect)
 
+        # Plot snapping points
+        for point in self.grid_spec.snapping_points:
+            _, (x_point, y_point) = Geometry.pop_axis(point, axis=axis)
+            x_point, y_point = (self._evaluate_inf(v) for v in (x_point, y_point))
+            ax.scatter(x_point, y_point, color=plot_params.edgecolor)
+
         ax = Scene._set_plot_bounds(
             bounds=self.simulation_bounds, ax=ax, x=x, y=y, z=z, hlim=hlim, vlim=vlim
         )