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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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