fix: make StructuredGrid3DGeometry agree with the structured support it duplicates

packages/loop_common/src/loop_common/geometry/_structured_grid_3d.py

@@ -172,20 +172,37 @@ def rotate(self, pos):
         return np.einsum("ijk,ik->ij", self.rotation_xy[None, :, :], pos)
 
     def position_to_cell_index(self, pos: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
-        inside = self.inside(pos)
+        """Get the (i, j, k) index of the cell that each point is in.
+
+        Points are expected in the grid's own frame; ``rotation_xy`` is not applied
+        (see :meth:`rotate`).
+
+        Parameters
+        ----------
+        pos : np.ndarray
+            Nx3 array of xyz locations
+
+        Returns
+        -------
+        tuple[np.ndarray, np.ndarray]
+            N,3 integer cell indexes and a boolean mask of the points that are inside the
+            grid. Points on the maximum boundary belong to the last cell. The index of a point
+            outside the grid is (0, 0, 0): use the mask, not the index, to detect them.
+        """
         pos = self.check_position(pos)
+        inside = self.inside(pos)
         cell_indexes = np.zeros((pos.shape[0], 3), dtype=int)
-        cell_indexes[:, 0] = (pos[:, 0] - self.origin[0]) / self.step_vector[0]
-        cell_indexes[:, 1] = (pos[:, 1] - self.origin[1]) / self.step_vector[1]
-        cell_indexes[:, 2] = (pos[:, 2] - self.origin[2]) / self.step_vector[2]
-        return cell_indexes.astype(int), inside
+        for axis in range(3):
+            cell = np.floor((pos[inside, axis] - self.origin[axis]) / self.step_vector[axis])
+            cell_indexes[inside, axis] = np.clip(cell.astype(int), 0, self.nsteps_cells[axis] - 1)
+        return cell_indexes, inside
 
     def inside(self, pos: np.ndarray) -> np.ndarray:
-        inside = np.ones(pos.shape[0]).astype(bool)
-        for i in range(3):
-            inside *= pos[:, i] > self.origin[i]
-            inside *= pos[:, i] < self.maximum[i]
-        return inside
+        """Boolean mask of the points inside the grid, boundary included (to a small tolerance)."""
+        pos = self.check_position(pos)
+        extent = np.maximum(np.abs(self.maximum - self.origin), 1.0)
+        tol = np.finfo(float).eps * 32.0 * extent
+        return np.all((pos >= (self.origin - tol)) & (pos <= (self.maximum + tol)), axis=1)
 
     def check_position(self, pos: np.ndarray) -> np.ndarray:
         if len(pos.shape) == 1:
@@ -229,15 +246,36 @@ def global_index_to_cell_index(self, global_index):
         cell_indexes = np.zeros((global_index.shape[0], 3), dtype=np.int64)
         cell_indexes[:, 0] = global_index % self.nsteps_cells[0, None]
         cell_indexes[:, 1] = (global_index // self.nsteps_cells[0, None]) % self.nsteps_cells[1, None]
-        cell_indexes[:, 2] = (global_index // (self.nsteps_cells[0, None] * self.nsteps_cells[1, None])) % self.nsteps_cells[2, None]
+        cell_indexes[:, 2] = global_index // self.nsteps_cells[0, None] // self.nsteps_cells[1, None]
         return cell_indexes
 
     def global_index_to_node_index(self, global_index):
-        cell_indexes = np.zeros((global_index.shape[0], 3), dtype=np.int64)
-        cell_indexes[:, 0] = global_index % self.nsteps[0, None]
-        cell_indexes[:, 1] = (global_index // self.nsteps[0, None]) % self.nsteps[1, None]
-        cell_indexes[:, 2] = (global_index // (self.nsteps[0, None] * self.nsteps[1, None])) % self.nsteps[2, None]
-        return cell_indexes
+        """Node ijk indexes (shape ``(*global_index.shape, 3)``) from flat node indexes."""
+        global_index = np.asarray(global_index)
+        original_shape = global_index.shape
+        global_index = global_index.reshape(-1)
+        node_indexes = np.zeros((global_index.shape[0], 3), dtype=np.int64)
+        node_indexes[:, 0] = global_index % self.nsteps[0, None]
+        node_indexes[:, 1] = global_index // self.nsteps[0, None] % self.nsteps[1, None]
+        node_indexes[:, 2] = global_index // self.nsteps[0, None] // self.nsteps[1, None]
+        return node_indexes.reshape(*original_shape, 3)
+
+    def _global_indicies(self, indexes: np.ndarray, nsteps: np.ndarray) -> np.ndarray:
+        """Convert (..., 3) ijk indexes to flat indexes for a grid with ``nsteps`` per axis (x fastest)."""
+        indexes = np.asarray(indexes)
+        if indexes.ndim == 1:
+            raise ValueError("Cell indexes needs to be Nx3")
+        if indexes.shape[-1] != 3:
+            raise ValueError("Last dimensions should be ijk indexing")
+        original_shape = indexes.shape
+        indexes = indexes.reshape(-1, 3)
+        gi = indexes[:, 0] + nsteps[0] * indexes[:, 1] + nsteps[0] * nsteps[1] * indexes[:, 2]
+        return gi.reshape(original_shape[:-1])
 
     def global_node_indices(self, node_indexes):
-        return node_indexes
+        """Flat node indexes from (..., 3) node ijk indexes."""
+        return self._global_indicies(node_indexes, self.nsteps)
+
+    def global_cell_indices(self, cell_indexes):
+        """Flat cell indexes from (..., 3) cell ijk indexes; the order cell data is stored in."""
+        return self._global_indicies(cell_indexes, self.nsteps_cells)

packages/loop_common/tests/test_structured_grid_geometry.py

@@ -0,0 +1,122 @@
+"""StructuredGrid3DGeometry must agree with the interpolation support it duplicates.
+
+``loop_common.supports.StructuredGrid`` (via ``BaseStructuredSupport``) and
+``StructuredGrid3DGeometry`` implement the same grid indexing. The tests below pin the shared
+behaviour so the two cannot drift apart again; they compare the geometry class against the
+support on the same inputs, and spell out the boundary rules.
+"""
+import numpy as np
+import pytest
+from loop_common.geometry import StructuredGrid3DGeometry
+from loop_common.supports import StructuredGrid
+
+# (origin, nodes per axis, step vector)
+GRIDS = {
+    "unit": ([0.0, 0.0, 0.0], [5, 5, 5], [1.0, 1.0, 1.0]),
+    "offset_anisotropic": ([10.0, -3.0, 2.0], [7, 5, 4], [0.5, 2.0, 1.5]),
+    "minimal": ([0.0, 0.0, 0.0], [2, 2, 2], [3.0, 2.0, 1.0]),
+}
+
+
+@pytest.fixture(params=list(GRIDS))
+def pair(request):
+    origin, nsteps, step = GRIDS[request.param]
+    geometry = StructuredGrid3DGeometry(origin=origin, nsteps=nsteps, step_vector=step)
+    support = StructuredGrid(origin=origin, nsteps_cells=np.array(nsteps) - 1, step_vector=step)
+    return geometry, support
+
+
+def _points(geometry, n=300, seed=0):
+    rng = np.random.default_rng(seed)
+    low = geometry.origin - 1.0
+    high = geometry.maximum + 1.0
+    return rng.uniform(low, high, (n, 3))
+
+
+def test_counts_and_extent_agree(pair):
+    geometry, support = pair
+    np.testing.assert_array_equal(geometry.nsteps, support.nsteps)
+    np.testing.assert_array_equal(geometry.nsteps_cells, support.nsteps_cells)
+    np.testing.assert_allclose(geometry.maximum, support.maximum)
+    assert geometry.n_nodes == support.n_nodes
+    assert geometry.n_elements == support.n_elements
+    np.testing.assert_allclose(geometry.nodes, support.nodes)
+
+
+def test_position_to_cell_index_agrees(pair):
+    geometry, support = pair
+    points = _points(geometry)
+    g_index, g_inside = geometry.position_to_cell_index(points)
+    s_index, s_inside = support.position_to_cell_index(points)
+    np.testing.assert_array_equal(g_inside, s_inside)
+    np.testing.assert_array_equal(g_index, s_index)
+
+
+def test_inside_agrees_including_the_boundary(pair):
+    geometry, support = pair
+    corners = np.array([geometry.origin, geometry.maximum])
+    points = np.vstack([_points(geometry), corners])
+    np.testing.assert_array_equal(geometry.inside(points), support.inside(points))
+    assert geometry.inside(corners).all()
+
+
+def test_index_conversions_agree(pair):
+    geometry, support = pair
+    cells = np.arange(geometry.n_elements)
+    nodes = np.arange(geometry.n_nodes)
+    np.testing.assert_array_equal(
+        geometry.global_index_to_cell_index(cells), support.global_index_to_cell_index(cells)
+    )
+    np.testing.assert_array_equal(
+        geometry.global_index_to_node_index(nodes), support.global_index_to_node_index(nodes)
+    )
+    cell_ijk = geometry.global_index_to_cell_index(cells)
+    node_ijk = geometry.global_index_to_node_index(nodes)
+    np.testing.assert_array_equal(
+        geometry.global_cell_indices(cell_ijk), support.global_cell_indices(cell_ijk)
+    )
+    np.testing.assert_array_equal(
+        geometry.global_node_indices(node_ijk), support.global_node_indices(node_ijk)
+    )
+
+
+def test_elements_are_flat_node_indexes(pair):
+    geometry, support = pair
+    assert geometry.elements.shape == (geometry.n_elements, 8)
+    np.testing.assert_array_equal(geometry.elements, support.elements)
+
+
+def test_cell_index_round_trips_through_the_flat_index(pair):
+    geometry, _ = pair
+    cells = np.arange(geometry.n_elements)
+    ijk = geometry.global_index_to_cell_index(cells)
+    np.testing.assert_array_equal(geometry.global_cell_indices(ijk), cells)
+
+
+def test_point_on_the_maximum_boundary_is_in_the_last_cell():
+    geometry = StructuredGrid3DGeometry(origin=[0, 0, 0], nsteps=[5, 5, 5], step_vector=[1, 1, 1])
+    index, inside = geometry.position_to_cell_index(np.array([[4.0, 4.0, 4.0], [0.0, 0.0, 0.0]]))
+    assert inside.all()
+    np.testing.assert_array_equal(index, [[3, 3, 3], [0, 0, 0]])
+
+
+def test_point_outside_is_flagged_not_indexed():
+    geometry = StructuredGrid3DGeometry(origin=[0, 0, 0], nsteps=[5, 5, 5], step_vector=[1, 1, 1])
+    index, inside = geometry.position_to_cell_index(np.array([[6.0, 2.0, 2.0], [-0.5, 2.0, 2.0]]))
+    assert not inside.any()
+    np.testing.assert_array_equal(index, [[0, 0, 0], [0, 0, 0]])
+
+
+def test_single_point_is_accepted():
+    geometry = StructuredGrid3DGeometry(origin=[0, 0, 0], nsteps=[5, 5, 5], step_vector=[1, 1, 1])
+    index, inside = geometry.position_to_cell_index(np.array([1.5, 2.5, 0.5]))
+    np.testing.assert_array_equal(index, [[1, 2, 0]])
+    assert inside.tolist() == [True]
+
+
+def test_index_helpers_reject_bad_shapes():
+    geometry = StructuredGrid3DGeometry(origin=[0, 0, 0], nsteps=[3, 3, 3], step_vector=[1, 1, 1])
+    with pytest.raises(ValueError):
+        geometry.global_cell_indices(np.array([0, 0, 0]))
+    with pytest.raises(ValueError):
+        geometry.global_cell_indices(np.zeros((4, 2), dtype=int))