adding RegularMesh.get_indices_at_coords to match CylindricalMesh

openmc/mesh.py

@@ -542,6 +542,22 @@ def n_dimension(self):
     def _grids(self):
         pass
 
+    @abstractmethod
+    def get_indices_at_coords(self, coords: Sequence[float]) -> tuple:
+        """Finds the index of the mesh voxel at the specified coordinates.
+
+        Parameters
+        ----------
+        coords : Sequence[float]
+            The x, y, z axis coordinates
+
+        Returns
+        -------
+        tuple
+            Mesh indices matching the dimensionality of the mesh
+
+        """
+
     @property
     def vertices(self):
         """Return coordinates of mesh vertices in Cartesian coordinates. Also
@@ -1416,6 +1432,44 @@ def build_cells(self, bc: str | None = None):
 
         return root_cell, cells
 
+    def get_indices_at_coords(self, coords: Sequence[float]) -> tuple:
+        """Finds the index of the mesh voxel at the specified x,y,z coordinates.
+
+        .. versionadded:: 0.15.4
+
+        Parameters
+        ----------
+        coords : Sequence[float]
+            The x, y, z axis coordinates
+
+        Returns
+        -------
+        tuple
+            Mesh indices matching the dimensionality of the mesh
+
+        """
+        ndim = self.n_dimension
+        if len(coords) < ndim:
+            raise ValueError(
+                f"coords must have at least {ndim} values for a "
+                f"{ndim}D mesh, got {len(coords)}"
+            )
+        lower_left = np.array(self.lower_left)
+        upper_right = np.array(self.upper_right)
+        dimension = np.array(self.dimension)
+
+        # Calculate spacing for each dimension
+        spacing = (upper_right - lower_left) / dimension
+
+        # Calculate indices for each coordinate
+        coords_array = np.array(coords[:ndim])
+        indices = np.floor((coords_array - lower_left) / spacing).astype(int)
+
+        # Clamp indices to valid range [0, dimension-1]
+        indices = np.clip(indices, 0, dimension - 1)
+
+        return tuple(int(i) for i in indices[:ndim])
+
 
 def Mesh(*args, **kwargs):
     warnings.warn("Mesh has been renamed RegularMesh. Future versions of "
@@ -1623,6 +1677,11 @@ 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"
+        )
+
 
 class CylindricalMesh(StructuredMesh):
     """A 3D cylindrical mesh
@@ -2444,6 +2503,11 @@ def _convert_to_cartesian(arr, origin: Sequence[float]):
         arr[..., 2] = z + origin[2]
         return arr
 
+    def get_indices_at_coords(self, coords: Sequence[float]) -> tuple:
+        raise NotImplementedError(
+            "get_indices_at_coords is not yet implemented for SphericalMesh"
+        )
+
 
 def require_statepoint_data(func):
     @wraps(func)

tests/unit_tests/test_mesh.py

@@ -920,3 +920,73 @@ def test_filter_time_mesh(run_in_tmpdir):
         f"Collision vs tracklength tallies disagree: chi2={chi2_stat:.2f} "
         f">= {crit=:.2f} ({dof=}, {alpha=})"
     )
+
+
+def test_regular_mesh_get_indices_at_coords():
+    """Test get_indices_at_coords method for RegularMesh"""
+    # Create a 10x10x10 mesh from (0,0,0) to (1,1,1)
+    # Each voxel is 0.1 x 0.1 x 0.1
+    mesh = openmc.RegularMesh()
+    mesh.lower_left = (0, 0, 0)
+    mesh.upper_right = (1, 1, 1)
+    mesh.dimension = [10, 10, 10]
+
+    # Test lower-left corner maps to first voxel (0, 0, 0)
+    assert mesh.get_indices_at_coords([0.0, 0.0, 0.0]) == (0, 0, 0)
+
+    # Test upper-right corner maps to last voxel (9, 9, 9) - clamped
+    assert mesh.get_indices_at_coords([1.0, 1.0, 1.0]) == (9, 9, 9)
+
+    # Test centroid of first voxel
+    # Voxel 0 spans [0.0, 0.1], so centroid is at 0.05
+    assert mesh.get_indices_at_coords([0.05, 0.05, 0.05]) == (0, 0, 0)
+
+    # Test centroid of last voxel maps correctly
+    # Voxel 9 spans [0.9, 1.0], so centroid is at 0.95
+    assert mesh.get_indices_at_coords([0.95, 0.95, 0.95]) == (9, 9, 9)
+
+    # Test a middle voxel
+    # Voxel 4 spans [0.4, 0.5], so 0.45 should map to it
+    assert mesh.get_indices_at_coords([0.45, 0.45, 0.45]) == (4, 4, 4)
+
+    # Test mixed indices
+    assert mesh.get_indices_at_coords([0.05, 0.45, 0.95]) == (0, 4, 9)
+    assert mesh.get_indices_at_coords([0.95, 0.05, 0.45]) == (9, 0, 4)
+
+    # Test coordinates outside mesh bounds get clamped to edge voxels
+    assert mesh.get_indices_at_coords([-0.5, 0.5, 0.5]) == (0, 5, 5)
+    assert mesh.get_indices_at_coords([1.5, 0.5, 0.5]) == (9, 5, 5)
+    assert mesh.get_indices_at_coords([0.5, -0.5, 0.5]) == (5, 0, 5)
+    assert mesh.get_indices_at_coords([0.5, 1.5, 0.5]) == (5, 9, 5)
+    assert mesh.get_indices_at_coords([0.5, 0.5, -0.5]) == (5, 5, 0)
+    assert mesh.get_indices_at_coords([0.5, 0.5, 1.5]) == (5, 5, 9)
+
+    # Test that results match expected dimensionality (3D mesh returns 3-tuple)
+    result = mesh.get_indices_at_coords([0.5, 0.5, 0.5])
+    assert isinstance(result, tuple)
+    assert len(result) == 3
+
+    # Test that indices can be used directly with centroids array
+    idx = mesh.get_indices_at_coords([0.95, 0.95, 0.95])
+    centroid = mesh.centroids[idx]
+    np.testing.assert_array_almost_equal(centroid, [0.95, 0.95, 0.95])
+
+    # Test with a 2D mesh
+    mesh_2d = openmc.RegularMesh()
+    mesh_2d.lower_left = (0, 0)
+    mesh_2d.upper_right = (1, 1)
+    mesh_2d.dimension = [10, 10]
+    result_2d = mesh_2d.get_indices_at_coords([0.5, 0.5, 999.0])
+    assert isinstance(result_2d, tuple)
+    assert len(result_2d) == 2
+    assert result_2d == (5, 5)
+
+    # Test with a 1D mesh
+    mesh_1d = openmc.RegularMesh()
+    mesh_1d.lower_left = [0]
+    mesh_1d.upper_right = [1]
+    mesh_1d.dimension = [10]
+    result_1d = mesh_1d.get_indices_at_coords([0.5, 999.0, 999.0])
+    assert isinstance(result_1d, tuple)
+    assert len(result_1d) == 1
+    assert result_1d == (5,)