Skip to content

Optimized connectivity: Simultaneous face_edges and edge_nodes #1560

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
cmdupuis3 wants to merge 9 commits into
base: main
Choose a base branch
from
+85 −104
Draft

Optimized connectivity: Simultaneous face_edges and edge_nodes #1560

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
2fe0a5c
work on optimizing connectivity construction, cleanup Grid
philipc2 Apr 3, 2025
418adbc
update docstrings
philipc2 Apr 4, 2025
9b0e284
update face_face_connectivity
philipc2 Apr 7, 2025
8a31a4d
update face_face_connectivity
philipc2 Apr 7, 2025
39326ba
add derived geometries
philipc2 Apr 7, 2025
ffdd4bb
add geometry module
philipc2 Apr 8, 2025
55b3601
Merge branch 'main' into optimize-face-edges
philipc2 Apr 8, 2025
2a2eee9
Refresh optimize-face-edges branch
cmdupuis3 Jul 10, 2026
19bed33
Rework centroid triangle test
cmdupuis3 Jul 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions test/grid/geometry/test_centroids.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
174 changes: 76 additions & 98 deletions uxarray/grid/connectivity.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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)
Expand All @@ -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
)

Expand All @@ -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.
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down
Loading