refactor: PropertyGrid takes its geometry from StructuredGrid3DGeometry

packages/loop_common/src/loop_common/geometry/_structured_grid.py

@@ -4,18 +4,25 @@
 
 from loop_common.logging import get_logger as getLogger
 
+from ._structured_grid_3d import StructuredGrid3DGeometry
+
 logger = getLogger(__name__)
 
 
 @dataclass
 class PropertyGrid:
     """A regular 3D grid for storing property values (the result container).
 
-    This is a plain data container, not an interpolation support: it holds
+    This is a data container, not an interpolation support: it holds
     node-based ``properties`` and ``cell_properties`` on a regular grid and can
     export them (``vtk``, ``save``). For the class used to interpolate on a
     grid see :class:`loop_common.supports.StructuredGrid`.
 
+    The grid geometry (node and cell positions, point-in-cell lookup) is defined once, by
+    :class:`StructuredGrid3DGeometry`, which is built from ``origin``, ``step_vector`` and
+    ``nsteps`` (the number of **nodes**) whenever it is needed; see :attr:`geometry`. Points
+    passed to the lookup methods are in the grid's own frame (the grid is not rotated).
+
     It was previously named ``StructuredGrid``; that name remains as an alias.
 
     This class represents a regular 3D grid with properties and cell properties
@@ -62,6 +69,16 @@ def to_dict(self):
             "name": self.name,
         }
 
+    @property
+    def geometry(self) -> StructuredGrid3DGeometry:
+        """The grid geometry, built from the current ``origin``, ``step_vector`` and ``nsteps``.
+
+        It is a fresh object each time, so it always reflects the current fields.
+        """
+        return StructuredGrid3DGeometry(
+            origin=self.origin, nsteps=self.nsteps, step_vector=self.step_vector
+        )
+
     @property
     def maximum(self):
         """Calculate the maximum coordinates of the grid.
@@ -71,7 +88,70 @@ def maximum(self):
         np.ndarray
             Maximum coordinates (origin + (nsteps - 1) * step_vector)
         """
-        return self.origin + (self.nsteps - 1) * self.step_vector
+        return self.geometry.maximum
+
+    @property
+    def nsteps_cells(self):
+        """Number of cells in each direction (``nsteps - 1``)."""
+        return self.geometry.nsteps_cells
+
+    @property
+    def n_nodes(self):
+        """Total number of nodes."""
+        return self.geometry.n_nodes
+
+    @property
+    def n_cells(self):
+        """Total number of cells."""
+        return self.geometry.n_elements
+
+    def inside(self, points):
+        """Boolean mask of the points inside the grid (boundary included)."""
+        return self.geometry.inside(np.asarray(points, dtype=float))
+
+    def cell_index_at(self, points):
+        """Flat index of the cell each point falls in.
+
+        The flat index addresses ``cell_properties`` arrays and the rows of ``cell_centres``
+        (x varies fastest). Points on the maximum boundary belong to the last cell.
+
+        Parameters
+        ----------
+        points : array-like
+            N x 3 locations (or one location) in the grid's own frame
+
+        Returns
+        -------
+        np.ndarray
+            N integer indexes, ``-1`` for points outside the grid
+        """
+        geometry = self.geometry
+        cell_ijk, inside = geometry.position_to_cell_index(np.asarray(points, dtype=float))
+        index = np.full(len(inside), -1, dtype=np.int64)
+        index[inside] = geometry.global_cell_indices(cell_ijk[inside])
+        return index
+
+    def cell_value_at(self, name, points):
+        """Value of the cell property ``name`` in the cell each point falls in.
+
+        Parameters
+        ----------
+        name : str
+            key of ``cell_properties``
+        points : array-like
+            N x 3 locations (or one location) in the grid's own frame
+
+        Returns
+        -------
+        np.ndarray
+            one entry per point; NaN for points outside the grid
+        """
+        values = np.asarray(self.cell_properties[name])
+        index = self.cell_index_at(points)
+        found = index >= 0
+        result = np.full((len(index),) + values.shape[1:], np.nan)
+        result[found] = values[index[found]]
+        return result
 
     def vtk(self):
         """Convert the structured grid to a PyVista RectilinearGrid.
@@ -122,39 +202,21 @@ def plot(self, pyvista_kwargs=None):
 
     @property
     def cell_centres(self):
-        """Calculate the coordinates of cell centres.
+        """Coordinates of the centres of the ``nsteps - 1`` cells in each direction.
 
         Returns
         -------
-        tuple of np.ndarray
-            X, Y, Z coordinates of all cell centres
+        np.ndarray
+            (n_cells, 3) array in the same order as ``cell_properties`` (x varies fastest)
         """
-        x = np.linspace(
-            self.origin[0] + self.step_vector[0] * 0.5,
-            self.maximum[0] - self.step_vector[0] * 0.5,
-            self.nsteps[0] - 1,
-        )
-        y = np.linspace(
-            self.origin[1] + self.step_vector[1] * 0.5,
-            self.maximum[1] - self.step_vector[1] * 0.5,
-            self.nsteps[1] - 1,
-        )
-        z = np.linspace(
-            self.origin[2] + self.step_vector[2] * 0.5,
-            self.maximum[2] - self.step_vector[2] * 0.5,
-            self.nsteps[2] - 1,
-        )
-        x, y, z = np.meshgrid(x, y, z, indexing="ij")
-        return np.vstack([x.flatten(order="f"), y.flatten(order="f"), z.flatten(order="f")]).T
+        geometry = self.geometry
+        cell_ijk = geometry.global_index_to_cell_index(np.arange(geometry.n_elements))
+        return geometry.origin + (cell_ijk + 0.5) * geometry.step_vector
 
     @property
     def nodes(self):
-
-        x = np.linspace(self.origin[0], self.maximum[0], self.nsteps[0])
-        y = np.linspace(self.origin[1], self.maximum[1], self.nsteps[1])
-        z = np.linspace(self.origin[2], self.maximum[2], self.nsteps[2])
-        x, y, z = np.meshgrid(x, y, z, indexing="ij")
-        return np.vstack([x.flatten(order="f"), y.flatten(order="f"), z.flatten(order="f")]).T
+        """Coordinates of the nodes, (n_nodes, 3), in the same order as ``properties``."""
+        return self.geometry.nodes
 
     def merge(self, other):
         if not np.all(np.isclose(self.origin, other.origin)):

packages/loop_common/tests/test_property_grid.py

@@ -0,0 +1,91 @@
+"""PropertyGrid: geometry comes from StructuredGrid3DGeometry, and cell lookups index cell data."""
+import numpy as np
+import pytest
+from loop_common.geometry import PropertyGrid, StructuredGrid3DGeometry
+
+GRIDS = {
+    "unit_int": dict(origin=np.array([0, 0, 0]), step_vector=np.array([1, 1, 1]), nsteps=np.array([4, 4, 4])),
+    "offset_anisotropic": dict(
+        origin=np.array([10.0, -3.0, 2.0]),
+        step_vector=np.array([0.5, 2.0, 1.5]),
+        nsteps=np.array([7, 5, 4]),
+    ),
+    "minimal": dict(origin=np.zeros(3), step_vector=np.array([3.0, 2.0, 1.0]), nsteps=np.array([2, 2, 2])),
+}
+
+
+@pytest.fixture(params=list(GRIDS))
+def grid(request):
+    return PropertyGrid(**GRIDS[request.param])
+
+
+def _axis_nodes(grid, axis):
+    return grid.origin[axis] + grid.step_vector[axis] * np.arange(grid.nsteps[axis])
+
+
+def test_geometry_is_built_from_the_current_fields(grid):
+    assert isinstance(grid.geometry, StructuredGrid3DGeometry)
+    np.testing.assert_array_equal(grid.geometry.nsteps, grid.nsteps)
+    grid.nsteps = grid.nsteps + 1
+    np.testing.assert_array_equal(grid.geometry.nsteps, grid.nsteps)
+
+
+def test_maximum_and_counts(grid):
+    np.testing.assert_allclose(grid.maximum, grid.origin + (grid.nsteps - 1) * grid.step_vector)
+    np.testing.assert_array_equal(grid.nsteps_cells, grid.nsteps - 1)
+    assert grid.n_nodes == np.prod(grid.nsteps)
+    assert grid.n_cells == np.prod(grid.nsteps - 1)
+
+
+def test_nodes_match_the_independent_construction(grid):
+    xs, ys, zs = (_axis_nodes(grid, a) for a in range(3))
+    expected = np.array(np.meshgrid(xs, ys, zs, indexing="ij")).reshape(3, -1, order="F").T
+    np.testing.assert_allclose(grid.nodes, expected)
+
+
+def test_cell_centres_match_the_independent_construction(grid):
+    mids = [0.5 * (n[:-1] + n[1:]) for n in (_axis_nodes(grid, a) for a in range(3))]
+    expected = np.array(np.meshgrid(*mids, indexing="ij")).reshape(3, -1, order="F").T
+    assert grid.cell_centres.shape == (grid.n_cells, 3)
+    np.testing.assert_allclose(grid.cell_centres, expected)
+
+
+def test_cell_centres_are_inside_the_grid(grid):
+    assert grid.inside(grid.cell_centres).all()
+    assert (grid.cell_centres >= grid.origin).all() and (grid.cell_centres <= grid.maximum).all()
+
+
+def test_each_cell_centre_is_in_its_own_cell(grid):
+    np.testing.assert_array_equal(grid.cell_index_at(grid.cell_centres), np.arange(grid.n_cells))
+
+
+def test_cell_value_at_reads_cell_data_in_storage_order(grid):
+    # cell data defined as the x coordinate of the cell centre
+    grid.cell_properties["x"] = grid.cell_centres[:, 0].copy()
+    rng = np.random.default_rng(0)
+    points = rng.uniform(grid.origin, grid.maximum, (200, 3))
+    index = grid.cell_index_at(points)
+    assert (index >= 0).all()
+    # a point lies within half a step of the centre of the cell it is reported to be in
+    centres = grid.cell_centres[index]
+    assert (np.abs(points - centres) <= 0.5 * grid.step_vector + 1e-9).all()
+    np.testing.assert_allclose(grid.cell_value_at("x", points), centres[:, 0])
+
+
+def test_points_outside_give_minus_one_and_nan(grid):
+    outside = np.array([grid.maximum + grid.step_vector, grid.origin - grid.step_vector])
+    np.testing.assert_array_equal(grid.cell_index_at(outside), [-1, -1])
+    grid.cell_properties["x"] = np.arange(grid.n_cells, dtype=float)
+    assert np.isnan(grid.cell_value_at("x", outside)).all()
+    assert not grid.inside(outside).any()
+
+
+def test_boundary_points_belong_to_the_last_cell(grid):
+    np.testing.assert_array_equal(grid.cell_index_at(grid.maximum), [grid.n_cells - 1])
+    np.testing.assert_array_equal(grid.cell_index_at(grid.origin), [0])
+
+
+def test_vector_valued_cell_data(grid):
+    grid.cell_properties["v"] = np.column_stack([np.arange(grid.n_cells), -np.arange(grid.n_cells)]).astype(float)
+    values = grid.cell_value_at("v", grid.cell_centres[[0, -1]])
+    np.testing.assert_allclose(values, [[0, 0], [grid.n_cells - 1, -(grid.n_cells - 1)]])