Implement SphericalMesh.get_indices_at_coords

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

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)