Merge pull request #1311 from UXARRAY/construct-node-edge-conn

Construct `node_edge_connectivity`

Author: rajeeja

test/meshfiles/ugrid/quad-hexagon/grid.nc

@@ -0,0 +1,197 @@
+import pytest
+from pathlib import Path
+import uxarray as ux
+import numpy as np
+
+from uxarray.constants import INT_DTYPE, INT_FILL_VALUE as fv
+
+
+class TestQuadHexagon:
+    """Test suite for the quad-hexagon mesh connectivity construction"""
+
+    @pytest.fixture()
+    def uxgrid(self):
+        grid_path = (
+            Path(__file__).parent / "meshfiles" / "ugrid" / "quad-hexagon" / "grid.nc"
+        )
+        return ux.open_grid(grid_path)
+
+    # ─── Node Connectivity ────────────────────────────────────────────────────────────────────────────────────────────
+
+    def test_node_edge_connectivity(self, uxgrid):
+        node_edges = uxgrid.node_edge_connectivity.values
+        expected_node_edges = np.array(
+            [
+                [0, 1, 2],
+                [0, 3, 4],
+                [3, 5, fv],
+                [5, 6, fv],
+                [6, 7, 8],
+                [1, 7, 9],
+                [10, 11, fv],
+                [10, 12, fv],
+                [13, 14, fv],
+                [13, 15, fv],
+                [16, 17, fv],
+                [16, 18, fv],
+                [2, 11, 14],
+                [4, 15, fv],
+                [8, 18, fv],
+                [9, 12, 17],
+            ],
+            dtype=INT_DTYPE,
+        )
+        assert "face_node_connectivity" in uxgrid.connectivity
+        assert "edge_node_connectivity" in uxgrid.connectivity
+        assert "node_edge_connectivity" in uxgrid.connectivity
+        np.testing.assert_array_equal(node_edges, expected_node_edges)
+
+    def test_node_face_connectivity(self, uxgrid):
+        node_faces = uxgrid.node_face_connectivity.values
+        expected_node_faces = np.array(
+            [
+                [0, 1, 2],
+                [0, 2, fv],
+                [0, fv, fv],
+                [0, fv, fv],
+                [0, 3, fv],
+                [0, 1, 3],
+                [1, fv, fv],
+                [1, fv, fv],
+                [2, fv, fv],
+                [2, fv, fv],
+                [3, fv, fv],
+                [3, fv, fv],
+                [1, 2, fv],
+                [2, fv, fv],
+                [3, fv, fv],
+                [1, 3, fv],
+            ],
+            dtype=INT_DTYPE,
+        )
+        assert "face_node_connectivity" in uxgrid.connectivity
+        assert "node_face_connectivity" in uxgrid.connectivity
+        np.testing.assert_array_equal(node_faces, expected_node_faces)
+
+    def test_node_node_connectivity(self, uxgrid):
+        with pytest.raises(NotImplementedError):
+            _ = uxgrid.node_node_connectivity.values
+
+    # ─── Edge Connectivity ───────────────────────────────────────────────────────────────────────────────────────────
+
+    def test_edge_node_connectivity(self, uxgrid):
+        edge_nodes = uxgrid.edge_node_connectivity.values
+        expected_edge_nodes = np.array(
+            [
+                [0, 1],
+                [0, 5],
+                [0, 12],
+                [1, 2],
+                [1, 13],
+                [2, 3],
+                [3, 4],
+                [4, 5],
+                [4, 14],
+                [5, 15],
+                [6, 7],
+                [6, 12],
+                [7, 15],
+                [8, 9],
+                [8, 12],
+                [9, 13],
+                [10, 11],
+                [10, 15],
+                [11, 14],
+            ],
+            dtype=INT_DTYPE,
+        )
+        assert "face_node_connectivity" in uxgrid.connectivity
+        assert "edge_node_connectivity" in uxgrid.connectivity
+        np.testing.assert_array_equal(edge_nodes, expected_edge_nodes)
+
+    def test_edge_edge_connectivity(self, uxgrid):
+        with pytest.raises(NotImplementedError):
+            _ = uxgrid.edge_edge_connectivity.values
+
+    def test_edge_face_connectivity(self, uxgrid):
+        edge_faces = uxgrid.edge_face_connectivity.values
+        expected_edge_faces = np.array(
+            [
+                [0, 2],
+                [0, 1],
+                [1, 2],
+                [0, fv],
+                [2, fv],
+                [0, fv],
+                [0, fv],
+                [0, 3],
+                [3, fv],
+                [1, 3],
+                [1, fv],
+                [1, fv],
+                [1, fv],
+                [2, fv],
+                [2, fv],
+                [2, fv],
+                [3, fv],
+                [3, fv],
+                [3, fv],
+            ],
+            dtype=INT_DTYPE,
+        )
+        assert "face_node_connectivity" in uxgrid.connectivity
+        assert "edge_node_connectivity" in uxgrid.connectivity
+        assert "edge_face_connectivity" in uxgrid.connectivity
+        assert "face_edge_connectivity" in uxgrid.connectivity
+        np.testing.assert_array_equal(edge_faces, expected_edge_faces)
+
+    # ─── Face Connectivity ───────────────────────────────────────────────────────────────────────────────────────────
+
+    def test_face_node_connectivity(self, uxgrid):
+        expected_face_nodes = np.array(
+            [
+                [0, 1, 2, 3, 4, 5],
+                [15, 7, 6, 12, 0, 5],
+                [0, 12, 8, 9, 13, 1],
+                [4, 14, 11, 10, 15, 5],
+            ],
+            dtype=INT_DTYPE,
+        )
+        face_nodes = uxgrid.face_node_connectivity.values
+        assert "face_node_connectivity" in uxgrid.connectivity
+        assert len(uxgrid.connectivity) == 1
+        np.testing.assert_array_equal(face_nodes, expected_face_nodes)
+
+    def test_face_edge_connectivity(self, uxgrid):
+        expected_face_edges = np.array(
+            [
+                [0, 3, 5, 6, 7, 1],
+                [12, 10, 11, 2, 1, 9],
+                [2, 14, 13, 15, 4, 0],
+                [8, 18, 16, 17, 9, 7],
+            ],
+            dtype=INT_DTYPE,
+        )
+        face_edges = uxgrid.face_edge_connectivity.values
+        assert "face_node_connectivity" in uxgrid.connectivity
+        assert "face_edge_connectivity" in uxgrid.connectivity
+        assert "edge_node_connectivity" in uxgrid.connectivity
+        np.testing.assert_array_equal(face_edges, expected_face_edges)
+
+    def test_face_face_connectivity(self, uxgrid):
+        face_faces = uxgrid.face_face_connectivity.values
+        expected_face_faces = np.array(
+            [
+                [2, 1, 3, fv, fv, fv],
+                [0, 2, 3, fv, fv, fv],
+                [0, 1, fv, fv, fv, fv],
+                [0, 1, fv, fv, fv, fv],
+            ],
+            dtype=INT_DTYPE,
+        )
+        assert "face_node_connectivity" in uxgrid.connectivity
+        assert "edge_node_connectivity" in uxgrid.connectivity
+        assert "edge_face_connectivity" in uxgrid.connectivity
+        assert "face_edge_connectivity" in uxgrid.connectivity
+        assert "face_face_connectivity" in uxgrid.connectivity
+        np.testing.assert_array_equal(face_faces, expected_face_faces)

test/meshfiles/ugrid/quad-hexagon/index_layout.png

@@ -771,10 +771,6 @@ def test_normalize_existing_coordinates_norm_initial():
 
 
 
-
-
-
-
 def test_from_topology():
     node_lon = np.array([-20.0, 0.0, 20.0, -20, -40])
     node_lat = np.array([-10.0, 10.0, -10.0, 10, -10])

test/test_connectivity.py

@@ -453,3 +453,54 @@ def _build_face_face_connectivity(grid):
     ]
 
     return face_face_connectivity
+
+
+def _populate_node_edge_connectivity(grid):
+    """Constructs the UGRID connectivity variable (``edge_node_connectivity``)
+    and stores it within the internal (``Grid._ds``) and through the attribute
+    (``Grid.edge_node_connectivity``)."""
+    node_edge_connectivity = _build_node_edge_connectivity(
+        grid.edge_node_connectivity.values, grid.n_node
+    )
+
+    grid._ds["node_edge_connectivity"] = xr.DataArray(
+        data=node_edge_connectivity,
+        dims=ugrid.NODE_EDGE_CONNECTIVITY_DIMS,
+        attrs=ugrid.NODE_EDGE_CONNECTIVITY_ATTRS,
+    )
+
+
+@njit
+def _build_node_edge_connectivity(edge_nodes, n_node):
+    """Constructs the Node Edge Connectivity, which stores the indices of the edges that are shared by each node."""
+    n_edge, nodes_per_edge = edge_nodes.shape
+
+    # count how many edges touch each node
+    counts = np.zeros(n_node, dtype=INT_DTYPE)
+    for e in range(n_edge):
+        for j in range(nodes_per_edge):
+            node = edge_nodes[e, j]
+            if node != INT_FILL_VALUE:
+                counts[node] += 1
+
+    # find the maximum
+    max_edges = 0
+    for i in range(n_node):
+        if counts[i] > max_edges:
+            max_edges = counts[i]
+
+    # allocate output, pad with fill
+    node_edge = np.full((n_node, max_edges), INT_FILL_VALUE, dtype=INT_DTYPE)
+
+    ptr = np.zeros(n_node, dtype=INT_DTYPE)
+
+    # fill in
+    for e in range(n_edge):
+        for j in range(nodes_per_edge):
+            node = edge_nodes[e, j]
+            if node != INT_FILL_VALUE:
+                idx = ptr[node]
+                node_edge[node, idx] = e
+                ptr[node] += 1
+
+    return node_edge

test/test_grid.py

@@ -28,6 +28,7 @@
     _populate_face_edge_connectivity,
     _populate_face_face_connectivity,
     _populate_n_nodes_per_face,
+    _populate_node_edge_connectivity,
     _populate_node_face_connectivity,
 )
 from uxarray.grid.coordinates import (
@@ -1286,9 +1287,7 @@ def edge_edge_connectivity(self, value):
     def node_edge_connectivity(self) -> xr.DataArray:
         """Indices of the edges that surround each node."""
         if "node_edge_connectivity" not in self._ds:
-            raise NotImplementedError(
-                "Construction of `node_edge_connectivity` not yet supported."
-            )
+            _populate_node_edge_connectivity(self)
 
         return self._ds["node_edge_connectivity"]