diff --git a/test/grid/grid/test_core.py b/test/grid/grid/test_core.py index da5882bea..e2cc34840 100644 --- a/test/grid/grid/test_core.py +++ b/test/grid/grid/test_core.py @@ -5,6 +5,11 @@ import uxarray as ux from uxarray.constants import ERROR_TOLERANCE, INT_FILL_VALUE +from uxarray.grid.validation import ( + _check_duplicate_nodes_indices, + _find_duplicate_nodes, +) +from uxarray.errors import GridInvalidError def test_grid_with_holes(gridpath): @@ -129,7 +134,170 @@ def test_dual_mesh_mpas(gridpath): def test_dual_duplicate(gridpath): - """Test dual mesh creation with duplicate grids.""" - dataset = ux.open_dataset(gridpath("ugrid", "geoflow-small", "grid.nc"), gridpath("ugrid", "geoflow-small", "grid.nc")) - with pytest.raises(ux.errors.GridInvalidError): - dataset.get_dual() + """Test dual mesh creation on a grid whose source file has duplicate + (coincident) node indices, merged at construction time.""" + grid_path = gridpath("ugrid", "geoflow-small", "grid.nc") + grid = ux.open_grid(grid_path) + + # The source file really does contain duplicates: 6000 node coordinates for + # 3850 distinct locations, so 2150 indices are coincident with an earlier one. + duplicates = _find_duplicate_nodes(grid) + assert grid.n_node == 6000 + assert len(duplicates) == 2150 + + # Connectivity is canonicalized to a single index per coincident group, so no + # face references any of those 2150 duplicate indices. + assert not _check_duplicate_nodes_indices(grid) + # duplicate coordinates are left in place by design, but connectivity is + # fully canonicalized, so validation passes + assert grid.validate() + + dual = grid.get_dual() + + assert dual.n_node == grid.n_face + + # One dual face per node that is a corner of at least three faces. After the + # merge, 3850 distinct nodes remain, ten of which are touched by a single face + # only and so produce no dual cell, leaving 3840. + face_nodes = grid.face_node_connectivity.values + faces_per_node = np.bincount( + face_nodes[face_nodes != INT_FILL_VALUE], minlength=grid.n_node + ) + assert grid.n_node - len(duplicates) == 3850 + assert (faces_per_node >= 3).sum() == 3840 + assert dual.n_face == 3840 + + dataset = ux.open_dataset(grid_path, grid_path) + dual_ds = dataset.get_dual() + assert dual_ds.uxgrid.n_face == dual.n_face + + +def test_dual_duplicate_geos_cs(gridpath): + """Test dual mesh creation on a cube-sphere grid with duplicate node + indices (issue #865).""" + grid_path = gridpath("geos-cs", "c12", "test-c12.native.nc4") + grid = ux.open_grid(grid_path) + + assert len(_find_duplicate_nodes(grid)) > 0 + assert not _check_duplicate_nodes_indices(grid) + + dual = grid.get_dual() + assert dual.n_node == grid.n_face + assert dual.n_face > 0 + + +def test_duplicate_nodes_minimal_example(): + """Two quads that share an edge, but whose shared corners are stored twice. + + Nodes 2 and 3 are repeated as nodes 6 and 7, so the file describes 8 nodes at + 6 distinct locations. Node 6 must canonicalize to node 2 and node 7 to node 3, + leaving the second face pointing at the first face's corners. + + 3---2---7 lat 1 nodes 2,3 are the shared edge + | | | nodes 7,6 are their duplicates + 0---1---6 lat 0 + """ + node_lon = np.array([0.0, 1.0, 1.0, 0.0, 2.0, 2.0, 1.0, 1.0]) + node_lat = np.array([0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0]) + # left quad right quad, via the duplicates + face_node_connectivity = np.array([[0, 1, 2, 3], [6, 4, 5, 7]]) + + grid = ux.Grid.from_topology(node_lon, node_lat, face_node_connectivity) + + duplicates = _find_duplicate_nodes(grid) + assert duplicates == {6: 1, 7: 2} + + # No face may still reference a duplicate index. + assert not _check_duplicate_nodes_indices(grid) + nt.assert_equal( + grid.face_node_connectivity.values, np.array([[0, 1, 2, 3], [1, 4, 5, 2]]) + ) + + +def test_get_dual_rejects_faces_referencing_duplicate_nodes(): + """``construct_dual`` reads ``node_face_connectivity`` with no duplicate + handling, so a face still pointing at a dead duplicate index would yield a + degenerate dual face instead of an error. Merging at construction makes this + unreachable today; the guard keeps it that way.""" + node_lon = np.array([0.0, 1.0, 1.0, 0.0, 2.0, 2.0, 1.0, 1.0]) + node_lat = np.array([0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0]) + unmerged = np.array([[0, 1, 2, 3], [6, 4, 5, 7]]) + + grid = ux.Grid.from_topology(node_lon, node_lat, unmerged) + # Construction canonicalized the connectivity; put the duplicates back. + grid.face_node_connectivity = xr.DataArray( + unmerged, dims=grid.face_node_connectivity.dims + ) + + assert _check_duplicate_nodes_indices(grid) + with pytest.raises(GridInvalidError): + grid.get_dual() + + +def test_pole_exception_uses_a_chord_tolerance(): + """The pole carve-out must be a chord radius, not a raw ``|z|`` deviation. + + ``np.isclose(|z|, 1.0, atol=tolerance)`` also carries numpy's default + ``rtol=1e-5``, so the carve-out spanned ``1 - |z| <= 1.001e-5`` -- a chord of + 4.5e-3, or ~28 km on Earth. Every node within that cap was exempted from + merging. Only nodes at the pole itself may be exempt. + """ + from uxarray.grid.validation import _coincident_node_canonical_indices + + # Colatitude chosen so 1 - z = 1e-6: well inside the old carve-out, and far + # outside a chord of ERROR_TOLERANCE (whose cap is 1 - z <= 5e-17). + z = 1.0 - 1e-6 + x = np.sqrt(1.0 - z * z) + + points_xyz = np.array( + [ + [0.0, 0.0, 1.0], # north pole, kept distinct from the next node + [0.0, 0.0, 1.0], # same location, its own face-specific longitude + [x, 0.0, z], # near the pole, genuinely coincident with the next + [x, 0.0, z], + ] + ) + + canonical = _coincident_node_canonical_indices(points_xyz) + + # Nodes at a pole are still never merged with one another. + nt.assert_equal(canonical[:2], np.array([0, 1])) + # Near-pole coincident nodes now merge; before the fix they were exempt. + nt.assert_equal(canonical[2:], np.array([2, 2])) + + +def test_coincident_prescreen_keeps_both_ends_of_a_run(): + """The x-sorted prescreen must mark both members of a close pair. + + It flags a point when the gap to its predecessor *or* its successor is within + tolerance. Dropping either half of that OR silently loses one node of every + coincident pair, so this places coincident pairs at both ends of the sorted + order, where only one of the two neighbour tests fires. + """ + from uxarray.grid.validation import _coincident_node_canonical_indices + + # x strictly increasing and far apart, except for the duplicated first and + # last points, which have no predecessor / no successor respectively. + points_xyz = np.array( + [ + [0.0, 0.0, 0.0], # 0 + [0.0, 0.0, 0.0], # 1, coincident with 0 -> first in sorted x + [0.25, 0.5, 0.0], # 2 + [0.5, 0.5, 0.0], # 3 + [1.0, 0.0, 0.0], # 4 + [1.0, 0.0, 0.0], # 5, coincident with 4 -> last in sorted x + ] + ) + + canonical = _coincident_node_canonical_indices(points_xyz) + + nt.assert_equal(canonical, np.array([0, 0, 2, 3, 4, 4])) + + +def test_no_duplicate_nodes_ne30pg3(gridpath): + """``esmf/ne30/ne30pg3.grid.nc`` no longer reproduces issue #865's + duplicate-node bug; this only checks the general fix is a safe no-op.""" + grid_path = gridpath("esmf", "ne30", "ne30pg3.grid.nc") + grid = ux.open_grid(grid_path) + + assert len(_find_duplicate_nodes(grid)) == 0 diff --git a/test/test_subset.py b/test/test_subset.py index 29830559f..e04321331 100644 --- a/test/test_subset.py +++ b/test/test_subset.py @@ -78,8 +78,12 @@ def test_grid_nn_subset(gridpath): for grid_path in GRID_PATHS: grid = ux.open_grid(grid_path) - # corner-nodes - ks = [1, 2, grid.n_node - 1] + # corner-nodes -- k is bounded by the number of *live* (non-duplicate) + # nodes, since the node search tree excludes dead coincident indices + from uxarray.grid.validation import _live_node_indices + + n_live_nodes = len(_live_node_indices(grid)) + ks = [1, 2, n_live_nodes - 1] for coord in coord_locs: for k in ks: grid_subset = grid.subset.nearest_neighbor(coord, diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 1f397349b..5ccda041d 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -2331,8 +2331,12 @@ def get_dual(self): -------- dual : uxda Dual Mesh `uxda` constructed - """ + Raises + ------ + GridInvalidError + If any face still references a coincident duplicate node. + """ if _check_duplicate_nodes_indices(self.uxgrid): raise GridInvalidError("Duplicate nodes found, cannot construct dual") diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index e81431707..22d42eeed 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -899,8 +899,12 @@ def get_dual(self): -------- dual : uxds Dual Mesh `uxds` constructed - """ + Raises + ------ + GridInvalidError + If any face still references a coincident duplicate node. + """ if _check_duplicate_nodes_indices(self.uxgrid): raise GridInvalidError("Duplicate nodes found, cannot construct dual") diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index 54da63763..e45200fdb 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -4,7 +4,7 @@ import xarray as xr from numba import njit, prange -from uxarray.constants import INT_DTYPE, INT_FILL_VALUE +from uxarray.constants import ERROR_TOLERANCE, INT_DTYPE, INT_FILL_VALUE from uxarray.conventions import ugrid from uxarray.grid.utils import ( _build_pair_index, @@ -510,6 +510,161 @@ def _build_face_edge_connectivity( return face_edge_connectivity +def _remap_node_connectivity(connectivity, duplicate_node_map, n_node): + """Return a copy of connectivity with duplicate node indices canonicalized.""" + if not duplicate_node_map: + return connectivity + + lookup = np.arange(n_node, dtype=INT_DTYPE) + keys = np.fromiter( + duplicate_node_map.keys(), dtype=INT_DTYPE, count=len(duplicate_node_map) + ) + vals = np.fromiter( + duplicate_node_map.values(), dtype=INT_DTYPE, count=len(duplicate_node_map) + ) + lookup[keys] = vals + + remapped_connectivity = connectivity.copy() + valid = connectivity != INT_FILL_VALUE + remapped_connectivity[valid] = lookup[connectivity[valid]] + return remapped_connectivity + + +def _collapse_repeated_face_corners(face_node_connectivity, canonical_values): + """Collapse consecutive (cyclically) repeated corners in each face row. + + Remapping two originally-distinct, now-coincident corners of the same + face to a single canonical node can leave that node referenced twice in + a row, e.g. a quad (A, P, P, B) at a merged pole -- a triangle stored as + a 4-column row with one corner repeated. This pads it back down to + (A, P, B, FILL) so it is treated as the triangle it actually is. + + Only rows containing a value in ``canonical_values`` (nodes that + absorbed at least one duplicate) are inspected, since no other row can + have gained a repeat from the remap. + """ + if len(canonical_values) == 0: + return face_node_connectivity + + affected_rows = np.flatnonzero( + np.isin(face_node_connectivity, canonical_values).any(axis=1) + ) + if len(affected_rows) == 0: + return face_node_connectivity + + face_node_connectivity = face_node_connectivity.copy() + for row_index in affected_rows: + row = face_node_connectivity[row_index] + valid = row != INT_FILL_VALUE + n_valid = int(valid.sum()) + if n_valid <= 1: + continue + + corners = row[:n_valid] + keep = corners != np.roll(corners, 1) + if keep.all(): + continue + + compacted = corners[keep] + new_row = np.full_like(row, INT_FILL_VALUE) + new_row[: len(compacted)] = compacted + face_node_connectivity[row_index] = new_row + + return face_node_connectivity + + +# node-index-valued connectivity: safe to remap element-wise in place +_NODE_INDEX_CONNECTIVITY_TO_REMAP = ("face_node_connectivity", "node_node_connectivity") + +# connectivity derived from (and referencing) node indices, but whose rows must stay +# unique (e.g. edge_node_connectivity) -- dropped rather than remapped, so the +# existing lazy `@property` getters rebuild them cleanly from the corrected +# face_node_connectivity instead of leaving phantom duplicate rows behind. +_DERIVED_CONNECTIVITY_TO_INVALIDATE = ( + "edge_node_connectivity", + "face_edge_connectivity", + "edge_face_connectivity", + "face_face_connectivity", + "node_edge_connectivity", + "node_face_connectivity", +) + + +def _merge_coincident_grid_ds_nodes(grid_ds, tolerance=ERROR_TOLERANCE): + """Canonicalize coincident (within ``tolerance``) node indices in a raw grid + dataset's connectivity, before it is wrapped in a ``Grid``. + + Per issue #865, node coordinate/data arrays are left untouched -- only + connectivity references to coincident nodes are remapped to a single canonical + (lowest-indexed) node. Note this differs from TempestRemap and MOAB, which + delete the redundant nodes and renumber; keeping them preserves round-trip + fidelity and leaves node-centered data index-aligned, at the cost of leaving + unreferenced coordinates behind (see ``_live_node_indices``). + """ + from uxarray.grid.coordinates import _lonlat_rad_to_xyz + from uxarray.grid.validation import _coincident_node_canonical_indices + + if "face_node_connectivity" not in grid_ds: + return grid_ds + + if {"node_x", "node_y", "node_z"} <= set(grid_ds.variables): + points_xyz = np.column_stack( + ( + grid_ds["node_x"].values, + grid_ds["node_y"].values, + grid_ds["node_z"].values, + ) + ) + elif "node_lon" in grid_ds and "node_lat" in grid_ds: + points_xyz = np.column_stack( + _lonlat_rad_to_xyz( + np.deg2rad(grid_ds["node_lon"].values), + np.deg2rad(grid_ds["node_lat"].values), + ) + ) + else: + return grid_ds + + n_node = points_xyz.shape[0] + canonical = _coincident_node_canonical_indices(points_xyz, tolerance) + duplicate_node_map = { + INT_DTYPE(index): INT_DTYPE(canonical[index]) + for index in np.flatnonzero(canonical != np.arange(n_node, dtype=INT_DTYPE)) + } + if not duplicate_node_map: + return grid_ds + + grid_ds = grid_ds.copy() + + for name in _NODE_INDEX_CONNECTIVITY_TO_REMAP: + if name in grid_ds: + grid_ds[name] = grid_ds[name].copy( + data=_remap_node_connectivity( + grid_ds[name].values, duplicate_node_map, n_node + ) + ) + + if "face_node_connectivity" in grid_ds: + canonical_values = np.unique( + np.fromiter( + duplicate_node_map.values(), + dtype=INT_DTYPE, + count=len(duplicate_node_map), + ) + ) + grid_ds["face_node_connectivity"] = grid_ds["face_node_connectivity"].copy( + data=_collapse_repeated_face_corners( + grid_ds["face_node_connectivity"].values, canonical_values + ) + ) + + for name in _DERIVED_CONNECTIVITY_TO_INVALIDATE: + if name in grid_ds: + grid_ds = grid_ds.drop_vars(name) + + return grid_ds + + def _populate_node_face_connectivity(grid): """Constructs the UGRID connectivity variable (``node_face_connectivity``) and stores it within the internal (``Grid._ds``) and through the attribute diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 425e2a092..5a7c30b6b 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -26,6 +26,7 @@ from uxarray.grid.area import _get_all_face_area_from_coords from uxarray.grid.bounds import _populate_face_bounds from uxarray.grid.connectivity import ( + _merge_coincident_grid_ds_nodes, _populate_edge_face_connectivity, _populate_edge_node_connectivity, _populate_face_edge_connectivity, @@ -189,6 +190,11 @@ def __init__( ) # TODO: more checks for validate grid (lat/lon coords, etc) + # canonicalize coincident node indices in connectivity before this dataset + # is wrapped in a Grid, so every construction path benefits and no + # lazily-computed connectivity is ever built from stale indices. + grid_ds = _merge_coincident_grid_ds_nodes(grid_ds) + # mapping of ugrid dimensions and variables to source dataset's conventions self._source_dims_dict = source_dims_dict or {} @@ -2722,22 +2728,40 @@ def to_linecollection( return copy.deepcopy(line_collection) - def get_dual(self, check_duplicate_nodes: bool = False): + def get_dual(self, check_duplicate_nodes: bool | None = None): """Compute the dual for a grid, which constructs a new grid centered around the nodes, where the nodes of the primal become the face centers of the dual, and the face centers of the primal become the nodes of the dual. Returns a new `Grid` object. + Parameters + ---------- + check_duplicate_nodes : bool, optional + Deprecated and ignored. Coincident nodes are merged at grid + construction, so the check below always runs and always passes. + Returns -------- dual : Grid Dual Mesh Grid constructed + + Raises + ------ + GridInvalidError + If any face still references a coincident duplicate node. The dual + reads ``node_face_connectivity`` directly, so a dead duplicate index + yields a degenerate dual face rather than an error. """ + if check_duplicate_nodes is not None: + warnings.warn( + "`check_duplicate_nodes` is deprecated and ignored; coincident " + "nodes are merged at grid construction and always checked here.", + DeprecationWarning, + stacklevel=2, + ) - if check_duplicate_nodes: - if _check_duplicate_nodes_indices(self): - # TODO: This is very slow - raise GridInvalidError("Duplicate nodes found, cannot construct dual") + if _check_duplicate_nodes_indices(self): + raise GridInvalidError("Duplicate nodes found, cannot construct dual") # Get dual mesh node face connectivity dual_node_face_conn = construct_dual(grid=self) diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index eaf1f9b6a..be3da99b7 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -14,6 +14,7 @@ INT_FILL_VALUE, ) from uxarray.errors import DimensionError +from uxarray.grid.validation import _live_node_indices class KDTree: @@ -63,11 +64,18 @@ def __init__( self._tree_from_nodes = None self._tree_from_face_centers = None self._tree_from_edge_centers = None + # maps node-tree-local index -> original grid node index, set only when + # the node tree excludes dead duplicate node indices (see _build_from_nodes) + self._node_index_map = None # Build the tree based on nodes, face centers, or edge centers if coordinates == "nodes": self._tree_from_nodes = self._build_from_nodes() - self._n_elements = self._source_grid.n_node + self._n_elements = ( + len(self._node_index_map) + if self._node_index_map is not None + else self._source_grid.n_node + ) elif coordinates == "face centers": self._tree_from_face_centers = self._build_from_face_centers() self._n_elements = self._source_grid.n_face @@ -111,6 +119,13 @@ def _build_from_nodes(self): f"'spherical'" ) + live_indices = _live_node_indices(self._source_grid) + if len(live_indices) < len(coords): + self._node_index_map = live_indices + coords = coords[live_indices] + else: + self._node_index_map = None + self._tree_from_nodes = SKKDTree(coords, metric=self.distance_metric) return self._tree_from_nodes @@ -271,6 +286,9 @@ def query( ind = np.asarray(ind, dtype=INT_DTYPE) + if self._coordinates == "nodes" and self._node_index_map is not None: + ind = self._node_index_map[ind] + if coords.shape[0] == 1: ind = ind.squeeze() @@ -291,6 +309,9 @@ def query( ind = np.asarray(ind, dtype=INT_DTYPE) + if self._coordinates == "nodes" and self._node_index_map is not None: + ind = self._node_index_map[ind] + if coords.shape[0] == 1: ind = ind.squeeze() return ind @@ -360,6 +381,8 @@ def query_radius( ) ind = [np.asarray(cur_ind, dtype=INT_DTYPE) for cur_ind in ind] + if self._coordinates == "nodes" and self._node_index_map is not None: + ind = [self._node_index_map[cur_ind] for cur_ind in ind] d = [np.asarray(cur_d) for cur_d in d] if coords.shape[0] == 1: @@ -376,6 +399,8 @@ def query_radius( ) ind = [np.asarray(cur_ind, dtype=INT_DTYPE) for cur_ind in ind] + if self._coordinates == "nodes" and self._node_index_map is not None: + ind = [self._node_index_map[cur_ind] for cur_ind in ind] if coords.shape[0] == 1: ind = ind[0] @@ -394,7 +419,11 @@ def coordinates(self, value): if self._coordinates == "nodes": if self._tree_from_nodes is None or self.reconstruct: self._tree_from_nodes = self._build_from_nodes() - self._n_elements = self._source_grid.n_node + self._n_elements = ( + len(self._node_index_map) + if self._node_index_map is not None + else self._source_grid.n_node + ) elif self._coordinates == "face centers": if self._tree_from_face_centers is None or self.reconstruct: self._tree_from_face_centers = self._build_from_face_centers() @@ -455,11 +484,18 @@ def __init__( self._tree_from_nodes = None self._tree_from_face_centers = None self._tree_from_edge_centers = None + # maps node-tree-local index -> original grid node index, set only when + # the node tree excludes dead duplicate node indices (see _build_from_nodes) + self._node_index_map = None # set up appropriate reference to tree if coordinates == "nodes": self._tree_from_nodes = self._build_from_nodes() - self._n_elements = self._source_grid.n_node + self._n_elements = ( + len(self._node_index_map) + if self._node_index_map is not None + else self._source_grid.n_node + ) elif coordinates == "face centers": self._tree_from_face_centers = self._build_from_face_centers() self._n_elements = self._source_grid.n_face @@ -532,6 +568,14 @@ def _build_from_nodes(self): ), axis=-1, ) + + live_indices = _live_node_indices(self._source_grid) + if len(live_indices) < len(coords): + self._node_index_map = live_indices + coords = coords[live_indices] + else: + self._node_index_map = None + self._tree_from_nodes = SKBallTree(coords, metric=self.distance_metric) return self._tree_from_nodes @@ -655,6 +699,9 @@ def query( ind = np.asarray(ind, dtype=INT_DTYPE) + if self._coordinates == "nodes" and self._node_index_map is not None: + ind = self._node_index_map[ind] + if coords.shape[0] == 1: ind = ind.squeeze() @@ -675,6 +722,9 @@ def query( ind = np.asarray(ind, dtype=INT_DTYPE) + if self._coordinates == "nodes" and self._node_index_map is not None: + ind = self._node_index_map[ind] + if coords.shape[0] == 1: ind = ind.squeeze() @@ -742,6 +792,8 @@ def query_radius( ) ind = [np.asarray(cur_ind, dtype=INT_DTYPE) for cur_ind in ind] + if self._coordinates == "nodes" and self._node_index_map is not None: + ind = [self._node_index_map[cur_ind] for cur_ind in ind] d = [np.asarray(cur_d) for cur_d in d] if coords.shape[0] == 1: @@ -758,6 +810,8 @@ def query_radius( ) ind = [np.asarray(cur_ind, dtype=INT_DTYPE) for cur_ind in ind] + if self._coordinates == "nodes" and self._node_index_map is not None: + ind = [self._node_index_map[cur_ind] for cur_ind in ind] if coords.shape[0] == 1: ind = ind[0] @@ -776,7 +830,11 @@ def coordinates(self, value): if self._coordinates == "nodes": if self._tree_from_nodes is None or self.reconstruct: self._tree_from_nodes = self._build_from_nodes() - self._n_elements = self._source_grid.n_node + self._n_elements = ( + len(self._node_index_map) + if self._node_index_map is not None + else self._source_grid.n_node + ) elif self._coordinates == "face centers": if self._tree_from_face_centers is None or self.reconstruct: self._tree_from_face_centers = self._build_from_face_centers() diff --git a/uxarray/grid/validation.py b/uxarray/grid/validation.py index 36c05b119..faede2a29 100644 --- a/uxarray/grid/validation.py +++ b/uxarray/grid/validation.py @@ -4,10 +4,18 @@ import polars as pl from uxarray.constants import ERROR_TOLERANCE, INT_DTYPE +from uxarray.grid.coordinates import _lonlat_rad_to_xyz def _check_connectivity(grid): - """Check if all nodes are referenced by at least one element.""" + """Check if all nodes are referenced by at least one element. + + Node indices that are coincident duplicates of a node that *is* + referenced are expected to be unreferenced -- connectivity is + canonicalized to point at a single index per coincident group, while + the duplicate coordinates themselves are left in place (see + ``_find_duplicate_nodes``). + """ # Convert face_node_connectivity to a Polars Series and get unique values nodes_in_conn = pl.Series(grid.face_node_connectivity.values.flatten()).unique() @@ -15,12 +23,15 @@ def _check_connectivity(grid): # Filter out negative values nodes_in_conn = nodes_in_conn.filter(nodes_in_conn >= 0) - # Check if the size of unique nodes in connectivity is equal to the number of nodes - if len(nodes_in_conn) == grid.n_node: + n_duplicate_nodes = len(_find_duplicate_nodes(grid)) + + # Check if the size of unique nodes in connectivity is equal to the number of + # non-duplicate nodes + if len(nodes_in_conn) == grid.n_node - n_duplicate_nodes: return True else: warn( - f"Some nodes may not be referenced by any element. {len(nodes_in_conn)} and {grid.n_node}", + f"Some nodes may not be referenced by any element. {len(nodes_in_conn)} and {grid.n_node - n_duplicate_nodes}", RuntimeWarning, ) return False @@ -49,17 +60,17 @@ def _check_duplicate_nodes(grid): def _check_duplicate_nodes_indices(grid): - """Check if there are duplicate node indices, returns True if there are.""" - - # Create a duplication dictionary - duplicate_node_dict = _find_duplicate_nodes(grid) + """Check if any face still references a duplicate node index, returns True if + it does.""" - for face_nodes in grid.face_node_connectivity.values: - for node in face_nodes: - if node in duplicate_node_dict.keys(): - return True + duplicate_node_map = _find_duplicate_nodes(grid) + if not duplicate_node_map: + return False - return False + duplicate_indices = np.fromiter( + duplicate_node_map.keys(), dtype=INT_DTYPE, count=len(duplicate_node_map) + ) + return bool(np.isin(grid.face_node_connectivity.values, duplicate_indices).any()) def _check_area(grid): @@ -76,31 +87,113 @@ def _check_area(grid): return True -def _find_duplicate_nodes(grid): - # list of tuple indices - lonlat_t = [ - (lon, lat) for lon, lat in zip(grid.node_lon.values, grid.node_lat.values) +def _coincident_node_canonical_indices(points_xyz, tolerance=ERROR_TOLERANCE): + """For each point, find the lowest-indexed point within ``tolerance`` chordal + distance on the unit sphere (a point with no coincident neighbor maps to itself). + + Points at the geographic poles are never merged with one another: longitude is + singular there, and grid files (e.g. SCRIP cube-sphere) commonly give each face + touching a pole its own arbitrary-but-meaningful longitude for that corner, which + downstream lat/lon bounds and zonal-weight code relies on staying distinct per + face even though the xyz location is identical. + """ + from scipy.sparse import coo_matrix + from scipy.sparse.csgraph import connected_components + from scipy.spatial import KDTree + + n_points = len(points_xyz) + canonical = np.arange(n_points, dtype=INT_DTYPE) + + # ``tolerance`` is a chord radius (see the ``query_pairs`` call below), so it + # cannot be used directly as a deviation of |z| from 1. For a point at + # colatitude t from the pole, 1 - |z| = 1 - cos(t) = 2*sin(t/2)**2 = chord**2/2. + pole_mask = np.isclose( + np.abs(points_xyz[:, 2]), 1.0, rtol=0.0, atol=tolerance**2 / 2 + ) + mergeable_indices = np.flatnonzero(~pole_mask) + + if len(mergeable_indices) < 2: + return canonical + + mergeable_xyz = points_xyz[mergeable_indices] + + # Prescreen on x before paying for a KDTree. Two points within a chord of + # ``tolerance`` differ by at most ``tolerance`` in x, so in x-sorted order every + # consecutive gap between them is also at most ``tolerance``. A point whose + # sorted neighbours are both further than that in x therefore cannot be + # coincident with anything and is dropped. This is exact -- no candidate pair is + # lost -- and on the common case of a grid with no coincident nodes it replaces + # the tree build entirely with one sort. + order = np.argsort(mergeable_xyz[:, 0], kind="stable") + gap_is_small = np.diff(mergeable_xyz[order, 0]) <= tolerance + is_candidate = np.zeros(len(order), dtype=bool) + is_candidate[:-1] |= gap_is_small + is_candidate[1:] |= gap_is_small + # kept ascending so that the first member of a connected component below is + # still the lowest-numbered node in it + candidates = np.sort(order[is_candidate]) + + if len(candidates) < 2: + return canonical + + tree = KDTree(mergeable_xyz[candidates]) + pairs = tree.query_pairs(r=tolerance, output_type="ndarray") + + if len(pairs) == 0: + return canonical + + rows = np.concatenate([pairs[:, 0], pairs[:, 1]]) + cols = np.concatenate([pairs[:, 1], pairs[:, 0]]) + n_candidates = len(candidates) + adj_matrix = coo_matrix( + (np.ones(len(rows)), (rows, cols)), shape=(n_candidates, n_candidates) + ) + _, labels = connected_components(csgraph=adj_matrix, directed=False) + + unique_labels, first_indices = np.unique(labels, return_index=True) + sub_canonical = first_indices[np.searchsorted(unique_labels, labels)] + canonical[mergeable_indices[candidates]] = mergeable_indices[ + candidates[sub_canonical] ] + return canonical - # # Dictionary to track first occurrence and subsequent indices - occurrences = {} - # Iterate through the list and track occurrences - for index, tpl in enumerate(lonlat_t): - if tpl in occurrences: - occurrences[tpl].append((INT_DTYPE(index))) - else: - occurrences[tpl] = [INT_DTYPE(index)] +def _find_duplicate_node_map(node_lon, node_lat, tolerance=ERROR_TOLERANCE): + """Map duplicate (within ``tolerance`` on the unit sphere) node indices to the + lowest-indexed node sharing their location.""" + points_xyz = np.column_stack( + _lonlat_rad_to_xyz(np.deg2rad(node_lon), np.deg2rad(node_lat)) + ) + canonical = _coincident_node_canonical_indices(points_xyz, tolerance) - duplicate_dict = {} + n_node = len(node_lon) + duplicate_indices = np.flatnonzero(canonical != np.arange(n_node, dtype=INT_DTYPE)) + return { + INT_DTYPE(index): INT_DTYPE(canonical[index]) for index in duplicate_indices + } - for tpl, indices in occurrences.items(): - if len(indices) > 1: - source_idx = indices[0] - for duplicate_idx in indices[1:]: - duplicate_dict[duplicate_idx] = source_idx - return duplicate_dict +def _find_duplicate_nodes(grid): + """Map duplicate node indices to the canonical (lowest-indexed) node sharing + their coordinates.""" + return _find_duplicate_node_map(grid.node_lon.values, grid.node_lat.values) + + +def _live_node_indices(grid): + """Node indices still referenced after connectivity is canonicalized. + + Duplicate node coordinates are left in the node arrays by design (see + ``_find_duplicate_nodes``), but a raw coordinate-space search (e.g. a + node KDTree/BallTree) can otherwise select a dead duplicate index that no + face references. Callers building such trees should restrict to this set. + """ + duplicate_map = _find_duplicate_nodes(grid) + if not duplicate_map: + return np.arange(grid.n_node, dtype=INT_DTYPE) + dead = np.fromiter(duplicate_map.keys(), dtype=INT_DTYPE, count=len(duplicate_map)) + return np.setdiff1d( + np.arange(grid.n_node, dtype=INT_DTYPE), dead, assume_unique=True + ) def _check_normalization(grid): diff --git a/uxarray/io/_scrip.py b/uxarray/io/_scrip.py index 676c3b35a..9428fa1a1 100644 --- a/uxarray/io/_scrip.py +++ b/uxarray/io/_scrip.py @@ -60,8 +60,11 @@ def _to_ugrid(in_ds, out_ds): "original_index" ) - # Get unique rows (first occurrence). This preserves the order in which they appear. - unique_df = df.unique(subset=["lon", "lat"], keep="first") + # Get unique rows (first occurrence). maintain_order is required for this to be + # deterministic across runs -- polars' default unique() may otherwise reorder rows, + # which would make the resulting node index assignment (and downstream duplicate-node + # canonicalization) non-reproducible. + unique_df = df.unique(subset=["lon", "lat"], keep="first", maintain_order=True) # unq_ind: The indices of the unique rows in the original array unq_ind = unique_df["original_index"].to_numpy().astype(INT_DTYPE)