Merge branch 'add_grid_cells_connectivity' into 'main'

Grid: add cells_connectivity and cells_definition

See merge request !275

CHANGELOG.md

@@ -12,6 +12,14 @@
 ### Features
 
 * Components and adapters automatically provide default metadata that can be extended by implementations (!274, !276)
+* Grid class now have attributes providing connectivity information for the contained cells (!275)
+  * `cells_connectivity`: connectivity array as used by ESMF and VTK
+  * `cells_definition`: cell definition as used by PyVista and legacy VTK
+  * `cells_offset`: location of the start of each cell in `cells_connectivity`
+* added convenience functions and constants to `grid_tools`  (!275)
+  * `get_cells_matrix`: convert `cells_connectivity` or `cells_definition` back to the default cells matrix used in the Grid class (can be used to convert VTK-grids into FINAM-grids)
+  * `INV_VTK_TYPE_MAP`: inverse mapping to `VTK_TYPE_MAP` - FINAM cell type to VTK cell type
+  * `VTK_CELL_DIM`: parametric dimension for each VTK cell type
 
 ### Documentation
 
@@ -33,7 +41,7 @@
 * Automatic conversion between compatible grids (!255)
 * Adds methods `to_canonical`, `from_canonical` and `get_transform_to` to grid classes (!255)
 * Adds support for masked grids using `numpy.ma.MaskedArray` (!258, !260)
-* Adds convenience functions for dealing with masked arrays in `data.tools` (!260):  
+* Adds convenience functions for dealing with masked arrays in `data.tools` (!260):
   `is_masked_array`, `has_masked_values`, `filled`, `to_compressed`, `from_compressed`, `check_data_covers_domain`
 
 ### Documentation
@@ -56,7 +64,7 @@
 * FINAM uses a new scheduling algorithm that allows components to use future data instead of only past/current (!157, !159)
 * New adapters to resolve circular coupling through the use of delayed data (!187)
 * It is now possible to set up static couplings that run only once and have no explicit time or stepping (!166)
-* FINAM can handle different starting times of components by pushing initial data twice (!206):  
+* FINAM can handle different starting times of components by pushing initial data twice (!206):
   Once for the common starting time, and once for the actual component time
 * Components are no longer required to push all outputs on every step (!208)
 

src/finam/data/grid_base.py

@@ -74,13 +74,36 @@ class Grid(GridBase):
     @property
     @abstractmethod
     def cells(self):
-        """np.ndarray: Cell nodes in ESMF format."""
+        """np.ndarray: Cell nodes as 2D array."""
 
     @property
     @abstractmethod
     def cell_types(self):
         """np.ndarray: Cell types."""
 
+    @property
+    def cells_connectivity(self):
+        """np.ndarray: Cells connectivity in ESMF format (list of node IDs)."""
+        return flatten_cells(self.cells)
+
+    @property
+    def cells_definition(self):
+        """np.ndarray: Cell definition in VTK format (list of number of nodes with node IDs)."""
+        return flatten_cells(
+            np.squeeze(
+                np.hstack(
+                    (np.atleast_2d(self.cell_node_counts).T, np.atleast_2d(self.cells))
+                )
+            )
+        )
+
+    @property
+    def cells_offset(self):
+        """np.ndarray: The location of the start of each cell in cells_connectivity."""
+        return np.concatenate(
+            (np.array([0], dtype=int), np.cumsum(self.cell_node_counts))
+        )
+
     @property
     def cell_centers(self):
         """np.ndarray: Grid cell centers."""
@@ -215,8 +238,9 @@ class Grid(GridBase):
             x = np.ascontiguousarray(points[:, 0])
             y = np.ascontiguousarray(points[:, 1] if self.dim > 1 else np.zeros_like(x))
             z = np.ascontiguousarray(points[:, 2] if self.dim > 2 else np.zeros_like(x))
-            con = flatten_cells(self.cells)
-            off = np.cumsum(NODE_COUNT[self.cell_types])
+            con = self.cells_connectivity
+            # pyevtk only needs the ends of the cell definition
+            off = self.cells_offset[1:]
             typ = VTK_TYPE_MAP[self.cell_types]
             unstructuredGridToVTK(path, x, y, z, con, off, typ, **kw)
         else:

src/finam/data/grid_tools.py

@@ -399,6 +399,61 @@ def flatten_cells(cells):
     return np.ma.masked_values(cells, -1).compressed()
 
 
+def get_cells_matrix(cell_types, cells, connectivity=False):
+    """
+    Create the cells matrix as used in the Grid class.
+
+    Parameters
+    ----------
+    cell_types : np.ndarray
+        Cell types.
+
+    cells : np.ndarray
+        Either cell definitions given as list of number of nodes with node IDs:
+        ``[n0, p0_0, p0_1, ..., p0_n, n1, p1_0, p1_1, ..., p1_n, ...]``
+
+        Or cell connectivity given as list of node IDs:
+        ``[p0_0, p0_1, ..., p0_n, p1_0, p1_1, ..., p1_n, ...]``
+
+    connectivity : bool, optional
+        Indicate that cells are given by connectivity. Default: False
+
+    Returns
+    -------
+    np.ndarray
+        Cell nodes as 2D array.
+    """
+    unique_cell_types = np.unique(cell_types)
+    max_nodes = np.max(NODE_COUNT[unique_cell_types])
+    cells_matrix = np.full((len(cell_types), max_nodes), fill_value=-1, dtype=int)
+    cell_sizes = NODE_COUNT[cell_types]
+
+    if connectivity:
+        cell_ends = np.cumsum(cell_sizes)
+        cell_starts = np.concatenate(
+            [np.array([0], dtype=cell_ends.dtype), cell_ends[:-1]]
+        )
+    else:
+        # adding one to skip cell size entry in cell definitions array
+        # [(n0), p0_0, p0_1, ..., p0_n, (n1), p1_0, p1_1, ..., p1_n, ...]
+        cell_ends = np.cumsum(cell_sizes + 1)
+        cell_starts = (
+            np.concatenate([np.array([0], dtype=cell_ends.dtype), cell_ends[:-1]]) + 1
+        )
+
+    for cell_type in unique_cell_types:
+        cell_size = NODE_COUNT[cell_type]
+        mask = cell_types == cell_type
+        current_cell_starts = cell_starts[mask]
+
+        cells_inds = current_cell_starts[..., np.newaxis] + np.arange(cell_size)[
+            np.newaxis
+        ].astype(cell_starts.dtype)
+        cells_matrix[:, :cell_size][mask] = cells[cells_inds]
+
+    return cells_matrix
+
+
 class Location(Enum):
     """Data location in the grid."""
 
@@ -433,9 +488,37 @@ CELL_DIM = np.array([0, 1, 2, 2, 3, 3], dtype=int)
 """np.ndarray: Cell dimension per CellType."""
 
 
-VTK_TYPE_MAP = np.array([1, 3, 5, 9, 10, 12], dtype=int)
+ESMF_TYPE_MAP = np.array([-1, -1, 3, 4, 10, 12], dtype=int)
+"""np.ndarray: ESMF type per CellType (-1 for unsupported)."""
+
+
+VTK_TYPE_MAP = np.array((vtk_t := [1, 3, 5, 9, 10, 12]), dtype=int)
 """np.ndarray: VTK type per CellType."""
 
 
-ESMF_TYPE_MAP = np.array([-1, -1, 3, 4, 10, 12], dtype=int)
-"""np.ndarray: ESMF type per CellType (-1 for unsupported)."""
+INV_VTK_TYPE_MAP = np.array(
+    [vtk_t.index(i) if i in vtk_t else -1 for i in range(82)],
+    dtype=int,
+)
+"""np.ndarray: CellType per VTK type."""
+
+
+# VTK v9.3
+VTK_CELL_DIM = np.array(
+    # Linear cells (0-16)
+    [0, 0, 0, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3] + 4 * [-1]
+    # Quadratic / Cubic, isoparametric cells (21-37)
+    + [1, 2, 2, 3, 3, 3, 3, 2, 3, 2, 3, 3, 3, 2, 1, 2, 3] + 3 * [-1]
+    # Special class and Polyhedron cell (41-42)
+    + [1, 3] + 8 * [-1]
+    # Higher order cells in parametric form (51-56)
+    + [1, 2, 2, 2, 3, 3] + 3 * [-1]
+    # Higher order cells (60-67)
+    + [1, 2, 2, 2, 3, 3, 3, 3]
+    # Arbitrary order Lagrange elements (68-74)
+    + [1, 2, 2, 3, 3, 3, 3]
+    # Arbitrary order Bezier elements (75-81)
+    + [1, 2, 2, 3, 3, 3, 3],
+    dtype=int,
+)
+"""np.ndarray: Cell dimension per VTK type."""

tests/data/test_grid_tools.py

@@ -5,6 +5,10 @@ from numpy.testing import assert_allclose, assert_array_equal
 
 from finam import CellType, EsriGrid, NoGrid, UniformGrid, UnstructuredGrid
 from finam.data.grid_tools import (
+    CELL_DIM,
+    INV_VTK_TYPE_MAP,
+    VTK_CELL_DIM,
+    VTK_TYPE_MAP,
     check_axes_monotonicity,
     check_axes_uniformity,
     flatten_cells,
@@ -12,6 +16,7 @@ from finam.data.grid_tools import (
     gen_cells,
     gen_node_centers,
     gen_points,
+    get_cells_matrix,
     order_map,
     point_order,
 )
@@ -189,3 +194,64 @@ class TestGridTools(unittest.TestCase):
         # layer height info results in 3D grid
         grid2 = UniformGrid((4, 3, 1))
         self.assertFalse(grid1.compatible_with(grid2))
+
+    def test_cells_definition(self):
+        # number of nodes with node IDs
+        c_def = np.array(
+            [3, 0, 6, 5]
+            + [3, 0, 1, 6]
+            + [4, 1, 2, 7, 6]
+            + [3, 2, 8, 7]
+            + [3, 2, 3, 8]
+            + [4, 3, 4, 9, 8]
+        )
+        # node IDs
+        c_con = np.array(
+            [0, 6, 5] + [0, 1, 6] + [1, 2, 7, 6] + [2, 8, 7] + [2, 3, 8] + [3, 4, 9, 8]
+        )
+        # cell offsets: start of cell definition in connectivity plus size of connectivity
+        c_off = np.array([0, 3, 6, 10, 13, 16, 20])
+        # node IDs as matrix
+        c_mat = np.array(
+            [
+                [0, 6, 5, -1],
+                [0, 1, 6, -1],
+                [1, 2, 7, 6],
+                [2, 8, 7, -1],
+                [2, 3, 8, -1],
+                [3, 4, 9, 8],
+            ]
+        )
+        cell_types = np.array([2, 2, 3, 2, 2, 3])
+        pnts = np.array(
+            [
+                [0.0, 0.0],
+                [1.0, 0.0],
+                [2.0, 0.0],
+                [3.0, 0.0],
+                [4.0, 0.0],
+                [0.0, 1.0],
+                [1.0, 1.0],
+                [2.0, 1.0],
+                [3.0, 1.0],
+                [4.0, 1.0],
+            ]
+        )
+        grid = UnstructuredGrid(points=pnts, cells=c_mat, cell_types=cell_types)
+
+        assert_array_equal(grid.cells_connectivity, c_con)
+        assert_array_equal(grid.cells_definition, c_def)
+        assert_array_equal(grid.cells_offset, c_off)
+        self.assertEqual(grid.cells_offset[0], 0)
+        self.assertEqual(grid.cells_offset[-1], len(c_con))
+        assert_array_equal(get_cells_matrix(grid.cell_types, c_def), c_mat)
+        assert_array_equal(
+            get_cells_matrix(grid.cell_types, c_con, connectivity=True), c_mat
+        )
+
+    def test_vtk_maps(self):
+        test_dims = VTK_CELL_DIM[VTK_TYPE_MAP]
+        assert_array_equal(test_dims, CELL_DIM)
+
+        test_type = INV_VTK_TYPE_MAP[VTK_TYPE_MAP]
+        assert_array_equal(test_type, range(6))