Implement SphericalMesh.get_indices_at_coords (#3919)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

Author: EdenRochmanSharabi

openmc/mesh.py

@@ -3,20 +3,18 @@
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Sequence, Mapping
 from functools import wraps
-from math import pi, sqrt, atan2
+import math
 from numbers import Integral, Real
 from pathlib import Path
 from typing import Protocol
 
 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
 
@@ -2028,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
@@ -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.
@@ -2645,10 +2643,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 = math.hypot(dx, dy, dz)
+
+        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 = math.acos(dz / r_value)
+            phi_value = math.atan2(dy, dx)
+            if phi_value < 0:
+                phi_value += 2 * math.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):

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
@@ -1065,3 +1065,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)