fix: consistency between mode source and mode frame

Author: dbochkov-flexcompute

tests/test_components/test_grid.py

@@ -402,3 +402,37 @@ def test_get_geo_inds_with_span():
 
     inds = g._get_geo_inds(geo_full, span_inds=span_inds)
     assert np.array_equal(inds, expected)
+
+
+def test_discretize_inds_relax_precision():
+    """Test that relax_precision handles floating-point precision issues at cell boundaries."""
+    g = make_grid()
+
+    # Create a box where boundaries are exactly on cell boundaries
+    box_exact = td.Box(center=(0, 0, 0), size=(2, 2, 2))
+    inds_exact = g.discretize_inds(box=box_exact, extend=False, relax_precision=False)
+
+    # Create a box where min boundary is slightly below the cell boundary (precision issue)
+    # The grid has boundaries at x = -1, 0, 1; y = -2, -1, 0, 1, 2; z = -3, -2, -1, 0, 1, 2, 3
+    # A box from -1 to 1 should span indices [0, 2] in x
+    eps = 1e-14  # Small epsilon to simulate floating-point precision issues
+    box_min_below = td.Box(center=(0, 0, 0), size=(2 + eps, 2 + eps, 2 + eps))
+
+    # Without relax_precision, the slightly larger box may include extra cells
+    inds_no_relax = g.discretize_inds(box=box_min_below, extend=False, relax_precision=False)
+
+    # With relax_precision, the indices should match the exact case since boundaries are "close enough"
+    inds_with_relax = g.discretize_inds(box=box_min_below, extend=False, relax_precision=True)
+
+    # The relaxed precision should produce the same result as the exact box
+    assert inds_exact == inds_with_relax
+    assert inds_no_relax != inds_with_relax
+
+    # Test case where box max is slightly below a cell boundary
+    box_max_below = td.Box(center=(0, 0, 0), size=(2 - eps, 2 - eps, 2 - eps))
+    inds_max_no_relax = g.discretize_inds(box=box_max_below, extend=False, relax_precision=False)
+    inds_max_with_relax = g.discretize_inds(box=box_max_below, extend=False, relax_precision=True)
+
+    # With relaxed precision, the boundaries close to cell boundaries should be treated as equal
+    assert inds_exact == inds_max_with_relax
+    assert inds_max_no_relax != inds_max_with_relax

tidy3d/components/grid/grid.py

@@ -605,7 +605,9 @@ def _yee_h(self, axis: Axis):
 
         return Coords(**yee_coords)
 
-    def discretize_inds(self, box: Box, extend: bool = False) -> list[tuple[int, int]]:
+    def discretize_inds(
+        self, box: Box, extend: bool = False, relax_precision: bool = False
+    ) -> list[tuple[int, int]]:
         """Start and stopping indexes for the cells that intersect with a :class:`Box`.
 
         Parameters
@@ -616,6 +618,9 @@ def discretize_inds(self, box: Box, extend: bool = False) -> list[tuple[int, int
             If ``True``, ensure that the returned indexes extend sufficiently in every direction to
             be able to interpolate any field component at any point within the ``box``, for field
             components sampled on the Yee grid.
+        relax_precision : bool = False
+            If ``True``, relax the precision of the discretization to allow for small numerical
+            differences between the box boundaries and the cell boundaries.
 
         Returns
         -------
@@ -646,17 +651,36 @@ def discretize_inds(self, box: Box, extend: bool = False) -> list[tuple[int, int
             # handle extensions
             if ind_max > ind_min and extend:
                 # Left side
-                if box.bounds[0][axis] < self.centers.to_list[axis][ind_min]:
+                if pts_min[axis] < self.centers.to_list[axis][ind_min]:
                     # Box bounds on the left side are to the left of the closest grid center
                     ind_min -= 1
 
                 # We always need an extra pixel on the right for the tangential components
                 ind_max += 1
 
             # store indexes
-            inds_list.append((ind_min, ind_max))
-
-        return inds_list
+            inds_list.append([ind_min, ind_max])
+
+        if relax_precision:
+            for dim in range(3):
+                # Fix some corner cases when the box boundary is very close to
+                # cell boundaries but due to finite precision is slightly smaller or larger
+                cell_bounds = np.array(boundaries.to_list[dim])
+                num_cells = len(cell_bounds) - 1
+                min_ind = inds_list[dim][0]
+                max_ind = inds_list[dim][1]
+                box_min = pts_min[dim]
+                box_max = pts_max[dim]
+                # Check if the cell boundary after current min is close to the box min,
+                # if it is close enough then choose that to be the new minimum cell index
+                if min_ind + 1 < num_cells and np.isclose(box_min, cell_bounds[min_ind + 1]):
+                    inds_list[dim][0] += 1
+                # Same but for the max cell boundary. If the current cell boundary is close to the box bounds,
+                # then it is considered equal and the stop index should be incremented
+                if max_ind < num_cells and np.isclose(box_max, cell_bounds[max_ind]):
+                    inds_list[dim][1] += 1
+
+        return [(ind_min, ind_max) for ind_min, ind_max in inds_list]
 
     def extended_subspace(
         self,

tidy3d/components/simulation.py

@@ -2212,7 +2212,7 @@ def _pec_frame_box(
     ) -> tuple[Box, int, str]:
         """Return pec bounding box, frame axis and object's direction"""
 
-        span_inds = np.array(self.grid.discretize_inds(obj))
+        span_inds = np.array(self.grid.discretize_inds(obj, relax_precision=True))
         coords = self.grid.boundaries.to_list
         direction = obj.direction
         if isinstance(obj, ModeSource):