diff --git a/test/grid/grid/test_core.py b/test/grid/grid/test_core.py index da5882bea..73d687bc7 100644 --- a/test/grid/grid/test_core.py +++ b/test/grid/grid/test_core.py @@ -129,7 +129,53 @@ 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.""" + from uxarray.grid.validation import _check_duplicate_nodes_indices, _find_duplicate_nodes + + grid_path = gridpath("ugrid", "geoflow-small", "grid.nc") + grid = ux.open_grid(grid_path) + + # source file has duplicate node coordinates; connectivity should already + # be canonicalized to a single index per coincident group + assert len(_find_duplicate_nodes(grid)) > 0 + 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 + 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).""" + from uxarray.grid.validation import _check_duplicate_nodes_indices, _find_duplicate_nodes + + 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_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.""" + from uxarray.grid.validation import _find_duplicate_nodes + + 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/io/test_structured.py b/test/io/test_structured.py index 54e735e4b..166387a18 100644 --- a/test/io/test_structured.py +++ b/test/io/test_structured.py @@ -74,3 +74,52 @@ def test_from_xarray_with_grid_from_latlon(ds_name): subset = uxds["air"].isel(time=0).subset.bounding_circle((-100.0, 40.0), 5) assert "n_face" in subset.dims assert subset.sizes["n_face"] > 0 + + +def test_global_structured_grid_merges_poles_and_seam(): + """Nodes coincident on the sphere must be merged, even though their + (lon, lat) pairs differ. Regression test for issue #1689.""" + import numpy as np + + n_lon, n_lat = 36, 18 + d_lat = 180.0 / n_lat + lon = np.linspace(-180, 180, n_lon, endpoint=False) + lat = np.linspace(-90 + d_lat / 2, 90 - d_lat / 2, n_lat) + + uxgrid = ux.Grid.from_structured(lon=lon, lat=lat) + + # Every duplicated pole node and antimeridian node must be gone. + assert uxgrid.n_node < (n_lon + 1) * (n_lat + 1) + assert np.isclose(uxgrid.node_lat.values, 90.0).sum() == 1 + assert np.isclose(uxgrid.node_lat.values, -90.0).sum() == 1 + + # A closed sphere: V - E + F == 2. + assert uxgrid.n_node - uxgrid.n_edge + uxgrid.n_face == 2 + + # The pole is now a real singularity touching every longitude column, and + # its faces are triangles rather than quads with a repeated corner. + face_nodes = uxgrid.face_node_connectivity.values + n_nodes_per_face = uxgrid.n_nodes_per_face.values + assert (n_nodes_per_face == 3).sum() == 2 * n_lon + + for face, n_nodes in zip(face_nodes, n_nodes_per_face): + nodes = face.tolist()[:n_nodes] + assert len(set(nodes)) == n_nodes + + pole = int(np.flatnonzero(np.isclose(uxgrid.node_lat.values, 90.0))[0]) + assert (face_nodes == pole).any(axis=1).sum() == n_lon + + +def test_regional_structured_grid_is_unchanged(): + """A grid that touches neither pole nor the antimeridian must keep every + node and stay entirely quadrilateral.""" + import numpy as np + + lon = np.linspace(-50, -10, 20) + lat = np.linspace(10, 40, 15) + + uxgrid = ux.Grid.from_structured(lon=lon, lat=lat) + + assert uxgrid.n_node == 21 * 16 + assert uxgrid.n_face == 20 * 15 + assert (uxgrid.n_nodes_per_face.values == 4).all() 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 b62b7f1c9..9587050e9 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -34,7 +34,6 @@ from uxarray.formatting_html import array_repr from uxarray.grid import Grid from uxarray.grid.dual import construct_dual -from uxarray.grid.validation import _check_duplicate_nodes_indices from uxarray.io._healpix import get_zoom_from_cells from uxarray.plot.accessor import UxDataArrayPlotAccessor from uxarray.remap.accessor import RemapAccessor @@ -2177,9 +2176,6 @@ def get_dual(self): Dual Mesh `uxda` constructed """ - if _check_duplicate_nodes_indices(self.uxgrid): - raise GridInvalidError("Duplicate nodes found, cannot construct dual") - if self.uxgrid.partial_sphere_coverage: warn( "This mesh is partial, which could cause inconsistent results and data will be lost", diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 129466589..07b4bc027 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -18,7 +18,6 @@ from uxarray.formatting_html import dataset_repr from uxarray.grid import Grid from uxarray.grid.dual import construct_dual -from uxarray.grid.validation import _check_duplicate_nodes_indices from uxarray.io._healpix import get_zoom_from_cells from uxarray.plot.accessor import UxDatasetPlotAccessor from uxarray.remap.accessor import RemapAccessor @@ -706,9 +705,6 @@ def get_dual(self): Dual Mesh `uxds` constructed """ - if _check_duplicate_nodes_indices(self.uxgrid): - raise GridInvalidError("Duplicate nodes found, cannot construct dual") - if self.uxgrid.partial_sphere_coverage: warn( "This mesh is partial, which could cause inconsistent results and data will be lost", diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index ac9658979..a6854e0e8 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -2,7 +2,7 @@ import xarray as xr from numba import njit -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 @@ -299,6 +299,158 @@ def _build_face_edge_connectivity(inverse_indices, n_face, n_max_face_nodes): return inverse_indices +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 _dedupe_grid_ds_nodes(grid_ds, tolerance=ERROR_TOLERANCE): + """Canonicalize duplicate (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 duplicate nodes are remapped to a single canonical + (lowest-indexed) node. + """ + 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 711b30389..f4bd120c5 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -23,6 +23,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 ( + _dedupe_grid_ds_nodes, _populate_edge_face_connectivity, _populate_edge_node_connectivity, _populate_face_edge_connectivity, @@ -71,7 +72,6 @@ _check_area, _check_connectivity, _check_duplicate_nodes, - _check_duplicate_nodes_indices, _check_normalization, ) from uxarray.io._delaunay import ( @@ -89,7 +89,10 @@ from uxarray.io._icon import _read_icon from uxarray.io._mpas import _read_mpas from uxarray.io._scrip import _encode_scrip, _read_scrip -from uxarray.io._structured import _read_structured_grid +from uxarray.io._structured import ( + _DEFAULT_STRUCTURED_TOL_DEG, + _read_structured_grid, +) from uxarray.io._topology import _read_topology from uxarray.io._ugrid import ( _encode_ugrid, @@ -185,6 +188,11 @@ def __init__( ) # TODO: more checks for validate grid (lat/lon coords, etc) + # canonicalize duplicate (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 = _dedupe_grid_ds_nodes(grid_ds) + # mapping of ugrid dimensions and variables to source dataset's conventions self._source_dims_dict = source_dims_dict or {} @@ -501,7 +509,11 @@ def from_topology( @classmethod def from_structured( - cls, ds: xr.Dataset = None, lon=None, lat=None, tol: float | None = 1e-10 + cls, + ds: xr.Dataset = None, + lon=None, + lat=None, + tol: float | None = _DEFAULT_STRUCTURED_TOL_DEG, ): """ Converts a structured ``xarray.Dataset`` or longitude and latitude coordinates into an unstructured ``uxarray.Grid``. @@ -525,8 +537,9 @@ def from_structured( Should be a one-dimensional or two-dimensional array following CF conventions. tol : float, optional - Tolerance for considering nodes as identical when constructing the grid from longitude and latitude. - Default is `1e-10`. + Tolerance in degrees for considering nodes as identical when constructing the grid from + longitude and latitude. Defaults to the angle whose chord length on the unit sphere equals + ``uxarray.constants.ERROR_TOLERANCE``. Returns ------- @@ -2606,22 +2619,30 @@ 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. Duplicate nodes are now merged while + constructing the dual, so they no longer prevent it. + Returns -------- dual : Grid Dual Mesh Grid constructed """ - - 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 is not None: + warnings.warn( + "`check_duplicate_nodes` is deprecated and ignored; duplicate nodes " + "are now merged when constructing the dual.", + DeprecationWarning, + stacklevel=2, + ) # 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 195bf4138..e5e56dde7 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -5,6 +5,7 @@ from uxarray.constants import ERROR_TOLERANCE, INT_DTYPE, INT_FILL_VALUE from uxarray.errors import DimensionError +from uxarray.grid.validation import _live_node_indices class KDTree: @@ -54,11 +55,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 @@ -102,6 +110,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 @@ -262,6 +277,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() @@ -282,6 +300,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 @@ -351,6 +372,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: @@ -367,6 +390,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] @@ -385,7 +410,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() @@ -446,11 +475,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 @@ -523,6 +559,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 @@ -646,6 +690,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() @@ -666,6 +713,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() @@ -733,6 +783,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: @@ -749,6 +801,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] @@ -767,7 +821,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..0885d1282 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 @@ -76,31 +87,85 @@ def _check_area(grid): return True +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) + + pole_mask = np.isclose(np.abs(points_xyz[:, 2]), 1.0, atol=tolerance) + mergeable_indices = np.flatnonzero(~pole_mask) + + if len(mergeable_indices) < 2: + return canonical + + tree = KDTree(points_xyz[mergeable_indices]) + 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_mergeable = len(mergeable_indices) + adj_matrix = coo_matrix( + (np.ones(len(rows)), (rows, cols)), shape=(n_mergeable, n_mergeable) + ) + _, 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] = mergeable_indices[sub_canonical] + return canonical + + +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) + + 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 + } + + 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) - ] - - # # 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)] - - duplicate_dict = {} - - 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 + """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) diff --git a/uxarray/io/_structured.py b/uxarray/io/_structured.py index e3b8cab1e..cced26e20 100644 --- a/uxarray/io/_structured.py +++ b/uxarray/io/_structured.py @@ -1,11 +1,17 @@ import numpy as np import xarray as xr -from uxarray.constants import INT_DTYPE +from uxarray.constants import ERROR_TOLERANCE, INT_DTYPE, INT_FILL_VALUE from uxarray.conventions import ugrid +from uxarray.grid.coordinates import _lonlat_rad_to_xyz +# ``ERROR_TOLERANCE`` is defined as a Cartesian distance on the unit sphere; convert +# it to the angular tolerance `tol` is documented in so the default tracks the same +# precision assumption used everywhere else in the codebase. +_DEFAULT_STRUCTURED_TOL_DEG = np.rad2deg(2.0 * np.arcsin(ERROR_TOLERANCE / 2.0)) -def _read_structured_grid(lon, lat, tol=1e-10): + +def _read_structured_grid(lon, lat, tol=_DEFAULT_STRUCTURED_TOL_DEG): """ Constructs an unstructured grid dataset from structured longitude and latitude coordinates. @@ -21,7 +27,9 @@ def _read_structured_grid(lon, lat, tol=1e-10): lat : array_like 1D array of latitude coordinates in degrees. tol : float, optional - Tolerance for considering nodes as identical (default is `1e-10`). + Tolerance in degrees for considering nodes as identical. Defaults to the angle + whose chord length on the unit sphere equals ``uxarray.constants.ERROR_TOLERANCE``, + matching the precision assumption used elsewhere in the codebase. Returns ------- @@ -42,6 +50,13 @@ def _read_structured_grid(lon, lat, tol=1e-10): out_ds = xr.Dataset() + # Coincidence detection below relies on float64 precision (~1e-16); real-world + # datasets often store lon/lat as float32 (~1e-7), which silently propagates + # through this pipeline and causes pole/antimeridian merges to fail or merge + # only partially, regardless of ``tol``. + lon = np.asarray(lon, dtype=np.float64) + lat = np.asarray(lat, dtype=np.float64) + sorted_indices = np.argsort(lon) lon = lon[sorted_indices] @@ -79,11 +94,22 @@ def _read_structured_grid(lon, lat, tol=1e-10): # Stack longitude and latitude for processing nodes = np.column_stack((node_lon, node_lat)) + # Match nodes on the sphere rather than in the lon/lat plane, so that the poles + # (many lon values, one point) and the antimeridian seam (lon differing by 360) + # are recognized as coincident. + lon_rad = np.deg2rad(node_lon) + lat_rad = np.deg2rad(node_lat) + node_xyz = np.column_stack(_lonlat_rad_to_xyz(lon_rad, lat_rad)) + # Build KDTree - tree = KDTree(nodes) + tree = KDTree(node_xyz) + + # ``tol`` is an angle in degrees; on the unit sphere the matching radius is the + # chord subtended by that angle, so the threshold keeps its documented meaning. + chord_tol = 2.0 * np.sin(np.deg2rad(tol) / 2.0) # Find all pairs of nodes within the tolerance - pairs = tree.query_pairs(r=tol) + pairs = tree.query_pairs(r=chord_tol) n_nodes = len(nodes) if pairs: @@ -138,6 +164,34 @@ def _read_structured_grid(lon, lat, tol=1e-10): # Stack the node indices to form face_node_connectivity face_node_conn = np.vstack((n1, n2, n3, n4), dtype=INT_DTYPE).T + # No new faces are created here -- this only shrinks the width of existing rows + # in face_node_conn for faces that became degenerate after the pole merge above. + # + # A face touching the pole is built from 2 distinct edge-longitudes at the pole + # latitude, e.g. corners (n1, n2, n3, n4) = (A, P, P, B), where P is the single + # merged pole node that both pole-row corners now point to (n2 == n3). That is a + # triangle A-P-B stored as a 4-column quad with one corner repeated, so: + # 1. `keep` marks, per face, which corners differ from their cyclic predecessor + # (n2 == n3 above means the P at position 2 is dropped from that row). + # 2. The kept corners are pushed to the front of each row (`order`), giving + # (A, P, B, B) instead of (A, P, P, B) -- still 4 columns, but the last + # column is now the padding slot for a 3-node face. + # 3. `n_max_face_nodes` is the largest node count any face still needs (3 here, + # unless some other face in the grid still has 4 distinct corners, in which + # case nothing is trimmed and this is a no-op). Columns beyond each face's + # own count are set to `INT_FILL_VALUE`, giving (A, P, B, FILL). + keep = face_node_conn != np.roll(face_node_conn, 1, axis=1) + if not keep.all(): + n_nodes_per_face = keep.sum(axis=1) + order = np.argsort(~keep, axis=1, kind="stable") + compacted = np.take_along_axis(face_node_conn, order, axis=1) + n_max_face_nodes = n_nodes_per_face.max() + compacted = compacted[:, :n_max_face_nodes] + compacted[np.arange(n_max_face_nodes) >= n_nodes_per_face[:, None]] = ( + INT_FILL_VALUE + ) + face_node_conn = compacted + out_ds["node_lon"] = xr.DataArray( data=unique_node_lon, dims=ugrid.NODE_DIM, attrs=ugrid.NODE_LON_ATTRS )