From 2fe0a5c38b636848aa830329192de96def4900fa Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec <67855069+philipc2@users.noreply.github.com> Date: Thu, 3 Apr 2025 01:22:47 -0500 Subject: [PATCH 01/19] work on optimizing connectivity construction, cleanup Grid --- test/test_gradient.py | 1 + test/test_grid.py | 12 +- uxarray/grid/connectivity.py | 159 ++++++++--------- uxarray/grid/grid.py | 326 +++++++++++------------------------ uxarray/grid/utils.py | 17 ++ 5 files changed, 198 insertions(+), 317 deletions(-) diff --git a/test/test_gradient.py b/test/test_gradient.py index 4d1a7c7ef..1f196ffee 100644 --- a/test/test_gradient.py +++ b/test/test_gradient.py @@ -53,6 +53,7 @@ def test_quad_hex(): else: assert grad.values[i] != 0 + # TODO: expected_values = np.array([27.95, 20.79, 28.96, 0, 0, 0, 0, 60.64, 0, 86.45, 0, 0, 0, 0, 0, 0, 0, 0, 0]) nt.assert_almost_equal(grad.values, expected_values, 1e-2) diff --git a/test/test_grid.py b/test/test_grid.py index 89e11fcf2..73c802a14 100644 --- a/test/test_grid.py +++ b/test/test_grid.py @@ -500,15 +500,17 @@ def test_connectivity_build_face_edges_connectivity_mpas(): edge_nodes_expected.sort(axis=1) edge_nodes_expected = np.unique(edge_nodes_expected, axis=0) - edge_nodes_output, _, _ = _build_edge_node_connectivity(mpas_grid_ux.face_node_connectivity.values, - mpas_grid_ux.n_face, - mpas_grid_ux.n_max_face_nodes) - assert np.array_equal(edge_nodes_expected, edge_nodes_output) + + edge_node_connectivity, _ = _build_edge_node_connectivity(mpas_grid_ux.face_node_connectivity.values, + mpas_grid_ux.n_nodes_per_face.values) + + print() + assert np.array_equal(edge_nodes_expected, edge_node_connectivity) n_face = mpas_grid_ux.n_node n_node = mpas_grid_ux.n_face - n_edge = edge_nodes_output.shape[0] + n_edge = edge_node_connectivity.shape[0] assert (n_face == n_edge - n_node + 2) diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index 41382ed50..b12d28a68 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -159,85 +159,58 @@ def _build_n_nodes_per_face(face_nodes, n_face, n_max_face_nodes): return n_nodes_per_face -def _populate_edge_node_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``).""" - - edge_nodes, inverse_indices, fill_value_mask = _build_edge_node_connectivity( - grid.face_node_connectivity.values, grid.n_face, grid.n_max_face_nodes - ) - - edge_node_attrs = ugrid.EDGE_NODE_CONNECTIVITY_ATTRS - edge_node_attrs["inverse_indices"] = inverse_indices - edge_node_attrs["fill_value_mask"] = fill_value_mask +@njit(cache=True) +def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face): + edge_idx = 0 + edge_dict = {} - # 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 + # Keep track of face_edge_connectivity + face_edge_connectivity = np.full_like( + face_node_connectivity, INT_FILL_VALUE, dtype=INT_DTYPE ) + 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] -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``). + if current_node == n_edges - 1: + end_node = face_node_connectivity[i, 0] + else: + end_node = face_node_connectivity[i, current_node + 1] - Additionally, the attributes (``inverse_indices``) and - (``fill_value_mask``) are stored for constructing other - connectivity variables. + # TODO: Maybe store direction here? + edge = (min(start_node, end_node), max(start_node, end_node)) - 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) + if edge not in edge_dict: + edge_dict[edge] = edge_idx + edge_idx += 1 - # 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) + face_edge_connectivity[i, current_node] = edge_dict[edge] - # first index includes starting node up to non-padded value - edge_nodes[:, 0] = padded_face_nodes[:, :-1].ravel() + edge_node_connectivity = np.asarray(list(edge_dict.keys()), dtype=INT_DTYPE) - # second index includes second node up to padded value - edge_nodes[:, 1] = padded_face_nodes[:, 1:].ravel() + return edge_node_connectivity, face_edge_connectivity - # sorted edge nodes - edge_nodes.sort(axis=1) - # 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, +def _populate_edge_node_connectivity(grid): + edge_node_connectivity, face_edge_connectivity = _build_edge_node_connectivity( + grid.face_node_connectivity.values, grid.n_nodes_per_face.values ) - # 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] - - # Update inverse_indices accordingly - indices_to_update = np.where(fill_value_mask)[0] - - remove_mask = np.isin(inverse_indices, indices_to_update) - inverse_indices[remove_mask] = INT_FILL_VALUE - - # 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] + grid._ds["edge_node_connectivity"] = xr.DataArray( + edge_node_connectivity, + dims=ugrid.EDGE_NODE_CONNECTIVITY_DIMS, + attrs=ugrid.EDGE_NODE_CONNECTIVITY_ATTRS, + ) - return edge_nodes_unique, inverse_indices, fill_value_mask + grid._ds["face_edge_connectivity"] = xr.DataArray( + face_edge_connectivity, + dims=ugrid.FACE_EDGE_CONNECTIVITY_DIMS, + attrs=ugrid.FACE_EDGE_CONNECTIVITY_ATTRS, + ) +# def _populate_edge_face_connectivity(grid): """Constructs the UGRID connectivity variable (``edge_node_connectivity``) and stores it within the internal (``Grid._ds``) and through the attribute @@ -256,7 +229,7 @@ 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 + edge_face_connectivity = 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) @@ -264,12 +237,12 @@ def _build_edge_face_connectivity(face_edges, n_nodes_per_face, n_edge): # obtain all the edges that make up a face (excluding fill values) edges = cur_face_edges[:n_edges] for edge_idx in edges: - if edge_faces[edge_idx, 0] == INT_FILL_VALUE: - edge_faces[edge_idx, 0] = face_idx + if edge_face_connectivity[edge_idx, 0] == INT_FILL_VALUE: + edge_face_connectivity[edge_idx, 0] = face_idx else: - edge_faces[edge_idx, 1] = face_idx + edge_face_connectivity[edge_idx, 1] = face_idx - return edge_faces + return edge_face_connectivity def _populate_face_edge_connectivity(grid): @@ -277,29 +250,33 @@ 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, - ) - - -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 + # TODO: Check if "edge_edge_connectivity" is already present + + _populate_edge_node_connectivity(grid) + + # 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, + # ) + + +# 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 def _populate_node_face_connectivity(grid): diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 9cae36bac..cc6a95b59 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -14,7 +14,7 @@ Tuple, ) -from uxarray.grid.utils import _get_cartesian_face_edge_nodes +from uxarray.grid.utils import _get_cartesian_face_edge_nodes, make_setter # reader and writer imports from uxarray.io._exodus import _read_exodus, _encode_exodus @@ -179,7 +179,7 @@ def __init__( self, grid_ds: xr.Dataset, source_grid_spec: Optional[str] = None, - source_dims_dict: Optional[dict] = {}, + source_dims_dict: Optional[dict] = None, is_subset: bool = False, inverse_indices: Optional[xr.Dataset] = None, ): @@ -204,7 +204,9 @@ def __init__( # TODO: more checks for validate grid (lat/lon coords, etc) # mapping of ugrid dimensions and variables to source dataset's conventions - self._source_dims_dict = source_dims_dict + self._source_dims_dict = ( + source_dims_dict if source_dims_dict is not None else {} + ) # source grid specification (i.e. UGRID, MPAS, SCRIP, etc.) self.source_grid_spec = source_grid_spec @@ -679,7 +681,7 @@ def __repr__(self): prefix = "\n" original_grid_str = f"Original Grid Type: {self.source_grid_spec}\n" - dims_heading = "Grid Dimensions:\n" + dims_heading = "Grid Shape:\n" dims_str = "" for dim_name in ugrid.DIM_NAMES: @@ -859,7 +861,12 @@ def n_face(self) -> int: @property def n_max_face_nodes(self) -> int: - """The maximum number of nodes that can make up a single face.""" + """The maximum number of nodes that can make up a single face. + + For example, if a grid is composed entirely of triangular faces, the value would be 3. If a grid is composed + of a mix of triangles and hexagons, the value would be 6. + + """ return self.face_node_connectivity.shape[1] @property @@ -894,24 +901,30 @@ def n_max_node_edges(self) -> int: def n_nodes_per_face(self) -> xr.DataArray: """The number of nodes that make up each face. - Dimensions: ``(n_node, )`` + Shape: ``(n_face, )`` """ if "n_nodes_per_face" not in self._ds: _populate_n_nodes_per_face(self) return self._ds["n_nodes_per_face"] - @n_nodes_per_face.setter - def n_nodes_per_face(self, value): - """Setter for ``n_nodes_per_face``""" - assert isinstance(value, xr.DataArray) - self._ds["n_nodes_per_face"] = value + n_nodes_per_face = n_nodes_per_face.setter(make_setter("n_nodes_per_face")) + + @property + def n_edges_per_face(self) -> xr.DataArray: + """The number of edges that make up each face. Equivalent to ``n_nodes_per_face``. + + Shape: ``(n_face, )`` + """ + return self.n_nodes_per_face @property def node_lon(self) -> xr.DataArray: """Longitude of each node in degrees. - Dimensions: ``(n_node, )`` + Values are expected to be in the range ``[-180.0, 180.0]``. + + Shape: ``(n_node, )`` """ if "node_lon" not in self._ds: if self.source_grid_spec == "HEALPix": @@ -921,17 +934,13 @@ def node_lon(self) -> xr.DataArray: _populate_node_latlon(self) return self._ds["node_lon"] - @node_lon.setter - def node_lon(self, value): - """Setter for ``node_lon``""" - assert isinstance(value, xr.DataArray) - self._ds["node_lon"] = value + node_lon = node_lon.setter(make_setter("node_lon")) @property def node_lat(self) -> xr.DataArray: """Latitude of each node in degrees. - Dimensions: ``(n_node, )`` + Shape: ``(n_node, )`` """ if "node_lat" not in self._ds: if self.source_grid_spec == "HEALPix": @@ -941,166 +950,130 @@ def node_lat(self) -> xr.DataArray: _populate_node_latlon(self) return self._ds["node_lat"] - @node_lat.setter - def node_lat(self, value): - """Setter for ``node_lat``""" - assert isinstance(value, xr.DataArray) - self._ds["node_lat"] = value + node_lat = node_lat.setter(make_setter("node_lat")) @property def node_x(self) -> xr.DataArray: """Cartesian x location of each node in meters. - Dimensions: ``(n_node, )`` + Shape: ``(n_node, )`` """ if "node_x" not in self._ds: _populate_node_xyz(self) return self._ds["node_x"] - @node_x.setter - def node_x(self, value): - """Setter for ``node_x``""" - assert isinstance(value, xr.DataArray) - self._ds["node_x"] = value + node_x = node_x.setter(make_setter("node_x")) @property def node_y(self) -> xr.DataArray: """Cartesian y location of each node in meters. - Dimensions: ``(n_node, )`` + Shape: ``(n_node, )`` """ if "node_y" not in self._ds: _populate_node_xyz(self) return self._ds["node_y"] - @node_y.setter - def node_y(self, value): - """Setter for ``node_y``""" - assert isinstance(value, xr.DataArray) - self._ds["node_y"] = value + node_y = node_y.setter(make_setter("node_y")) @property def node_z(self) -> xr.DataArray: """Cartesian z location of each node in meters. - Dimensions: ``(n_node, )`` + Shape: ``(n_node, )`` """ if "node_z" not in self._ds: _populate_node_xyz(self) return self._ds["node_z"] - @node_z.setter - def node_z(self, value): - """Setter for ``node_z``""" - assert isinstance(value, xr.DataArray) - self._ds["node_z"] = value + node_z = node_z.setter(make_setter("node_z")) @property def edge_lon(self) -> xr.DataArray: """Longitude of the center of each edge in degrees. - Dimensions: ``(n_edge, )`` + Values are expected to be in the range ``[-180.0, 180.0]``. + + Shape: ``(n_edge, )`` """ if "edge_lon" not in self._ds: _populate_edge_centroids(self) _set_desired_longitude_range(self) return self._ds["edge_lon"] - @edge_lon.setter - def edge_lon(self, value): - """Setter for ``edge_lon``""" - assert isinstance(value, xr.DataArray) - self._ds["edge_lon"] = value + edge_lon = edge_lon.setter(make_setter("edge_lon")) @property def edge_lat(self) -> xr.DataArray: """Latitude of the center of each edge in degrees. - Dimensions: ``(n_edge, )`` + Shape: ``(n_edge, )`` """ if "edge_lat" not in self._ds: _populate_edge_centroids(self) _set_desired_longitude_range(self) return self._ds["edge_lat"] - @edge_lat.setter - def edge_lat(self, value): - """Setter for ``edge_lat``""" - assert isinstance(value, xr.DataArray) - self._ds["edge_lat"] = value + edge_lat = edge_lat.setter(make_setter("edge_lat")) @property def edge_x(self) -> xr.DataArray: """Cartesian x location of the center of each edge in meters. - Dimensions: ``(n_edge, )`` + Shape: ``(n_edge, )`` """ if "edge_x" not in self._ds: _populate_edge_centroids(self) return self._ds["edge_x"] - @edge_x.setter - def edge_x(self, value): - """Setter for ``edge_x``""" - assert isinstance(value, xr.DataArray) - self._ds["edge_x"] = value + edge_x = edge_x.setter(make_setter("edge_x")) @property def edge_y(self) -> xr.DataArray: """Cartesian y location of the center of each edge in meters. - Dimensions: ``(n_edge, )`` + Shape: ``(n_edge, )`` """ if "edge_y" not in self._ds: _populate_edge_centroids(self) return self._ds["edge_y"] - @edge_y.setter - def edge_y(self, value): - """Setter for ``edge_y``""" - assert isinstance(value, xr.DataArray) - self._ds["edge_y"] = value + edge_y = edge_y.setter(make_setter("edge_y")) @property def edge_z(self) -> xr.DataArray: """Cartesian z location of the center of each edge in meters. - Dimensions: ``(n_edge, )`` + Shape: ``(n_edge, )`` """ if "edge_z" not in self._ds: _populate_edge_centroids(self) return self._ds["edge_z"] - @edge_z.setter - def edge_z(self, value): - """Setter for ``edge_z``""" - assert isinstance(value, xr.DataArray) - self._ds["edge_z"] = value + edge_z = edge_z.setter(make_setter("edge_z")) @property def face_lon(self) -> xr.DataArray: """Longitude of the center of each face in degrees. - Dimensions: ``(n_face, )`` + Values are expected to be in the range ``[-180.0, 180.0]``. + + Shape: ``(n_face, )`` """ if "face_lon" not in self._ds: _populate_face_centroids(self) _set_desired_longitude_range(self) return self._ds["face_lon"] - @face_lon.setter - def face_lon(self, value): - """Setter for ``face_lon``""" - assert isinstance(value, xr.DataArray) - self._ds["face_lon"] = value + face_lon = face_lon.setter(make_setter("face_lon")) @property def face_lat(self) -> xr.DataArray: """Latitude of the center of each face in degrees. - Dimensions: ``(n_face, )`` + Shape: ``(n_face, )`` """ if "face_lat" not in self._ds: _populate_face_centroids(self) @@ -1108,66 +1081,51 @@ def face_lat(self) -> xr.DataArray: return self._ds["face_lat"] - @face_lat.setter - def face_lat(self, value): - """Setter for ``face_lat``""" - assert isinstance(value, xr.DataArray) - self._ds["face_lat"] = value + face_lat = face_lat.setter(make_setter("face_lat")) @property def face_x(self) -> xr.DataArray: """Cartesian x location of the center of each face in meters. - Dimensions: ``(n_face, )`` + Shape: ``(n_face, )`` """ if "face_x" not in self._ds: _populate_face_centroids(self) return self._ds["face_x"] - @face_x.setter - def face_x(self, value): - """Setter for ``face_x``""" - assert isinstance(value, xr.DataArray) - self._ds["face_x"] = value + face_x = face_x.setter(make_setter("face_x")) @property def face_y(self) -> xr.DataArray: """Cartesian y location of the center of each face in meters. - Dimensions: ``(n_face, )`` + Shape: ``(n_face, )`` """ if "face_y" not in self._ds: _populate_face_centroids(self) return self._ds["face_y"] - @face_y.setter - def face_y(self, value): - """Setter for ``face_x``""" - assert isinstance(value, xr.DataArray) - self._ds["face_y"] = value + face_y = face_y.setter(make_setter("face_y")) @property def face_z(self) -> xr.DataArray: """Cartesian z location of the center of each face in meters. - Dimensions: ``(n_face, )`` + Shape: ``(n_face, )`` """ if "face_z" not in self._ds: _populate_face_centroids(self) return self._ds["face_z"] - @face_z.setter - def face_z(self, value): - """Setter for ``face_z``""" - assert isinstance(value, xr.DataArray) - self._ds["face_z"] = value + face_z = face_z.setter(make_setter("face_z")) @property def face_node_connectivity(self) -> xr.DataArray: - """Indices of the nodes that make up each face. + """ + Indices of the nodes that make up each face. - Dimensions: ``(n_face, n_max_face_nodes)`` + Shape: ``(n_face, n_max_face_nodes)`` Nodes are in counter-clockwise order. """ @@ -1191,17 +1149,15 @@ def face_node_connectivity(self) -> xr.DataArray: return self._ds["face_node_connectivity"] - @face_node_connectivity.setter - def face_node_connectivity(self, value): - """Setter for ``face_node_connectivity``""" - assert isinstance(value, xr.DataArray) - self._ds["face_node_connectivity"] = value + face_node_connectivity = face_node_connectivity.setter( + make_setter("face_node_connectivity") + ) @property def edge_node_connectivity(self) -> xr.DataArray: """Indices of the two nodes that make up each edge. - Dimensions: ``(n_edge, two)`` + Shape: ``(n_edge, two)`` Nodes are in arbitrary order. """ @@ -1210,17 +1166,15 @@ def edge_node_connectivity(self) -> xr.DataArray: return self._ds["edge_node_connectivity"] - @edge_node_connectivity.setter - def edge_node_connectivity(self, value): - """Setter for ``edge_node_connectivity``""" - assert isinstance(value, xr.DataArray) - self._ds["edge_node_connectivity"] = value + edge_node_connectivity = edge_node_connectivity.setter( + make_setter("edge_node_connectivity") + ) @property def edge_node_x(self) -> xr.DataArray: """Cartesian x location for the two nodes that make up every edge. - Dimensions: ``(n_edge, two)`` + Shape: ``(n_edge, two)`` """ if "edge_node_x" not in self._ds: @@ -1233,40 +1187,6 @@ def edge_node_x(self) -> xr.DataArray: return self._ds["edge_node_x"] - @property - def edge_node_y(self) -> xr.DataArray: - """Cartesian y location for the two nodes that make up every edge. - - Dimensions: ``(n_edge, two)`` - """ - - if "edge_node_y" not in self._ds: - _edge_node_y = self.node_y[self.edge_node_connectivity] - - self._ds["edge_node_y"] = xr.DataArray( - data=_edge_node_y, - dims=["n_edge", "two"], - ) - - return self._ds["edge_node_y"] - - @property - def edge_node_z(self) -> xr.DataArray: - """Cartesian z location for the two nodes that make up every edge. - - Dimensions: ``(n_edge, two)`` - """ - - if "edge_node_z" not in self._ds: - _edge_node_z = self.node_z[self.edge_node_connectivity] - - self._ds["edge_node_z"] = xr.DataArray( - data=_edge_node_z, - dims=["n_edge", "two"], - ) - - return self._ds["edge_node_z"] - @property def node_node_connectivity(self) -> xr.DataArray: """Indices of the nodes that surround each node.""" @@ -1276,34 +1196,30 @@ def node_node_connectivity(self) -> xr.DataArray: ) return self._ds["node_node_connectivity"] - @node_node_connectivity.setter - def node_node_connectivity(self, value): - """Setter for ``node_node_connectivity``""" - assert isinstance(value, xr.DataArray) - self._ds["node_node_connectivity"] = value + node_node_connectivity = node_node_connectivity.setter( + make_setter("node_node_connectivity") + ) @property def face_edge_connectivity(self) -> xr.DataArray: """Indices of the edges that surround each face. - Dimensions: ``(n_face, n_max_face_edges)`` + Shape: ``(n_face, n_max_face_edges)`` """ if "face_edge_connectivity" not in self._ds: _populate_face_edge_connectivity(self) return self._ds["face_edge_connectivity"] - @face_edge_connectivity.setter - def face_edge_connectivity(self, value): - """Setter for ``face_edge_connectivity``""" - assert isinstance(value, xr.DataArray) - self._ds["face_edge_connectivity"] = value + face_edge_connectivity = face_edge_connectivity.setter( + make_setter("face_edge_connectivity") + ) @property def edge_edge_connectivity(self) -> xr.DataArray: """Indices of the edges that surround each edge. - Dimensions: ``(n_face, n_max_edge_edges)`` + Shape: ``(n_face, n_max_edge_edges)`` """ if "edge_edge_connectivity" not in self._ds: raise NotImplementedError( @@ -1312,11 +1228,9 @@ def edge_edge_connectivity(self) -> xr.DataArray: return self._ds["edge_edge_connectivity"] - @edge_edge_connectivity.setter - def edge_edge_connectivity(self, value): - """Setter for ``edge_edge_connectivity``""" - assert isinstance(value, xr.DataArray) - self._ds["edge_edge_connectivity"] = value + edge_edge_connectivity = edge_edge_connectivity.setter( + make_setter("edge_edge_connectivity") + ) @property def node_edge_connectivity(self) -> xr.DataArray: @@ -1328,11 +1242,9 @@ def node_edge_connectivity(self) -> xr.DataArray: return self._ds["node_edge_connectivity"] - @node_edge_connectivity.setter - def node_edge_connectivity(self, value): - """Setter for ``node_edge_connectivity``""" - assert isinstance(value, xr.DataArray) - self._ds["node_edge_connectivity"] = value + node_edge_connectivity = node_edge_connectivity.setter( + make_setter("node_edge_connectivity") + ) @property def face_face_connectivity(self) -> xr.DataArray: @@ -1345,11 +1257,9 @@ def face_face_connectivity(self) -> xr.DataArray: return self._ds["face_face_connectivity"] - @face_face_connectivity.setter - def face_face_connectivity(self, value): - """Setter for ``face_face_connectivity``""" - assert isinstance(value, xr.DataArray) - self._ds["face_face_connectivity"] = value + face_face_connectivity = face_face_connectivity.setter( + make_setter("face_face_connectivity") + ) @property def edge_face_connectivity(self) -> xr.DataArray: @@ -1362,11 +1272,9 @@ def edge_face_connectivity(self) -> xr.DataArray: return self._ds["edge_face_connectivity"] - @edge_face_connectivity.setter - def edge_face_connectivity(self, value): - """Setter for ``edge_face_connectivity``""" - assert isinstance(value, xr.DataArray) - self._ds["edge_face_connectivity"] = value + edge_face_connectivity = edge_face_connectivity.setter( + make_setter("edge_face_connectivity") + ) @property def node_face_connectivity(self) -> xr.DataArray: @@ -1379,11 +1287,9 @@ def node_face_connectivity(self) -> xr.DataArray: return self._ds["node_face_connectivity"] - @node_face_connectivity.setter - def node_face_connectivity(self, value): - """Setter for ``node_face_connectivity``""" - assert isinstance(value, xr.DataArray) - self._ds["node_face_connectivity"] = value + node_face_connectivity = node_face_connectivity.setter( + make_setter("node_face_connectivity") + ) @property def edge_node_distances(self): @@ -1395,11 +1301,7 @@ def edge_node_distances(self): _populate_edge_node_distances(self) return self._ds["edge_node_distances"] - @edge_node_distances.setter - def edge_node_distances(self, value): - """Setter for ``edge_node_distances``""" - assert isinstance(value, xr.DataArray) - self._ds["edge_node_distances"] = value + edge_node_distances = edge_node_distances.setter(make_setter("edge_node_distances")) @property def edge_face_distances(self): @@ -1412,11 +1314,7 @@ def edge_face_distances(self): _populate_edge_face_distances(self) return self._ds["edge_face_distances"] - @edge_face_distances.setter - def edge_face_distances(self, value): - """Setter for ``edge_face_distances``""" - assert isinstance(value, xr.DataArray) - self._ds["edge_face_distances"] = value + edge_face_distances = edge_face_distances.setter(make_setter("edge_face_distances")) @property def antimeridian_face_indices(self) -> np.ndarray: @@ -1437,11 +1335,7 @@ def face_areas(self) -> xr.DataArray: ) return self._ds["face_areas"] - @face_areas.setter - def face_areas(self, value): - """Setter for ``face_areas``""" - assert isinstance(value, xr.DataArray) - self._ds["face_areas"] = value + face_areas = face_areas.setter(make_setter("face_areas")) @property def bounds(self): @@ -1459,11 +1353,7 @@ def bounds(self): _populate_bounds(self) return self._ds["bounds"] - @bounds.setter - def bounds(self, value): - """Setter for ``bounds``""" - assert isinstance(value, xr.DataArray) - self._ds["bounds"] = value + bounds = bounds.setter(make_setter("bounds")) @property def face_bounds_lon(self): @@ -1515,11 +1405,9 @@ def boundary_edge_indices(self): ) return self._ds["boundary_edge_indices"] - @boundary_edge_indices.setter - def boundary_edge_indices(self, value): - """Setter for ``boundary_edge_indices``""" - assert isinstance(value, xr.DataArray) - self._ds["boundary_edge_indices"] = value + boundary_edge_indices = boundary_edge_indices.setter( + make_setter("boundary_edge_indices") + ) @property def boundary_node_indices(self): @@ -1530,11 +1418,9 @@ def boundary_node_indices(self): return self._ds["boundary_node_indices"] - @boundary_node_indices.setter - def boundary_node_indices(self, value): - """Setter for ``boundary_node_indices``""" - assert isinstance(value, xr.DataArray) - self._ds["boundary_node_indices"] = value + boundary_node_indices = boundary_node_indices.setter( + make_setter("boundary_node_indices") + ) @property def boundary_face_indices(self): @@ -1551,11 +1437,9 @@ def boundary_face_indices(self): return self._ds["boundary_face_indices"] - @boundary_face_indices.setter - def boundary_face_indices(self, value): - """Setter for ``boundary_face_indices``""" - assert isinstance(value, xr.DataArray) - self._ds["boundary_face_indices"] = value + boundary_face_indices = boundary_face_indices.setter( + make_setter("boundary_face_indices") + ) @property def triangular(self): diff --git a/uxarray/grid/utils.py b/uxarray/grid/utils.py index d9a6621c9..4e93b62e2 100644 --- a/uxarray/grid/utils.py +++ b/uxarray/grid/utils.py @@ -1,9 +1,21 @@ import numpy as np +import xarray as xr from uxarray.constants import INT_FILL_VALUE from numba import njit +def make_setter(key: str): + """Return a setter that assigns the value to self._ds[key] after type-checking.""" + + def setter(self, value): + if not isinstance(value, xr.DataArray): + raise ValueError(f"{key} must be an xr.DataArray") + self._ds[key] = value + + return setter + + @njit(cache=True) def _small_angle_of_2_vectors(u, v): """ @@ -229,6 +241,11 @@ def _get_cartesian_face_edge_nodes( [[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE], [INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]]]]) """ + + # face_edge_connectivity (n_face, n_edge) + + # each edge should have a shape (2, 3) + # Shift node connections to create edge connections face_node_conn_shift = np.roll(face_node_conn, -1, axis=1) From 418adbcedef8df568a4c15a2f1f9390a100809f3 Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec Date: Thu, 3 Apr 2025 21:22:26 -0500 Subject: [PATCH 02/19] update docstrings --- docs/api.rst | 10 + docs/getting-started/overview.rst | 37 ++- test/test_centroids.py | 7 +- test/test_grid.py | 1 - uxarray/grid/connectivity.py | 394 ++++++++++++----------- uxarray/grid/coordinates.py | 16 - uxarray/grid/grid.py | 497 ++++++++++++++++++++---------- uxarray/grid/utils.py | 116 ++++++- 8 files changed, 711 insertions(+), 367 deletions(-) diff --git a/docs/api.rst b/docs/api.rst index 5e76ba319..a2573cb29 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -480,3 +480,13 @@ Accurate Computing utils.computing.cross_fma utils.computing.dot_fma + + +Constants +------------------ + +.. autosummary:: + :toctree: generated/ + + constants.INT_FILL_VALUE + constants.INT_DTYPE diff --git a/docs/getting-started/overview.rst b/docs/getting-started/overview.rst index fc69d8635..fe0870f61 100644 --- a/docs/getting-started/overview.rst +++ b/docs/getting-started/overview.rst @@ -36,19 +36,21 @@ other geometric faces. Core Data Structures ==================== -The functionality of UXarray is built around three core data structures which provide -an Unstructured Grid aware implementation of many Xarray functions and use cases. +The functionality of UXarray is built around three core data structures: + +* :class:`Grid` + Used to represent an Unstructured Grid, housing grid-specific methods and + topology variables. + +* :class:`UxDataset` + Inherits from :py:class:`xarray.Dataset`, providing the same functionality + but extended to operate directly on Unstructured Grids. An :class:`UxDataset` + is linked to a :class:`Grid` object via the :attr:`UxDataset.uxgrid` property. + +* :class:`UxDataArray` + Similarly inherits from :py:class:`xarray.DataArray` and contains a + :attr:`UxDataArray.uxgrid` property just like :class:`UxDataset`. -* ``Grid`` is used to represent our Unstructured Grid, housing grid-specific methods - and topology variables. -* ``UxDataset`` inherits from the ``xarray.Dataset`` class, providing much of the same - functionality but extended to operate on Unstructured Grids. Other than new and - overloaded methods, it is linked to a ``Grid`` object through the use of a class - property (``UxDataset.uxgrid``) to provide a grid-aware implementation. An instance - of ``UxDataset`` can be thought of as a collection of Data Variables that reside on - some Unstructured Grid as defined in the ``uxgrid`` property. -* ``UxDataArray`` similarly inherits from the ``xarray.DataArray`` class and contains - a ``Grid`` property (``UxDataArray.uxgrid``) just like ``UxDataset``. Core Functionality ================== @@ -56,10 +58,13 @@ Core Functionality In addition to providing a way to load in and interface with Unstructured Grids, we also aim to provide computational and analysis operators that directly operate on Unstructured Grids. Some of these include: -* Visualization -* Remapping -* Subsetting & Selection -* Aggregations +- Visualization +- Remapping +- Subsetting & Selection +- Cross Sections +- Aggregations +- Calculus Operations +- Zonal Averaging A more detailed overview of supported functionality can be found in our `API Reference `_ and `User Guide `_ sections. diff --git a/test/test_centroids.py b/test/test_centroids.py index 7fdfc1dec..4347ac093 100644 --- a/test/test_centroids.py +++ b/test/test_centroids.py @@ -69,9 +69,10 @@ 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:]]) + + centroid_x = grid.node_x[grid.edge_node_connectivity].mean(axis=1) + centroid_y = grid.node_y[grid.edge_node_connectivity].mean(axis=1) + centroid_z = grid.node_z[grid.edge_node_connectivity].mean(axis=1) assert centroid_x == grid.edge_x[0] assert centroid_y == grid.edge_y[0] diff --git a/test/test_grid.py b/test/test_grid.py index 73c802a14..ca89d0583 100644 --- a/test/test_grid.py +++ b/test/test_grid.py @@ -501,7 +501,6 @@ def test_connectivity_build_face_edges_connectivity_mpas(): edge_nodes_expected = np.unique(edge_nodes_expected, axis=0) - edge_node_connectivity, _ = _build_edge_node_connectivity(mpas_grid_ux.face_node_connectivity.values, mpas_grid_ux.n_nodes_per_face.values) diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index b12d28a68..47805f6ef 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -1,131 +1,18 @@ import numpy as np import xarray as xr - from uxarray.constants import INT_DTYPE, INT_FILL_VALUE from uxarray.conventions import ugrid from numba import njit - -def close_face_nodes(face_node_connectivity, n_face, n_max_face_nodes): - """Closes (``face_node_connectivity``) by inserting the first node index - after the last non-fill-value node. - - Parameters - ---------- - face_node_connectivity : np.ndarray - Connectivity array for constructing a face from its nodes - n_face : constant - Number of faces - n_max_face_nodes : constant - Max number of nodes that compose a face - - Returns - ---------- - closed : ndarray - Closed (padded) face_node_connectivity - - Example - ---------- - Given face nodes with shape [2 x 5] - [0, 1, 2, 3, FILL_VALUE] - [4, 5, 6, 7, 8] - Pads them to the following with shape [2 x 6] - [0, 1, 2, 3, 0, FILL_VALUE] - [4, 5, 6, 7, 8, 4] - """ - - # padding to shape [n_face, n_max_face_nodes + 1] - closed = np.ones((n_face, n_max_face_nodes + 1), dtype=INT_DTYPE) * INT_FILL_VALUE - - # set all non-paded values to original face nodee values - closed[:, :-1] = face_node_connectivity.copy() - - # instance of first fill value - first_fv_idx_2d = np.argmax(closed == INT_FILL_VALUE, axis=1) - - # 2d to 1d index for np.put() - first_fv_idx_1d = first_fv_idx_2d + ((n_max_face_nodes + 1) * np.arange(0, n_face)) - - # column of first node values - first_node_value = face_node_connectivity[:, 0].copy() - - # insert first node column at occurrence of first fill value - np.put(closed.ravel(), first_fv_idx_1d, first_node_value) - - return closed - - -def _replace_fill_values(grid_var, original_fill, new_fill, new_dtype=None): - """Replaces all instances of the current fill value (``original_fill``) in - (``grid_var``) with (``new_fill``) and converts to the dtype defined by - (``new_dtype``) - - Parameters - ---------- - grid_var : xr.DataArray - Grid variable to be modified - original_fill : constant - Original fill value used in (``grid_var``) - new_fill : constant - New fill value to be used in (``grid_var``) - new_dtype : np.dtype, optional - New data type to convert (``grid_var``) to - - Returns - ------- - grid_var : xr.DataArray - Modified DataArray with updated fill values and dtype - """ - - # Identify fill value locations - if original_fill is not None and np.isnan(original_fill): - # For NaN fill values - fill_val_idx = grid_var.isnull() - # Temporarily replace NaNs with a placeholder if dtype conversion is needed - if new_dtype is not None and np.issubdtype(new_dtype, np.floating): - grid_var = grid_var.fillna(0.0) - else: - # Choose an appropriate placeholder for non-floating types - grid_var = grid_var.fillna(new_fill) - else: - # For non-NaN fill values - fill_val_idx = grid_var == original_fill - - # Convert to the new data type if specified - if new_dtype is not None and new_dtype != grid_var.dtype: - grid_var = grid_var.astype(new_dtype) - - # Validate that the new_fill can be represented in the new_dtype - if new_dtype is not None: - if np.issubdtype(new_dtype, np.integer): - int_min = np.iinfo(new_dtype).min - int_max = np.iinfo(new_dtype).max - if not (int_min <= new_fill <= int_max): - raise ValueError( - f"New fill value: {new_fill} not representable by integer dtype: {new_dtype}" - ) - elif np.issubdtype(new_dtype, np.floating): - if not ( - np.isnan(new_fill) - or (np.finfo(new_dtype).min <= new_fill <= np.finfo(new_dtype).max) - ): - raise ValueError( - f"New fill value: {new_fill} not representable by float dtype: {new_dtype}" - ) - else: - raise ValueError(f"Data type {new_dtype} not supported for grid variables") - - grid_var = grid_var.where(~fill_val_idx, new_fill) - - return grid_var +# ====================================================================================================================== +# n_nodes_per_face: Number of non-fill-value nodes/edges per face +# ====================================================================================================================== def _populate_n_nodes_per_face(grid): - """Constructs the connectivity variable (``n_nodes_per_face``) and stores - it within the internal (``Grid._ds``) and through the attribute - (``Grid.n_nodes_per_face``).""" + """Populates the ``n_nodes_per_face`` variable for a ``ux.Grid`` instance.""" n_nodes_per_face = _build_n_nodes_per_face( grid.face_node_connectivity.values, grid.n_face, grid.n_max_face_nodes @@ -148,7 +35,6 @@ def _build_n_nodes_per_face(face_nodes, n_face, n_max_face_nodes): """Constructs ``n_nodes_per_face``, which contains the number of non-fill- value nodes for each face in ``face_node_connectivity``""" - n_face, n_max_face_nodes = face_nodes.shape n_nodes_per_face = np.empty(n_face, dtype=INT_DTYPE) for i in range(n_face): c = 0 @@ -159,8 +45,55 @@ def _build_n_nodes_per_face(face_nodes, n_face, n_max_face_nodes): return n_nodes_per_face +# ====================================================================================================================== +# edge_node_connectivity: Indices of the two nodes that make up each edge +# ====================================================================================================================== + + +def _populate_edge_node_connectivity(grid): + """Populates the ``edge_node_connectivity`` and ``face_node_connectivity`` variables for a ``ux.Grid`` instance.""" + + # Check edge coordinates already exist, if they do this might cause issues + + edge_node_connectivity, face_edge_connectivity = _build_edge_node_connectivity( + grid.face_node_connectivity.values, grid.n_nodes_per_face.values + ) + + grid._ds["edge_node_connectivity"] = xr.DataArray( + edge_node_connectivity, + dims=ugrid.EDGE_NODE_CONNECTIVITY_DIMS, + attrs=ugrid.EDGE_NODE_CONNECTIVITY_ATTRS, + ) + + grid._ds["face_edge_connectivity"] = xr.DataArray( + face_edge_connectivity, + dims=ugrid.FACE_EDGE_CONNECTIVITY_DIMS, + attrs=ugrid.FACE_EDGE_CONNECTIVITY_ATTRS, + ) + + @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 + ---------- + face_node_connectivity : np.ndarray + Face Node Connectivity + n_nodes_per_face : np.ndarray + Number of nodes/edges per face + + 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) + + """ edge_idx = 0 edge_dict = {} @@ -172,45 +105,28 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face): 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] - if current_node == n_edges - 1: - end_node = face_node_connectivity[i, 0] - else: - end_node = face_node_connectivity[i, current_node + 1] - - # TODO: Maybe store direction here? edge = (min(start_node, end_node), max(start_node, end_node)) if edge not in edge_dict: + # Only store unique edges edge_dict[edge] = edge_idx edge_idx += 1 face_edge_connectivity[i, current_node] = edge_dict[edge] + # TODO: maybe sort these, but I don't think it's necessary edge_node_connectivity = np.asarray(list(edge_dict.keys()), dtype=INT_DTYPE) return edge_node_connectivity, face_edge_connectivity -def _populate_edge_node_connectivity(grid): - edge_node_connectivity, face_edge_connectivity = _build_edge_node_connectivity( - grid.face_node_connectivity.values, grid.n_nodes_per_face.values - ) +# ====================================================================================================================== +# edge_face_connectivity: Indices of the faces that saddle each edge +# ====================================================================================================================== - grid._ds["edge_node_connectivity"] = xr.DataArray( - edge_node_connectivity, - dims=ugrid.EDGE_NODE_CONNECTIVITY_DIMS, - attrs=ugrid.EDGE_NODE_CONNECTIVITY_ATTRS, - ) - - grid._ds["face_edge_connectivity"] = xr.DataArray( - face_edge_connectivity, - dims=ugrid.FACE_EDGE_CONNECTIVITY_DIMS, - attrs=ugrid.FACE_EDGE_CONNECTIVITY_ATTRS, - ) - -# def _populate_edge_face_connectivity(grid): """Constructs the UGRID connectivity variable (``edge_node_connectivity``) and stores it within the internal (``Grid._ds``) and through the attribute @@ -245,6 +161,11 @@ def _build_edge_face_connectivity(face_edges, n_nodes_per_face, n_edge): return edge_face_connectivity +# ====================================================================================================================== +# face_edge_connectivity: Indicies of the edges that make up each face +# ====================================================================================================================== + + def _populate_face_edge_connectivity(grid): """Constructs the UGRID connectivity variable (``face_edge_connectivity``) and stores it within the internal (``Grid._ds``) and through the attribute @@ -252,7 +173,8 @@ def _populate_face_edge_connectivity(grid): # TODO: Check if "edge_edge_connectivity" is already present - _populate_edge_node_connectivity(grid) + if "edge_node_connectivity" not in grid._ds: + _populate_edge_node_connectivity(grid) # if ( # "edge_node_connectivity" not in grid._ds @@ -278,6 +200,10 @@ def _populate_face_edge_connectivity(grid): # inverse_indices = inverse_indices.reshape(n_face, n_max_face_nodes) # return inverse_indices +# ====================================================================================================================== +# node_face_connectivity: Indices of the faces that share each node +# ====================================================================================================================== + def _populate_node_face_connectivity(grid): """Constructs the UGRID connectivity variable (``node_face_connectivity``) @@ -331,6 +257,55 @@ def _build_node_faces_connectivity(face_nodes, n_node): return node_face_connectivity, n_max_node_faces +# ====================================================================================================================== +# face_face_connectivity: Indices of the faces that neighbor each face +# ====================================================================================================================== + + +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) + + grid._ds["face_face_connectivity"] = xr.DataArray( + data=face_face, + dims=ugrid.FACE_FACE_CONNECTIVITY_DIMS, + attrs=ugrid.FACE_FACE_CONNECTIVITY_ATTRS, + ) + + +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)} + + # 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) + + # 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 + ] + + return face_face_connectivity + + +# ====================================================================================================================== +# Utils +# ====================================================================================================================== + + def _face_nodes_to_sparse_matrix(dense_matrix: np.ndarray) -> tuple: """Converts a given dense matrix connectivity to a sparse matrix format where the locations of non fill-value entries are stored using COO @@ -395,40 +370,115 @@ def get_face_node_partitions(n_nodes_per_face): return change_ind, n_nodes_per_face_sorted_ind, element_sizes, size_counts -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) +def close_face_nodes(face_node_connectivity, n_face, n_max_face_nodes): + """Closes (``face_node_connectivity``) by inserting the first node index + after the last non-fill-value node. - grid._ds["face_face_connectivity"] = xr.DataArray( - data=face_face, - dims=ugrid.FACE_FACE_CONNECTIVITY_DIMS, - attrs=ugrid.FACE_FACE_CONNECTIVITY_ATTRS, - ) + Parameters + ---------- + face_node_connectivity : np.ndarray + Connectivity array for constructing a face from its nodes + n_face : constant + Number of faces + n_max_face_nodes : constant + Max number of nodes that compose a face + Returns + ---------- + closed : ndarray + Closed (padded) face_node_connectivity -def _build_face_face_connectivity(grid): - """Returns face-face connectivity.""" + Example + ---------- + Given face nodes with shape [2 x 5] + [0, 1, 2, 3, FILL_VALUE] + [4, 5, 6, 7, 8] + Pads them to the following with shape [2 x 6] + [0, 1, 2, 3, 0, FILL_VALUE] + [4, 5, 6, 7, 8, 4] + """ - # Dictionary to store each faces adjacent faces - face_neighbors = {i: [] for i in range(grid.n_face)} + # padding to shape [n_face, n_max_face_nodes + 1] + closed = np.ones((n_face, n_max_face_nodes + 1), dtype=INT_DTYPE) * INT_FILL_VALUE - # 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) + # set all non-paded values to original face nodee values + closed[:, :-1] = face_node_connectivity.copy() - # 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 - ] + # instance of first fill value + first_fv_idx_2d = np.argmax(closed == INT_FILL_VALUE, axis=1) - return face_face_connectivity + # 2d to 1d index for np.put() + first_fv_idx_1d = first_fv_idx_2d + ((n_max_face_nodes + 1) * np.arange(0, n_face)) + + # column of first node values + first_node_value = face_node_connectivity[:, 0].copy() + + # insert first node column at occurrence of first fill value + np.put(closed.ravel(), first_fv_idx_1d, first_node_value) + + return closed + + +def _replace_fill_values(grid_var, original_fill, new_fill, new_dtype=None): + """Replaces all instances of the current fill value (``original_fill``) in + (``grid_var``) with (``new_fill``) and converts to the dtype defined by + (``new_dtype``) + + Parameters + ---------- + grid_var : xr.DataArray + Grid variable to be modified + original_fill : constant + Original fill value used in (``grid_var``) + new_fill : constant + New fill value to be used in (``grid_var``) + new_dtype : np.dtype, optional + New data type to convert (``grid_var``) to + + Returns + ------- + grid_var : xr.DataArray + Modified DataArray with updated fill values and dtype + """ + + # Identify fill value locations + if original_fill is not None and np.isnan(original_fill): + # For NaN fill values + fill_val_idx = grid_var.isnull() + # Temporarily replace NaNs with a placeholder if dtype conversion is needed + if new_dtype is not None and np.issubdtype(new_dtype, np.floating): + grid_var = grid_var.fillna(0.0) + else: + # Choose an appropriate placeholder for non-floating types + grid_var = grid_var.fillna(new_fill) + else: + # For non-NaN fill values + fill_val_idx = grid_var == original_fill + + # Convert to the new data type if specified + if new_dtype is not None and new_dtype != grid_var.dtype: + grid_var = grid_var.astype(new_dtype) + + # Validate that the new_fill can be represented in the new_dtype + if new_dtype is not None: + if np.issubdtype(new_dtype, np.integer): + int_min = np.iinfo(new_dtype).min + int_max = np.iinfo(new_dtype).max + if not (int_min <= new_fill <= int_max): + raise ValueError( + f"New fill value: {new_fill} not representable by integer dtype: {new_dtype}" + ) + elif np.issubdtype(new_dtype, np.floating): + if not ( + np.isnan(new_fill) + or (np.finfo(new_dtype).min <= new_fill <= np.finfo(new_dtype).max) + ): + raise ValueError( + f"New fill value: {new_fill} not representable by float dtype: {new_dtype}" + ) + else: + raise ValueError(f"Data type {new_dtype} not supported for grid variables") + + grid_var = grid_var.where(~fill_val_idx, new_fill) + + return grid_var diff --git a/uxarray/grid/coordinates.py b/uxarray/grid/coordinates.py index ffe5437d5..213093e47 100644 --- a/uxarray/grid/coordinates.py +++ b/uxarray/grid/coordinates.py @@ -792,22 +792,6 @@ def _xyz_to_lonlat_rad_no_norm( return lon, lat -def _normalize_xyz( - x: Union[np.ndarray, float], - y: Union[np.ndarray, float], - z: Union[np.ndarray, float], -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Normalizes a set of Cartesiain coordinates.""" - denom = np.linalg.norm( - np.asarray(np.array([x, y, z]), dtype=np.float64), ord=2, axis=0 - ) - - x_norm = x / denom - y_norm = y / denom - z_norm = z / denom - return x_norm, y_norm, z_norm - - @njit(cache=True) def _lonlat_rad_to_xyz( lon: Union[np.ndarray, float], diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index cc6a95b59..4583e7d29 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -16,7 +16,6 @@ from uxarray.grid.utils import _get_cartesian_face_edge_nodes, make_setter -# reader and writer imports from uxarray.io._exodus import _read_exodus, _encode_exodus from uxarray.io._mpas import _read_mpas from uxarray.io._geopandas import _read_geodataframe @@ -133,16 +132,17 @@ class Grid: """Represents a two-dimensional unstructured grid encoded following the UGRID conventions and provides grid-specific functionality. - Can be used standalone to work with unstructured grids, or can be paired with either a ``ux.UxDataArray`` or - ``ux.UxDataset`` and accessed through the ``.uxgrid`` attribute. + Can be used standalone to work with unstructured grids, or can be paired with either a `:py:class:`~uxarray.UxDataArray` or + :py:class:`~uxarray.UxDataset`and accessed through the `:py:attr:`~uxarray.UxDataArray.uxgrid` or `:py:attr:`~uxarray.UxDataset.uxgrid` + attributes. For constructing a grid from non-UGRID datasets or other types of supported data, see our ``ux.open_grid`` method or - specific class methods (``Grid.from_dataset``, ``Grid.from_face_verticies``, etc.) + specific class methods (py:meth:`.from_dataset`, py:meth:`.from_topology`, etc.) Parameters ---------- - grid_ds : xr.Dataset + grid_ds : :py:class:`xarray.Dataset` ``xarray.Dataset`` encoded in the UGRID conventions source_grid_spec : str, default="UGRID" @@ -154,7 +154,7 @@ class Grid: is_subset : bool, default=False Flag to mark if the grid is a subset or not - inverse_indices: xr.Dataset, default=None + inverse_indices: :py:class:`xarray.Dataset`, default=None A dataset of indices that correspond to the original grid, if the grid being constructed is a subset Examples @@ -277,11 +277,11 @@ def __init__( @classmethod def from_dataset(cls, dataset, use_dual: Optional[bool] = False, **kwargs): - """Constructs a ``Grid`` object from a dataset. + """Constructs a py:class:`~uxarray.Grid` object from a dataset. Parameters ---------- - dataset : xr.Dataset or path-like + dataset : py:class:`xarray.Dataset` or path-like ``xarray.Dataset`` containing unstructured grid coordinates and connectivity variables or a directory containing ASCII files represents a FESOM2 grid. use_dual : bool, default=False @@ -350,8 +350,7 @@ def from_file( backend: Optional[str] = "geopandas", **kwargs, ): - """Constructs a ``Grid`` object from a using the read_file method with - a specified backend. + """Constructs a py:class:`~uxarray.Grid` from a file using a specific backend. Parameters ---------- @@ -399,7 +398,7 @@ def from_points( boundary_points=None, **kwargs, ): - """Create a grid from unstructured points. + """Create a py:class:`~uxarray.Grid` from unstructured points. This class method generates connectivity information based on the provided points. Depending on the chosen `method`, it constructs either a spherical Voronoi diagram @@ -467,7 +466,7 @@ def from_topology( dims_dict: Optional[dict] = None, **kwargs, ): - """Constructs a ``Grid`` object from user-defined topology variables + """Constructs a py:class:`~uxarray.Grid` from user-defined topology variables provided in the UGRID conventions. Note @@ -518,7 +517,7 @@ def from_structured( cls, ds: xr.Dataset = None, lon=None, lat=None, tol: Optional[float] = 1e-10 ): """ - Converts a structured ``xarray.Dataset`` or longitude and latitude coordinates into an unstructured ``uxarray.Grid``. + Converts a structured py:class:`xarray.Dataset` or longitude and latitude coordinates into an unstructured py:class:`~uxarray.Grid`. This class method provides flexibility in converting structured grid data into an unstructured `uxarray.UxDataset`. Users can either supply an existing structured `xarray.Dataset` or provide longitude and latitude coordinates @@ -565,7 +564,7 @@ def from_face_vertices( face_vertices: Union[list, tuple, np.ndarray], latlon: Optional[bool] = True, ): - """Constructs a ``Grid`` object from user-defined face vertices. + """Constructs a py:class:`~uxarray.Grid` from user-defined face vertices. Parameters ---------- @@ -596,7 +595,7 @@ def from_face_vertices( @classmethod def from_healpix(cls, zoom: int, pixels_only: bool = True, nest: bool = True): - """Constructs a ``Grid`` object representing a given HEALPix zoom level. + """Constructs a py:class:`~uxarray.Grid` object representing a given HEALPix zoom level. Parameters ---------- @@ -607,8 +606,8 @@ def from_healpix(cls, zoom: int, pixels_only: bool = True, nest: bool = True): Returns ------- - Grid - An instance of ``uxarray.Grid`` + Grid: py:class:`~uxarray.Grid` + A py:class:`~uxarray.Grid` representing a HEALPix grid """ grid_ds = _pixels_to_ugrid(zoom, nest) @@ -618,7 +617,7 @@ def from_healpix(cls, zoom: int, pixels_only: bool = True, nest: bool = True): return cls.from_dataset(grid_ds, source_grid_spec="HEALPix") def validate(self, check_duplicates=True): - """Validates the current ``Grid``, checking for Duplicate Nodes, + """Validates the current py:class:`~uxarray.Grid`, checking for Duplicate Nodes, Present Connectivity, and Non-Zero Face Areas. Raises @@ -675,7 +674,7 @@ def construct_face_centers(self, method="cartesian average"): ) def __repr__(self): - """Constructs a string representation of the contents of a ``Grid``.""" + """Constructs a string representation of the contents of a py:class:`~uxarray.Grid`.""" from uxarray.conventions import descriptors @@ -780,7 +779,7 @@ def __ne__(self, other) -> bool: Parameters ---------- - other : uxarray.Grid + other : py:class:`~uxarray.Grid` The second grid object to be compared with `self` Returns @@ -789,6 +788,10 @@ def __ne__(self, other) -> bool: """ return not self.__eq__(other) + # ================================================================================================================== + # Grid Information Properties + # ================================================================================================================== + @property def dims(self) -> set: """Names of all unstructured grid dimensions.""" @@ -827,104 +830,167 @@ def descriptors(self) -> set: return set([desc for desc in DESCRIPTOR_NAMES if desc in self._ds]) - @property - def parsed_attrs(self) -> dict: - """Dictionary of parsed attributes from the source grid.""" - warn( - "Grid.parsed_attrs will be deprecated in a future release. Please use Grid.attrs instead.", - DeprecationWarning, - ) - return self._ds.attrs - @property def attrs(self) -> dict: """Dictionary of parsed attributes from the source grid.""" return self._ds.attrs + # ================================================================================================================== + # Dimension Properties + # ================================================================================================================== + @property def n_node(self) -> int: - """Total number of nodes.""" + """Total number of nodes. + + Returns + ------- + n_node : int + The total number of nodes. + """ return self._ds.sizes["n_node"] @property def n_edge(self) -> int: - """Total number of edges.""" + """Total number of edges. + + Returns + ------- + n_edge : int + The total number of edges. + """ if "edge_node_connectivity" not in self._ds: _populate_edge_node_connectivity(self) - return self._ds.sizes["n_edge"] @property def n_face(self) -> int: - """Total number of faces.""" + """Total number of faces. + + Returns + ------- + n_face : int + The total number of faces. + """ return self._ds.sizes["n_face"] @property def n_max_face_nodes(self) -> int: - """The maximum number of nodes that can make up a single face. + """Maximum number of nodes defining a single face. - For example, if a grid is composed entirely of triangular faces, the value would be 3. If a grid is composed - of a mix of triangles and hexagons, the value would be 6. + For example, if the grid is composed entirely of triangular faces, the value would be 3. + If the grid is composed of a mix of triangles and hexagons, the value would be 6. + Returns + ------- + n_max_face_nodes : int + The maximum number of nodes that can define a face. """ return self.face_node_connectivity.shape[1] @property def n_max_face_edges(self) -> int: - """The maximum number of edges that surround a single face. + """Maximum number of edges surrounding a single face. + + This is equivalent to :py:attr:`~uxarray.Grid.n_max_face_nodes`. - Equivalent to ``n_max_face_nodes`` + Returns + ------- + n_max_face_edges : int + The maximum number of edges that can surround a face. """ return self.face_edge_connectivity.shape[1] @property def n_max_face_faces(self) -> int: - """The maximum number of faces that surround a single face.""" + """Maximum number of neighboring faces surrounding a single face. + + Returns + ------- + n_max_face_faces : int + The maximum number of faces that can surround a face. + """ return self.face_face_connectivity.shape[1] @property def n_max_edge_edges(self) -> int: - """The maximum number of edges that surround a single edge.""" + """Maximum number of edges surrounding a single edge. + + Returns + ------- + n_max_edge_edges : int + The maximum number of edges that can surround an edge. + """ return self.edge_edge_connectivity.shape[1] @property def n_max_node_faces(self) -> int: - """The maximum number of faces that surround a single node.""" + """Maximum number of faces surrounding a single node. + + Returns + ------- + n_max_node_faces : int + The maximum number of faces that can surround a node. + """ return self.node_face_connectivity.shape[1] @property def n_max_node_edges(self) -> int: - """The maximum number of edges that surround a single node.""" + """Maximum number of edges surrounding a single node. + + Returns + ------- + n_max_node_edges : int + The maximum number of edges that can surround a node. + """ return self.node_edge_connectivity.shape[1] @property def n_nodes_per_face(self) -> xr.DataArray: - """The number of nodes that make up each face. + """Number of nodes defining each face. + + Shape: (:py:attr:`~uxarray.Grid.n_face`,) - Shape: ``(n_face, )`` + Returns + ------- + n_nodes_per_face : :py:class:`xarray.DataArray` + An array containing the number of nodes per face. """ if "n_nodes_per_face" not in self._ds: _populate_n_nodes_per_face(self) - return self._ds["n_nodes_per_face"] n_nodes_per_face = n_nodes_per_face.setter(make_setter("n_nodes_per_face")) @property def n_edges_per_face(self) -> xr.DataArray: - """The number of edges that make up each face. Equivalent to ``n_nodes_per_face``. + """Number of edges defining each face. - Shape: ``(n_face, )`` + This is equivalent to :py:attr:`~uxarray.Grid.n_nodes_per_face`. + + Shape: (:py:attr:`~uxarray.Grid.n_face`,) + + Returns + ------- + n_edges_per_face : :py:class:`xarray.DataArray` + An array containing the number of edges per face. """ return self.n_nodes_per_face + # ================================================================================================================== + # Coordinate Properties + # ================================================================================================================== + @property def node_lon(self) -> xr.DataArray: - """Longitude of each node in degrees. + """Longitude coordinate of each node (in degrees). - Values are expected to be in the range ``[-180.0, 180.0]``. + Values are expected to be in the range [-180.0, 180.0]. - Shape: ``(n_node, )`` + Returns + ------- + node_lon : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_node`,) """ if "node_lon" not in self._ds: if self.source_grid_spec == "HEALPix": @@ -938,9 +1004,12 @@ def node_lon(self) -> xr.DataArray: @property def node_lat(self) -> xr.DataArray: - """Latitude of each node in degrees. + """Latitude coordinate of each node (in degrees). - Shape: ``(n_node, )`` + Returns + ------- + node_lat : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_node`,) """ if "node_lat" not in self._ds: if self.source_grid_spec == "HEALPix": @@ -954,22 +1023,27 @@ def node_lat(self) -> xr.DataArray: @property def node_x(self) -> xr.DataArray: - """Cartesian x location of each node in meters. + """Cartesian x coordinate of each node (in meters). - Shape: ``(n_node, )`` + Returns + ------- + node_x : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_node`,) """ if "node_x" not in self._ds: _populate_node_xyz(self) - return self._ds["node_x"] node_x = node_x.setter(make_setter("node_x")) @property def node_y(self) -> xr.DataArray: - """Cartesian y location of each node in meters. + """Cartesian y coordinate of each node (in meters). - Shape: ``(n_node, )`` + Returns + ------- + node_y : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_node`,) """ if "node_y" not in self._ds: _populate_node_xyz(self) @@ -979,9 +1053,12 @@ def node_y(self) -> xr.DataArray: @property def node_z(self) -> xr.DataArray: - """Cartesian z location of each node in meters. + """Cartesian z coordinate of each node (in meters). - Shape: ``(n_node, )`` + Returns + ------- + node_z : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_node`,) """ if "node_z" not in self._ds: _populate_node_xyz(self) @@ -991,11 +1068,14 @@ def node_z(self) -> xr.DataArray: @property def edge_lon(self) -> xr.DataArray: - """Longitude of the center of each edge in degrees. + """Longitude coordinate of the center of each edge (in degrees). - Values are expected to be in the range ``[-180.0, 180.0]``. + Values are expected to be in the range [-180.0, 180.0]. - Shape: ``(n_edge, )`` + Returns + ------- + edge_lon : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_edge`,) """ if "edge_lon" not in self._ds: _populate_edge_centroids(self) @@ -1006,35 +1086,43 @@ def edge_lon(self) -> xr.DataArray: @property def edge_lat(self) -> xr.DataArray: - """Latitude of the center of each edge in degrees. + """Latitude coordinate of the center of each edge (in degrees). - Shape: ``(n_edge, )`` + Returns + ------- + edge_lat : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_edge`,) """ if "edge_lat" not in self._ds: _populate_edge_centroids(self) - _set_desired_longitude_range(self) + _set_desired_longitude_range(self) return self._ds["edge_lat"] edge_lat = edge_lat.setter(make_setter("edge_lat")) @property def edge_x(self) -> xr.DataArray: - """Cartesian x location of the center of each edge in meters. + """Cartesian x coordinate of the center of each edge (in meters). - Shape: ``(n_edge, )`` + Returns + ------- + edge_x : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_edge`,) """ if "edge_x" not in self._ds: _populate_edge_centroids(self) - return self._ds["edge_x"] edge_x = edge_x.setter(make_setter("edge_x")) @property def edge_y(self) -> xr.DataArray: - """Cartesian y location of the center of each edge in meters. + """Cartesian y coordinate of the center of each edge (in meters). - Shape: ``(n_edge, )`` + Returns + ------- + edge_y : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_edge`,) """ if "edge_y" not in self._ds: _populate_edge_centroids(self) @@ -1044,9 +1132,12 @@ def edge_y(self) -> xr.DataArray: @property def edge_z(self) -> xr.DataArray: - """Cartesian z location of the center of each edge in meters. + """Cartesian z coordinate of the center of each edge (in meters). - Shape: ``(n_edge, )`` + Returns + ------- + edge_z : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_edge`,) """ if "edge_z" not in self._ds: _populate_edge_centroids(self) @@ -1056,11 +1147,14 @@ def edge_z(self) -> xr.DataArray: @property def face_lon(self) -> xr.DataArray: - """Longitude of the center of each face in degrees. + """Longitude coordinate of the center of each face (in degrees). - Values are expected to be in the range ``[-180.0, 180.0]``. + Values are expected to be in the range [-180.0, 180.0]. - Shape: ``(n_face, )`` + Returns + ------- + face_lon : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_face`,) """ if "face_lon" not in self._ds: _populate_face_centroids(self) @@ -1071,36 +1165,43 @@ def face_lon(self) -> xr.DataArray: @property def face_lat(self) -> xr.DataArray: - """Latitude of the center of each face in degrees. + """Latitude coordinate of the center of each face (in degrees). - Shape: ``(n_face, )`` + Returns + ------- + face_lat : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_face`,) """ if "face_lat" not in self._ds: _populate_face_centroids(self) _set_desired_longitude_range(self) - return self._ds["face_lat"] face_lat = face_lat.setter(make_setter("face_lat")) @property def face_x(self) -> xr.DataArray: - """Cartesian x location of the center of each face in meters. + """Cartesian x coordinate of the center of each face (in meters). - Shape: ``(n_face, )`` + Returns + ------- + face_x : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_face`,) """ if "face_x" not in self._ds: _populate_face_centroids(self) - return self._ds["face_x"] face_x = face_x.setter(make_setter("face_x")) @property def face_y(self) -> xr.DataArray: - """Cartesian y location of the center of each face in meters. + """Cartesian y coordinate of the center of each face (in meters). - Shape: ``(n_face, )`` + Returns + ------- + face_y : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_face`,) """ if "face_y" not in self._ds: _populate_face_centroids(self) @@ -1110,9 +1211,12 @@ def face_y(self) -> xr.DataArray: @property def face_z(self) -> xr.DataArray: - """Cartesian z location of the center of each face in meters. + """Cartesian z coordinate of the center of each face (in meters). - Shape: ``(n_face, )`` + Returns + ------- + face_z : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_face`,) """ if "face_z" not in self._ds: _populate_face_centroids(self) @@ -1120,16 +1224,26 @@ def face_z(self) -> xr.DataArray: face_z = face_z.setter(make_setter("face_z")) + # ================================================================================================================== + # Connectivity Properties + # ================================================================================================================== + @property def face_node_connectivity(self) -> xr.DataArray: """ - Indices of the nodes that make up each face. + Connectivity variable representing the indices of nodes (mesh vertices) that define each face. - Shape: ``(n_face, n_max_face_nodes)`` + Each row (i.e., each face) contains at least three node indices and up to a maximum of + :py:attr:`~uxarray.Grid.n_max_face_nodes`. In grids with a mix of geometries (e.g., triangles and hexagons), + rows containing fewer than :py:attr:`~uxarray.Grid.n_max_face_nodes` indices are padded with the fill value defined in + :py:attr:`~uxarray.constants.INT_FILL_VALUE`. The node indices are stored in counter-clockwise order. - Nodes are in counter-clockwise order. + Returns + ------- + face_node_connectivity : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_face`, :py:attr:`~uxarray.Grid.n_max_face_nodes`) + representing the connectivity. """ - if ( "face_node_connectivity" not in self._ds and self.source_grid_spec == "HEALPix" @@ -1146,7 +1260,6 @@ def face_node_connectivity(self) -> xr.DataArray: dims=["n_face", "n_max_face_nodes"], attrs=self._ds["face_node_connectivity"].attrs, ) - return self._ds["face_node_connectivity"] face_node_connectivity = face_node_connectivity.setter( @@ -1155,15 +1268,20 @@ def face_node_connectivity(self) -> xr.DataArray: @property def edge_node_connectivity(self) -> xr.DataArray: - """Indices of the two nodes that make up each edge. + """ + Connectivity variable representing the indices of nodes (mesh vertices) that define each edge. - Shape: ``(n_edge, two)`` + Each row (i.e., each edge) contains exactly two node indices that define the start and end points of the edge. + The nodes are stored in an arbitrary order. - Nodes are in arbitrary order. + Returns + ------- + edge_node_connectivity : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_edge`, 2) + representing the connectivity. """ if "edge_node_connectivity" not in self._ds: _populate_edge_node_connectivity(self) - return self._ds["edge_node_connectivity"] edge_node_connectivity = edge_node_connectivity.setter( @@ -1171,25 +1289,16 @@ def edge_node_connectivity(self) -> xr.DataArray: ) @property - def edge_node_x(self) -> xr.DataArray: - """Cartesian x location for the two nodes that make up every edge. - - Shape: ``(n_edge, two)`` + def node_node_connectivity(self) -> xr.DataArray: """ + Connectivity variable representing the indices of nodes (mesh vertices) that surround each node. - if "edge_node_x" not in self._ds: - _edge_node_x = self.node_x[self.edge_node_connectivity] - - self._ds["edge_node_x"] = xr.DataArray( - data=_edge_node_x, - dims=["n_edge", "two"], - ) - - return self._ds["edge_node_x"] - - @property - def node_node_connectivity(self) -> xr.DataArray: - """Indices of the nodes that surround each node.""" + Returns + ------- + node_node_connectivity : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_node`, n_max_node_nodes) + representing the connectivity. + """ if "node_node_connectivity" not in self._ds: raise NotImplementedError( "Construction of `node_node_connectivity` not yet supported." @@ -1202,13 +1311,22 @@ def node_node_connectivity(self) -> xr.DataArray: @property def face_edge_connectivity(self) -> xr.DataArray: - """Indices of the edges that surround each face. + """ + Connectivity variable representing the indices of edges that define each face. - Shape: ``(n_face, n_max_face_edges)`` + Each row (i.e., each face) contains at least three edge indices and up to a maximum of + :py:attr:`~uxarray.Grid.n_max_face_edges`. In grids with a mix of geometries (e.g., triangles and hexagons), + rows containing fewer than :py:attr:`~uxarray.Grid.n_max_face_edges` indices are padded with the fill value defined in + :py:attr:`~uxarray.constants.INT_FILL_VALUE`. + + Returns + ------- + face_edge_connectivity : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_face`, :py:attr:`~uxarray.Grid.n_max_face_edges`) + representing the connectivity. """ if "face_edge_connectivity" not in self._ds: _populate_face_edge_connectivity(self) - return self._ds["face_edge_connectivity"] face_edge_connectivity = face_edge_connectivity.setter( @@ -1217,15 +1335,22 @@ def face_edge_connectivity(self) -> xr.DataArray: @property def edge_edge_connectivity(self) -> xr.DataArray: - """Indices of the edges that surround each edge. + """ + Connectivity variable representing the indices of edges that share at least one node. + + In grids with a mix of geometries (e.g., triangles and hexagons), rows containing fewer than the maximum number + of edge indices are padded with the fill value defined in :py:attr:`~uxarray.constants.INT_FILL_VALUE`. - Shape: ``(n_face, n_max_edge_edges)`` + Returns + ------- + edge_edge_connectivity : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_edge`, :py:attr:`~uxarray.Grid.n_max_edge_edges`) + representing the connectivity. """ if "edge_edge_connectivity" not in self._ds: raise NotImplementedError( "Construction of `edge_edge_connectivity` not yet supported." ) - return self._ds["edge_edge_connectivity"] edge_edge_connectivity = edge_edge_connectivity.setter( @@ -1234,12 +1359,22 @@ def edge_edge_connectivity(self) -> xr.DataArray: @property def node_edge_connectivity(self) -> xr.DataArray: - """Indices of the edges that surround each node.""" + """ + Connectivity variable representing the indices of edges that contain each node. + + In grids with a mix of geometries (e.g., triangles and hexagons), rows containing fewer than the maximum number + of edge indices are padded with the fill value defined in :py:attr:`~uxarray.constants.INT_FILL_VALUE`. + + Returns + ------- + node_edge_connectivity : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_node`, :py:attr:`~uxarray.Grid.n_max_node_edges`) + representing the connectivity. + """ if "node_edge_connectivity" not in self._ds: raise NotImplementedError( "Construction of `node_edge_connectivity` not yet supported." ) - return self._ds["node_edge_connectivity"] node_edge_connectivity = node_edge_connectivity.setter( @@ -1248,13 +1383,21 @@ def node_edge_connectivity(self) -> xr.DataArray: @property def face_face_connectivity(self) -> xr.DataArray: - """Indices of the faces that surround each face. + """ + Connectivity variable representing the indices of faces that share edges. - Dimensions ``(n_face, n_max_face_faces)`` + In grids with a mix of geometries (e.g., triangles and hexagons), rows containing fewer than + :py:attr:`~uxarray.Grid.n_max_face_faces` indices are padded with the fill value defined in + :py:attr:`~uxarray.constants.INT_FILL_VALUE`. + + Returns + ------- + face_face_connectivity : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_face`, :py:attr:`~uxarray.Grid.n_max_face_faces`) + representing the connectivity. """ if "face_face_connectivity" not in self._ds: _populate_face_face_connectivity(self) - return self._ds["face_face_connectivity"] face_face_connectivity = face_face_connectivity.setter( @@ -1263,13 +1406,21 @@ def face_face_connectivity(self) -> xr.DataArray: @property def edge_face_connectivity(self) -> xr.DataArray: - """Indices of the faces that saddle each edge. + """ + Connectivity variable representing the indices of faces that saddle each edge. + + Each row (i.e., each edge) contains either one or two face indices. A single face indicates that there + exists an empty region not covered by any geometry (e.g., a coastline). If an edge neighbors only one face, + the second value is padded with :py:attr:`~uxarray.constants.INT_FILL_VALUE`. - Dimensions ``(n_edge, two)`` + Returns + ------- + edge_face_connectivity : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_edge`, 2) + representing the connectivity. """ if "edge_face_connectivity" not in self._ds: _populate_edge_face_connectivity(self) - return self._ds["edge_face_connectivity"] edge_face_connectivity = edge_face_connectivity.setter( @@ -1278,25 +1429,41 @@ def edge_face_connectivity(self) -> xr.DataArray: @property def node_face_connectivity(self) -> xr.DataArray: - """Indices of the faces that surround each node. + """ + Connectivity variable representing the indices of faces that share a given node. + + In grids with a mix of geometries (e.g., triangles and hexagons), rows containing fewer than + :py:attr:`~uxarray.Grid.n_max_node_faces` indices are padded with the fill value defined in + :py:attr:`~uxarray.constants.INT_FILL_VALUE`. - Dimensions ``(n_node, n_max_node_faces)`` + Returns + ------- + node_face_connectivity : :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_node`, :py:attr:`~uxarray.Grid.n_max_node_faces`) + representing the connectivity. """ if "node_face_connectivity" not in self._ds: _populate_node_face_connectivity(self) - return self._ds["node_face_connectivity"] node_face_connectivity = node_face_connectivity.setter( make_setter("node_face_connectivity") ) + # ================================================================================================================== + # Descriptor Properties + # ================================================================================================================== + @property def edge_node_distances(self): - """Distances between the two nodes that surround each edge in radians. + """Arc distance between the two nodes that make up each edge (in radians). - Dimensions ``(n_edge, )`` + Returns + ------- + edge_node_distances: :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_edge`,) """ + if "edge_node_distances" not in self._ds: _populate_edge_node_distances(self) return self._ds["edge_node_distances"] @@ -1305,11 +1472,14 @@ def edge_node_distances(self): @property def edge_face_distances(self): - """Distances between the centers of the faces that saddle each edge in - radians. + """Arc distance between the faces that saddle each edge (in radians). - Dimensions ``(n_edge, )`` + Returns + ------- + edge_face_distances: :py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_edge`,) """ + if "edge_face_distances" not in self._ds: _populate_edge_face_distances(self) return self._ds["edge_face_distances"] @@ -1443,19 +1613,19 @@ def boundary_face_indices(self): @property def triangular(self): - """Boolean indicated whether the Grid is strictly composed of + """Boolean flag indicating whether the Grid is strictly composed of triangular faces.""" return self.n_max_face_nodes == 3 @property def partial_sphere_coverage(self): - """Boolean indicated whether the Grid partial covers the unit sphere + """Boolean flag indicating whether the Grid partial covers the unit sphere (i.e. contains holes)""" return self.boundary_edge_indices.size != 0 @property def global_sphere_coverage(self): - """Boolean indicated whether the Grid completely covers the unit sphere + """Boolean flag indicating whether the Grid completely covers the unit sphere (i.e. contains no holes)""" return not self.partial_sphere_coverage @@ -1471,16 +1641,26 @@ def inverse_indices(self) -> xr.Dataset: @property def is_subset(self): - """Returns `True` if the Grid is a subset, 'False' otherwise.""" + """Boolean flag indicating whether the Grid is a subset.""" return self._is_subset @property def max_face_radius(self): - """Maximum face radius of the grid in degrees""" + """Maximum face radius of the grid (in degrees)""" if "max_face_radius" not in self._ds: self._ds["max_face_radius"] = _populate_max_face_radius(self) return self._ds["max_face_radius"] + # ================================================================================================================== + # Convenience Properties + # ================================================================================================================== + + # TODO: + + # ================================================================================================================== + # Grid Methods + # ================================================================================================================== + def chunk(self, n_node="auto", n_edge="auto", n_face="auto"): """Converts all arrays to dask arrays with given chunks across grid dimensions in-place. @@ -1556,7 +1736,7 @@ def get_ball_tree( distance_metric: Optional[str] = "haversine", reconstruct: bool = False, ): - """Get the BallTree data structure of this Grid that allows for nearest + """Get the `~uxarray.grid.neighbors.BallTree` data structure of this Grid that allows for nearest neighbor queries (k nearest or within some radius) on either the (``node_x``, ``node_y``, ``node_z``) and (``node_lon``, ``node_lat``), edge (``edge_x``, ``edge_y``, ``edge_z``) and (``edge_lon``, @@ -1606,7 +1786,7 @@ def get_kd_tree( distance_metric: Optional[str] = "minkowski", reconstruct: bool = False, ): - """Get the KDTree data structure of this Grid that allows for nearest + """Get the `~uxarray.grid.neighbors.KDTree` data structure of this Grid that allows for nearest neighbor queries (k nearest or within some radius) on either the (``node_x``, ``node_y``, ``node_z``) and (``node_lon``, ``node_lat``), edge (``edge_x``, ``edge_y``, ``edge_z``) and (``edge_lon``, @@ -1652,7 +1832,7 @@ def get_spatial_hash( self, reconstruct: bool = False, ): - """Get the SpatialHash data structure of this Grid that allows for + """Obtain the py:class:`~uxarray.grid.neighbors.SpatialHash` that allows for fast face search queries. Face searches are used to find the faces that a list of points, in spherical coordinates, are contained within. @@ -1663,7 +1843,7 @@ def get_spatial_hash( Returns ------- - self._spatialhash : grid.Neighbors.SpatialHash + self._spatialhash : `~uxarray.grid.neighbors.SpatialHash` SpatialHash instance Note @@ -1700,7 +1880,7 @@ def copy(self): ) def encode_as(self, grid_type: str) -> xr.Dataset: - """Encodes the grid as a new `xarray.Dataset` per grid format supplied + """Encodes the grid as a new py:class:`xarray.Dataset` per grid format supplied in the `grid_type` argument. Parameters @@ -1711,8 +1891,8 @@ def encode_as(self, grid_type: str) -> xr.Dataset: Returns ------- - out_ds : xarray.Dataset - The output `xarray.Dataset` that is encoded from the this grid. + out_ds : py:class:`xarray.Dataset` + The output dataset that is encoded from the this grid. Raises ------ @@ -1758,7 +1938,7 @@ def calculate_total_face_area( order : int, optional Order of quadrature rule. Defaults to 4. latitude_adjusted_area : bool, optional - If True, corrects the area of the faces accounting for lines of constant lattitude. Defaults to False. + If True, corrects the area of the faces accounting for lines of constant latitude. Defaults to False. Returns ------- @@ -1779,8 +1959,7 @@ def compute_face_areas( latlon: Optional[bool] = True, latitude_adjusted_area: Optional[bool] = False, ): - """Face areas calculation function for grid class, calculates area of - all faces in the grid. + """Computes the area of all faces in the grid. Parameters ---------- @@ -1897,7 +2076,7 @@ def normalize_cartesian_coordinates(self): self.face_z.data = face_z def to_xarray(self, grid_format: Optional[str] = "ugrid"): - """Returns an ``xarray.Dataset`` with the variables stored under the + """Returns an py:class:`xarray.Dataset` with the variables stored under the ``Grid`` encoded in a specific grid format. Parameters @@ -1908,7 +2087,7 @@ def to_xarray(self, grid_format: Optional[str] = "ugrid"): Returns ------- - out_ds: xarray.Dataset + out_ds: py:class:`xarray.Dataset` Dataset representing the unstructured grid in a given grid format """ @@ -1945,8 +2124,8 @@ def to_geodataframe( exclude_nan_polygons: Optional[bool] = True, **kwargs, ): - """Constructs a ``GeoDataFrame`` consisting of polygons representing - the faces of the current ``Grid`` + """Constructs a py:class:`spatialpandas.GeoDataFrame` or py:class:`geopandas.GeoDataFrame`consisting of polygons representing + the faces of the current py:class:`~uxarray.Grid` Periodic polygons (i.e. those that cross the antimeridian) can be handled using the ``periodic_elements`` parameter. Setting ``periodic_elements='split'`` will split each periodic polygon along the antimeridian. @@ -2070,8 +2249,7 @@ def to_polycollection( return_non_nan_polygon_indices: Optional[bool] = False, **kwargs, ): - """Constructs a ``matplotlib.collections.PolyCollection``` consisting - of polygons representing the faces of the current ``Grid`` + """Constructs a py:class:`matplotlib.collections.PolyCollection`consisting of polygons representing the faces of the current py:class:`~uxarray.Grid` Parameters ---------- @@ -2153,8 +2331,8 @@ def to_linecollection( override: Optional[bool] = False, **kwargs, ): - """Constructs a ``matplotlib.collections.LineCollection``` consisting - of lines representing the edges of the current ``Grid`` + """Constructs a py:class:`matplotlib.collections.LineCollection` consisting + of lines representing the edges of the current py:class:`~uxarray.Grid` Parameters ---------- @@ -2214,7 +2392,7 @@ def get_dual(self): Returns -------- - dual : Grid + dual : py:class:`~uxarray.Grid` Dual Mesh Grid constructed """ @@ -2284,6 +2462,9 @@ def isel( "Indexing must be along a grid dimension: ('n_node', 'n_edge', 'n_face')" ) + # ================================================================================================================== + # Geometry Methods + # ================================================================================================================== def get_edges_at_constant_latitude(self, lat: float, use_face_bounds: bool = False): """Identifies the indices of edges that intersect with a line of constant latitude. @@ -2448,7 +2629,7 @@ def get_faces_containing_point( Parameters ---------- point_xyz : numpy.ndarray - A point in cartesian coordinates. Best performance if + A point in cartesian coordinates. point_lonlat : numpy.ndarray A point in spherical coordinates. tolerance : numpy.ndarray diff --git a/uxarray/grid/utils.py b/uxarray/grid/utils.py index 4e93b62e2..bb6f5a53f 100644 --- a/uxarray/grid/utils.py +++ b/uxarray/grid/utils.py @@ -1,6 +1,6 @@ import numpy as np import xarray as xr -from uxarray.constants import INT_FILL_VALUE +from uxarray.constants import INT_FILL_VALUE, INT_DTYPE from numba import njit @@ -337,3 +337,117 @@ def _get_lonlat_rad_face_edge_nodes( face_edges_lonlat_rad[valid_mask, 1] = node_lat_rad[valid_edges] return face_edges_lonlat_rad.reshape(n_face, n_max_face_edges, 2, 2) + + +def close_face_nodes(face_node_connectivity, n_face, n_max_face_nodes): + """Closes (``face_node_connectivity``) by inserting the first node index + after the last non-fill-value node. + + Parameters + ---------- + face_node_connectivity : np.ndarray + Connectivity array for constructing a face from its nodes + n_face : constant + Number of faces + n_max_face_nodes : constant + Max number of nodes that compose a face + + Returns + ---------- + closed : ndarray + Closed (padded) face_node_connectivity + + Example + ---------- + Given face nodes with shape [2 x 5] + [0, 1, 2, 3, FILL_VALUE] + [4, 5, 6, 7, 8] + Pads them to the following with shape [2 x 6] + [0, 1, 2, 3, 0, FILL_VALUE] + [4, 5, 6, 7, 8, 4] + """ + + # padding to shape [n_face, n_max_face_nodes + 1] + closed = np.ones((n_face, n_max_face_nodes + 1), dtype=INT_DTYPE) * INT_FILL_VALUE + + # set all non-paded values to original face nodee values + closed[:, :-1] = face_node_connectivity.copy() + + # instance of first fill value + first_fv_idx_2d = np.argmax(closed == INT_FILL_VALUE, axis=1) + + # 2d to 1d index for np.put() + first_fv_idx_1d = first_fv_idx_2d + ((n_max_face_nodes + 1) * np.arange(0, n_face)) + + # column of first node values + first_node_value = face_node_connectivity[:, 0].copy() + + # insert first node column at occurrence of first fill value + np.put(closed.ravel(), first_fv_idx_1d, first_node_value) + + return closed + + +def _replace_fill_values(grid_var, original_fill, new_fill, new_dtype=None): + """Replaces all instances of the current fill value (``original_fill``) in + (``grid_var``) with (``new_fill``) and converts to the dtype defined by + (``new_dtype``) + + Parameters + ---------- + grid_var : xr.DataArray + Grid variable to be modified + original_fill : constant + Original fill value used in (``grid_var``) + new_fill : constant + New fill value to be used in (``grid_var``) + new_dtype : np.dtype, optional + New data type to convert (``grid_var``) to + + Returns + ------- + grid_var : xr.DataArray + Modified DataArray with updated fill values and dtype + """ + + # Identify fill value locations + if original_fill is not None and np.isnan(original_fill): + # For NaN fill values + fill_val_idx = grid_var.isnull() + # Temporarily replace NaNs with a placeholder if dtype conversion is needed + if new_dtype is not None and np.issubdtype(new_dtype, np.floating): + grid_var = grid_var.fillna(0.0) + else: + # Choose an appropriate placeholder for non-floating types + grid_var = grid_var.fillna(new_fill) + else: + # For non-NaN fill values + fill_val_idx = grid_var == original_fill + + # Convert to the new data type if specified + if new_dtype is not None and new_dtype != grid_var.dtype: + grid_var = grid_var.astype(new_dtype) + + # Validate that the new_fill can be represented in the new_dtype + if new_dtype is not None: + if np.issubdtype(new_dtype, np.integer): + int_min = np.iinfo(new_dtype).min + int_max = np.iinfo(new_dtype).max + if not (int_min <= new_fill <= int_max): + raise ValueError( + f"New fill value: {new_fill} not representable by integer dtype: {new_dtype}" + ) + elif np.issubdtype(new_dtype, np.floating): + if not ( + np.isnan(new_fill) + or (np.finfo(new_dtype).min <= new_fill <= np.finfo(new_dtype).max) + ): + raise ValueError( + f"New fill value: {new_fill} not representable by float dtype: {new_dtype}" + ) + else: + raise ValueError(f"Data type {new_dtype} not supported for grid variables") + + grid_var = grid_var.where(~fill_val_idx, new_fill) + + return grid_var From 9b0e2840805ac13a21b7a8b3a598a2dddad5d890 Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec <67855069+philipc2@users.noreply.github.com> Date: Mon, 7 Apr 2025 17:27:56 -0500 Subject: [PATCH 03/19] update face_face_connectivity --- uxarray/grid/connectivity.py | 58 +++++++++++++++++------------------- 1 file changed, 28 insertions(+), 30 deletions(-) diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index 47805f6ef..6ac5401b6 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -55,6 +55,10 @@ def _populate_edge_node_connectivity(grid): # Check edge coordinates already exist, if they do this might cause issues + if "n_edge" in grid.sizes: + # TODO: raise a warning or exception? + pass + edge_node_connectivity, face_edge_connectivity = _build_edge_node_connectivity( grid.face_node_connectivity.values, grid.n_nodes_per_face.values ) @@ -94,8 +98,11 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face): Face Edge Connectivity with shape (n_face, n_max_face_edges) """ + + # Dictionary to keep track of unique edges + unique_edge_dict = {} + edge_idx = 0 - edge_dict = {} # Keep track of face_edge_connectivity face_edge_connectivity = np.full_like( @@ -109,15 +116,15 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face): edge = (min(start_node, end_node), max(start_node, end_node)) - if edge not in edge_dict: + if edge not in unique_edge_dict: # Only store unique edges - edge_dict[edge] = edge_idx + unique_edge_dict[edge] = edge_idx edge_idx += 1 - face_edge_connectivity[i, current_node] = edge_dict[edge] + face_edge_connectivity[i, current_node] = unique_edge_dict[edge] # TODO: maybe sort these, but I don't think it's necessary - edge_node_connectivity = np.asarray(list(edge_dict.keys()), dtype=INT_DTYPE) + edge_node_connectivity = np.asarray(list(unique_edge_dict.keys()), dtype=INT_DTYPE) return edge_node_connectivity, face_edge_connectivity @@ -210,7 +217,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 ) @@ -221,7 +228,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. @@ -266,7 +273,7 @@ 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, @@ -275,28 +282,19 @@ 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)} - - # 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) - - # 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 - ] +@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) + + 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 + + face_face_connectivity[face_b, face_index_position[face_b]] = face_a + face_index_position[face_b] += 1 return face_face_connectivity From 8a31a4de9aeafa7eea2c59a6155e082975182c42 Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec <67855069+philipc2@users.noreply.github.com> Date: Mon, 7 Apr 2025 17:29:21 -0500 Subject: [PATCH 04/19] update face_face_connectivity --- uxarray/grid/connectivity.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index 6ac5401b6..711ce17d8 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -273,7 +273,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.edge_face_connectivity.values, grid.n_face, grid.n_max_face_nodes) + 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, @@ -284,7 +286,9 @@ def _populate_face_face_connectivity(grid): @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_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) for edge_faces in edge_face_connectivity: From 39326ba3f0b66c7950049d986ce1dbb06fd9ce4b Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec <67855069+philipc2@users.noreply.github.com> Date: Mon, 7 Apr 2025 18:13:07 -0500 Subject: [PATCH 05/19] add derived geometries --- uxarray/core/zonal.py | 19 +-- uxarray/grid/connectivity.py | 2 +- uxarray/grid/geometry.py | 28 ++-- uxarray/grid/grid.py | 124 ++++++++++++++++-- uxarray/grid/utils.py | 240 +---------------------------------- 5 files changed, 131 insertions(+), 282 deletions(-) diff --git a/uxarray/core/zonal.py b/uxarray/core/zonal.py index e57975b87..eb5c0f0a3 100644 --- a/uxarray/core/zonal.py +++ b/uxarray/core/zonal.py @@ -3,7 +3,7 @@ from uxarray.grid.integrate import _zonal_face_weights, _zonal_face_weights_robust -from uxarray.grid.utils import _get_cartesian_face_edge_nodes +# from uxarray.grid.utils import _get_cartesian_face_edge_nodes def _compute_non_conservative_zonal_mean(uxda, latitudes, use_robust_weights=False): @@ -18,14 +18,7 @@ def _compute_non_conservative_zonal_mean(uxda, latitudes, use_robust_weights=Fal # Create a NumPy array for storing results result = np.zeros(shape, dtype=uxda.dtype) - faces_edge_nodes_xyz = _get_cartesian_face_edge_nodes( - uxgrid.face_node_connectivity.values, - uxgrid.n_face, - uxgrid.n_max_face_nodes, - uxgrid.node_x.values, - uxgrid.node_y.values, - uxgrid.node_z.values, - ) + face_edge_nodes_cartesian = uxda.uxgrid.face_edge_nodes_cartesian bounds = uxgrid.bounds.values @@ -34,7 +27,9 @@ def _compute_non_conservative_zonal_mean(uxda, latitudes, use_robust_weights=Fal z = np.sin(np.deg2rad(lat)) - faces_edge_nodes_xyz_candidate = faces_edge_nodes_xyz[face_indices, :, :, :] + face_edge_nodes_cartesian_candidate = face_edge_nodes_cartesian[ + face_indices, :, :, : + ] n_nodes_per_face_candidate = n_nodes_per_face[face_indices] @@ -42,11 +37,11 @@ def _compute_non_conservative_zonal_mean(uxda, latitudes, use_robust_weights=Fal if use_robust_weights: weights = _zonal_face_weights_robust( - faces_edge_nodes_xyz_candidate, z, bounds_candidate + face_edge_nodes_cartesian_candidate, z, bounds_candidate )["weight"].to_numpy() else: weights = _zonal_face_weights( - faces_edge_nodes_xyz_candidate, + face_edge_nodes_cartesian_candidate, bounds_candidate, n_nodes_per_face_candidate, z, diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index 711ce17d8..60f90a1a1 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -169,7 +169,7 @@ def _build_edge_face_connectivity(face_edges, n_nodes_per_face, n_edge): # ====================================================================================================================== -# face_edge_connectivity: Indicies of the edges that make up each face +# face_edge_connectivity: Indices of the edges that make up each face # ====================================================================================================================== diff --git a/uxarray/grid/geometry.py b/uxarray/grid/geometry.py index 1c5660da6..41140f307 100644 --- a/uxarray/grid/geometry.py +++ b/uxarray/grid/geometry.py @@ -28,10 +28,12 @@ from uxarray.grid.intersections import ( gca_gca_intersection, ) -from uxarray.grid.utils import ( - _get_cartesian_face_edge_nodes, - _get_lonlat_rad_face_edge_nodes, -) +# from uxarray.grid.utils import ( +# _get_cartesian_face_edge_nodes, +# _get_lonlat_rad_face_edge_nodes, +# ) + + from uxarray.utils.computing import allclose, isclose POLE_POINTS_XYZ = { @@ -1410,22 +1412,10 @@ def _populate_bounds( grid.normalize_cartesian_coordinates() # Prepare data for Numba functions - faces_edges_cartesian = _get_cartesian_face_edge_nodes( - grid.face_node_connectivity.values, - grid.n_face, - grid.n_max_face_edges, - grid.node_x.values, - grid.node_y.values, - grid.node_z.values, - ) + faces_edges_cartesian = grid.face_edge_nodes_cartesian - faces_edges_lonlat_rad = _get_lonlat_rad_face_edge_nodes( - grid.face_node_connectivity.values, - grid.n_face, - grid.n_max_face_edges, - grid.node_lon.values, - grid.node_lat.values, - ) + # TODO: update variable names + faces_edges_lonlat_rad = grid.face_edge_nodes_spherical n_nodes_per_face = grid.n_nodes_per_face.values diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 4583e7d29..0b9da2188 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -14,7 +14,12 @@ Tuple, ) -from uxarray.grid.utils import _get_cartesian_face_edge_nodes, make_setter +from uxarray.grid.utils import make_setter + +from uxarray.geometry.face_edges import ( + _construct_face_edge_nodes_cartesian, + _construct_face_edge_nodes_spherical, +) from uxarray.io._exodus import _read_exodus, _encode_exodus from uxarray.io._mpas import _read_mpas @@ -263,6 +268,8 @@ def __init__( # flag to track if coordinates are normalized self._normalized = None + self._cache_geometry = False + # set desired longitude range to [-180, 180] _set_desired_longitude_range(self) @@ -1652,10 +1659,112 @@ def max_face_radius(self): return self._ds["max_face_radius"] # ================================================================================================================== - # Convenience Properties + # Derived Geometry Arrays # ================================================================================================================== - # TODO: + @property + def cache_geometry(self): + """Boolean flag indicating whether to cache intermediary geometry arrays used within internal computations. + + For example, if face_edges_cartesian and face_edges_spherical are constructed during the face bounds construction, + they will be cached for later use in other methods, such as zonal averaging. + + The value is set to False by default to reduce memory usage. + + """ + return self._cache_geometry + + @cache_geometry.setter + def cache_geometry(self, value: bool): + assert isinstance(value, bool) + self._cache_geometry = value + + @property + def face_edge_nodes_cartesian(self): + """ + Geometry variable containing the Cartesian coordinates of the edges that make up each face. + + Returns + ------- + face_edge_nodes_cartesian : py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_face`, :py:attr:`~uxarray.Grid.n_max_face_edges`, two, three) + """ + if self.cache_geometry and "face_edge_nodes_cartesian" in self._ds: + return self._ds["face_edges_cartesian"] + + face_edge_nodes_cartesian = _construct_face_edge_nodes_cartesian( + self.face_node_connectivity.values, + self.n_face, + self.n_max_face_edges, + self.node_x.values, + self.node_y.values, + self.node_z.values, + ) + + if self.cache_geometry: + self._ds["face_edge_nodes_cartesian"] = xr.DataArray( + data=face_edge_nodes_cartesian, + dims=["n_face", "n_max_face_edges", "two", "three"], + ) + return self._ds["face_edge_nodes_cartesian"] + else: + return face_edge_nodes_cartesian + + @property + def face_edge_nodes_spherical(self): + """ + Geometry variable containing the Spherical coordinates of the edges that make up each face. + + Returns + ------- + face_edge_nodes_cartesian : py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_face`, :py:attr:`~uxarray.Grid.n_max_face_edges`, two, two) + """ + if self.cache_geometry and "face_edge_nodes_spherical" in self._ds: + return self._ds["face_edge_nodes_spherical"] + + face_edge_nodes_spherical = _construct_face_edge_nodes_spherical( + self.face_node_connectivity.values, + self.n_face, + self.n_max_face_edges, + self.node_lon.values, + self.node_lat.values, + ) + + if self.cache_geometry: + self._ds["face_edge_nodes_spherical"] = xr.DataArray( + data=face_edge_nodes_spherical, + dims=["n_face", "n_max_face_edges", "two", "two"], + ) + return self._ds["face_edge_nodes_spherical"] + else: + return face_edge_nodes_spherical + + # TODO: Polygon Coordinates (face_nodes_spherical) + + @property + def face_nodes_cartesian(self): + """ + Geometry variable containing the closed Cartesian coordinates of the nodes that make up each face. + + Returns + ------- + face_nodes_cartesian : py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_face`, :py:attr:`~uxarray.Grid.n_max_face_nodes` + 1) + """ + return None + + @property + def face_nodes_spherical(self): + """ + Geometry variable containing the closed Spherical coordinates of the nodes that make up each face. + + Returns + ------- + face_nodes_cartesian : py:class:`xarray.DataArray` + An array of shape (:py:attr:`~uxarray.Grid.n_face`, :py:attr:`~uxarray.Grid.n_max_face_nodes` + 1) + """ + return None # ================================================================================================================== # Grid Methods @@ -2703,14 +2812,7 @@ def get_faces_containing_point( return np.empty(0, dtype=np.int64) # Get the faces in terms of their edges - face_edge_nodes_xyz = _get_cartesian_face_edge_nodes( - subset.face_node_connectivity.values, - subset.n_face, - subset.n_max_face_nodes, - subset.node_x.values, - subset.node_y.values, - subset.node_z.values, - ) + face_edge_nodes_xyz = self.face_edge_nodes_cartesian # Get the original face indices from the subset inverse_indices = subset.inverse_indices.face.values diff --git a/uxarray/grid/utils.py b/uxarray/grid/utils.py index bb6f5a53f..9514a7271 100644 --- a/uxarray/grid/utils.py +++ b/uxarray/grid/utils.py @@ -1,6 +1,6 @@ import numpy as np import xarray as xr -from uxarray.constants import INT_FILL_VALUE, INT_DTYPE +from uxarray.constants import INT_FILL_VALUE from numba import njit @@ -150,244 +150,6 @@ def _swap_first_fill_value_with_last(arr): return arr -def _get_cartesian_face_edge_nodes( - face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z -): - """Construct an array to hold the edge Cartesian coordinates connectivity - for multiple faces in a grid. - - Parameters - ---------- - face_node_conn : np.ndarray - An array of shape (n_face, n_max_face_edges) containing the node indices for each face. Accessed through `grid.face_node_connectivity.value`. - n_face : int - The number of faces in the grid. Accessed through `grid.n_face`. - n_max_face_edges : int - The maximum number of edges for any face in the grid. Accessed through `grid.n_max_face_edges`. - node_x : np.ndarray - An array of shape (n_nodes,) containing the x-coordinate values of the nodes. Accessed through `grid.node_x`. - node_y : np.ndarray - An array of shape (n_nodes,) containing the y-coordinate values of the nodes. Accessed through `grid.node_y`. - node_z : np.ndarray - An array of shape (n_nodes,) containing the z-coordinate values of the nodes. Accessed through `grid.node_z`. - - Returns - ------- - face_edges_cartesian : np.ndarray - An array of shape (n_face, n_max_face_edges, 2, 3) containing the Cartesian coordinates of the edges - for each face. It might contain dummy values if the grid has holes. - - Examples - -------- - >>> face_node_conn = np.array( - ... [ - ... [0, 1, 2, 3, 4], - ... [0, 1, 3, 4, INT_FILL_VALUE], - ... [0, 1, 3, INT_FILL_VALUE, INT_FILL_VALUE], - ... ] - ... ) - >>> n_face = 3 - >>> n_max_face_edges = 5 - >>> node_x = np.array([0, 1, 1, 0, 1, 0]) - >>> node_y = np.array([0, 0, 1, 1, 2, 2]) - >>> node_z = np.array([0, 0, 0, 0, 1, 1]) - >>> _get_cartesian_face_edge_nodes( - ... face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z - ... ) - array([[[[ 0, 0, 0], - [ 1, 0, 0]], - - [[ 1, 0, 0], - [ 1, 1, 0]], - - [[ 1, 1, 0], - [ 0, 1, 0]], - - [[ 0, 1, 0], - [ 1, 2, 1]], - - [[ 1, 2, 1], - [ 0, 0, 0]]], - - - [[[ 0, 0, 0], - [ 1, 0, 0]], - - [[ 1, 0, 0], - [ 0, 1, 0]], - - [[ 0, 1, 0], - [ 1, 2, 1]], - - [[ 1, 2, 1], - [ 0, 0, 0]], - - [[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE], - [INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]]], - - - [[[ 0, 0, 0], - [ 1, 0, 0]], - - [[ 1, 0, 0], - [ 0, 1, 0]], - - [[ 0, 1, 0], - [ 0, 0, 0]], - - [[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE], - [INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]], - - [[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE], - [INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]]]]) - """ - - # face_edge_connectivity (n_face, n_edge) - - # each edge should have a shape (2, 3) - - # Shift node connections to create edge connections - face_node_conn_shift = np.roll(face_node_conn, -1, axis=1) - - # Construct edge connections by combining original and shifted node connections - face_edge_conn = np.array([face_node_conn, face_node_conn_shift]).T.swapaxes(0, 1) - - # swap the first occurrence of INT_FILL_VALUE with the last value in each sub-array - face_edge_conn = _swap_first_fill_value_with_last(face_edge_conn) - - # Get the indices of the nodes from face_edge_conn - face_edge_conn_flat = face_edge_conn.reshape(-1) - - valid_mask = face_edge_conn_flat != INT_FILL_VALUE - - # Get the valid node indices - valid_edges = face_edge_conn_flat[valid_mask] - - # Create an array to hold the Cartesian coordinates of the edges - face_edges_cartesian = np.full( - (len(face_edge_conn_flat), 3), INT_FILL_VALUE, dtype=float - ) - - # Fill the array with the Cartesian coordinates of the edges - face_edges_cartesian[valid_mask, 0] = node_x[valid_edges] - face_edges_cartesian[valid_mask, 1] = node_y[valid_edges] - face_edges_cartesian[valid_mask, 2] = node_z[valid_edges] - - return face_edges_cartesian.reshape(n_face, n_max_face_edges, 2, 3) - - -def _get_lonlat_rad_face_edge_nodes( - face_node_conn, n_face, n_max_face_edges, node_lon, node_lat -): - """Construct an array to hold the edge latitude and longitude in radians - connectivity for multiple faces in a grid. - - Parameters - ---------- - face_node_conn : np.ndarray - An array of shape (n_face, n_max_face_edges) containing the node indices for each face. Accessed through `grid.face_node_connectivity.value`. - n_face : int - The number of faces in the grid. Accessed through `grid.n_face`. - n_max_face_edges : int - The maximum number of edges for any face in the grid. Accessed through `grid.n_max_face_edges`. - node_lon : np.ndarray - An array of shape (n_nodes,) containing the longitude values of the nodes in degrees. Accessed through `grid.node_lon`. - node_lat : np.ndarray - An array of shape (n_nodes,) containing the latitude values of the nodes in degrees. Accessed through `grid.node_lat`. - - Returns - ------- - face_edges_lonlat_rad : np.ndarray - An array of shape (n_face, n_max_face_edges, 2, 2) containing the latitude and longitude coordinates - in radians for the edges of each face. It might contain dummy values if the grid has holes. - - Notes - ----- - If the grid has holes, the function will return an entry of dummy value faces_edges_coordinates[i] filled with INT_FILL_VALUE. - """ - - # Convert node coordinates to radians - node_lon_rad = np.deg2rad(node_lon) - node_lat_rad = np.deg2rad(node_lat) - - # Shift node connections to create edge connections - face_node_conn_shift = np.roll(face_node_conn, -1, axis=1) - - # Construct edge connections by combining original and shifted node connections - face_edge_conn = np.array([face_node_conn, face_node_conn_shift]).T.swapaxes(0, 1) - - # swap the first occurrence of INT_FILL_VALUE with the last value in each sub-array - face_edge_conn = _swap_first_fill_value_with_last(face_edge_conn) - - # Get the indices of the nodes from face_edge_conn - face_edge_conn_flat = face_edge_conn.reshape(-1) - - valid_mask = face_edge_conn_flat != INT_FILL_VALUE - - # Get the valid node indices - valid_edges = face_edge_conn_flat[valid_mask] - - # Create an array to hold the latitude and longitude in radians for the edges - face_edges_lonlat_rad = np.full( - (len(face_edge_conn_flat), 2), INT_FILL_VALUE, dtype=float - ) - - # Fill the array with the latitude and longitude in radians for the edges - face_edges_lonlat_rad[valid_mask, 0] = node_lon_rad[valid_edges] - face_edges_lonlat_rad[valid_mask, 1] = node_lat_rad[valid_edges] - - return face_edges_lonlat_rad.reshape(n_face, n_max_face_edges, 2, 2) - - -def close_face_nodes(face_node_connectivity, n_face, n_max_face_nodes): - """Closes (``face_node_connectivity``) by inserting the first node index - after the last non-fill-value node. - - Parameters - ---------- - face_node_connectivity : np.ndarray - Connectivity array for constructing a face from its nodes - n_face : constant - Number of faces - n_max_face_nodes : constant - Max number of nodes that compose a face - - Returns - ---------- - closed : ndarray - Closed (padded) face_node_connectivity - - Example - ---------- - Given face nodes with shape [2 x 5] - [0, 1, 2, 3, FILL_VALUE] - [4, 5, 6, 7, 8] - Pads them to the following with shape [2 x 6] - [0, 1, 2, 3, 0, FILL_VALUE] - [4, 5, 6, 7, 8, 4] - """ - - # padding to shape [n_face, n_max_face_nodes + 1] - closed = np.ones((n_face, n_max_face_nodes + 1), dtype=INT_DTYPE) * INT_FILL_VALUE - - # set all non-paded values to original face nodee values - closed[:, :-1] = face_node_connectivity.copy() - - # instance of first fill value - first_fv_idx_2d = np.argmax(closed == INT_FILL_VALUE, axis=1) - - # 2d to 1d index for np.put() - first_fv_idx_1d = first_fv_idx_2d + ((n_max_face_nodes + 1) * np.arange(0, n_face)) - - # column of first node values - first_node_value = face_node_connectivity[:, 0].copy() - - # insert first node column at occurrence of first fill value - np.put(closed.ravel(), first_fv_idx_1d, first_node_value) - - return closed - - def _replace_fill_values(grid_var, original_fill, new_fill, new_dtype=None): """Replaces all instances of the current fill value (``original_fill``) in (``grid_var``) with (``new_fill``) and converts to the dtype defined by From ffdd4bb2ba8d2c3a09814a9221350c3bec219723 Mon Sep 17 00:00:00 2001 From: Philip Chmielowiec <67855069+philipc2@users.noreply.github.com> Date: Mon, 7 Apr 2025 22:33:29 -0500 Subject: [PATCH 06/19] add geometry module --- uxarray/geometry/__init__.py | 0 uxarray/geometry/face_edges.py | 193 +++++++++++++++++++++++++++++++++ 2 files changed, 193 insertions(+) create mode 100644 uxarray/geometry/__init__.py create mode 100644 uxarray/geometry/face_edges.py diff --git a/uxarray/geometry/__init__.py b/uxarray/geometry/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/uxarray/geometry/face_edges.py b/uxarray/geometry/face_edges.py new file mode 100644 index 000000000..847137d9b --- /dev/null +++ b/uxarray/geometry/face_edges.py @@ -0,0 +1,193 @@ +import numpy as np + +from uxarray.constants import INT_FILL_VALUE +from uxarray.grid.utils import _swap_first_fill_value_with_last + + +def _construct_face_edge_nodes_cartesian( + face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z +): + """Construct an array to hold the edge Cartesian coordinates connectivity + for multiple faces in a grid. + + Parameters + ---------- + face_node_conn : np.ndarray + An array of shape (n_face, n_max_face_edges) containing the node indices for each face. Accessed through `grid.face_node_connectivity.value`. + n_face : int + The number of faces in the grid. Accessed through `grid.n_face`. + n_max_face_edges : int + The maximum number of edges for any face in the grid. Accessed through `grid.n_max_face_edges`. + node_x : np.ndarray + An array of shape (n_nodes,) containing the x-coordinate values of the nodes. Accessed through `grid.node_x`. + node_y : np.ndarray + An array of shape (n_nodes,) containing the y-coordinate values of the nodes. Accessed through `grid.node_y`. + node_z : np.ndarray + An array of shape (n_nodes,) containing the z-coordinate values of the nodes. Accessed through `grid.node_z`. + + Returns + ------- + face_edges_cartesian : np.ndarray + An array of shape (n_face, n_max_face_edges, 2, 3) containing the Cartesian coordinates of the edges + for each face. It might contain dummy values if the grid has holes. + + Examples + -------- + >>> face_node_conn = np.array( + ... [ + ... [0, 1, 2, 3, 4], + ... [0, 1, 3, 4, INT_FILL_VALUE], + ... [0, 1, 3, INT_FILL_VALUE, INT_FILL_VALUE], + ... ] + ... ) + >>> n_face = 3 + >>> n_max_face_edges = 5 + >>> node_x = np.array([0, 1, 1, 0, 1, 0]) + >>> node_y = np.array([0, 0, 1, 1, 2, 2]) + >>> node_z = np.array([0, 0, 0, 0, 1, 1]) + >>> _get_cartesian_face_edge_nodes( + ... face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z + ... ) + array([[[[ 0, 0, 0], + [ 1, 0, 0]], + + [[ 1, 0, 0], + [ 1, 1, 0]], + + [[ 1, 1, 0], + [ 0, 1, 0]], + + [[ 0, 1, 0], + [ 1, 2, 1]], + + [[ 1, 2, 1], + [ 0, 0, 0]]], + + + [[[ 0, 0, 0], + [ 1, 0, 0]], + + [[ 1, 0, 0], + [ 0, 1, 0]], + + [[ 0, 1, 0], + [ 1, 2, 1]], + + [[ 1, 2, 1], + [ 0, 0, 0]], + + [[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE], + [INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]]], + + + [[[ 0, 0, 0], + [ 1, 0, 0]], + + [[ 1, 0, 0], + [ 0, 1, 0]], + + [[ 0, 1, 0], + [ 0, 0, 0]], + + [[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE], + [INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]], + + [[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE], + [INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]]]]) + """ + + # face_edge_connectivity (n_face, n_edge) + + # each edge should have a shape (2, 3) + + # Shift node connections to create edge connections + face_node_conn_shift = np.roll(face_node_conn, -1, axis=1) + + # Construct edge connections by combining original and shifted node connections + face_edge_conn = np.array([face_node_conn, face_node_conn_shift]).T.swapaxes(0, 1) + + # swap the first occurrence of INT_FILL_VALUE with the last value in each sub-array + face_edge_conn = _swap_first_fill_value_with_last(face_edge_conn) + + # Get the indices of the nodes from face_edge_conn + face_edge_conn_flat = face_edge_conn.reshape(-1) + + valid_mask = face_edge_conn_flat != INT_FILL_VALUE + + # Get the valid node indices + valid_edges = face_edge_conn_flat[valid_mask] + + # Create an array to hold the Cartesian coordinates of the edges + face_edges_cartesian = np.full( + (len(face_edge_conn_flat), 3), INT_FILL_VALUE, dtype=float + ) + + # Fill the array with the Cartesian coordinates of the edges + face_edges_cartesian[valid_mask, 0] = node_x[valid_edges] + face_edges_cartesian[valid_mask, 1] = node_y[valid_edges] + face_edges_cartesian[valid_mask, 2] = node_z[valid_edges] + + return face_edges_cartesian.reshape(n_face, n_max_face_edges, 2, 3) + + +def _construct_face_edge_nodes_spherical( + face_node_conn, n_face, n_max_face_edges, node_lon, node_lat +): + """Construct an array to hold the edge latitude and longitude in radians + connectivity for multiple faces in a grid. + + Parameters + ---------- + face_node_conn : np.ndarray + An array of shape (n_face, n_max_face_edges) containing the node indices for each face. Accessed through `grid.face_node_connectivity.value`. + n_face : int + The number of faces in the grid. Accessed through `grid.n_face`. + n_max_face_edges : int + The maximum number of edges for any face in the grid. Accessed through `grid.n_max_face_edges`. + node_lon : np.ndarray + An array of shape (n_nodes,) containing the longitude values of the nodes in degrees. Accessed through `grid.node_lon`. + node_lat : np.ndarray + An array of shape (n_nodes,) containing the latitude values of the nodes in degrees. Accessed through `grid.node_lat`. + + Returns + ------- + face_edges_lonlat_rad : np.ndarray + An array of shape (n_face, n_max_face_edges, 2, 2) containing the latitude and longitude coordinates + in radians for the edges of each face. It might contain dummy values if the grid has holes. + + Notes + ----- + If the grid has holes, the function will return an entry of dummy value faces_edges_coordinates[i] filled with INT_FILL_VALUE. + """ + + # Convert node coordinates to radians + node_lon_rad = np.deg2rad(node_lon) + node_lat_rad = np.deg2rad(node_lat) + + # Shift node connections to create edge connections + face_node_conn_shift = np.roll(face_node_conn, -1, axis=1) + + # Construct edge connections by combining original and shifted node connections + face_edge_conn = np.array([face_node_conn, face_node_conn_shift]).T.swapaxes(0, 1) + + # swap the first occurrence of INT_FILL_VALUE with the last value in each sub-array + face_edge_conn = _swap_first_fill_value_with_last(face_edge_conn) + + # Get the indices of the nodes from face_edge_conn + face_edge_conn_flat = face_edge_conn.reshape(-1) + + valid_mask = face_edge_conn_flat != INT_FILL_VALUE + + # Get the valid node indices + valid_edges = face_edge_conn_flat[valid_mask] + + # Create an array to hold the latitude and longitude in radians for the edges + face_edges_lonlat_rad = np.full( + (len(face_edge_conn_flat), 2), INT_FILL_VALUE, dtype=float + ) + + # Fill the array with the latitude and longitude in radians for the edges + face_edges_lonlat_rad[valid_mask, 0] = node_lon_rad[valid_edges] + face_edges_lonlat_rad[valid_mask, 1] = node_lat_rad[valid_edges] + + return face_edges_lonlat_rad.reshape(n_face, n_max_face_edges, 2, 2) From 19bed33678f164244212d7e3644410df44cdbfe1 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 10 Jul 2026 16:14:17 -0500 Subject: [PATCH 07/19] Rework centroid triangle test --- test/grid/geometry/test_centroids.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) 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.""" From b69b47f0cd26dd9d5638040e92ab8588d5084e68 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 17 Jul 2026 15:56:13 -0500 Subject: [PATCH 08/19] Add n_nodes_per_face benchmark --- benchmarks/bench_connectivity.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/benchmarks/bench_connectivity.py b/benchmarks/bench_connectivity.py index de1bf3895..f4c05c637 100644 --- a/benchmarks/bench_connectivity.py +++ b/benchmarks/bench_connectivity.py @@ -54,6 +54,9 @@ def teardown(self, resolution, *args, **kwargs): class Connectivity(GridBenchmark): + def time_n_nodes_per_face(self, resolution): + _ = self.uxgrid.n_nodes_per_face + def time_face_node(self, resolution): _ = self.uxgrid.face_node_connectivity From b00a124734052bfe7cdfa72e14f085a16619ac12 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 17 Jul 2026 16:11:16 -0500 Subject: [PATCH 09/19] lazy n_nodes_per_face --- uxarray/grid/connectivity.py | 20 +------------------- 1 file changed, 1 insertion(+), 19 deletions(-) diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index f50d9b6c1..0230e2a9d 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -125,9 +125,7 @@ def _populate_n_nodes_per_face(grid): it within the internal (``Grid._ds``) and through the attribute (``Grid.n_nodes_per_face``).""" - n_nodes_per_face = _build_n_nodes_per_face( - grid.face_node_connectivity.values, grid.n_face, grid.n_max_face_nodes - ) + n_nodes_per_face = (grid.face_node_connectivity != INT_FILL_VALUE).sum(axis=1) if n_nodes_per_face.ndim == 0: # convert scalar value into a [1, 1] array @@ -141,22 +139,6 @@ def _populate_n_nodes_per_face(grid): ) -@njit(cache=True) -def _build_n_nodes_per_face(face_nodes, n_face, n_max_face_nodes): - """Constructs ``n_nodes_per_face``, which contains the number of non-fill- - value nodes for each face in ``face_node_connectivity``""" - - n_face, n_max_face_nodes = face_nodes.shape - n_nodes_per_face = np.empty(n_face, dtype=INT_DTYPE) - for i in range(n_face): - c = 0 - for j in range(n_max_face_nodes): - if face_nodes[i, j] != INT_FILL_VALUE: - c += 1 - n_nodes_per_face[i] = c - return n_nodes_per_face - - def _populate_edge_node_connectivity(grid): """Constructs the UGRID connectivity variable (``edge_node_connectivity``) and stores it within the internal (``Grid._ds``) and through the attribute From bd1498585255cba92cb4f1ba76252f1eb2a61601 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Tue, 28 Jul 2026 18:39:07 -0500 Subject: [PATCH 10/19] Restore canonical edge ordering in _build_edge_node_connectivity The optimized edge builder deduped half edges with a numba hash map, which numbered edges in first-encounter order. Edges had previously been numbered lexicographically by their (min_node, max_node) pair, as a side effect of the np.unique(..., axis=0) the hash map replaced. Global edge index is a public identity: it indexes edge_lon/edge_lat, edge centered data variables, and edge_node_distances, so renumbering silently re-pairs user data with different physical edges. It also broke the five TestQuadHexagon connectivity tests, which assert on edge_node, face_edge, node_edge, edge_face and face_face -- all the same renumbering cascading through the derived connectivities. Sort as the dedup mechanism instead of hashing. Node indices are dense integers in [0, n_node), so a counting sort buckets the half edges by their first node without any comparisons, and sorting each bucket by its second node leaves the duplicates adjacent -- the dedup then falls out of the same walk. Buckets hold one entry per edge incident to a node, so on a real mesh they are tiny (node degree, typically under ten) and an insertion sort finishes them. A bucket above MAX_INSERTION_SORT_SIZE is heap sorted so that a degenerate mesh cannot degrade the build quadratically; np.argsort is deliberately not used there, as numba's implementation degrades badly on structured input. Half edges are identified throughout by their flat face_node_connectivity index, which is also the face_edge_connectivity slot they are written back to, so the sort needs a single permutation array and no mapping back. This is faster and leaner than the hash map it replaces. On a synthetic one million face quad mesh, measured by peak RSS rather than tracemalloc, which does not observe numba's typed dict allocations: dict build 397.3 ms 239.5 MB bucket sort 85.1 ms 91.6 MB The five legacy tests now pass unchanged. Adds order invariant coverage for the canonical ordering and the face_edge positional contract, plus high degree nodes either side of the insertion sort threshold. Also casts n_nodes_per_face back to INT_DTYPE, so that the builder is not compiled a second time for int64, and restores a blank line dropped between two top level functions. Co-Authored-By: Claude Opus 5 --- test/grid/grid/test_connectivity.py | 72 ++++++++- uxarray/grid/connectivity.py | 225 +++++++++++++++++++++++++--- 2 files changed, 278 insertions(+), 19 deletions(-) diff --git a/test/grid/grid/test_connectivity.py b/test/grid/grid/test_connectivity.py index 7ae6087a5..4faa82eb3 100644 --- a/test/grid/grid/test_connectivity.py +++ b/test/grid/grid/test_connectivity.py @@ -3,10 +3,10 @@ import pytest import uxarray as ux -from uxarray.constants import INT_FILL_VALUE, ERROR_TOLERANCE +from uxarray.constants import INT_DTYPE, INT_FILL_VALUE, ERROR_TOLERANCE from uxarray.grid.connectivity import (_populate_face_edge_connectivity, _build_edge_face_connectivity, _build_edge_node_connectivity, _build_face_face_connectivity, - _populate_face_face_connectivity) + _populate_face_face_connectivity, MAX_INSERTION_SORT_SIZE) def test_connectivity_build_n_nodes_per_face(gridpath): @@ -63,6 +63,74 @@ def test_connectivity_build_face_edges_connectivity(gridpath): assert np.all(valid_edges >= 0) assert np.all(valid_edges < uxgrid.n_edge) +@pytest.mark.parametrize("grid_parts", [("ugrid", "outCSne30", "outCSne30.ug"), + ("ugrid", "quad-hexagon", "grid.nc"), + ("ugrid", "geoflow-small", "grid.nc")]) +def test_connectivity_edge_node_canonical_order(gridpath, grid_parts): + """Test that constructed edges are numbered in lexicographic node order.""" + uxgrid = ux.open_grid(gridpath(*grid_parts)) + edge_nodes = uxgrid.edge_node_connectivity.values + + # Each edge is stored as an ascending node pair + assert np.all(edge_nodes[:, 0] < edge_nodes[:, 1]) + + # Edges are numbered lexicographically by that pair, with no duplicates + lexicographic_order = np.lexsort((edge_nodes[:, 1], edge_nodes[:, 0])) + nt.assert_array_equal(lexicographic_order, np.arange(uxgrid.n_edge)) + assert len(np.unique(edge_nodes, axis=0)) == uxgrid.n_edge + +@pytest.mark.parametrize("n_spoke", [MAX_INSERTION_SORT_SIZE - 1, MAX_INSERTION_SORT_SIZE + 1, 500]) +def test_connectivity_edge_node_high_degree_node(n_spoke): + """Test edge construction for a node shared by more faces than the bucket sort will + insertion sort, which takes the heap sort path.""" + # A fan of triangles around node 0, with the spokes numbered in descending order so + # that they reach the sort already reversed + spokes = np.arange(n_spoke, 0, -1, dtype=INT_DTYPE) + face_node_connectivity = np.stack( + [np.zeros(n_spoke, dtype=INT_DTYPE), spokes, np.roll(spokes, -1)], axis=1 + ) + + edge_nodes, face_edges = _build_edge_node_connectivity( + face_node_connectivity, np.full(n_spoke, 3, dtype=INT_DTYPE), n_spoke + 1 + ) + + # Same invariants as any other mesh: ascending pairs, lexicographic numbering + assert np.all(edge_nodes[:, 0] < edge_nodes[:, 1]) + nt.assert_array_equal( + np.lexsort((edge_nodes[:, 1], edge_nodes[:, 0])), np.arange(len(edge_nodes)) + ) + assert len(np.unique(edge_nodes, axis=0)) == len(edge_nodes) + + # The hub is shared by every face, so it has one edge per spoke + assert np.count_nonzero(edge_nodes == 0) == n_spoke + + # And face_edge_connectivity still points at the right node pairs + for face_idx in range(n_spoke): + for cur in range(3): + expected = sorted((face_node_connectivity[face_idx, cur], + face_node_connectivity[face_idx, (cur + 1) % 3])) + assert sorted(edge_nodes[face_edges[face_idx, cur]]) == expected + +def test_connectivity_face_edge_positional_alignment(gridpath): + """Test that face_edge_connectivity[i, j] is the edge between face nodes j and j+1.""" + uxgrid = ux.open_grid(gridpath("ugrid", "outCSne30", "outCSne30.ug")) + + face_nodes = uxgrid.face_node_connectivity.values + face_edges = uxgrid.face_edge_connectivity.values + edge_nodes = uxgrid.edge_node_connectivity.values + + for face_idx, n_edges in enumerate(uxgrid.n_nodes_per_face.values): + for cur in range(n_edges): + start_node = face_nodes[face_idx, cur] + end_node = face_nodes[face_idx, (cur + 1) % n_edges] + + expected = sorted((start_node, end_node)) + actual = sorted(edge_nodes[face_edges[face_idx, cur]]) + assert actual == expected + + # Remaining slots stay padded + assert np.all(face_edges[face_idx, n_edges:] == INT_FILL_VALUE) + def test_connectivity_build_face_edges_connectivity_fillvalues(): """Test face-edge connectivity with fill values.""" # Create a simple grid with mixed face types diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index 0230e2a9d..ad99930e4 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -5,6 +5,10 @@ from uxarray.constants import INT_DTYPE, INT_FILL_VALUE from uxarray.conventions import ugrid +# Largest number of edges incident to a single node that ``_build_edge_node_connectivity`` +# will insertion sort before falling back to a comparison sort +MAX_INSERTION_SORT_SIZE = 64 + def close_face_nodes(face_node_connectivity, n_face, n_max_face_nodes): """Closes (``face_node_connectivity``) by inserting the first node index @@ -125,7 +129,9 @@ def _populate_n_nodes_per_face(grid): it within the internal (``Grid._ds``) and through the attribute (``Grid.n_nodes_per_face``).""" - n_nodes_per_face = (grid.face_node_connectivity != INT_FILL_VALUE).sum(axis=1) + n_nodes_per_face = ( + (grid.face_node_connectivity != INT_FILL_VALUE).sum(axis=1).astype(INT_DTYPE) + ) if n_nodes_per_face.ndim == 0: # convert scalar value into a [1, 1] array @@ -151,7 +157,7 @@ def _populate_edge_node_connectivity(grid): pass edge_nodes, face_edges = _build_edge_node_connectivity( - grid.face_node_connectivity.values, grid.n_nodes_per_face.values + grid.face_node_connectivity.values, grid.n_nodes_per_face.values, grid.n_node ) grid._ds["edge_node_connectivity"] = xr.DataArray( @@ -167,12 +173,102 @@ def _populate_edge_node_connectivity(grid): ) +@njit(cache=True, inline="always") +def _other_node(face_node_connectivity, n_nodes_per_face, n_max_face_nodes, flat_idx): + """Returns the larger of the two nodes of the half edge stored at a flat + ``face_node_connectivity`` index.""" + face_idx = flat_idx // n_max_face_nodes + current_node = flat_idx % n_max_face_nodes + n_edges = n_nodes_per_face[face_idx] + + start_node = face_node_connectivity[face_idx, current_node] + end_node = face_node_connectivity[face_idx, (current_node + 1) % n_edges] + + return max(start_node, end_node) + + +@njit(cache=True) +def _sift_down( + order, bucket_start, root, size, face_node_connectivity, n_nodes_per_face, n_max +): + """Restores the max-heap property at ``root`` for a bucket keyed on ``node_b``.""" + while True: + child = 2 * root + 1 + if child >= size: + break + + if child + 1 < size and _other_node( + face_node_connectivity, n_nodes_per_face, n_max, order[bucket_start + child] + ) < _other_node( + face_node_connectivity, + n_nodes_per_face, + n_max, + order[bucket_start + child + 1], + ): + child += 1 + + if _other_node( + face_node_connectivity, n_nodes_per_face, n_max, order[bucket_start + root] + ) >= _other_node( + face_node_connectivity, n_nodes_per_face, n_max, order[bucket_start + child] + ): + break + + tmp = order[bucket_start + root] + order[bucket_start + root] = order[bucket_start + child] + order[bucket_start + child] = tmp + root = child + + +@njit(cache=True) +def _heap_sort_bucket( + order, bucket_start, size, face_node_connectivity, n_nodes_per_face, n_max +): + """Sorts a bucket by ``node_b`` in place, in ``O(size * log(size))`` and without + scratch space, for the rare bucket too large to insertion sort.""" + for root in range(size // 2 - 1, -1, -1): + _sift_down( + order, + bucket_start, + root, + size, + face_node_connectivity, + n_nodes_per_face, + n_max, + ) + + for end in range(size - 1, 0, -1): + tmp = order[bucket_start] + order[bucket_start] = order[bucket_start + end] + order[bucket_start + end] = tmp + _sift_down( + order, + bucket_start, + 0, + end, + face_node_connectivity, + n_nodes_per_face, + n_max, + ) + + @njit(cache=True) -def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face): +def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_node): """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. + Each edge is stored as an ascending ``(node_a, node_b)`` pair, and the edges are numbered in lexicographic + order of that pair. Since node indices are dense integers in ``[0, n_node)``, that ordering is obtained + with a counting sort that buckets the half edges by ``node_a``, after which each bucket is sorted by + ``node_b``. Buckets hold one entry per edge incident to a node, so for a real mesh they are tiny (node + degree, typically under ten) and an insertion sort is the cheapest way to finish them; the rare bucket + above ``MAX_INSERTION_SORT_SIZE`` is heap sorted instead so that a degenerate mesh cannot degrade the + build quadratically. Sorting also groups the duplicate half edges, so the dedup falls out of the same walk. + + Half edges are identified throughout by their flat ``face_node_connectivity`` index, which is also the + ``face_edge_connectivity`` slot they are written back to. That keeps the scratch space to a single + permutation array, and node pairs are recomputed on demand rather than materialized. Parameters ---------- @@ -180,6 +276,8 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face): Face Node Connectivity n_nodes_per_face : np.ndarray Number of nodes/edges per face + n_node : int + Total number of nodes, used as the number of buckets for the counting sort Returns ------- @@ -190,32 +288,124 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face): """ - # Dictionary to keep track of unique edges - unique_edge_dict = {} - - edge_idx = 0 + n_face, n_max_face_nodes = face_node_connectivity.shape # Keep track of face_edge_connectivity face_edge_connectivity = np.full_like( face_node_connectivity, INT_FILL_VALUE, dtype=INT_DTYPE ) - 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] + n_half_edge = 0 + for i in range(n_face): + n_half_edge += n_nodes_per_face[i] - edge = (min(start_node, end_node), max(start_node, end_node)) + if n_half_edge == 0: + return np.empty((0, 2), dtype=INT_DTYPE), face_edge_connectivity - if edge not in unique_edge_dict: + # Count how many half edges fall into each ``node_a`` bucket, then prefix sum so that + # ``bucket_bounds[a]`` is where bucket ``a`` starts + bucket_bounds = np.zeros(n_node + 1, dtype=INT_DTYPE) + for face_idx in range(n_face): + n_edges = n_nodes_per_face[face_idx] + for current_node in range(n_edges): + start_node = face_node_connectivity[face_idx, current_node] + end_node = face_node_connectivity[face_idx, (current_node + 1) % n_edges] + bucket_bounds[min(start_node, end_node) + 1] += 1 + for i in range(n_node): + bucket_bounds[i + 1] += bucket_bounds[i] + + # Scatter the half edges into their buckets. This advances each entry of + # ``bucket_bounds`` to the *end* of its bucket, so afterwards bucket ``a`` spans + # ``bucket_bounds[a - 1]`` up to ``bucket_bounds[a]``, with bucket 0 starting at 0 + order = np.empty(n_half_edge, dtype=INT_DTYPE) + for face_idx in range(n_face): + n_edges = n_nodes_per_face[face_idx] + for current_node in range(n_edges): + start_node = face_node_connectivity[face_idx, current_node] + end_node = face_node_connectivity[face_idx, (current_node + 1) % n_edges] + + node_a = min(start_node, end_node) + order[bucket_bounds[node_a]] = face_idx * n_max_face_nodes + current_node + bucket_bounds[node_a] += 1 + + # Sort each bucket by ``node_b`` and count the unique edges while the bucket is in + # cache, which gives the exact allocation size for the walk below + n_edge = 0 + bucket_start = 0 + for node_a in range(n_node): + bucket_end = bucket_bounds[node_a] + + if bucket_end - bucket_start > MAX_INSERTION_SORT_SIZE: + # A node this well connected is not expected of a real mesh, but insertion + # sort degrades quadratically, so fall back to a heap sort + _heap_sort_bucket( + order, + bucket_start, + bucket_end - bucket_start, + face_node_connectivity, + n_nodes_per_face, + n_max_face_nodes, + ) + else: + for i in range(bucket_start + 1, bucket_end): + flat_idx = order[i] + node_b = _other_node( + face_node_connectivity, n_nodes_per_face, n_max_face_nodes, flat_idx + ) + + j = i - 1 + while j >= bucket_start and ( + _other_node( + face_node_connectivity, + n_nodes_per_face, + n_max_face_nodes, + order[j], + ) + > node_b + ): + order[j + 1] = order[j] + j -= 1 + order[j + 1] = flat_idx + + prev_b = INT_FILL_VALUE + for i in range(bucket_start, bucket_end): + node_b = _other_node( + face_node_connectivity, n_nodes_per_face, n_max_face_nodes, order[i] + ) + if node_b != prev_b: + n_edge += 1 + prev_b = node_b + + bucket_start = bucket_end + + # Duplicate half edges are now adjacent, so a single walk assigns each unique edge its + # index and populates the face edge connectivity + edge_node_connectivity = np.empty((n_edge, 2), dtype=INT_DTYPE) + edge_idx = -1 + bucket_start = 0 + + for node_a in range(n_node): + bucket_end = bucket_bounds[node_a] + prev_b = INT_FILL_VALUE + + for i in range(bucket_start, bucket_end): + flat_idx = order[i] + node_b = _other_node( + face_node_connectivity, n_nodes_per_face, n_max_face_nodes, flat_idx + ) + + if node_b != prev_b: # Only store unique edges - unique_edge_dict[edge] = edge_idx edge_idx += 1 + edge_node_connectivity[edge_idx, 0] = node_a + edge_node_connectivity[edge_idx, 1] = node_b + prev_b = node_b - face_edge_connectivity[i, current_node] = unique_edge_dict[edge] + face_edge_connectivity[ + flat_idx // n_max_face_nodes, flat_idx % n_max_face_nodes + ] = edge_idx - # 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) + bucket_start = bucket_end return edge_node_connectivity, face_edge_connectivity @@ -414,6 +604,7 @@ def _build_face_face_connectivity(edge_face_connectivity, n_face, n_max_face_nod 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 From e7910dc1eae57600474b8112a111379cc70a2527 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 29 Jul 2026 13:26:34 -0500 Subject: [PATCH 11/19] Sorting algo cleanup --- test/grid/grid/test_connectivity.py | 67 ++++++----- uxarray/grid/connectivity.py | 175 +++++----------------------- uxarray/grid/utils.py | 104 +++++++++++++++++ 3 files changed, 175 insertions(+), 171 deletions(-) diff --git a/test/grid/grid/test_connectivity.py b/test/grid/grid/test_connectivity.py index 4faa82eb3..afaccd5b2 100644 --- a/test/grid/grid/test_connectivity.py +++ b/test/grid/grid/test_connectivity.py @@ -6,7 +6,9 @@ from uxarray.constants import INT_DTYPE, INT_FILL_VALUE, ERROR_TOLERANCE from uxarray.grid.connectivity import (_populate_face_edge_connectivity, _build_edge_face_connectivity, _build_edge_node_connectivity, _build_face_face_connectivity, - _populate_face_face_connectivity, MAX_INSERTION_SORT_SIZE) + _populate_face_face_connectivity) +from uxarray.grid.utils import (_adaptive_sort_bucket, _insertion_sort_bucket, + MIN_ADAPTIVE_SORT_SIZE) def test_connectivity_build_n_nodes_per_face(gridpath): @@ -79,37 +81,46 @@ def test_connectivity_edge_node_canonical_order(gridpath, grid_parts): nt.assert_array_equal(lexicographic_order, np.arange(uxgrid.n_edge)) assert len(np.unique(edge_nodes, axis=0)) == uxgrid.n_edge -@pytest.mark.parametrize("n_spoke", [MAX_INSERTION_SORT_SIZE - 1, MAX_INSERTION_SORT_SIZE + 1, 500]) -def test_connectivity_edge_node_high_degree_node(n_spoke): - """Test edge construction for a node shared by more faces than the bucket sort will - insertion sort, which takes the heap sort path.""" - # A fan of triangles around node 0, with the spokes numbered in descending order so - # that they reach the sort already reversed - spokes = np.arange(n_spoke, 0, -1, dtype=INT_DTYPE) - face_node_connectivity = np.stack( - [np.zeros(n_spoke, dtype=INT_DTYPE), spokes, np.roll(spokes, -1)], axis=1 - ) +@pytest.mark.parametrize("sort", [_insertion_sort_bucket, _adaptive_sort_bucket], + ids=["insertion", "adaptive"]) +def test_connectivity_bucket_sort(sort): + """Test that each bucket sort orders its own slice and nothing else. - edge_nodes, face_edges = _build_edge_node_connectivity( - face_node_connectivity, np.full(n_spoke, 3, dtype=INT_DTYPE), n_spoke + 1 - ) + The bucket sizes straddle ``MIN_ADAPTIVE_SORT_SIZE``: the small ones cannot accumulate + enough shifts to exhaust the budget, so the metered sort stays on its insertion path, + while the 500 element bucket is shuffled far past the budget and falls back to the heap + sort. Keys repeat, since an interior edge reaches its bucket once per adjacent face. + """ + rng = np.random.default_rng(0) - # Same invariants as any other mesh: ascending pairs, lexicographic numbering - assert np.all(edge_nodes[:, 0] < edge_nodes[:, 1]) - nt.assert_array_equal( - np.lexsort((edge_nodes[:, 1], edge_nodes[:, 0])), np.arange(len(edge_nodes)) - ) - assert len(np.unique(edge_nodes, axis=0)) == len(edge_nodes) + sizes = [5, MIN_ADAPTIVE_SORT_SIZE, MIN_ADAPTIVE_SORT_SIZE + 1, 500] + bounds = np.cumsum([0] + sizes) + n_half_edge = int(bounds[-1]) + buckets = list(zip(bounds[:-1], bounds[1:])) + + keys = rng.integers(0, 40, n_half_edge).astype(INT_DTYPE) + order = rng.permutation(n_half_edge).astype(INT_DTYPE) + + # the key each half edge must still be paired with once the permutation has moved it + key_for = np.empty(n_half_edge, dtype=INT_DTYPE) + key_for[order] = keys + + expected_keys = np.concatenate([np.sort(keys[start:end]) for start, end in buckets]) + + got_keys, got_order = keys.copy(), order.copy() + for start, end in buckets: + shuffle = rng.permutation(end - start) + got_keys[start:end] = got_keys[start:end][shuffle] + got_order[start:end] = got_order[start:end][shuffle] + + sort(got_keys, got_order, start, end - start) + + nt.assert_array_equal(got_keys, expected_keys) - # The hub is shared by every face, so it has one edge per spoke - assert np.count_nonzero(edge_nodes == 0) == n_spoke + # sorted keys alone would pass even if the permutation had been scrambled independently + nt.assert_array_equal(key_for[got_order], got_keys) + nt.assert_array_equal(np.sort(got_order), np.arange(n_half_edge)) - # And face_edge_connectivity still points at the right node pairs - for face_idx in range(n_spoke): - for cur in range(3): - expected = sorted((face_node_connectivity[face_idx, cur], - face_node_connectivity[face_idx, (cur + 1) % 3])) - assert sorted(edge_nodes[face_edges[face_idx, cur]]) == expected def test_connectivity_face_edge_positional_alignment(gridpath): """Test that face_edge_connectivity[i, j] is the edge between face nodes j and j+1.""" diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index ad99930e4..2b192a877 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -4,10 +4,11 @@ from uxarray.constants import INT_DTYPE, INT_FILL_VALUE from uxarray.conventions import ugrid - -# Largest number of edges incident to a single node that ``_build_edge_node_connectivity`` -# will insertion sort before falling back to a comparison sort -MAX_INSERTION_SORT_SIZE = 64 +from uxarray.grid.utils import ( + MIN_ADAPTIVE_SORT_SIZE, + _adaptive_sort_bucket, + _insertion_sort_bucket, +) def close_face_nodes(face_node_connectivity, n_face, n_max_face_nodes): @@ -173,85 +174,6 @@ def _populate_edge_node_connectivity(grid): ) -@njit(cache=True, inline="always") -def _other_node(face_node_connectivity, n_nodes_per_face, n_max_face_nodes, flat_idx): - """Returns the larger of the two nodes of the half edge stored at a flat - ``face_node_connectivity`` index.""" - face_idx = flat_idx // n_max_face_nodes - current_node = flat_idx % n_max_face_nodes - n_edges = n_nodes_per_face[face_idx] - - start_node = face_node_connectivity[face_idx, current_node] - end_node = face_node_connectivity[face_idx, (current_node + 1) % n_edges] - - return max(start_node, end_node) - - -@njit(cache=True) -def _sift_down( - order, bucket_start, root, size, face_node_connectivity, n_nodes_per_face, n_max -): - """Restores the max-heap property at ``root`` for a bucket keyed on ``node_b``.""" - while True: - child = 2 * root + 1 - if child >= size: - break - - if child + 1 < size and _other_node( - face_node_connectivity, n_nodes_per_face, n_max, order[bucket_start + child] - ) < _other_node( - face_node_connectivity, - n_nodes_per_face, - n_max, - order[bucket_start + child + 1], - ): - child += 1 - - if _other_node( - face_node_connectivity, n_nodes_per_face, n_max, order[bucket_start + root] - ) >= _other_node( - face_node_connectivity, n_nodes_per_face, n_max, order[bucket_start + child] - ): - break - - tmp = order[bucket_start + root] - order[bucket_start + root] = order[bucket_start + child] - order[bucket_start + child] = tmp - root = child - - -@njit(cache=True) -def _heap_sort_bucket( - order, bucket_start, size, face_node_connectivity, n_nodes_per_face, n_max -): - """Sorts a bucket by ``node_b`` in place, in ``O(size * log(size))`` and without - scratch space, for the rare bucket too large to insertion sort.""" - for root in range(size // 2 - 1, -1, -1): - _sift_down( - order, - bucket_start, - root, - size, - face_node_connectivity, - n_nodes_per_face, - n_max, - ) - - for end in range(size - 1, 0, -1): - tmp = order[bucket_start] - order[bucket_start] = order[bucket_start + end] - order[bucket_start + end] = tmp - _sift_down( - order, - bucket_start, - 0, - end, - face_node_connectivity, - n_nodes_per_face, - n_max, - ) - - @njit(cache=True) def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_node): """Constructs the ``edge_node_connectivity`` variable, which represents the indices of the two nodes that make up @@ -259,16 +181,7 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_no indices of the edges that make up each face. Each edge is stored as an ascending ``(node_a, node_b)`` pair, and the edges are numbered in lexicographic - order of that pair. Since node indices are dense integers in ``[0, n_node)``, that ordering is obtained - with a counting sort that buckets the half edges by ``node_a``, after which each bucket is sorted by - ``node_b``. Buckets hold one entry per edge incident to a node, so for a real mesh they are tiny (node - degree, typically under ten) and an insertion sort is the cheapest way to finish them; the rare bucket - above ``MAX_INSERTION_SORT_SIZE`` is heap sorted instead so that a degenerate mesh cannot degrade the - build quadratically. Sorting also groups the duplicate half edges, so the dedup falls out of the same walk. - - Half edges are identified throughout by their flat ``face_node_connectivity`` index, which is also the - ``face_edge_connectivity`` slot they are written back to. That keeps the scratch space to a single - permutation array, and node pairs are recomputed on demand rather than materialized. + order of that pair. Parameters ---------- @@ -318,63 +231,41 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_no # ``bucket_bounds`` to the *end* of its bucket, so afterwards bucket ``a`` spans # ``bucket_bounds[a - 1]`` up to ``bucket_bounds[a]``, with bucket 0 starting at 0 order = np.empty(n_half_edge, dtype=INT_DTYPE) + end_node_keys = np.empty(n_half_edge, dtype=INT_DTYPE) for face_idx in range(n_face): n_edges = n_nodes_per_face[face_idx] for current_node in range(n_edges): start_node = face_node_connectivity[face_idx, current_node] end_node = face_node_connectivity[face_idx, (current_node + 1) % n_edges] - node_a = min(start_node, end_node) - order[bucket_bounds[node_a]] = face_idx * n_max_face_nodes + current_node - bucket_bounds[node_a] += 1 + if start_node > end_node: + end_node, start_node = start_node, end_node + + slot = bucket_bounds[start_node] + order[slot] = face_idx * n_max_face_nodes + current_node + end_node_keys[slot] = end_node + bucket_bounds[start_node] = slot + 1 # Sort each bucket by ``node_b`` and count the unique edges while the bucket is in # cache, which gives the exact allocation size for the walk below n_edge = 0 bucket_start = 0 - for node_a in range(n_node): - bucket_end = bucket_bounds[node_a] - - if bucket_end - bucket_start > MAX_INSERTION_SORT_SIZE: - # A node this well connected is not expected of a real mesh, but insertion - # sort degrades quadratically, so fall back to a heap sort - _heap_sort_bucket( - order, - bucket_start, - bucket_end - bucket_start, - face_node_connectivity, - n_nodes_per_face, - n_max_face_nodes, - ) - else: - for i in range(bucket_start + 1, bucket_end): - flat_idx = order[i] - node_b = _other_node( - face_node_connectivity, n_nodes_per_face, n_max_face_nodes, flat_idx - ) + for n in range(n_node): + bucket_end = bucket_bounds[n] - j = i - 1 - while j >= bucket_start and ( - _other_node( - face_node_connectivity, - n_nodes_per_face, - n_max_face_nodes, - order[j], - ) - > node_b - ): - order[j + 1] = order[j] - j -= 1 - order[j + 1] = flat_idx + size = bucket_end - bucket_start + if size > MIN_ADAPTIVE_SORT_SIZE: + # Large enough that a bad ordering would be worth catching, which only a + # collapsed pole or a similarly degenerate node reaches + _adaptive_sort_bucket(end_node_keys, order, bucket_start, size) + elif size > 1: + _insertion_sort_bucket(end_node_keys, order, bucket_start, size) prev_b = INT_FILL_VALUE for i in range(bucket_start, bucket_end): - node_b = _other_node( - face_node_connectivity, n_nodes_per_face, n_max_face_nodes, order[i] - ) - if node_b != prev_b: + if end_node_keys[i] != prev_b: n_edge += 1 - prev_b = node_b + prev_b = end_node_keys[i] bucket_start = bucket_end @@ -384,22 +275,20 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_no edge_idx = -1 bucket_start = 0 - for node_a in range(n_node): - bucket_end = bucket_bounds[node_a] + for n in range(n_node): + bucket_end = bucket_bounds[n] prev_b = INT_FILL_VALUE for i in range(bucket_start, bucket_end): flat_idx = order[i] - node_b = _other_node( - face_node_connectivity, n_nodes_per_face, n_max_face_nodes, flat_idx - ) + end_node = end_node_keys[i] - if node_b != prev_b: + if end_node != prev_b: # Only store unique edges edge_idx += 1 - edge_node_connectivity[edge_idx, 0] = node_a - edge_node_connectivity[edge_idx, 1] = node_b - prev_b = node_b + edge_node_connectivity[edge_idx, 0] = n + edge_node_connectivity[edge_idx, 1] = end_node + prev_b = end_node face_edge_connectivity[ flat_idx // n_max_face_nodes, flat_idx % n_max_face_nodes diff --git a/uxarray/grid/utils.py b/uxarray/grid/utils.py index 9287bc325..f45288eab 100644 --- a/uxarray/grid/utils.py +++ b/uxarray/grid/utils.py @@ -443,3 +443,107 @@ def setter(self, value): self._ds[key] = value return setter + + +# Bucket sorts for the counting sort in ``connectivity._build_edge_node_connectivity``. Each +# sorts one contiguous ``[bucket_start, bucket_start + size)`` slice of ``node_b`` in place, +# applying the same permutation to ``order`` so the two stay aligned. + +# Smallest bucket worth watching for pathological input. A bucket of ``size`` holds at most +# ``size * (size - 1) / 2`` inversions, so at or below this size it cannot exceed the shift +# budget below and the bookkeeping would never pay for itself +MIN_ADAPTIVE_SORT_SIZE = 16 + +# Shifts per edge an insertion sort may spend on a bucket before it is abandoned for a heap +# sort. Insertion sort costs ``O(size + shifts)``, so a constant budget per edge keeps the +# adaptive path linear while leaving ample room for the near-sorted input it is chosen for +MAX_SHIFTS_PER_EDGE = 8 + + +@njit(cache=True) +def _sift_down(node_b, order, bucket_start, root, size): + """Restores the max-heap property at ``root`` for a bucket keyed on ``node_b``.""" + while True: + child = 2 * root + 1 + if child >= size: + break + + if ( + child + 1 < size + and node_b[bucket_start + child] < node_b[bucket_start + child + 1] + ): + child += 1 + + if node_b[bucket_start + root] >= node_b[bucket_start + child]: + break + + node_b[bucket_start + root], node_b[bucket_start + child] = ( + node_b[bucket_start + child], + node_b[bucket_start + root], + ) + order[bucket_start + root], order[bucket_start + child] = ( + order[bucket_start + child], + order[bucket_start + root], + ) + root = child + + +@njit(cache=True) +def _heap_sort_bucket(node_b, order, bucket_start, size): + """Sorts a bucket by ``node_b`` in place, in ``O(size * log(size))`` and without + scratch space, for the rare bucket an insertion sort cannot finish cheaply.""" + for root in range(size // 2 - 1, -1, -1): + _sift_down(node_b, order, bucket_start, root, size) + + for end in range(size - 1, 0, -1): + node_b[bucket_start], node_b[bucket_start + end] = ( + node_b[bucket_start + end], + node_b[bucket_start], + ) + order[bucket_start], order[bucket_start + end] = ( + order[bucket_start + end], + order[bucket_start], + ) + _sift_down(node_b, order, bucket_start, 0, end) + + +@njit(cache=True) +def _insertion_sort_bucket(node_b, order, bucket_start, size): + """Sorts a bucket by ``node_b`` in place, in ``O(size + inversions)``.""" + for i in range(bucket_start + 1, bucket_start + size): + key = node_b[i] + flat_idx = order[i] + + j = i - 1 + while j >= bucket_start and node_b[j] > key: + node_b[j + 1] = node_b[j] + order[j + 1] = order[j] + j -= 1 + node_b[j + 1] = key + order[j + 1] = flat_idx + + +@njit(cache=True) +def _adaptive_sort_bucket(node_b, order, bucket_start, size): + """Sorts a large bucket by ``node_b`` in place, insertion sorting it unless it turns out + to be badly ordered, in which case the partial work is abandoned for a heap sort. + """ + budget = MAX_SHIFTS_PER_EDGE * size + shifts = 0 + + for i in range(bucket_start + 1, bucket_start + size): + key = node_b[i] + flat_idx = order[i] + + j = i - 1 + while j >= bucket_start and node_b[j] > key: + node_b[j + 1] = node_b[j] + order[j + 1] = order[j] + j -= 1 + node_b[j + 1] = key + order[j + 1] = flat_idx + + shifts += i - 1 - j + if shifts > budget: + _heap_sort_bucket(node_b, order, bucket_start, size) + return From 5f0e40140de097e47f1a7a4f9b5b8ca0abb94587 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 29 Jul 2026 17:04:48 -0500 Subject: [PATCH 12/19] OFE: better parallelism --- uxarray/grid/connectivity.py | 53 ++++++++++++++++++++++-------------- 1 file changed, 33 insertions(+), 20 deletions(-) diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index 2b192a877..b28606754 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -1,6 +1,6 @@ import numpy as np import xarray as xr -from numba import njit +from numba import njit, prange from uxarray.constants import INT_DTYPE, INT_FILL_VALUE from uxarray.conventions import ugrid @@ -157,8 +157,16 @@ def _populate_edge_node_connectivity(grid): # TODO: raise a warning or exception? pass + # HACK: this is lieu of an xarray equivalent to `da.compute(a, b)` + computed = xr.Dataset( + { + "face_nodes": grid.face_node_connectivity.variable, + "n_nodes_per_face": grid.n_nodes_per_face.variable, + } + ).compute() + edge_nodes, face_edges = _build_edge_node_connectivity( - grid.face_node_connectivity.values, grid.n_nodes_per_face.values, grid.n_node + computed.face_nodes.data, computed.n_nodes_per_face.data, grid.n_node ) grid._ds["edge_node_connectivity"] = xr.DataArray( @@ -174,7 +182,7 @@ def _populate_edge_node_connectivity(grid): ) -@njit(cache=True) +@njit(cache=True, parallel=True) def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_node): """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 @@ -208,14 +216,12 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_no face_node_connectivity, INT_FILL_VALUE, dtype=INT_DTYPE ) - n_half_edge = 0 - for i in range(n_face): - n_half_edge += n_nodes_per_face[i] + n_half_edge = np.sum(n_nodes_per_face) if n_half_edge == 0: return np.empty((0, 2), dtype=INT_DTYPE), face_edge_connectivity - # Count how many half edges fall into each ``node_a`` bucket, then prefix sum so that + # Count how many half edges fall into each ``start_node`` bucket, then prefix sum so that # ``bucket_bounds[a]`` is where bucket ``a`` starts bucket_bounds = np.zeros(n_node + 1, dtype=INT_DTYPE) for face_idx in range(n_face): @@ -246,11 +252,11 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_no end_node_keys[slot] = end_node bucket_bounds[start_node] = slot + 1 - # Sort each bucket by ``node_b`` and count the unique edges while the bucket is in - # cache, which gives the exact allocation size for the walk below - n_edge = 0 - bucket_start = 0 - for n in range(n_node): + # Sort each bucket by ``node_b`` and count its unique edges while the bucket is in + # cache. Buckets are disjoint, so this runs one bucket per thread. + unique_per_bucket = np.empty(n_node, dtype=INT_DTYPE) + for n in prange(n_node): + bucket_start = bucket_bounds[n - 1] if n > 0 else 0 bucket_end = bucket_bounds[n] size = bucket_end - bucket_start @@ -261,22 +267,31 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_no elif size > 1: _insertion_sort_bucket(end_node_keys, order, bucket_start, size) + n_unique = 0 prev_b = INT_FILL_VALUE for i in range(bucket_start, bucket_end): if end_node_keys[i] != prev_b: - n_edge += 1 + n_unique += 1 prev_b = end_node_keys[i] - bucket_start = bucket_end + unique_per_bucket[n] = n_unique + + edge_offset = np.empty(n_node + 1, dtype=INT_DTYPE) + n_edge = 0 + for n in range(n_node): + edge_offset[n] = n_edge + n_edge += unique_per_bucket[n] + edge_offset[n_node] = n_edge # Duplicate half edges are now adjacent, so a single walk assigns each unique edge its - # index and populates the face edge connectivity + # index and populates the face edge connectivity. edge_node_connectivity = np.empty((n_edge, 2), dtype=INT_DTYPE) - edge_idx = -1 - bucket_start = 0 - for n in range(n_node): + for n in prange(n_node): + bucket_start = bucket_bounds[n - 1] if n > 0 else 0 bucket_end = bucket_bounds[n] + + edge_idx = edge_offset[n] - 1 prev_b = INT_FILL_VALUE for i in range(bucket_start, bucket_end): @@ -294,8 +309,6 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_no flat_idx // n_max_face_nodes, flat_idx % n_max_face_nodes ] = edge_idx - bucket_start = bucket_end - return edge_node_connectivity, face_edge_connectivity From 0a623d736b8d7f62ee4ebee82d0e164fea5567f3 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Tue, 4 Aug 2026 19:38:27 -0500 Subject: [PATCH 13/19] OFE: refactor connectivity and sorting for readability --- uxarray/grid/connectivity.py | 296 +++++++++++++++++++++++++---------- uxarray/grid/grid.py | 3 +- uxarray/grid/utils.py | 234 ++++++++++++++++++++++----- 3 files changed, 405 insertions(+), 128 deletions(-) diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index b28606754..dc64065c1 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -5,9 +5,10 @@ from uxarray.constants import INT_DTYPE, INT_FILL_VALUE from uxarray.conventions import ugrid from uxarray.grid.utils import ( - MIN_ADAPTIVE_SORT_SIZE, - _adaptive_sort_bucket, - _insertion_sort_bucket, + _build_pair_index, + _count_unique_in_bucket, + _search_bucket, + _sort_bucket, ) @@ -182,6 +183,100 @@ def _populate_edge_node_connectivity(grid): ) +@njit(cache=True, inline="always") +def _canonical_half_edge(face_node_connectivity, face_idx, local_idx, n_edges): + """The ``(low, high)`` node pair of the half edge leaving face slot ``local_idx``, + wrapping back to slot 0 after ``n_edges``.""" + start_node = face_node_connectivity[face_idx, local_idx] + end_node = face_node_connectivity[face_idx, (local_idx + 1) % n_edges] + + if start_node > end_node: + return end_node, start_node + return start_node, end_node + + +@njit(cache=True) +def _count_half_edges_per_node(face_node_connectivity, n_nodes_per_face, n_node): + """Bucket offsets keyed on each half edge's lower node: bucket ``a`` will occupy + ``[bucket_offset[a], bucket_offset[a + 1])``.""" + bucket_offset = np.zeros(n_node + 1, dtype=INT_DTYPE) + + for face_idx in range(face_node_connectivity.shape[0]): + n_edges = n_nodes_per_face[face_idx] + for local_idx in range(n_edges): + node_a, _ = _canonical_half_edge( + face_node_connectivity, face_idx, local_idx, n_edges + ) + bucket_offset[node_a + 1] += 1 + + for n in range(n_node): + bucket_offset[n + 1] += bucket_offset[n] + + return bucket_offset + + +@njit(cache=True) +def _scatter_half_edges( + face_node_connectivity, n_nodes_per_face, bucket_offset, n_half_edge +): + """Fills every bucket with its half edges, leaving ``bucket_offset`` as it found it. + + Each half edge is identified by ``half_edge_slot``, its flattened position + ``face_idx * n_max_face_nodes + local_idx`` in the face node connectivity, and keyed on + ``end_node``, the higher of its two nodes.""" + n_max_face_nodes = face_node_connectivity.shape[1] + + half_edge_slot = np.empty(n_half_edge, dtype=INT_DTYPE) + end_node = np.empty(n_half_edge, dtype=INT_DTYPE) + + for face_idx in range(face_node_connectivity.shape[0]): + n_edges = n_nodes_per_face[face_idx] + for local_idx in range(n_edges): + node_a, node_b = _canonical_half_edge( + face_node_connectivity, face_idx, local_idx, n_edges + ) + + slot = bucket_offset[node_a] + half_edge_slot[slot] = face_idx * n_max_face_nodes + local_idx + end_node[slot] = node_b + bucket_offset[node_a] = slot + 1 + + # The scatter left each entry at its bucket's end, i.e. one slot right of where the + # convention above wants it. One backward pass puts it back. + for n in range(bucket_offset.shape[0] - 1, 0, -1): + bucket_offset[n] = bucket_offset[n - 1] + bucket_offset[0] = 0 + + return half_edge_slot, end_node + + +@njit(cache=True) +def _emit_bucket_edges( + end_node, + half_edge_slot, + bucket_start, + bucket_end, + node_a, + first_edge_idx, + edge_node_connectivity, + face_edge_flat, +): + """Numbers a sorted bucket's unique edges from ``first_edge_idx`` and points each of its + half edges at the edge it belongs to.""" + edge_idx = first_edge_idx - 1 + previous_end_node = INT_FILL_VALUE + + for i in range(bucket_start, bucket_end): + if end_node[i] != previous_end_node: + # Duplicate half edges are adjacent, so a new key starts a new edge + edge_idx += 1 + edge_node_connectivity[edge_idx, 0] = node_a + edge_node_connectivity[edge_idx, 1] = end_node[i] + previous_end_node = end_node[i] + + face_edge_flat[half_edge_slot[i]] = edge_idx + + @njit(cache=True, parallel=True) def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_node): """Constructs the ``edge_node_connectivity`` variable, which represents the indices of the two nodes that make up @@ -191,6 +286,9 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_no Each edge is stored as an ascending ``(node_a, node_b)`` pair, and the edges are numbered in lexicographic order of that pair. + Every half edge is bucketed on its lower node, each bucket is sorted on its higher node, and the duplicates + that this makes adjacent are then collapsed into one edge apiece. + Parameters ---------- face_node_connectivity : np.ndarray @@ -208,12 +306,10 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_no Face Edge Connectivity with shape (n_face, n_max_face_edges) """ - - n_face, n_max_face_nodes = face_node_connectivity.shape - - # Keep track of face_edge_connectivity - face_edge_connectivity = np.full_like( - face_node_connectivity, INT_FILL_VALUE, dtype=INT_DTYPE + # ``np.full`` rather than ``np.full_like``, which would inherit a Fortran-ordered + # prototype's layout and make the flat view below unobtainable + face_edge_connectivity = np.full( + face_node_connectivity.shape, INT_FILL_VALUE, dtype=INT_DTYPE ) n_half_edge = np.sum(n_nodes_per_face) @@ -221,61 +317,27 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_no if n_half_edge == 0: return np.empty((0, 2), dtype=INT_DTYPE), face_edge_connectivity - # Count how many half edges fall into each ``start_node`` bucket, then prefix sum so that - # ``bucket_bounds[a]`` is where bucket ``a`` starts - bucket_bounds = np.zeros(n_node + 1, dtype=INT_DTYPE) - for face_idx in range(n_face): - n_edges = n_nodes_per_face[face_idx] - for current_node in range(n_edges): - start_node = face_node_connectivity[face_idx, current_node] - end_node = face_node_connectivity[face_idx, (current_node + 1) % n_edges] - bucket_bounds[min(start_node, end_node) + 1] += 1 - for i in range(n_node): - bucket_bounds[i + 1] += bucket_bounds[i] - - # Scatter the half edges into their buckets. This advances each entry of - # ``bucket_bounds`` to the *end* of its bucket, so afterwards bucket ``a`` spans - # ``bucket_bounds[a - 1]`` up to ``bucket_bounds[a]``, with bucket 0 starting at 0 - order = np.empty(n_half_edge, dtype=INT_DTYPE) - end_node_keys = np.empty(n_half_edge, dtype=INT_DTYPE) - for face_idx in range(n_face): - n_edges = n_nodes_per_face[face_idx] - for current_node in range(n_edges): - start_node = face_node_connectivity[face_idx, current_node] - end_node = face_node_connectivity[face_idx, (current_node + 1) % n_edges] - - if start_node > end_node: - end_node, start_node = start_node, end_node - - slot = bucket_bounds[start_node] - order[slot] = face_idx * n_max_face_nodes + current_node - end_node_keys[slot] = end_node - bucket_bounds[start_node] = slot + 1 + bucket_offset = _count_half_edges_per_node( + face_node_connectivity, n_nodes_per_face, n_node + ) + half_edge_slot, end_node = _scatter_half_edges( + face_node_connectivity, n_nodes_per_face, bucket_offset, n_half_edge + ) - # Sort each bucket by ``node_b`` and count its unique edges while the bucket is in - # cache. Buckets are disjoint, so this runs one bucket per thread. + # Sort each bucket and count its unique edges while the bucket is in cache. Buckets are + # disjoint, so this runs one bucket per thread. unique_per_bucket = np.empty(n_node, dtype=INT_DTYPE) for n in prange(n_node): - bucket_start = bucket_bounds[n - 1] if n > 0 else 0 - bucket_end = bucket_bounds[n] - - size = bucket_end - bucket_start - if size > MIN_ADAPTIVE_SORT_SIZE: - # Large enough that a bad ordering would be worth catching, which only a - # collapsed pole or a similarly degenerate node reaches - _adaptive_sort_bucket(end_node_keys, order, bucket_start, size) - elif size > 1: - _insertion_sort_bucket(end_node_keys, order, bucket_start, size) - - n_unique = 0 - prev_b = INT_FILL_VALUE - for i in range(bucket_start, bucket_end): - if end_node_keys[i] != prev_b: - n_unique += 1 - prev_b = end_node_keys[i] - - unique_per_bucket[n] = n_unique + bucket_start = bucket_offset[n] + bucket_end = bucket_offset[n + 1] + + _sort_bucket(end_node, half_edge_slot, bucket_start, bucket_end - bucket_start) + unique_per_bucket[n] = _count_unique_in_bucket( + end_node, bucket_start, bucket_end + ) + # Hand each bucket the edge index its first unique edge takes, so the emit below can run + # one bucket per thread as well edge_offset = np.empty(n_node + 1, dtype=INT_DTYPE) n_edge = 0 for n in range(n_node): @@ -283,31 +345,20 @@ def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_no n_edge += unique_per_bucket[n] edge_offset[n_node] = n_edge - # Duplicate half edges are now adjacent, so a single walk assigns each unique edge its - # index and populates the face edge connectivity. edge_node_connectivity = np.empty((n_edge, 2), dtype=INT_DTYPE) + face_edge_flat = face_edge_connectivity.reshape(-1) for n in prange(n_node): - bucket_start = bucket_bounds[n - 1] if n > 0 else 0 - bucket_end = bucket_bounds[n] - - edge_idx = edge_offset[n] - 1 - prev_b = INT_FILL_VALUE - - for i in range(bucket_start, bucket_end): - flat_idx = order[i] - end_node = end_node_keys[i] - - if end_node != prev_b: - # Only store unique edges - edge_idx += 1 - edge_node_connectivity[edge_idx, 0] = n - edge_node_connectivity[edge_idx, 1] = end_node - prev_b = end_node - - face_edge_connectivity[ - flat_idx // n_max_face_nodes, flat_idx % n_max_face_nodes - ] = edge_idx + _emit_bucket_edges( + end_node, + half_edge_slot, + bucket_offset[n], + bucket_offset[n + 1], + n, + edge_offset[n], + edge_node_connectivity, + face_edge_flat, + ) return edge_node_connectivity, face_edge_connectivity @@ -351,10 +402,87 @@ def _populate_face_edge_connectivity(grid): and stores it within the internal (``Grid._ds``) and through the attribute (``Grid.face_edge_connectivity``).""" - # TODO: Check if "edge_edge_connectivity" is already present - if "edge_node_connectivity" not in grid._ds: + # Constructing the edges derives this variable in the same pass _populate_edge_node_connectivity(grid) + return + + # In lieu of an xarray equivalent to `da.compute(a, b)`, we can batch these variables as + # an xarray Dataset and re-extract after graph traversal. + computed = xr.Dataset( + { + "face_nodes": grid.face_node_connectivity.variable, + "n_nodes_per_face": grid.n_nodes_per_face.variable, + "edge_nodes": grid.edge_node_connectivity.variable, + } + ).compute() + + face_edges = _build_face_edge_connectivity( + computed.face_nodes.data, + computed.n_nodes_per_face.data, + computed.edge_nodes.data, + grid.n_node, + ) + + grid._ds["face_edge_connectivity"] = xr.DataArray( + face_edges, + dims=ugrid.FACE_EDGE_CONNECTIVITY_DIMS, + attrs=ugrid.FACE_EDGE_CONNECTIVITY_ATTRS, + ) + + +@njit(cache=True, parallel=True) +def _build_face_edge_connectivity( + face_node_connectivity, n_nodes_per_face, edge_node_connectivity, n_node +): + """Constructs the ``face_edge_connectivity`` variable, which represents the indices of the edges that make up + each face, by looking each face's edges up in an existing ``edge_node_connectivity``. The edges keep the + numbering they arrived with. + + Edges are bucketed on their lower node so that each of a face's edges can be found by a binary search of one + bucket. Edges already in the canonical order that :func:`_build_edge_node_connectivity` emits are bucketed + without being sorted again. + + Parameters + ---------- + face_node_connectivity : np.ndarray + Face Node Connectivity + n_nodes_per_face : np.ndarray + Number of nodes/edges per face + edge_node_connectivity : np.ndarray + Edge Node Connectivity with shape (n_edge, 2), in any order or orientation + n_node : int + Total number of nodes, used as the number of buckets + + Returns + ------- + face_edge_connectivity : np.ndarray + Face Edge Connectivity with shape (n_face, n_max_face_edges). Edges of a face that are absent from + ``edge_node_connectivity`` are left as ``INT_FILL_VALUE``, as are the padding slots of a face with + fewer than ``n_max_face_edges`` edges. + + """ + face_edge_connectivity = np.full( + face_node_connectivity.shape, INT_FILL_VALUE, dtype=INT_DTYPE + ) + + bucket_offset, end_node, edge_id = _build_pair_index(edge_node_connectivity, n_node) + + for face_idx in prange(face_node_connectivity.shape[0]): + n_edges = n_nodes_per_face[face_idx] + for local_idx in range(n_edges): + node_a, node_b = _canonical_half_edge( + face_node_connectivity, face_idx, local_idx, n_edges + ) + face_edge_connectivity[face_idx, local_idx] = _search_bucket( + end_node, + edge_id, + bucket_offset[node_a], + bucket_offset[node_a + 1], + node_b, + ) + + return face_edge_connectivity def _populate_node_face_connectivity(grid): diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 27c65e183..881cb14a9 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -1268,7 +1268,8 @@ def edge_node_connectivity(self) -> xr.DataArray: Connectivity variable representing the indices of nodes (mesh vertices) that define each edge. Each row (i.e., each edge) contains exactly two node indices that define the start and end points of the edge. - The nodes are stored in an arbitrary order. + Constructed edges are stored as ascending node pairs and numbered in lexicographic order of that pair; edges + read from a file keep the order and orientation they were stored in. Returns ------- diff --git a/uxarray/grid/utils.py b/uxarray/grid/utils.py index f45288eab..6510c85d0 100644 --- a/uxarray/grid/utils.py +++ b/uxarray/grid/utils.py @@ -2,7 +2,7 @@ import xarray as xr from numba import njit -from uxarray.constants import INT_FILL_VALUE +from uxarray.constants import INT_DTYPE, INT_FILL_VALUE @njit(cache=True) @@ -445,9 +445,20 @@ def setter(self, value): return setter -# Bucket sorts for the counting sort in ``connectivity._build_edge_node_connectivity``. Each -# sorts one contiguous ``[bucket_start, bucket_start + size)`` slice of ``node_b`` in place, -# applying the same permutation to ``order`` so the two stay aligned. +# Bucket sorting and searching for the counting sorts in ``uxarray.grid.connectivity``, which +# bucket half edges or edges on one of their two nodes and then order each bucket by the other. +# Nothing below knows about meshes: a bucket is a contiguous ``[bucket_start, bucket_end)`` slice +# of a key array, carrying an equally long ``payload`` array that every reordering moves in step +# so the two stay aligned. +# +# ``_sort_bucket`` orders one bucket and is the entry point for the sorts; the kernels beneath it +# are chosen by bucket size and are exposed only for testing. ``_build_pair_index`` runs the whole +# count/scatter/sort sequence for a caller starting from an ``(n, 2)`` array, and +# ``_search_bucket`` is the lookup that index is built for. +# +# NOTE: these are inlined into ``cache=True`` kernels in another module, and numba stamps its +# cache against the defining file alone, so editing them does not invalidate a caller's cached +# object. Clear ``uxarray/grid/__pycache__/*.nbi *.nbc`` after changing anything here. # Smallest bucket worth watching for pathological input. A bucket of ``size`` holds at most # ``size * (size - 1) / 2`` inversions, so at or below this size it cannot exceed the shift @@ -461,8 +472,19 @@ def setter(self, value): @njit(cache=True) -def _sift_down(node_b, order, bucket_start, root, size): - """Restores the max-heap property at ``root`` for a bucket keyed on ``node_b``.""" +def _sort_bucket(end_node, payload, bucket_start, size): + """Orders one bucket by ``end_node``, picking the sort that suits its size.""" + if size > MIN_ADAPTIVE_SORT_SIZE: + # Large enough that a bad ordering would be worth catching, which only a + # collapsed pole or a similarly degenerate node reaches + _adaptive_sort_bucket(end_node, payload, bucket_start, size) + elif size > 1: + _insertion_sort_bucket(end_node, payload, bucket_start, size) + + +@njit(cache=True) +def _sift_down(end_node, payload, bucket_start, root, size): + """Restores the max-heap property at ``root`` for a bucket keyed on ``end_node``.""" while True: child = 2 * root + 1 if child >= size: @@ -470,80 +492,206 @@ def _sift_down(node_b, order, bucket_start, root, size): if ( child + 1 < size - and node_b[bucket_start + child] < node_b[bucket_start + child + 1] + and end_node[bucket_start + child] < end_node[bucket_start + child + 1] ): child += 1 - if node_b[bucket_start + root] >= node_b[bucket_start + child]: + if end_node[bucket_start + root] >= end_node[bucket_start + child]: break - node_b[bucket_start + root], node_b[bucket_start + child] = ( - node_b[bucket_start + child], - node_b[bucket_start + root], + end_node[bucket_start + root], end_node[bucket_start + child] = ( + end_node[bucket_start + child], + end_node[bucket_start + root], ) - order[bucket_start + root], order[bucket_start + child] = ( - order[bucket_start + child], - order[bucket_start + root], + payload[bucket_start + root], payload[bucket_start + child] = ( + payload[bucket_start + child], + payload[bucket_start + root], ) root = child @njit(cache=True) -def _heap_sort_bucket(node_b, order, bucket_start, size): - """Sorts a bucket by ``node_b`` in place, in ``O(size * log(size))`` and without +def _heap_sort_bucket(end_node, payload, bucket_start, size): + """Sorts a bucket by ``end_node`` in place, in ``O(size * log(size))`` and without scratch space, for the rare bucket an insertion sort cannot finish cheaply.""" for root in range(size // 2 - 1, -1, -1): - _sift_down(node_b, order, bucket_start, root, size) + _sift_down(end_node, payload, bucket_start, root, size) for end in range(size - 1, 0, -1): - node_b[bucket_start], node_b[bucket_start + end] = ( - node_b[bucket_start + end], - node_b[bucket_start], + end_node[bucket_start], end_node[bucket_start + end] = ( + end_node[bucket_start + end], + end_node[bucket_start], ) - order[bucket_start], order[bucket_start + end] = ( - order[bucket_start + end], - order[bucket_start], + payload[bucket_start], payload[bucket_start + end] = ( + payload[bucket_start + end], + payload[bucket_start], ) - _sift_down(node_b, order, bucket_start, 0, end) + _sift_down(end_node, payload, bucket_start, 0, end) @njit(cache=True) -def _insertion_sort_bucket(node_b, order, bucket_start, size): - """Sorts a bucket by ``node_b`` in place, in ``O(size + inversions)``.""" +def _insertion_sort_bucket(end_node, payload, bucket_start, size): + """Sorts a bucket by ``end_node`` in place, in ``O(size + inversions)``.""" for i in range(bucket_start + 1, bucket_start + size): - key = node_b[i] - flat_idx = order[i] + key = end_node[i] + key_payload = payload[i] j = i - 1 - while j >= bucket_start and node_b[j] > key: - node_b[j + 1] = node_b[j] - order[j + 1] = order[j] + while j >= bucket_start and end_node[j] > key: + end_node[j + 1] = end_node[j] + payload[j + 1] = payload[j] j -= 1 - node_b[j + 1] = key - order[j + 1] = flat_idx + end_node[j + 1] = key + payload[j + 1] = key_payload @njit(cache=True) -def _adaptive_sort_bucket(node_b, order, bucket_start, size): - """Sorts a large bucket by ``node_b`` in place, insertion sorting it unless it turns out +def _adaptive_sort_bucket(end_node, payload, bucket_start, size): + """Sorts a large bucket by ``end_node`` in place, insertion sorting it unless it turns out to be badly ordered, in which case the partial work is abandoned for a heap sort. + + This is ``_insertion_sort_bucket``'s loop with a shift meter around it. The duplication is + deliberate: metering every bucket instead of only the large ones measured ~5% slower + end-to-end, because typical buckets hold a handful of edges and the per-element bookkeeping + is a real fraction of that work. Keep the two in sync rather than merging them. """ budget = MAX_SHIFTS_PER_EDGE * size shifts = 0 for i in range(bucket_start + 1, bucket_start + size): - key = node_b[i] - flat_idx = order[i] + key = end_node[i] + key_payload = payload[i] j = i - 1 - while j >= bucket_start and node_b[j] > key: - node_b[j + 1] = node_b[j] - order[j + 1] = order[j] + while j >= bucket_start and end_node[j] > key: + end_node[j + 1] = end_node[j] + payload[j + 1] = payload[j] j -= 1 - node_b[j + 1] = key - order[j + 1] = flat_idx + end_node[j + 1] = key + payload[j + 1] = key_payload shifts += i - 1 - j if shifts > budget: - _heap_sort_bucket(node_b, order, bucket_start, size) + _heap_sort_bucket(end_node, payload, bucket_start, size) return + + +@njit(cache=True) +def _count_unique_in_bucket(end_node, bucket_start, bucket_end): + """Number of distinct keys in an already sorted bucket, where equal keys are adjacent.""" + n_unique = 0 + previous_end_node = INT_FILL_VALUE + + for i in range(bucket_start, bucket_end): + if end_node[i] != previous_end_node: + n_unique += 1 + previous_end_node = end_node[i] + + return n_unique + + +@njit(cache=True) +def _is_lexicographically_sorted(pairs): + """Whether every row of an ``(n, 2)`` array is ascending and the rows are themselves in + nondecreasing lexicographic order.""" + previous_low = INT_FILL_VALUE + previous_high = INT_FILL_VALUE + + for i in range(pairs.shape[0]): + low = pairs[i, 0] + high = pairs[i, 1] + + if low > high: + return False + if low < previous_low or (low == previous_low and high < previous_high): + return False + + previous_low = low + previous_high = high + + return True + + +@njit(cache=True) +def _count_pairs_per_bucket(pairs, n_bucket): + """Bucket offsets keyed on each row's lower value: bucket ``a`` will occupy + ``[bucket_offset[a], bucket_offset[a + 1])``.""" + bucket_offset = np.zeros(n_bucket + 1, dtype=INT_DTYPE) + + for i in range(pairs.shape[0]): + bucket_offset[min(pairs[i, 0], pairs[i, 1]) + 1] += 1 + + for a in range(n_bucket): + bucket_offset[a + 1] += bucket_offset[a] + + return bucket_offset + + +@njit(cache=True) +def _build_pair_index(pairs, n_bucket): + """Indexes an ``(n, 2)`` array of integer pairs so a pair can be looked up by value. + + Each row is bucketed on its lower value and each bucket ordered by its higher value, which is + what lets :func:`_search_bucket` find a row with a single binary search. Rows that are already + canonically ordered are indexed without being sorted again. + + Returns ``(bucket_offset, high, row)``: bucket ``a`` occupies + ``[bucket_offset[a], bucket_offset[a + 1])``, ``high`` holds each entry's higher value, and + ``row`` the index of ``pairs`` it came from. ``high`` is a copy rather than a column view, both + to keep the search off a strided array and so that either path below returns the same arrays. + """ + bucket_offset = _count_pairs_per_bucket(pairs, n_bucket) + + n_pair = pairs.shape[0] + high = np.empty(n_pair, dtype=INT_DTYPE) + row = np.empty(n_pair, dtype=INT_DTYPE) + + if _is_lexicographically_sorted(pairs): + # Already grouped by lower value, ascending within each group, so the buckets are the + # runs the counting pass just measured and no sorting is needed + for i in range(n_pair): + high[i] = pairs[i, 1] + row[i] = i + return bucket_offset, high, row + + for i in range(n_pair): + pair_low = pairs[i, 0] + pair_high = pairs[i, 1] + if pair_low > pair_high: + pair_low, pair_high = pair_high, pair_low + + slot = bucket_offset[pair_low] + high[slot] = pair_high + row[slot] = i + bucket_offset[pair_low] = slot + 1 + + # The scatter left each entry at its bucket's end, one slot right of where the convention + # above wants it. One backward pass puts it back. + for a in range(n_bucket, 0, -1): + bucket_offset[a] = bucket_offset[a - 1] + bucket_offset[0] = 0 + + for a in range(n_bucket): + bucket_start = bucket_offset[a] + _sort_bucket(high, row, bucket_start, bucket_offset[a + 1] - bucket_start) + + return bucket_offset, high, row + + +@njit(cache=True) +def _search_bucket(high, row, bucket_start, bucket_end, key): + """The ``row`` entry whose key is ``key`` within a sorted bucket, or ``INT_FILL_VALUE`` + when the bucket does not hold it.""" + low = bucket_start + stop = bucket_end + + while low < stop: + mid = (low + stop) // 2 + if high[mid] < key: + low = mid + 1 + else: + stop = mid + + if low < bucket_end and high[low] == key: + return row[low] + return INT_FILL_VALUE From 39dfc2d0a5c3d3011ad657fd8cd0484906c1f0b1 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 5 Aug 2026 18:29:32 -0500 Subject: [PATCH 14/19] OFE: edge corruption warning --- uxarray/grid/connectivity.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index dc64065c1..35d1b6c60 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -11,6 +11,8 @@ _sort_bucket, ) +from warnings import warn + def close_face_nodes(face_node_connectivity, n_face, n_max_face_nodes): """Closes (``face_node_connectivity``) by inserting the first node index @@ -155,10 +157,19 @@ def _populate_edge_node_connectivity(grid): # Check edge coordinates already exist, if they do this might cause issues if "n_edge" in grid.sizes: - # TODO: raise a warning or exception? - pass + stale = sorted(n for n in grid._ds if ugrid.EDGE_DIM in grid._ds[n].dims) + warn( + f"Constructing 'edge_node_connectivity' on a grid that already has " + f"edge-centered variables ({', '.join(stale)}). Constructed edges are " + f"numbered in lexicographic node-pair order, which need not match the " + f"numbering those variables were stored with; they may no longer refer " + f"to the same edges.", + RuntimeWarning, + stacklevel=2, + ) - # HACK: this is lieu of an xarray equivalent to `da.compute(a, b)` + # This is in lieu of an xarray equivalent to `da.compute(a, b)`. We traverse the + # grid once to gather both variables, possibly as chunks if dask is enabled computed = xr.Dataset( { "face_nodes": grid.face_node_connectivity.variable, From 51acffb1a3f04a416599ad827914a5258f169b81 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:32:28 +0000 Subject: [PATCH 15/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- uxarray/grid/connectivity.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index 35d1b6c60..400403a04 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -1,3 +1,5 @@ +from warnings import warn + import numpy as np import xarray as xr from numba import njit, prange @@ -11,8 +13,6 @@ _sort_bucket, ) -from warnings import warn - def close_face_nodes(face_node_connectivity, n_face, n_max_face_nodes): """Closes (``face_node_connectivity``) by inserting the first node index From ddfac01b2962ff5e0c58e8c631e5f392170e5914 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 5 Aug 2026 19:16:46 -0500 Subject: [PATCH 16/19] OFE: minor cleanup --- uxarray/grid/connectivity.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index 400403a04..0675ecc8e 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -1,4 +1,3 @@ -from warnings import warn import numpy as np import xarray as xr @@ -156,16 +155,14 @@ def _populate_edge_node_connectivity(grid): # Check edge coordinates already exist, if they do this might cause issues - if "n_edge" in grid.sizes: + if "n_edge" in grid.dims: stale = sorted(n for n in grid._ds if ugrid.EDGE_DIM in grid._ds[n].dims) - warn( + raise( f"Constructing 'edge_node_connectivity' on a grid that already has " f"edge-centered variables ({', '.join(stale)}). Constructed edges are " f"numbered in lexicographic node-pair order, which need not match the " f"numbering those variables were stored with; they may no longer refer " - f"to the same edges.", - RuntimeWarning, - stacklevel=2, + f"to the same edges." ) # This is in lieu of an xarray equivalent to `da.compute(a, b)`. We traverse the From ae04596ad0a48ab908f40021ed6e28d30d0ae3d2 Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Wed, 5 Aug 2026 19:30:36 -0500 Subject: [PATCH 17/19] OFE: typo --- uxarray/grid/connectivity.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index 0675ecc8e..e10705009 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -157,7 +157,7 @@ def _populate_edge_node_connectivity(grid): if "n_edge" in grid.dims: stale = sorted(n for n in grid._ds if ugrid.EDGE_DIM in grid._ds[n].dims) - raise( + raise ValueError( f"Constructing 'edge_node_connectivity' on a grid that already has " f"edge-centered variables ({', '.join(stale)}). Constructed edges are " f"numbered in lexicographic node-pair order, which need not match the " From 4e8c2d889ee10c913b2ed8df27d888dd0f5b234b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 00:33:01 +0000 Subject: [PATCH 18/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- uxarray/grid/connectivity.py | 1 - 1 file changed, 1 deletion(-) diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index e10705009..a35d1311f 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -1,4 +1,3 @@ - import numpy as np import xarray as xr from numba import njit, prange From 523dc6ae41a953587463d71859ad1a54503fc82b Mon Sep 17 00:00:00 2001 From: cmdupuis3 Date: Fri, 7 Aug 2026 12:03:45 -0500 Subject: [PATCH 19/19] Docstring for edge_node and face_edge combined behavior --- uxarray/grid/grid.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index db8e707c7..2dd7af5e3 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -1273,6 +1273,10 @@ def edge_node_connectivity(self) -> xr.DataArray: Constructed edges are stored as ascending node pairs and numbered in lexicographic order of that pair; edges read from a file keep the order and orientation they were stored in. + The result is cached after the first access; subsequent calls return the stored value without recomputing it. + Computing edge_node_connectivity always derives face_edge_connectivity as part of the same pass and + overwrites any existing face_edge_connectivity value, regardless of whether one was already present. + Returns ------- edge_node_connectivity : :py:class:`xarray.DataArray` @@ -1318,6 +1322,11 @@ def face_edge_connectivity(self) -> xr.DataArray: rows containing fewer than :py:attr:`~uxarray.Grid.n_max_face_edges` indices are padded with the fill value defined in :py:attr:`~uxarray.constants.INT_FILL_VALUE`. + The result is cached after the first access; subsequent calls return the stored value without recomputing it. + If edge_node_connectivity has not yet been computed, it is derived together with face_edge_connectivity in + the same pass. If edge_node_connectivity is already present, face_edge_connectivity is instead derived + independently from the existing connectivity data. + Returns ------- face_edge_connectivity : :py:class:`xarray.DataArray`