From 8081815b99bc1f13c7ea35da4bb11d82c6db4edd Mon Sep 17 00:00:00 2001 From: Jonathan Shimwell Date: Thu, 26 Feb 2026 03:19:02 +0100 Subject: [PATCH] Add RegularMesh.get_indices_at_coords method (#3824) Co-authored-by: Paul Romano --- openmc/mesh.py | 59 +++++++++++++++++++++++++++- tests/unit_tests/test_mesh.py | 73 +++++++++++++++++++++++++++++++++++ 2 files changed, 130 insertions(+), 2 deletions(-) diff --git a/openmc/mesh.py b/openmc/mesh.py index 3f0802d6b..030a57218 100644 --- a/openmc/mesh.py +++ b/openmc/mesh.py @@ -547,6 +547,10 @@ class StructuredMesh(MeshBase): 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 @@ class RegularMesh(StructuredMesh): 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 @@ class RectilinearMesh(StructuredMesh): 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 @@ class CylindricalMesh(StructuredMesh): 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 @@ class SphericalMesh(StructuredMesh): 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) diff --git a/tests/unit_tests/test_mesh.py b/tests/unit_tests/test_mesh.py index c5855a7b0..9d07eda0d 100644 --- a/tests/unit_tests/test_mesh.py +++ b/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,)