diff --git a/.github/workflows/qgis-compat.yml b/.github/workflows/qgis-compat.yml index 3a5078c17..287f7a92b 100644 --- a/.github/workflows/qgis-compat.yml +++ b/.github/workflows/qgis-compat.yml @@ -36,7 +36,7 @@ jobs: - name: Set up uv and Python uses: astral-sh/setup-uv@v5 with: - python-version: "3.9" + python-version: "3.12" enable-cache: true - name: Install plugin's non-QGIS test requirements diff --git a/packages/loop_common/src/loop_common/geometry/_structured_grid_2d.py b/packages/loop_common/src/loop_common/geometry/_structured_grid_2d.py index e4c8c8b74..7c48836b6 100644 --- a/packages/loop_common/src/loop_common/geometry/_structured_grid_2d.py +++ b/packages/loop_common/src/loop_common/geometry/_structured_grid_2d.py @@ -56,12 +56,10 @@ def elements(self) -> np.ndarray: return self.global_node_indices(self.cell_corner_indexes(cell_indexes)) def print_geometry(self): - logger.info(f"Origin: {self.origin[0]:f} {self.origin[1]:f} {self.origin[2]:f}") - logger.info( - f"Cell size: {self.step_vector[0]:f} {self.step_vector[1]:f} {self.step_vector[2]:f}" - ) + logger.info(f"Origin: {self.origin[0]:f} {self.origin[1]:f}") + logger.info(f"Cell size: {self.step_vector[0]:f} {self.step_vector[1]:f}") max = self.origin + self.nsteps_cells * self.step_vector - logger.info(f"Max extent: {max[0]:f} {max[1]:f} {max[2]:f}") + logger.info(f"Max extent: {max[0]:f} {max[1]:f}") def cell_centres(self, global_index: np.ndarray) -> np.ndarray: cell_indexes = self.global_index_to_cell_index(global_index) @@ -142,5 +140,157 @@ def global_index_to_node_index(self, global_index): cell_indexes[:, 1] = global_index // self.nsteps[0, None] % self.nsteps[1, None] return cell_indexes - def global_node_indices(self, node_indexes): - return node_indexes + def global_node_indices(self, node_indexes: np.ndarray) -> np.ndarray: + """Flat node indexes from (..., 2) node ij indexes.""" + return self._global_indices(node_indexes, self.nsteps) + + def bilinear(self, local_coords: np.ndarray) -> np.ndarray: + """ + returns the bilinear interpolation for the local coordinates + Parameters + ---------- + x - double, array of doubles + y - double, array of doubles + z - double, array of doubles + + Returns + ------- + array of interpolation coefficients + + """ + + return np.array( + [ + (1 - local_coords[:, 0]) * (1 - local_coords[:, 1]), + local_coords[:, 0] * (1 - local_coords[:, 1]), + (1 - local_coords[:, 0]) * local_coords[:, 1], + local_coords[:, 0] * local_coords[:, 1], + ] + ).T + + def position_to_local_coordinates(self, pos: np.ndarray) -> np.ndarray: + """ + Convert from global to local coordinates within a cel + Parameters + ---------- + pos - array of positions inside + + Returns + ------- + localx, localy, localz + + """ + # TODO check if inside mesh + + # calculate local coordinates for positions + local_coords = np.zeros(pos.shape) + local_coords[:, 0] = ( + (pos[:, 0] - self.origin[None, 0]) % self.step_vector[None, 0] + ) / self.step_vector[None, 0] + local_coords[:, 1] = ( + (pos[:, 1] - self.origin[None, 1]) % self.step_vector[None, 1] + ) / self.step_vector[None, 1] + + return local_coords + + def position_to_dof_coefs(self, pos: np.ndarray): + """ + global posotion to interpolation coefficients + Parameters + ---------- + pos + + Returns + ------- + + """ + local_coords = self.position_to_local_coordinates(pos) + weights = self.bilinear(local_coords) + return weights + + def _global_indices(self, indexes: np.ndarray, nsteps: np.ndarray) -> np.ndarray: + if len(indexes.shape) == 1: + raise ValueError("Indexes must be a 2D array") + if indexes.shape[-1] != 2: + raise ValueError("Last dimension of cell indexing needs to be ijk indexing") + original_shape = indexes.shape + indexes = indexes.reshape(-1, 2) + gi = indexes[:, 0] + nsteps[0] * indexes[:, 1] + return gi.reshape(original_shape[:-1]) + + def global_cell_indices(self, indexes: np.ndarray) -> np.ndarray: + return self._global_indices(indexes, self.nsteps_cells) + + def node_indexes_to_position(self, node_indexes: np.ndarray) -> np.ndarray: + + original_shape = node_indexes.shape + node_indexes = node_indexes.reshape((-1, 2)) + xy = np.zeros((node_indexes.shape[0], 2), dtype=float) + xy[:, 0] = self.origin[0] + self.step_vector[0] * node_indexes[:, 0] + xy[:, 1] = self.origin[1] + self.step_vector[1] * node_indexes[:, 1] + xy = xy.reshape(original_shape) + return xy + + def position_to_cell_corners(self, pos): + """Get the global indices of the vertices (corner) nodes of the cell containing each point. + + Parameters + ---------- + pos : np.array + (N, 2) array of xy coordinates representing the positions of N points. + + Returns + ------- + globalidx : np.array + (N, 4) array of global indices corresponding to the 4 corner nodes of the cell + each point lies in. If a point lies outside the support, its corresponding entry + will be set to -1. + inside : np.array + (N,) boolean array indicating whether each point is inside the support domain. + """ + corner_index, inside = self.position_to_cell_index(pos) + corners = self.cell_corner_indexes(corner_index) + globalidx = self.global_node_indices(corners) + # if global index is not inside the support set to -1 + globalidx[~inside] = -1 + return globalidx, inside + + def position_to_cell_vertices(self, pos): + """Get the vertices of the cell a point is in + + Parameters + ---------- + pos : np.array + Nx3 array of xyz locations + + Returns + ------- + np.array((N,3),dtype=float), np.array(N,dtype=int) + vertices, inside + """ + gi, inside = self.position_to_cell_corners(pos) + + node_indexes = self.global_index_to_node_index(gi.flatten()) + return self.node_indexes_to_position(node_indexes), inside + + def evaluate_value(self, evaluation_points: np.ndarray, property_array: np.ndarray): + """ + Evaluate the value of of the property at the locations. + Trilinear interpolation dot corner values + + Parameters + ---------- + evaluation_points np array of locations + property_name string of property name + + Returns + ------- + + """ + idc, inside = self.position_to_cell_corners(evaluation_points) + v = np.zeros(idc.shape) + v[:, :] = np.nan + + v[inside, :] = self.position_to_dof_coefs(evaluation_points[inside, :]) + v[inside, :] *= property_array[idc[inside, :]] + return np.sum(v, axis=1) diff --git a/packages/loop_common/src/loop_common/supports/_2d_structured_grid.py b/packages/loop_common/src/loop_common/supports/_2d_structured_grid.py index 974821b6d..85e977ef9 100644 --- a/packages/loop_common/src/loop_common/supports/_2d_structured_grid.py +++ b/packages/loop_common/src/loop_common/supports/_2d_structured_grid.py @@ -7,6 +7,8 @@ import numpy as np +from loop_common.geometry import StructuredGrid2DGeometry + from ..math.finite_difference_stencil import Operator from . import SupportType from ._base_support import BaseSupport @@ -14,8 +16,14 @@ logger = logging.getLogger(__name__) -class StructuredGrid2D(BaseSupport): - """ """ +class StructuredGrid2D(StructuredGrid2DGeometry, BaseSupport): + """A 2D structured grid used as an interpolation support. + + The grid geometry (node and cell positions, point location, index conversions, bilinear + weights and node-value interpolation) comes from + :class:`loop_common.geometry.StructuredGrid2DGeometry`; this class adds the interpolation + machinery. ``nsteps`` is the number of **nodes** in each direction. + """ dimension = 2 @@ -33,57 +41,15 @@ def __init__( nsteps - 2d list or numpy array of ints step_vector - 2d list or numpy array of int """ - if origin is None: - origin = np.zeros(2) - if nsteps is None: - nsteps = np.array([10, 10]) - if step_vector is None: - step_vector = np.ones(2) - + StructuredGrid2DGeometry.__init__(self, origin=origin, nsteps=nsteps, step_vector=step_vector) self.type = SupportType.StructuredGrid2D - self.nsteps = np.ceil(np.array(nsteps)).astype(int) - self.step_vector = np.array(step_vector) - self.origin = np.array(origin) - self.maximum = origin + self.nsteps * self.step_vector - - self.dim = 2 - self.nsteps_cells = self.nsteps - 1 - self.n_cell_x = self.nsteps[0] - 1 - self.n_cell_y = self.nsteps[1] - 1 self.properties = {} - - # calculate the node positions using numpy (this should probably not - # be stored as it defeats - # the purpose of a structured grid - - # self.barycentre = self.cell_centres(np.arange(self.n_elements)) - self.regions = {} self.regions["everywhere"] = np.ones(self.n_nodes).astype(bool) - @property - def nodes(self): - max = self.origin + self.nsteps_cells * self.step_vector - x = np.linspace(self.origin[0], max[0], self.nsteps[0]) - y = np.linspace(self.origin[1], max[1], self.nsteps[1]) - xx, yy = np.meshgrid(x, y, indexing="ij") - return np.array([xx.flatten(order="F"), yy.flatten(order="F")]).T - - @property - def n_nodes(self): - return self.nsteps[0] * self.nsteps[1] - def set_nelements(self, nelements) -> int: raise NotImplementedError("Cannot set number of elements for 2D structured grid") - @property - def n_elements(self): - return self.nsteps_cells[0] * self.nsteps_cells[1] - - @property - def element_size(self): - return np.prod(self.step_vector) - @property def barycentre(self): return self.cell_centres(np.arange(self.n_elements)) @@ -91,339 +57,6 @@ def barycentre(self): # @property # def barycentre(self): # return self.cell_centres(np.arange(self.n_elements)) - @property - def elements(self) -> np.ndarray: - global_index = np.arange(self.n_elements) - cell_indexes = self.global_index_to_cell_index(global_index) - - return self.global_node_indices(self.cell_corner_indexes(cell_indexes)) - - def print_geometry(self): - print(f"Origin: {self.origin[0]:f} {self.origin[1]:f} {self.origin[2]:f}") - print( - f"Cell size: {self.step_vector[0]:f} {self.step_vector[1]:f} {self.step_vector[2]:f}" - ) - max = self.origin + self.nsteps_cells * self.step_vector - print(f"Max extent: {max[0]:f} {max[1]:f} {max[2]:f}") - - def cell_centres(self, global_index: np.ndarray) -> np.ndarray: - """[summary] - - [extended_summary] - - Parameters - ---------- - global_index : [type] - [description] - - Returns - ------- - [type] - [description] - """ - cell_indexes = self.global_index_to_cell_index(global_index) - cell_centres = np.zeros((cell_indexes.shape[0], 2)) - - cell_centres[:, 0] = ( - self.origin[None, 0] - + self.step_vector[None, 0] * 0.5 - + self.step_vector[None, 0] * cell_indexes[:, 0] - ) - cell_centres[:, 1] = ( - self.origin[None, 1] - + self.step_vector[None, 1] * 0.5 - + self.step_vector[None, 1] * cell_indexes[:, 1] - ) - return cell_centres - - def position_to_cell_index(self, pos: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """[summary] - - [extended_summary] - - Parameters - ---------- - pos : [type] - [description] - - Returns - ------- - [type] - [description] - """ - inside = self.inside(pos) - cell_indexes = np.zeros((pos.shape[0], 2)) - cell_indexes[:, 0] = pos[:, 0] - self.origin[None, 0] - cell_indexes[:, 1] = pos[:, 1] - self.origin[None, 1] - cell_indexes /= self.step_vector[None, :] - return cell_indexes.astype(int), inside - - def inside(self, pos: np.ndarray) -> np.ndarray: - # check whether point is inside box - inside = np.ones(pos.shape[0]).astype(bool) - for i in range(self.dim): - inside *= pos[:, i] > self.origin[None, i] - inside *= ( - pos[:, i] - < self.origin[None, i] + self.step_vector[None, i] * self.nsteps_cells[None, i] - ) - return inside - - def check_position(self, pos: np.ndarray) -> np.ndarray: - """[summary] - - [extended_summary] - - Parameters - ---------- - pos : [type] - [description] - - Returns - ------- - [type] - [description] - """ - - if len(pos.shape) == 1: - pos = np.array([pos]) - if len(pos.shape) != 2: - raise ValueError("Position array needs to be a list of points or a point") - - return pos - - def bilinear(self, local_coords: np.ndarray) -> np.ndarray: - """ - returns the bilinear interpolation for the local coordinates - Parameters - ---------- - x - double, array of doubles - y - double, array of doubles - z - double, array of doubles - - Returns - ------- - array of interpolation coefficients - - """ - - return np.array( - [ - (1 - local_coords[:, 0]) * (1 - local_coords[:, 1]), - local_coords[:, 0] * (1 - local_coords[:, 1]), - (1 - local_coords[:, 0]) * local_coords[:, 1], - local_coords[:, 0] * local_coords[:, 1], - ] - ).T - - def position_to_local_coordinates(self, pos: np.ndarray) -> np.ndarray: - """ - Convert from global to local coordinates within a cel - Parameters - ---------- - pos - array of positions inside - - Returns - ------- - localx, localy, localz - - """ - # TODO check if inside mesh - - # calculate local coordinates for positions - local_coords = np.zeros(pos.shape) - local_coords[:, 0] = ( - (pos[:, 0] - self.origin[None, 0]) % self.step_vector[None, 0] - ) / self.step_vector[None, 0] - local_coords[:, 1] = ( - (pos[:, 1] - self.origin[None, 1]) % self.step_vector[None, 1] - ) / self.step_vector[None, 1] - - return local_coords - - def position_to_dof_coefs(self, pos: np.ndarray): - """ - global posotion to interpolation coefficients - Parameters - ---------- - pos - - Returns - ------- - - """ - local_coords = self.position_to_local_coordinates(pos) - weights = self.bilinear(local_coords) - return weights - - def neighbour_global_indexes(self, mask=None, **kwargs): - """ - Get neighbour indexes - - Parameters - ---------- - kwargs - indexes array specifying the cells to return neighbours - - Returns - ------- - - """ - indexes = None - if "indexes" in kwargs: - indexes = kwargs["indexes"] - if "indexes" not in kwargs: - gi = np.arange(self.n_nodes) - indexes = self.global_index_to_node_index(gi) - edge_mask = ( - (indexes[:, 0] > 0) - & (indexes[:, 0] < self.nsteps[0] - 1) - & (indexes[:, 1] > 0) - & (indexes[:, 1] < self.nsteps[1] - 1) - ) - indexes = indexes[edge_mask, :].T - # ii = [] - # jj = [] - # for i in range(1, self.nsteps[0] - 1): - # for j in range(1, self.nsteps[1] - 1): - # ii.append(i) - # jj.append(j) - # indexes = np.array([ii, jj]) - # indexes = np.array(indexes).T - if indexes.ndim != 2: - print(indexes.ndim) - return - # determine which neighbours to return default is diagonals included. - if mask is None: - mask = np.array([[-1, 0, 1, -1, 0, 1, -1, 0, 1], [1, 1, 1, 0, 0, 0, -1, -1, -1]]) - neighbours = indexes[:, None, :] + mask[:, :, None] - return (neighbours[0, :, :] + self.nsteps[0, None, None] * neighbours[1, :, :]).astype( - np.int64 - ) - - def cell_corner_indexes(self, cell_indexes: np.ndarray) -> np.ndarray: - """ - Returns the indexes of the corners of a cell given its location xi, - yi, zi - - Parameters - ---------- - x_cell_index - y_cell_index - z_cell_index - - Returns - ------- - - """ - corner_indexes = np.zeros((cell_indexes.shape[0], 4, 2), dtype=np.int64) - xcorner = np.array([0, 1, 0, 1]) - ycorner = np.array([0, 0, 1, 1]) - corner_indexes[:, :, 0] = ( - cell_indexes[:, None, 0] + corner_indexes[:, :, 0] + xcorner[None, :] - ) - corner_indexes[:, :, 1] = ( - cell_indexes[:, None, 1] + corner_indexes[:, :, 1] + ycorner[None, :] - ) - return corner_indexes - - def global_index_to_cell_index(self, global_index): - """ - Convert from global indexes to xi,yi,zi - - Parameters - ---------- - global_index - - Returns - ------- - - """ - # determine the ijk indices for the global index. - # remainder when dividing by nx = i - # remained when dividing modulus of nx by ny is j - cell_indexes = np.zeros((global_index.shape[0], 2), 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] - return cell_indexes - - def global_index_to_node_index(self, global_index): - cell_indexes = np.zeros((global_index.shape[0], 2), 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] - return cell_indexes - - def _global_indices(self, indexes: np.ndarray, nsteps: np.ndarray) -> np.ndarray: - if len(indexes.shape) == 1: - raise ValueError("Indexes must be a 2D array") - if indexes.shape[-1] != 2: - raise ValueError("Last dimension of cell indexing needs to be ijk indexing") - original_shape = indexes.shape - indexes = indexes.reshape(-1, 2) - gi = indexes[:, 0] + nsteps[0] * indexes[:, 1] - return gi.reshape(original_shape[:-1]) - - def global_cell_indices(self, indexes: np.ndarray) -> np.ndarray: - return self._global_indices(indexes, self.nsteps_cells) - - def global_node_indices(self, indexes: np.ndarray) -> np.ndarray: - return self._global_indices(indexes, self.nsteps) - - def node_indexes_to_position(self, node_indexes: np.ndarray) -> np.ndarray: - - original_shape = node_indexes.shape - node_indexes = node_indexes.reshape((-1, 2)) - xy = np.zeros((node_indexes.shape[0], 2), dtype=float) - xy[:, 0] = self.origin[0] + self.step_vector[0] * node_indexes[:, 0] - xy[:, 1] = self.origin[1] + self.step_vector[1] * node_indexes[:, 1] - xy = xy.reshape(original_shape) - return xy - - def position_to_cell_corners(self, pos): - """Get the global indices of the vertices (corner) nodes of the cell containing each point. - - Parameters - ---------- - pos : np.array - (N, 2) array of xy coordinates representing the positions of N points. - - Returns - ------- - globalidx : np.array - (N, 4) array of global indices corresponding to the 4 corner nodes of the cell - each point lies in. If a point lies outside the support, its corresponding entry - will be set to -1. - inside : np.array - (N,) boolean array indicating whether each point is inside the support domain. - """ - corner_index, inside = self.position_to_cell_index(pos) - corners = self.cell_corner_indexes(corner_index) - globalidx = self.global_node_indices(corners) - # if global index is not inside the support set to -1 - globalidx[~inside] = -1 - return globalidx, inside - - def evaluate_value(self, evaluation_points: np.ndarray, property_array: np.ndarray): - """ - Evaluate the value of of the property at the locations. - Trilinear interpolation dot corner values - - Parameters - ---------- - evaluation_points np array of locations - property_name string of property name - - Returns - ------- - - """ - idc, inside = self.position_to_cell_corners(evaluation_points) - v = np.zeros(idc.shape) - v[:, :] = np.nan - - v[inside, :] = self.position_to_dof_coefs(evaluation_points[inside, :]) - v[inside, :] *= property_array[idc[inside, :]] - return np.sum(v, axis=1) - def evaluate_gradient(self, evaluation_points, property_array): T = np.zeros((evaluation_points.shape[0], 2, 4)) _vertices, T, elements, inside = self.get_element_gradient_for_location(evaluation_points) @@ -478,24 +111,6 @@ def get_element_for_location( a = self.position_to_dof_coefs(pos) return vertices, a, elements, inside - def position_to_cell_vertices(self, pos): - """Get the vertices of the cell a point is in - - Parameters - ---------- - pos : np.array - Nx3 array of xyz locations - - Returns - ------- - np.array((N,3),dtype=float), np.array(N,dtype=int) - vertices, inside - """ - gi, inside = self.position_to_cell_corners(pos) - - node_indexes = self.global_index_to_node_index(gi.flatten()) - return self.node_indexes_to_position(node_indexes), inside - def onGeometryChange(self): pass diff --git a/packages/loop_common/tests/data/structured_grid_2d_reference.npz b/packages/loop_common/tests/data/structured_grid_2d_reference.npz new file mode 100644 index 000000000..1a272d143 Binary files /dev/null and b/packages/loop_common/tests/data/structured_grid_2d_reference.npz differ diff --git a/packages/loop_common/tests/test_structured_grid_2d_geometry.py b/packages/loop_common/tests/test_structured_grid_2d_geometry.py new file mode 100644 index 000000000..28b33f398 --- /dev/null +++ b/packages/loop_common/tests/test_structured_grid_2d_geometry.py @@ -0,0 +1,74 @@ +"""StructuredGrid2DGeometry: the geometry and node-value interpolation the 2D support builds on.""" +import logging + +import numpy as np +import pytest +from loop_common.geometry import StructuredGrid2DGeometry +from loop_common.supports import StructuredGrid2D + +GRIDS = { + "unit": dict(origin=[0.0, 0.0], nsteps=[5, 5], step_vector=[1.0, 1.0]), + "aniso": dict(origin=[10.0, -3.0], nsteps=[7, 4], step_vector=[0.5, 2.0]), +} + + +@pytest.fixture(params=list(GRIDS)) +def geometry(request): + return StructuredGrid2DGeometry(**GRIDS[request.param]) + + +def test_the_support_is_a_geometry(): + assert issubclass(StructuredGrid2D, StructuredGrid2DGeometry) + + +def test_elements_are_flat_node_indexes(geometry): + assert geometry.elements.shape == (geometry.n_elements, 4) + nodes = geometry.nodes + # every element's four corners are the corners of one cell + corners = nodes[geometry.elements] + spans = corners.max(axis=1) - corners.min(axis=1) + np.testing.assert_allclose(spans, np.broadcast_to(geometry.step_vector, spans.shape)) + + +def test_index_round_trip(geometry): + cells = np.arange(geometry.n_elements) + ij = geometry.global_index_to_cell_index(cells) + np.testing.assert_array_equal(geometry.global_cell_indices(ij), cells) + nodes = np.arange(geometry.n_nodes) + ij = geometry.global_index_to_node_index(nodes) + np.testing.assert_array_equal(geometry.global_node_indices(ij), nodes) + + +def test_index_helpers_reject_bad_shapes(geometry): + with pytest.raises(ValueError): + geometry.global_cell_indices(np.array([0, 0])) + with pytest.raises(ValueError): + geometry.global_cell_indices(np.zeros((3, 3), dtype=int)) + + +def test_bilinear_reproduces_a_linear_field(geometry): + coefficients = np.array([1.5, -2.0]) + values = geometry.nodes @ coefficients + 3.0 + low = geometry.origin + high = geometry.origin + (geometry.nsteps - 1) * geometry.step_vector + points = np.random.default_rng(0).uniform(low + 1e-6, high - 1e-6, (200, 2)) + np.testing.assert_allclose( + geometry.evaluate_value(points, values), points @ coefficients + 3.0, atol=1e-9 + ) + + +def test_bilinear_returns_node_values_at_interior_nodes(geometry): + values = np.random.default_rng(1).normal(size=geometry.n_nodes) + nodes = geometry.nodes + interior = ( + (nodes > geometry.origin + 1e-9).all(axis=1) + & (nodes < geometry.origin + (geometry.nsteps - 1) * geometry.step_vector - 1e-9).all(axis=1) + ) + np.testing.assert_allclose( + geometry.evaluate_value(nodes[interior], values)[...], values[interior], atol=1e-12 + ) + + +def test_print_geometry_works_on_a_2d_grid(geometry, caplog): + with caplog.at_level(logging.INFO): + geometry.print_geometry() diff --git a/packages/loop_common/tests/test_structured_grid_2d_reference.py b/packages/loop_common/tests/test_structured_grid_2d_reference.py new file mode 100644 index 000000000..993502c14 --- /dev/null +++ b/packages/loop_common/tests/test_structured_grid_2d_reference.py @@ -0,0 +1,182 @@ +"""Pin the behaviour of the 2D structured interpolation support. + +Records a wide set of ``StructuredGrid2D`` outputs (geometry, node/element arrays, point location +including exact boundary points, bilinear weights, value/gradient evaluation, element lookup, +regularisation operators, index conversions, ``set_nelements``) on a few grids and compares them +with ``data/structured_grid_2d_reference.npz``. + +The reference was generated from the code *before* ``StructuredGrid2D`` was changed to extend +``StructuredGrid2DGeometry``. Regenerate after an intended change, and review the diff: + + python packages/loop_common/tests/test_structured_grid_2d_reference.py --update +""" +import sys +import warnings +from pathlib import Path + +import numpy as np +import pytest +from loop_common.supports import StructuredGrid2D + +REFERENCE = Path(__file__).parent / "data" / "structured_grid_2d_reference.npz" + +# (origin, nodes per axis, step vector). StructuredGrid2D takes node counts. +GRIDS = { + "unit": ([0.0, 0.0], [5, 5], [1.0, 1.0]), + "aniso": ([10.0, -3.0], [7, 4], [0.5, 2.0]), + "minimal": ([0.0, 0.0], [2, 2], [3.0, 1.0]), +} +ATTRIBUTES = ( + "nsteps", + "nsteps_cells", + "maximum", + "origin", + "step_vector", + "n_nodes", + "n_elements", + "element_size", + "n_cell_x", + "n_cell_y", + "dim", + "dimension", + "nodes", + "elements", + "barycentre", +) +OPERATOR_WEIGHTS = {"dxx": 1.0, "dyy": 1.0, "dxy": 1.0} + + +def _flatten(value, prefix, out): + if isinstance(value, dict): + for key in sorted(value): + _flatten(value[key], f"{prefix}.{key}", out) + elif isinstance(value, (tuple, list)) and not all(np.isscalar(v) for v in value): + for i, item in enumerate(value): + _flatten(item, f"{prefix}.{i}", out) + else: + out[prefix] = np.asarray(value) + + +def _record(out, key, call): + try: + _flatten(call(), key, out) + except Exception as exc: # noqa: BLE001 - the exception type is part of the behaviour + out[key] = np.asarray(f"EXC:{type(exc).__name__}") + + +def collect() -> dict[str, np.ndarray]: + out: dict[str, np.ndarray] = {} + rng = np.random.default_rng(5) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + for name, (origin, nodes, step) in GRIDS.items(): + grid = StructuredGrid2D(origin=origin, nsteps=nodes, step_vector=step) + _probe(out, name, grid, origin, nodes, step, rng) + return out + + +def _probe(out, prefix, grid, origin, nodes, step, rng): + for attribute in ATTRIBUTES: + _record(out, f"{prefix}/{attribute}", lambda attribute=attribute: getattr(grid, attribute)) + _record(out, f"{prefix}/str_type", lambda: np.asarray(int(grid.type))) + + low = np.array(origin) - 1.0 + high = np.array(origin) + np.array(nodes) * np.array(step) + 1.0 + anywhere = rng.uniform(low, high, (120, 2)) + # points exactly on the boundary of the nodes' extent and of `maximum` + node_max = np.array(origin) + (np.array(nodes) - 1) * np.array(step) + corners = np.array(np.meshgrid(*zip(origin, node_max), indexing="ij")).reshape(2, -1).T + faces = [] + centre = 0.5 * (np.array(origin) + node_max) + for axis in range(2): + for bound in (origin[axis], node_max[axis]): + point = centre.copy() + point[axis] = bound + faces.append(point) + anywhere = np.vstack([anywhere, corners, np.array(faces)]) + interior = rng.uniform(np.array(origin) + 1e-3, node_max - 1e-3, (80, 2)) + + _record(out, f"{prefix}/inside", lambda: grid.inside(anywhere)) + _record(out, f"{prefix}/position_to_cell_index", lambda: grid.position_to_cell_index(anywhere)) + _record(out, f"{prefix}/check_position_1d", lambda: grid.check_position(np.array([1.0, 2.0]))) + _record(out, f"{prefix}/check_position_3d", lambda: grid.check_position(np.zeros((2, 2, 2)))) + _record(out, f"{prefix}/position_to_cell_corners", lambda: grid.position_to_cell_corners(anywhere)) + _record(out, f"{prefix}/position_to_cell_vertices", lambda: grid.position_to_cell_vertices(interior)) + _record(out, f"{prefix}/position_to_local_coordinates", lambda: grid.position_to_local_coordinates(interior)) + _record(out, f"{prefix}/position_to_dof_coefs", lambda: grid.position_to_dof_coefs(interior)) + _record(out, f"{prefix}/bilinear", lambda: grid.bilinear(rng.uniform(0, 1, (10, 2)))) + + cell_index = np.arange(grid.n_elements) + node_index = np.arange(grid.n_nodes) + _record(out, f"{prefix}/cell_centres", lambda: grid.cell_centres(cell_index)) + _record(out, f"{prefix}/global_index_to_cell_index", lambda: grid.global_index_to_cell_index(cell_index)) + node_ijk = grid.global_index_to_node_index(node_index) + _record(out, f"{prefix}/global_index_to_node_index", lambda: node_ijk) + cell_ijk = grid.global_index_to_cell_index(cell_index) + _record(out, f"{prefix}/cell_corner_indexes", lambda: grid.cell_corner_indexes(cell_ijk)) + _record(out, f"{prefix}/global_node_indices", lambda: grid.global_node_indices(node_ijk)) + _record(out, f"{prefix}/global_cell_indices", lambda: grid.global_cell_indices(cell_ijk)) + _record(out, f"{prefix}/node_indexes_to_position", lambda: grid.node_indexes_to_position(node_ijk[:10])) + _record(out, f"{prefix}/neighbour_global_indexes", lambda: grid.neighbour_global_indexes()) + + values = rng.normal(size=grid.n_nodes) + _record(out, f"{prefix}/evaluate_value", lambda: grid.evaluate_value(anywhere, values)) + _record(out, f"{prefix}/evaluate_gradient", lambda: grid.evaluate_gradient(interior, values)) + _record(out, f"{prefix}/get_element_for_location", lambda: grid.get_element_for_location(interior)) + _record( + out, + f"{prefix}/get_element_gradient_for_location", + lambda: grid.get_element_gradient_for_location(interior), + ) + _record( + out, + f"{prefix}/get_operators", + lambda: dict(grid.get_operators(OPERATOR_WEIGHTS)), + ) + _record(out, f"{prefix}/vtk_n_cells", lambda: np.array([grid.vtk(0.0).n_cells])) + _record(out, f"{prefix}/set_nelements", lambda: grid.set_nelements(100)) + + +def _assert_same(key, got, want): + if want.dtype.kind in "US": + assert got.dtype.kind in "US" and str(got) == str(want), f"{key}: {got} != {want}" + return + assert got.dtype.kind not in "US", f"{key}: now raises {got}" + assert got.shape == want.shape, f"{key}: shape {got.shape} != {want.shape}" + assert got.dtype.kind == want.dtype.kind, f"{key}: dtype {got.dtype} != {want.dtype}" + if want.dtype.kind == "f": + np.testing.assert_allclose(got, want, rtol=1e-9, atol=1e-12, err_msg=key, equal_nan=True) + else: + np.testing.assert_array_equal(got, want, err_msg=key) + + +@pytest.fixture(scope="module") +def current(): + return collect() + + +@pytest.fixture(scope="module") +def reference(): + with np.load(REFERENCE) as data: + return {key: data[key] for key in data.files} + + +def test_same_probes_are_recorded(current, reference): + assert sorted(current) == sorted(reference) + + +@pytest.mark.parametrize("grid_name", list(GRIDS)) +def test_structured_grid_2d_behaviour_is_unchanged(grid_name, current, reference): + keys = [k for k in reference if k.startswith(f"{grid_name}/")] + assert keys + for key in keys: + _assert_same(key, current[key], reference[key]) + + +if __name__ == "__main__": + if sys.argv[1:] != ["--update"]: + raise SystemExit(__doc__) + REFERENCE.parent.mkdir(exist_ok=True) + result = collect() + np.savez_compressed(REFERENCE, **result) + print(f"wrote {len(result)} arrays to {REFERENCE}")