From bc9c31e0f95e984f588f7143aac197ba4c2a3588 Mon Sep 17 00:00:00 2001 From: Marco De Pietri Date: Sat, 14 Mar 2026 05:19:51 +0100 Subject: [PATCH] get indices for rectilinear meshes (#3876) Co-authored-by: Paul Romano --- openmc/mesh.py | 45 +++++++++++++++++++++++++++++++---- tests/unit_tests/test_mesh.py | 34 ++++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 4 deletions(-) diff --git a/openmc/mesh.py b/openmc/mesh.py index 030a57218..79982d78a 100644 --- a/openmc/mesh.py +++ b/openmc/mesh.py @@ -1688,10 +1688,47 @@ 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" - ) + 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): diff --git a/tests/unit_tests/test_mesh.py b/tests/unit_tests/test_mesh.py index aa8bcae5f..0b28bdfbe 100644 --- a/tests/unit_tests/test_mesh.py +++ b/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.])