diff --git a/packages/loop_common/src/loop_common/geometry/_structured_grid.py b/packages/loop_common/src/loop_common/geometry/_structured_grid.py index a6622ce5e..3b2a8de2e 100644 --- a/packages/loop_common/src/loop_common/geometry/_structured_grid.py +++ b/packages/loop_common/src/loop_common/geometry/_structured_grid.py @@ -4,6 +4,8 @@ from loop_common.logging import get_logger as getLogger +from ._structured_grid_3d import StructuredGrid3DGeometry + logger = getLogger(__name__) @@ -11,11 +13,16 @@ 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)): diff --git a/packages/loop_common/tests/test_property_grid.py b/packages/loop_common/tests/test_property_grid.py new file mode 100644 index 000000000..797662981 --- /dev/null +++ b/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)]])