Optimized connectivity: Simultaneous face_edges and edge_nodes

test/grid/geometry/test_centroids.py

@@ -63,13 +63,16 @@ def test_edge_centroids_from_triangle():
     grid = ux.open_grid(test_triangle, latlon=False)
     _populate_edge_centroids(grid)
 
-    centroid_x = np.mean(grid.node_x[grid.edge_node_connectivity[0][0:]])
-    centroid_y = np.mean(grid.node_y[grid.edge_node_connectivity[0][0:]])
-    centroid_z = np.mean(grid.node_z[grid.edge_node_connectivity[0][0:]])
+    edge_nodes = grid.edge_node_connectivity.values
 
-    assert centroid_x == grid.edge_x[0]
-    assert centroid_y == grid.edge_y[0]
-    assert centroid_z == grid.edge_z[0]
+    centroid_x = grid.node_x.values[edge_nodes].mean(axis=1)
+    centroid_y = grid.node_y.values[edge_nodes].mean(axis=1)
+    centroid_z = grid.node_z.values[edge_nodes].mean(axis=1)
+    centroid_x, centroid_y, centroid_z = _normalize_xyz(centroid_x, centroid_y, centroid_z)
+
+    nt.assert_array_almost_equal(grid.edge_x.values, centroid_x)
+    nt.assert_array_almost_equal(grid.edge_y.values, centroid_y)
+    nt.assert_array_almost_equal(grid.edge_z.values, centroid_z)
 
 def test_edge_centroids_from_mpas(gridpath):
     """Test computed centroid values compared to values from a MPAS dataset."""

uxarray/grid/connectivity.py

@@ -162,77 +162,80 @@ def _populate_edge_node_connectivity(grid):
     and stores it within the internal (``Grid._ds``) and through the attribute
     (``Grid.edge_node_connectivity``)."""
 
-    edge_nodes, inverse_indices, fill_value_mask = _build_edge_node_connectivity(
-        grid.face_node_connectivity.values, grid.n_face, grid.n_max_face_nodes
-    )
+    # Check edge coordinates already exist, if they do this might cause issues
 
-    edge_node_attrs = ugrid.EDGE_NODE_CONNECTIVITY_ATTRS
-    edge_node_attrs["inverse_indices"] = inverse_indices
+    if "n_edge" in grid.sizes:
+        # TODO: raise a warning or exception?
+        pass
+
+    edge_nodes, face_edges = _build_edge_node_connectivity(
+        grid.face_node_connectivity.values, grid.n_nodes_per_face.values
+    )
 
-    # add edge_node_connectivity to internal dataset
     grid._ds["edge_node_connectivity"] = xr.DataArray(
-        edge_nodes, dims=ugrid.EDGE_NODE_CONNECTIVITY_DIMS, attrs=edge_node_attrs
+        edge_nodes,
+        dims=ugrid.EDGE_NODE_CONNECTIVITY_DIMS,
+        attrs=ugrid.EDGE_NODE_CONNECTIVITY_ATTRS,
     )
 
+    grid._ds["face_edge_connectivity"] = xr.DataArray(
+        face_edges,
+        dims=ugrid.FACE_EDGE_CONNECTIVITY_DIMS,
+        attrs=ugrid.FACE_EDGE_CONNECTIVITY_ATTRS,
+    )
 
-def _build_edge_node_connectivity(face_nodes, n_face, n_max_face_nodes):
-    """Constructs the UGRID connectivity variable (``edge_node_connectivity``)
-    and stores it within the internal (``Grid._ds``) and through the attribute
-    (``Grid.edge_node_connectivity``).
 
-    Additionally, the attributes (``inverse_indices``) and
-    (``fill_value_mask``) are stored for constructing other
-    connectivity variables.
+@njit(cache=True)
+def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face):
+    """Constructs the ``edge_node_connectivity`` variable, which represents the indices of the two nodes that make up
+    each edge. Additionally, the ``face_edge_connectivity`` is derived during construction,  which represents the
+    indices of the edges that make up each face.
+
 
     Parameters
     ----------
-    repopulate : bool, optional
-        Flag used to indicate if we want to overwrite the existed `edge_node_connectivity` and generate a new
-        inverse_indices, default is False
-    """
-
-    padded_face_nodes = close_face_nodes(face_nodes, n_face, n_max_face_nodes)
+    face_node_connectivity : np.ndarray
+        Face Node Connectivity
+    n_nodes_per_face : np.ndarray
+        Number of nodes/edges per face
 
-    # array of empty edge nodes where each entry is a pair of indices
-    edge_nodes = np.empty((n_face * n_max_face_nodes, 2), dtype=INT_DTYPE)
+    Returns
+    -------
+    edge_node_connectivity : np.ndarray
+        Edge Node Connectivity with shape (n_edge, 2)
+    face_edge_connectivity : np.ndarray
+        Face Edge Connectivity with shape (n_face, n_max_face_edges)
 
-    # first index includes starting node up to non-padded value
-    edge_nodes[:, 0] = padded_face_nodes[:, :-1].ravel()
+    """
 
-    # second index includes second node up to padded value
-    edge_nodes[:, 1] = padded_face_nodes[:, 1:].ravel()
+    # Dictionary to keep track of unique edges
+    unique_edge_dict = {}
 
-    # sorted edge nodes
-    edge_nodes.sort(axis=1)
+    edge_idx = 0
 
-    # unique edge nodes
-    edge_nodes_unique, inverse_indices = np.unique(
-        edge_nodes, return_inverse=True, axis=0
-    )
-    # find all edge nodes that contain a fill value
-    fill_value_mask = np.logical_or(
-        edge_nodes_unique[:, 0] == INT_FILL_VALUE,
-        edge_nodes_unique[:, 1] == INT_FILL_VALUE,
+    # Keep track of face_edge_connectivity
+    face_edge_connectivity = np.full_like(
+        face_node_connectivity, INT_FILL_VALUE, dtype=INT_DTYPE
     )
 
-    # all edge nodes that do not contain a fill value
-    non_fill_value_mask = np.logical_not(fill_value_mask)
-    edge_nodes_unique = edge_nodes_unique[non_fill_value_mask]
+    for i, n_edges in enumerate(n_nodes_per_face):
+        for current_node in range(n_edges):
+            start_node = face_node_connectivity[i, current_node]
+            end_node = face_node_connectivity[i, (current_node + 1) % n_edges]
+
+            edge = (min(start_node, end_node), max(start_node, end_node))
 
-    # Update inverse_indices accordingly
-    indices_to_update = np.where(fill_value_mask)[0]
+            if edge not in unique_edge_dict:
+                # Only store unique edges
+                unique_edge_dict[edge] = edge_idx
+                edge_idx += 1
 
-    remove_mask = np.isin(inverse_indices, indices_to_update)
-    inverse_indices[remove_mask] = INT_FILL_VALUE
+            face_edge_connectivity[i, current_node] = unique_edge_dict[edge]
 
-    # Compute the indices where inverse_indices exceeds the values in indices_to_update
-    indexes = np.searchsorted(indices_to_update, inverse_indices, side="right")
-    # subtract the corresponding indexes from `inverse_indices`
-    for i in range(len(inverse_indices)):
-        if inverse_indices[i] != INT_FILL_VALUE:
-            inverse_indices[i] -= indexes[i]
+    # TODO: maybe sort these, but I don't think it's necessary
+    edge_node_connectivity = np.asarray(list(unique_edge_dict.keys()), dtype=INT_DTYPE)
 
-    return edge_nodes_unique, inverse_indices, fill_value_mask
+    return edge_node_connectivity, face_edge_connectivity
 
 
 def _populate_edge_face_connectivity(grid):
@@ -252,8 +255,8 @@ def _populate_edge_face_connectivity(grid):
 
 @njit(cache=True)
 def _build_edge_face_connectivity(face_edges, n_nodes_per_face, n_edge):
-    """Helper for (``edge_face_connectivity``) construction."""
-    edge_faces = np.ones(shape=(n_edge, 2), dtype=face_edges.dtype) * INT_FILL_VALUE
+    """Helper for (``edge_faces``) construction."""
+    edge_faces = np.full((n_edge, 2), INT_FILL_VALUE, dtype=INT_DTYPE)
 
     for face_idx, (cur_face_edges, n_edges) in enumerate(
         zip(face_edges, n_nodes_per_face)
@@ -274,37 +277,18 @@ def _populate_face_edge_connectivity(grid):
     and stores it within the internal (``Grid._ds``) and through the attribute
     (``Grid.face_edge_connectivity``)."""
 
-    if (
-        "edge_node_connectivity" not in grid._ds
-        or "inverse_indices" not in grid._ds["edge_node_connectivity"].attrs
-    ):
-        _populate_edge_node_connectivity(grid)
-
-    face_edges = _build_face_edge_connectivity(
-        grid.edge_node_connectivity.attrs["inverse_indices"],
-        grid.n_face,
-        grid.n_max_face_nodes,
-    )
-
-    grid._ds["face_edge_connectivity"] = xr.DataArray(
-        data=face_edges,
-        dims=ugrid.FACE_EDGE_CONNECTIVITY_DIMS,
-        attrs=ugrid.FACE_EDGE_CONNECTIVITY_ATTRS,
-    )
-
+    # TODO: Check if "edge_edge_connectivity" is already present
 
-def _build_face_edge_connectivity(inverse_indices, n_face, n_max_face_nodes):
-    """Helper for (``face_edge_connectivity``) construction."""
-    inverse_indices = inverse_indices.reshape(n_face, n_max_face_nodes)
-    return inverse_indices
+    if "edge_node_connectivity" not in grid._ds:
+        _populate_edge_node_connectivity(grid)
 
 
 def _populate_node_face_connectivity(grid):
     """Constructs the UGRID connectivity variable (``node_face_connectivity``)
     and stores it within the internal (``Grid._ds``) and through the attribute
     (``Grid.node_face_connectivity``)."""
 
-    node_faces, n_max_faces_per_node = _build_node_faces_connectivity(
+    node_faces, n_max_faces_per_node = _build_node_face_connectivity(
         grid.face_node_connectivity.values, grid.n_node
     )
 
@@ -315,7 +299,7 @@ def _populate_node_face_connectivity(grid):
     )
 
 
-def _build_node_faces_connectivity(face_nodes, n_node):
+def _build_node_face_connectivity(face_nodes, n_node):
     """Builds the `Grid.node_faces_connectivity`: integer DataArray of size
     (n_node, n_max_faces_per_node) (optional) A DataArray of indices indicating
     faces that are neighboring each node.
@@ -419,7 +403,9 @@ def _populate_face_face_connectivity(grid):
     """Constructs the UGRID connectivity variable (``face_face_connectivity``)
     and stores it within the internal (``Grid._ds``) and through the attribute
     (``Grid.face_face_connectivity``)."""
-    face_face = _build_face_face_connectivity(grid)
+    face_face = _build_face_face_connectivity(
+        grid.edge_face_connectivity.values, grid.n_face, grid.n_max_face_nodes
+    )
 
     grid._ds["face_face_connectivity"] = xr.DataArray(
         data=face_face,
@@ -428,32 +414,24 @@ def _populate_face_face_connectivity(grid):
     )
 
 
-def _build_face_face_connectivity(grid):
-    """Returns face-face connectivity."""
-
-    # Dictionary to store each faces adjacent faces
-    face_neighbors = {i: [] for i in range(grid.n_face)}
+@njit(cache=True)
+def _build_face_face_connectivity(edge_face_connectivity, n_face, n_max_face_nodes):
+    face_face_connectivity = np.full(
+        (n_face, n_max_face_nodes), INT_FILL_VALUE, INT_DTYPE
+    )
+    face_index_position = np.zeros(n_face, dtype=INT_DTYPE)
 
-    # Loop through each edge_face and add to the dictionary every face that shares an edge
-    for edge_face in grid.edge_face_connectivity.values:
-        face1, face2 = edge_face
-        if face1 != INT_FILL_VALUE and face2 != INT_FILL_VALUE:
-            # Append to each face's dictionary index the opposite face index
-            face_neighbors[face1].append(face2)
-            face_neighbors[face2].append(face1)
+    for edge_faces in edge_face_connectivity:
+        face_a, face_b = edge_faces
+        if face_a != INT_FILL_VALUE and face_b != INT_FILL_VALUE:
+            face_face_connectivity[face_a, face_index_position[face_a]] = face_b
+            face_index_position[face_a] += 1
 
-    # Convert to an array and pad it with fill values
-    face_face_conn = list(face_neighbors.values())
-    face_face_connectivity = [
-        np.pad(
-            arr, (0, grid.n_max_face_edges - len(arr)), constant_values=INT_FILL_VALUE
-        )
-        for arr in face_face_conn
-    ]
+            face_face_connectivity[face_b, face_index_position[face_b]] = face_a
+            face_index_position[face_b] += 1
 
     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