From f311d5493766d1beb2738601121a22b8bcfa8a4f Mon Sep 17 00:00:00 2001 From: Eden Rochman Date: Thu, 9 Apr 2026 15:31:25 +0200 Subject: [PATCH 1/3] Implement SphericalMesh.get_indices_at_coords Resolves the remaining work from #3867. Converts Cartesian (x, y, z) input to spherical (r, theta, phi) relative to the mesh origin, validates against grid bounds, and returns bin indices via np.searchsorted. Handles the degenerate r=0 case where angular coordinates are undefined. Adds comprehensive unit tests covering basic lookups, multi-bin angular grids, non-default origin, boundary values, r=0, and out-of-bounds errors for each axis. --- openmc/mesh.py | 82 +++++++++++++++++++++++++++++-- tests/unit_tests/test_mesh.py | 91 ++++++++++++++++++++++++++++++++++- 2 files changed, 167 insertions(+), 6 deletions(-) diff --git a/openmc/mesh.py b/openmc/mesh.py index 3c3c0a1ac5d..da456c7e008 100644 --- a/openmc/mesh.py +++ b/openmc/mesh.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod from collections.abc import Iterable, Sequence, Mapping from functools import wraps -from math import pi, sqrt, atan2 +from math import acos, atan2, pi, sqrt from numbers import Integral, Real from pathlib import Path from typing import Protocol @@ -2644,10 +2644,82 @@ 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 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] + The r, theta, phi indices. + + Raises + ------ + ValueError + If the coordinates fall outside the mesh grid boundaries. + + """ + dx = coords[0] - self.origin[0] + dy = coords[1] - self.origin[1] + dz = coords[2] - self.origin[2] + + r_value = sqrt(dx**2 + dy**2 + dz**2) + + if r_value < self.r_grid[0] or r_value > self.r_grid[-1]: + raise ValueError( + f'The r value {r_value} computed from the specified ' + f'coordinates is outside the r grid range ' + f'[{self.r_grid[0]}, {self.r_grid[-1]}].' + ) + + r_index = int(min( + np.searchsorted(self.r_grid, r_value, side='right') - 1, + len(self.r_grid) - 2 + )) + + if r_value == 0.0: + theta_value = 0.0 + phi_value = 0.0 + else: + theta_value = acos(dz / r_value) + phi_value = atan2(dy, dx) + if phi_value < 0: + phi_value += 2 * pi + + if theta_value < self.theta_grid[0] or theta_value > self.theta_grid[-1]: + raise ValueError( + f'The theta value {theta_value} computed from the specified ' + f'coordinates is outside the theta grid range ' + f'[{self.theta_grid[0]}, {self.theta_grid[-1]}].' + ) + + theta_index = int(min( + np.searchsorted(self.theta_grid, theta_value, side='right') - 1, + len(self.theta_grid) - 2 + )) + + if phi_value < self.phi_grid[0] or phi_value > self.phi_grid[-1]: + raise ValueError( + f'The phi value {phi_value} computed from the specified ' + f'coordinates is outside the phi grid range ' + f'[{self.phi_grid[0]}, {self.phi_grid[-1]}].' + ) + + phi_index = int(min( + np.searchsorted(self.phi_grid, phi_value, side='right') - 1, + len(self.phi_grid) - 2 + )) + + return (r_index, theta_index, phi_index) def require_statepoint_data(func): diff --git a/tests/unit_tests/test_mesh.py b/tests/unit_tests/test_mesh.py index 0b28bdfbe44..77b64048927 100644 --- a/tests/unit_tests/test_mesh.py +++ b/tests/unit_tests/test_mesh.py @@ -1,4 +1,4 @@ -from math import pi +from math import pi, sqrt from tempfile import TemporaryDirectory from pathlib import Path import itertools @@ -1028,3 +1028,92 @@ def test_rectilinear_mesh_get_indices_at_coords(): mesh.get_indices_at_coords([0.5, -0.5, 110.]) with pytest.raises(ValueError): mesh.get_indices_at_coords([0.5, -20., 110.]) + + +def test_SphericalMesh_get_indices_at_coords(): + """Test get_indices_at_coords method for SphericalMesh""" + + # Basic mesh with default phi and theta grids (single angular bin) + mesh = openmc.SphericalMesh(r_grid=(0, 5, 10)) + + assert mesh.get_indices_at_coords([3, 0, 0]) == (0, 0, 0) + assert mesh.get_indices_at_coords([0, 0, 3]) == (0, 0, 0) + assert mesh.get_indices_at_coords([0, 0, -3]) == (0, 0, 0) + assert mesh.get_indices_at_coords([7, 0, 0]) == (1, 0, 0) + assert mesh.get_indices_at_coords([10, 0, 0]) == (1, 0, 0) + + # Out-of-bounds r + with pytest.raises(ValueError): + mesh.get_indices_at_coords([11, 0, 0]) + + mesh2 = openmc.SphericalMesh(r_grid=(2, 5, 10)) + with pytest.raises(ValueError): + mesh2.get_indices_at_coords([1, 0, 0]) + + # Multi-bin angular grids: use points clearly inside bins + mesh3 = openmc.SphericalMesh( + r_grid=(0, 5, 10), + theta_grid=(0, pi/4, pi/2, pi), + phi_grid=(0, pi/2, pi, 3*pi/2, 2*pi) + ) + + # Near z-axis: theta~0 -> bin 0 + assert mesh3.get_indices_at_coords([0.01, 0, 3]) == (0, 0, 0) + + # theta in (0, pi/4) -> bin 0: [1, 0, 2] theta=arccos(2/sqrt(5))~0.46 + assert mesh3.get_indices_at_coords([1, 0, 2]) == (0, 0, 0) + + # theta in (pi/4, pi/2) -> bin 1: [2, 0, 1] theta=arccos(1/sqrt(5))~1.107 + assert mesh3.get_indices_at_coords([2, 0, 1]) == (0, 1, 0) + + # theta in (pi/2, pi) -> bin 2: [1, 0, -2] theta=arccos(-2/sqrt(5))~2.034 + assert mesh3.get_indices_at_coords([1, 0, -2]) == (0, 2, 0) + + # phi in (pi/2, pi) -> bin 1: [-1, 1, 0.5] + assert mesh3.get_indices_at_coords([-1, 1, 0.5]) == (0, 1, 1) + + # phi in (pi, 3*pi/2) -> bin 2: [-1, -1, 0.5] + assert mesh3.get_indices_at_coords([-1, -1, 0.5]) == (0, 1, 2) + + # phi in (3*pi/2, 2*pi) -> bin 3: [1, -1, 0.5] + assert mesh3.get_indices_at_coords([1, -1, 0.5]) == (0, 1, 3) + + # Non-default origin + mesh4 = openmc.SphericalMesh( + r_grid=(0, 5, 10), + origin=(100, 200, 300) + ) + + assert mesh4.get_indices_at_coords([103, 200, 300]) == (0, 0, 0) + assert mesh4.get_indices_at_coords([100, 200, 307]) == (1, 0, 0) + + with pytest.raises(ValueError): + mesh4.get_indices_at_coords([111, 200, 300]) + + # Degenerate case: point at origin with r_grid starting at 0 + mesh5 = openmc.SphericalMesh(r_grid=(0, 5)) + assert mesh5.get_indices_at_coords([0, 0, 0]) == (0, 0, 0) + + # Out-of-bounds theta: restricted theta grid + mesh6 = openmc.SphericalMesh( + r_grid=(0, 10), + theta_grid=(0, pi/4) + ) + with pytest.raises(ValueError): + mesh6.get_indices_at_coords([5, 0, 0]) # theta=pi/2 > pi/4 + + # Out-of-bounds phi: restricted phi grid + mesh7 = openmc.SphericalMesh( + r_grid=(0, 10), + phi_grid=(0, pi/2) + ) + with pytest.raises(ValueError): + mesh7.get_indices_at_coords([-5, 0, 0]) # phi=pi > pi/2 + + # Diagonal point: verify r, theta, phi all computed correctly + r = 6.0 + val = r / sqrt(3) + result = mesh3.get_indices_at_coords([val, val, val]) + assert result[0] == 1 # r=6 in second bin [5, 10] + assert result[1] == 1 # theta=arccos(1/sqrt(3))~0.955, in (pi/4, pi/2) + assert result[2] == 0 # phi=pi/4, in [0, pi/2) From a25bae027b19abfd0190eb991a904aab6a64358e Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Fri, 29 May 2026 23:40:17 -0500 Subject: [PATCH 2/3] Update math imports --- openmc/mesh.py | 40 +++++++++++++++++++--------------------- 1 file changed, 19 insertions(+), 21 deletions(-) diff --git a/openmc/mesh.py b/openmc/mesh.py index 6007baa1c0f..c3f8022dd4d 100644 --- a/openmc/mesh.py +++ b/openmc/mesh.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod from collections.abc import Iterable, Sequence, Mapping from functools import wraps -from math import acos, atan2, pi, sqrt +import math from numbers import Integral, Real from pathlib import Path from typing import Protocol @@ -11,12 +11,10 @@ import h5py import lxml.etree as ET import numpy as np -from pathlib import Path import openmc import openmc.checkvalue as cv from openmc.checkvalue import PathLike -from openmc.utility_funcs import change_directory from .bounding_box import BoundingBox from ._xml import get_elem_list, get_text from .mixin import IDManagerMixin @@ -291,7 +289,7 @@ def from_hdf5(cls, group: h5py.Group): """ mesh_type = 'regular' if 'type' not in group.keys() else group['type'][()].decode() mesh_id = int(group.name.split('/')[-1].lstrip('mesh ')) - mesh_name = '' if not 'name' in group else group['name'][()].decode() + mesh_name = '' if 'name' not in group else group['name'][()].decode() if mesh_type == 'regular': return RegularMesh.from_hdf5(group, mesh_id, mesh_name) @@ -1903,7 +1901,7 @@ def __init__( self, r_grid: Sequence[float], z_grid: Sequence[float], - phi_grid: Sequence[float] = (0, 2*pi), + phi_grid: Sequence[float] = (0, 2*math.pi), origin: Sequence[float] = (0., 0., 0.), mesh_id: int | None = None, name: str = '', @@ -1960,7 +1958,7 @@ def phi_grid(self, grid): cv.check_length('mesh phi_grid', grid, 2) cv.check_increasing('mesh phi_grid', grid) grid = np.asarray(grid, dtype=float) - if np.any((grid < 0.0) | (grid > 2*pi)): + if np.any((grid < 0.0) | (grid > 2*math.pi)): raise ValueError("phi_grid values must be in [0, 2π].") self._phi_grid = grid @@ -2046,7 +2044,7 @@ def get_indices_at_coords( The r, phi, z indices """ - r_value_from_origin = sqrt((coords[0]-self.origin[0])**2 + (coords[1]-self.origin[1])**2) + r_value_from_origin = math.hypot(coords[0]-self.origin[0], coords[1]-self.origin[1]) if r_value_from_origin < self.r_grid[0] or r_value_from_origin > self.r_grid[-1]: raise ValueError( @@ -2070,13 +2068,13 @@ def get_indices_at_coords( delta_x = coords[0] - self.origin[0] delta_y = coords[1] - self.origin[1] # atan2 returns values in -pi to +pi range - phi_value = atan2(delta_y, delta_x) + phi_value = math.atan2(delta_y, delta_x) if delta_x < 0 and delta_y < 0: # returned phi_value anticlockwise and negative - phi_value += 2 * pi + phi_value += 2 * math.pi if delta_x > 0 and delta_y < 0: # returned phi_value anticlockwise and negative - phi_value += 2 * pi + phi_value += 2 * math.pi phi_grid_values = np.array(self.phi_grid) @@ -2111,7 +2109,7 @@ def from_bounding_box( dimension: Sequence[int] = (10, 10, 10), mesh_id: int | None = None, name: str = '', - phi_grid_bounds: Sequence[float] = (0.0, 2*pi), + phi_grid_bounds: Sequence[float] = (0.0, 2*math.pi), enclose_domain: bool = False, ) -> CylindricalMesh: """Create CylindricalMesh from a bounding box. @@ -2339,8 +2337,8 @@ class SphericalMesh(StructuredMesh): def __init__( self, r_grid: Sequence[float], - phi_grid: Sequence[float] = (0, 2*pi), - theta_grid: Sequence[float] = (0, pi), + phi_grid: Sequence[float] = (0, 2*math.pi), + theta_grid: Sequence[float] = (0, math.pi), origin: Sequence[float] = (0., 0., 0.), mesh_id: int | None = None, name: str = '', @@ -2397,7 +2395,7 @@ def theta_grid(self, grid): cv.check_length('mesh theta_grid', grid, 2) cv.check_increasing('mesh theta_grid', grid) grid = np.asarray(grid, dtype=float) - if np.any((grid < 0.0) | (grid > pi)): + if np.any((grid < 0.0) | (grid > math.pi)): raise ValueError("theta_grid values must be in [0, π].") self._theta_grid = grid @@ -2411,7 +2409,7 @@ def phi_grid(self, grid): cv.check_length('mesh phi_grid', grid, 2) cv.check_increasing('mesh phi_grid', grid) grid = np.asarray(grid, dtype=float) - if np.any((grid < 0.0) | (grid > 2*pi)): + if np.any((grid < 0.0) | (grid > 2*math.pi)): raise ValueError("phi_grid values must be in [0, 2π].") self._phi_grid = grid @@ -2483,8 +2481,8 @@ def from_bounding_box( dimension: Sequence[int] = (10, 10, 10), mesh_id: int | None = None, name: str = '', - phi_grid_bounds: Sequence[float] = (0.0, 2*pi), - theta_grid_bounds: Sequence[float] = (0.0, pi), + phi_grid_bounds: Sequence[float] = (0.0, 2*math.pi), + theta_grid_bounds: Sequence[float] = (0.0, math.pi), enclose_domain: bool = False, ) -> SphericalMesh: """Create SphericalMesh from a bounding box. @@ -2673,7 +2671,7 @@ def get_indices_at_coords( dy = coords[1] - self.origin[1] dz = coords[2] - self.origin[2] - r_value = sqrt(dx**2 + dy**2 + dz**2) + r_value = math.hypot(dx, dy, dz) if r_value < self.r_grid[0] or r_value > self.r_grid[-1]: raise ValueError( @@ -2691,10 +2689,10 @@ def get_indices_at_coords( theta_value = 0.0 phi_value = 0.0 else: - theta_value = acos(dz / r_value) - phi_value = atan2(dy, dx) + theta_value = math.acos(dz / r_value) + phi_value = math.atan2(dy, dx) if phi_value < 0: - phi_value += 2 * pi + phi_value += 2 * math.pi if theta_value < self.theta_grid[0] or theta_value > self.theta_grid[-1]: raise ValueError( From 49c3058919f925bf8f5d4814fffb0b298ed952df Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Fri, 29 May 2026 23:49:45 -0500 Subject: [PATCH 3/3] Style fixes --- openmc/mesh.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/openmc/mesh.py b/openmc/mesh.py index c3f8022dd4d..ff8e1457724 100644 --- a/openmc/mesh.py +++ b/openmc/mesh.py @@ -2026,9 +2026,9 @@ def __repr__(self): return string def get_indices_at_coords( - self, - coords: Sequence[float] - ) -> tuple[int, int, int]: + self, + coords: Sequence[float] + ) -> tuple[int, int, int]: """Finds the index of the mesh element at the specified coordinates. .. versionadded:: 0.15.0 @@ -2644,9 +2644,9 @@ def _convert_to_cartesian(arr, origin: Sequence[float]): return arr def get_indices_at_coords( - self, - coords: Sequence[float] - ) -> tuple[int, int, int]: + self, + coords: Sequence[float] + ) -> tuple[int, int, int]: """Find the mesh cell indices containing the specified coordinates. .. versionadded:: 0.15.4