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 1/6] 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 2/6] 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 3/6] 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 4/6] 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 5/6] 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 6/6] 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)