get indices for rectilinear meshes (#3876)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
This commit is contained in:
Marco De Pietri 2026-03-14 05:19:51 +01:00 committed by GitHub
parent dd6d31bfae
commit bc9c31e0f9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 75 additions and 4 deletions

View file

@ -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):

View file

@ -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.])