Add RegularMesh.get_indices_at_coords method (#3824)

Co-authored-by: Paul Romano <[email protected]>

Author: shimwell

openmc/mesh.py

@@ -547,6 +547,10 @@ def _axis_labels(self):
     def _grids(self):
         pass
 
+    @abstractmethod
+    def get_indices_at_coords(self, coords: Sequence[float]) -> tuple:
+        pass
+
     @property
     def vertices(self):
         """Return coordinates of mesh vertices in Cartesian coordinates. Also
@@ -1432,6 +1436,47 @@ 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 element at the specified coordinates.
+
+        .. versionadded:: 0.15.4
+
+        Parameters
+        ----------
+        coords : Sequence[float]
+            Cartesian coordinates of the point.
+
+        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)}"
+            )
+
+        coords_array = np.array(coords[:ndim])
+        lower_left = np.array(self.lower_left)
+        upper_right = np.array(self.upper_right)
+        dimension = np.array(self.dimension)
+
+        if np.any(coords_array < lower_left) or np.any(coords_array > upper_right):
+            raise ValueError(
+                f"coords {tuple(coords_array)} are outside mesh bounds "
+                f"[{tuple(lower_left)}, {tuple(upper_right)}]"
+            )
+
+        # Calculate spacing for each dimension
+        spacing = (upper_right - lower_left) / dimension
+
+        # Calculate indices for each coordinate
+        indices = np.floor((coords_array - lower_left) / spacing).astype(int)
+        return tuple(int(i) for i in indices[:ndim])
+
 
 def Mesh(*args, **kwargs):
     warnings.warn("Mesh has been renamed RegularMesh. Future versions of "
@@ -1643,6 +1688,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
@@ -1835,14 +1885,14 @@ def get_indices_at_coords(
             self,
             coords: Sequence[float]
         ) -> tuple[int, int, int]:
-        """Finds the index of the mesh voxel at the specified x,y,z coordinates.
+        """Finds the index of the mesh element at the specified coordinates.
 
         .. versionadded:: 0.15.0
 
         Parameters
         ----------
         coords : Sequence[float]
-            The x, y, z axis coordinates
+            Cartesian coordinates of the point.
 
         Returns
         -------
@@ -2478,6 +2528,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,76 @@ 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 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 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, 0.5])
+    with pytest.raises(ValueError):
+        mesh.get_indices_at_coords([0.5, 1.5, 0.5])
+    with pytest.raises(ValueError):
+        mesh.get_indices_at_coords([0.5, 0.5, -0.5])
+    with pytest.raises(ValueError):
+        mesh.get_indices_at_coords([0.5, 0.5, 1.5])
+
+    # 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,)