Adding get cylindrical mesh index at specified coordinates (#2782)

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

openmc/mesh.py

@@ -2,7 +2,7 @@
 import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
-from math import pi
+from math import pi, sqrt, atan2
 from numbers import Integral, Real
 from pathlib import Path
 from typing import Optional, Sequence, Tuple
@@ -1347,6 +1347,66 @@ def __repr__(self):
             string += fmt.format('\tZ Max:', '=\t', self._z_grid[-1])
         return string
 
+    def get_indices_at_coords(
+            self,
+            coords: Sequence[float]
+        ) -> Tuple[int, int, int]:
+        """Finds the index of the mesh voxel at the specified x,y,z coordinates.
+
+        Parameters
+        ----------
+        coords : Sequence[float]
+            The x, y, z axis coordinates
+
+        Returns
+        -------
+        Tuple[int, int, int]
+            The r, phi, z indices
+
+        """
+        r_value_from_origin = sqrt((coords[0]-self.origin[0])**2 + (coords[1]-self.origin[1])**2)
+
+        if r_value_from_origin < self.r_grid[0] or r_value_from_origin > self.r_grid[-1]:
+            raise ValueError(
+                f'The specified x, y ({coords[0]}, {coords[1]}) combine to give an r value of '
+                f'{r_value_from_origin} from the origin of {self.origin}.which '
+                f'is outside the origin absolute r grid values {self.r_grid}.'
+            )
+
+        r_index = np.searchsorted(self.r_grid, r_value_from_origin) - 1
+
+        z_grid_values = np.array(self.z_grid) + self.origin[2]
+
+        if coords[2] < z_grid_values[0] or coords[2] > z_grid_values[-1]:
+            raise ValueError(
+                f'The specified z value ({coords[2]}) from the z origin of '
+                f'{self.origin[-1]} is outside of the absolute z grid range {z_grid_values}.'
+            )
+
+        z_index = np.argmax(z_grid_values > coords[2]) - 1
+
+        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)
+        if delta_x < 0 and delta_y < 0:
+            # returned phi_value anticlockwise and negative
+            phi_value += 2 * pi
+        if delta_x > 0 and delta_y < 0:
+            # returned phi_value anticlockwise and negative
+            phi_value += 2 * pi
+
+        phi_grid_values = np.array(self.phi_grid)
+
+        if phi_value < phi_grid_values[0] or phi_value > phi_grid_values[-1]:
+            raise ValueError(
+                f'The phi value ({phi_value}) resulting from the specified x, y '
+                f'values is outside of the absolute  phi grid range {phi_grid_values}.'
+            )
+        phi_index = np.argmax(phi_grid_values > phi_value) - 1
+
+        return (r_index, phi_index, z_index)
+
     @classmethod
     def from_hdf5(cls, group: h5py.Group):
         mesh_id = int(group.name.split('/')[-1].lstrip('mesh '))

tests/unit_tests/test_mesh.py

@@ -77,6 +77,7 @@ def test_cylindrical_mesh_bounding_box():
     np.testing.assert_array_equal(mesh.lower_left, (2, 4, -3))
     np.testing.assert_array_equal(mesh.upper_right, (4, 6, 17))
 
+
 def test_spherical_mesh_bounding_box():
     # test with mesh at origin (0, 0, 0)
     mesh = openmc.SphericalMesh([0.1, 0.2, 0.5, 1.], origin=(0., 0., 0.))
@@ -239,4 +240,51 @@ def test_mesh_vertices(mesh_type):
         np.testing.assert_equal(mesh.vertices_spherical[ijk], exp_vert)
 
 
+def test_CylindricalMesh_get_indices_at_coords():
+    # default origin (0, 0, 0) and default phi grid (0, 2*pi)
+    mesh = openmc.CylindricalMesh(r_grid=(0, 5, 10), z_grid=(0, 5, 10))
+    assert mesh.get_indices_at_coords([1, 0, 1]) == (0, 0, 0)
+    assert mesh.get_indices_at_coords([6, 0, 1]) == (1, 0, 0)
+    assert mesh.get_indices_at_coords([9, 0, 1]) == (1, 0, 0)
+    assert mesh.get_indices_at_coords([0, 6, 0]) == (1, 0, 0)
+    assert mesh.get_indices_at_coords([0, 9, 6]) == (1, 0, 1)
+    assert mesh.get_indices_at_coords([-2, -2, 9]) == (0, 0, 1)
+
+    with pytest.raises(ValueError):
+        assert mesh.get_indices_at_coords([8, 8, 1])  # resulting r value to large
+    with pytest.raises(ValueError):
+        assert mesh.get_indices_at_coords([-8, -8, 1])  # resulting r value to large
+    with pytest.raises(ValueError):
+        assert mesh.get_indices_at_coords([1, 0, -1])  # z value below range
+    with pytest.raises(ValueError):
+        assert mesh.get_indices_at_coords([1, 0, 11])  # z value above range
+
+    assert mesh.get_indices_at_coords([1, 1, 1]) == (0, 0, 0)
+
+    # negative range on z grid
+    mesh = openmc.CylindricalMesh(
+        r_grid=(0, 5, 10),
+        phi_grid=(0, 0.5 * pi, pi, 1.5 * pi, 1.9 * pi),
+        z_grid=(-5, 0, 5, 10),
+    )
+    assert mesh.get_indices_at_coords([1, 1, 1]) == (0, 0, 1)  # first angle quadrant
+    assert mesh.get_indices_at_coords([2, 2, 6]) == (0, 0, 2)  # first angle quadrant
+    assert mesh.get_indices_at_coords([-2, 0.1, -1]) == (0, 1, 0)  # second angle quadrant
+    assert mesh.get_indices_at_coords([-2, -0.1, -1]) == (0, 2, 0)  # third angle quadrant
+    assert mesh.get_indices_at_coords([2, -0.9, -1]) == (0, 3, 0)  # forth angle quadrant
 
+    with pytest.raises(ValueError):
+        assert mesh.get_indices_at_coords([2, -0.1, 1])  # outside of phi range
+
+    # origin of mesh not default
+    mesh = openmc.CylindricalMesh(
+        r_grid=(0, 5, 10),
+        phi_grid=(0, 0.5 * pi, pi, 1.5 * pi, 1.9 * pi),
+        z_grid=(-5, 0, 5, 10),
+        origin=(100, 200, 300),
+    )
+    assert mesh.get_indices_at_coords([101, 201, 301]) == (0, 0, 1)  # first angle quadrant
+    assert mesh.get_indices_at_coords([102, 202, 306]) == (0, 0, 2)  # first angle quadrant
+    assert mesh.get_indices_at_coords([98, 200.1, 299]) == (0, 1, 0)  # second angle quadrant
+    assert mesh.get_indices_at_coords([98, 199.9, 299]) == (0, 2, 0)  # third angle quadrant
+    assert mesh.get_indices_at_coords([102, 199.1, 299]) == (0, 3, 0)  # forth angle quadrant