diff --git a/test/test_subset.py b/test/test_subset.py index 29830559f..ad4cf633b 100644 --- a/test/test_subset.py +++ b/test/test_subset.py @@ -5,6 +5,7 @@ from uxarray.grid.slice import _remap_dense, _remap_kernel, _remap_searchsorted import pytest +import xarray as xr def test_repr(gridpath, datasetpath): @@ -331,3 +332,63 @@ def open_with_edges(**kwargs): eager_subset._ds[name].values, err_msg=f"force_sparse={force_sparse}: {name}", ) + + +def test_bounding_box_with_stray_grid_time(gridpath): + """Subsetting must work when the grid was built from a source that carried a stray + scalar ``time`` coordinate. Ensures issue #1444 has been fixed. + """ + # grid built from a source carrying a stray scalar `time` + grid_ds = xr.open_dataset(gridpath("ugrid", "quad-hexagon", "grid.nc")) + uxgrid = ux.open_grid(grid_ds.assign_coords(time=np.datetime64("2016-01-01"))) + + # the stray coordinate is dropped during grid construction + assert "time" not in uxgrid._ds.coords + + # a face-centered variable carrying a *different* time coordinate + times = np.array(["2018-01-01", "2020-01-01"], dtype="datetime64[ns]") + uxda = ux.UxDataArray( + data=np.ones((times.size, uxgrid.n_face)), + dims=("time", "n_face"), + coords={"time": times}, + uxgrid=uxgrid, + ) + + # previously raised: "IndexError: dimension coordinate 'time' conflicts ..." + res = uxda.subset.bounding_box(lon_bounds=(-10, 10), lat_bounds=(-10, 10)) + + assert isinstance(res, ux.UxDataArray) + assert res.sizes["time"] == times.size + assert "n_face" in res.dims + + +def test_bounding_box_with_stray_grid_time_dimension(gridpath): + """Subsetting must also work when the grid source carried a full ``time`` *dimension* + (a variable along ``time``) rather than the scalar coordinate of issue #1444. + + A dimensional coordinate has dims ``(time,)``, which is not a subset of the + ``(n_face,)`` ``subgrid_face_indices`` indexer, so — unlike a 0-d coordinate — it + never attaches to that indexer and cannot collide during ``.isel``. The stray + dimension harmlessly rides along in ``Grid._ds`` but never reaches the indexing path. + """ + # grid built from a source carrying a variable along a stray `time` dimension + grid_ds = xr.open_dataset(gridpath("ugrid", "quad-hexagon", "grid.nc")) + uxgrid = ux.open_grid(grid_ds.assign(stray=("time", [0.0, 1.0]))) + + # the bare dimension is not a coordinate, so it is not dropped; that is harmless + assert "time" in uxgrid._ds.dims + + # a face-centered variable carrying its own time coordinate + times = np.array(["2018-01-01", "2020-01-01"], dtype="datetime64[ns]") + uxda = ux.UxDataArray( + data=np.ones((times.size, uxgrid.n_face)), + dims=("time", "n_face"), + coords={"time": times}, + uxgrid=uxgrid, + ) + + res = uxda.subset.bounding_box(lon_bounds=(-10, 10), lat_bounds=(-10, 10)) + + assert isinstance(res, ux.UxDataArray) + assert res.sizes["time"] == times.size + assert "n_face" in res.dims diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 711b30389..39e181877 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -66,7 +66,7 @@ _populate_edge_node_distances, ) from uxarray.grid.point_in_face import _point_in_face_query -from uxarray.grid.utils import make_setter +from uxarray.grid.utils import _drop_non_grid_coords, make_setter from uxarray.grid.validation import ( _check_area, _check_connectivity, @@ -191,8 +191,10 @@ def __init__( # source grid specification (i.e. UGRID, MPAS, SCRIP, etc.) self.source_grid_spec = source_grid_spec - # internal xarray dataset for storing grid variables - self._ds = grid_ds + # internal xarray dataset for storing grid variables. + # drop stray scalar coordinates (e.g. a `time` carried in from the source file) + # so they can't leak onto the grid and collide during subsetting (see #1444). + self._ds = _drop_non_grid_coords(grid_ds) # source grid specification (i.e. UGRID, MPAS, SCRIP, etc.) self.source_grid_spec = source_grid_spec diff --git a/uxarray/grid/utils.py b/uxarray/grid/utils.py index a33f50629..5b3706012 100644 --- a/uxarray/grid/utils.py +++ b/uxarray/grid/utils.py @@ -3,6 +3,7 @@ from numba import njit, prange from uxarray.constants import INT_FILL_VALUE +from uxarray.conventions import ugrid from uxarray.utils.numba_math import ( _numba_add3, _numba_mul3_scalar, @@ -11,6 +12,27 @@ ) +def _drop_non_grid_coords(ds): + """Drop scalar (0-d) coordinates that aren't recognized grid coordinates + (e.g. a stray ``time`` carried in from the source file). + + Only 0-d strays can cause the subset collision behind issue #1444: a scalar + coordinate attaches to every variable, including the ``(n_face,)`` + ``subgrid_*_indices`` indexer, where it can clash with a like-named dimension + on the data being indexed. A coordinate with its own dimension (e.g. ICON's + 1-d ``clon``/``clat`` or FESOM2's ``lon``/``lat``) never attaches to that + indexer, so it is left in place and those grids round-trip unchanged. + + Coordinate-only, so grid data variables — connectivity, descriptors, and the + subset's ``subgrid_*_indices`` — are always left intact. + """ + grid_coords = set(ugrid.SPHERICAL_COORD_NAMES) | set(ugrid.CARTESIAN_COORD_NAMES) + stray = [ + coord for coord in ds.coords if coord not in grid_coords and ds[coord].ndim == 0 + ] + return ds.drop_vars(stray, errors="ignore") + + @njit(cache=True) def _small_angle_of_2_vectors(u, v): """