diff --git a/test/grid/geometry/test_centroids.py b/test/grid/geometry/test_centroids.py index 1ee8c7b38..887dcc0f0 100644 --- a/test/grid/geometry/test_centroids.py +++ b/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.""" diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index ac9658979..f50d9b6c1 100644 --- a/uxarray/grid/connectivity.py +++ b/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,29 +277,10 @@ 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): @@ -304,7 +288,7 @@ def _populate_node_face_connectivity(grid): 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