get indices for rectilinear meshes (#3876)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

Author: marco-de-pietri

openmc/mesh.py

@@ -1688,10 +1688,47 @@ def to_xml_element(self):
 
         return element
 
-    def get_indices_at_coords(self, coords: Sequence[float]) -> tuple:
-        raise NotImplementedError(
-            "get_indices_at_coords is not yet implemented for RectilinearMesh"
-        )
+    def get_indices_at_coords(self, coords: Sequence[float]) -> tuple[int, int, int]:
+        """Find the mesh cell indices containing the specified coordinates.
+
+        .. versionadded:: 0.15.4
+
+        Parameters
+        ----------
+        coords : Sequence[float]
+            Cartesian coordinates of the point as (x, y, z).
+
+        Returns
+        -------
+        tuple[int, int, int]
+            Mesh indices (ix, iy, iz).
+
+        Raises
+        ------
+        ValueError
+            If coords does not contain exactly 3 values, or if a coordinate is
+            outside the mesh grid boundaries.
+        """
+        if len(coords) != 3:
+            raise ValueError(
+                f"coords must contain exactly 3 values for a rectilinear mesh, "
+                f"got {len(coords)}"
+            )
+
+        grids = (self.x_grid, self.y_grid, self.z_grid)
+        indices = []
+
+        for grid, value in zip(grids, coords):
+            if value < grid[0] or value > grid[-1]:
+                raise ValueError(
+                    f"Coordinate value {value} is outside the mesh grid boundaries: "
+                    f"[{grid[0]}, {grid[-1]}]"
+                )
+
+            idx = np.searchsorted(grid, value, side="right") - 1
+            indices.append(int(min(idx, len(grid) - 2)))
+
+        return tuple(indices)
 
 
 class CylindricalMesh(StructuredMesh):

tests/unit_tests/test_mesh.py

@@ -994,3 +994,37 @@ def test_regular_mesh_get_indices_at_coords():
     assert isinstance(result_1d, tuple)
     assert len(result_1d) == 1
     assert result_1d == (5,)
+
+
+def test_rectilinear_mesh_get_indices_at_coords():
+    """Test get_indices_at_coords method for RectilinearMesh"""
+    # Create a 3x2x2 rectilinear mesh with non-uniform spacing
+    mesh = openmc.RectilinearMesh()
+    mesh.x_grid = [0., 1., 5., 10.]
+    mesh.y_grid = [-10., -5., 0.]
+    mesh.z_grid = [-100., 0., 100.]
+
+    # Test lower-left corner maps to first voxel (0, 0, 0)
+    assert mesh.get_indices_at_coords([0.0, -10., -100.]) == (0, 0, 0)
+
+    # Test centroid of first voxel
+    assert mesh.get_indices_at_coords([0.5, -7.5, -50.]) == (0, 0, 0)
+
+    # Test centroid of last voxel maps correctly
+    assert mesh.get_indices_at_coords([7.5, -2.5, 50.]) == (2, 1, 1)
+
+    # Test upper_right corner maps to last voxel
+    assert mesh.get_indices_at_coords([10., 0., 100.]) == (2, 1, 1)
+
+    # Test a middle voxel
+    assert mesh.get_indices_at_coords([2., -5., 0.]) == (1, 1, 1)
+
+    # Test coordinates outside mesh bounds raise ValueError
+    with pytest.raises(ValueError):
+        mesh.get_indices_at_coords([-0.5, 0.5, 0.5])
+    with pytest.raises(ValueError):
+        mesh.get_indices_at_coords([1.5, 0.5, 0.5])
+    with pytest.raises(ValueError):
+        mesh.get_indices_at_coords([0.5, -0.5, 110.])
+    with pytest.raises(ValueError):
+        mesh.get_indices_at_coords([0.5, -20., 110.])