From 1d9cc01c76c84cf43ba393311ea6e0defed9e9c4 Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Mon, 14 Sep 2026 14:44:07 -0400 Subject: [PATCH 01/14] fix UxDataset.isel() bugs (#1713, #1714) #1713: UxDataset.isel() silently fails to slice when providing grid dim not in the dataset #1714: In UxDataset.isel() if "n_face" coords assigned: ValueError: dimension 'n_face' already exists as a scalar variable --- test/core/test_indexing.py | 61 +++++++++++++++++++------------------- uxarray/core/dataarray.py | 20 ++++++------- uxarray/core/dataset.py | 57 +++++++++++++---------------------- 3 files changed, 60 insertions(+), 78 deletions(-) diff --git a/test/core/test_indexing.py b/test/core/test_indexing.py index f70725267..34bf0930c 100644 --- a/test/core/test_indexing.py +++ b/test/core/test_indexing.py @@ -39,17 +39,17 @@ def test_sel_indexes_grid(): def test_sel_uses_grid_dim_labels(): """ensure obj.sel({grid_dim: ...}) actually utilizes coordinate labels on that grid dim, for UxDataArrays and UxDatasets. Regression test for #1641. - TODO: fix #1714 then uncomment the UxDataset tests below + Also contains regression test for #1714. """ # test corresponding to the workflow described in #1641, but for UxDataset uxds = ux.tutorial.open_dataset("outCSne30-vortex") - # (uncomment the next few lines after fixing #1714) - # uxds1 = uxds.assign_coords(n_face=np.arange(uxds.n_face.size)) - # uxds2 = uxds1.isel(n_face=range(0, 100, 5)) - # uxds3 = uxds2 + 7 - # # "check what the results look like on what were originally faces 20, 30, and 40" - # uxds4 = uxds3.sel(n_face=[20,30,40]) - # assert uxds4.sizes['n_face'] == uxds4.uxgrid.n_face == 3 + # (the next few lines, through the `assert`, also serve as a regression test for #1714) + uxds1 = uxds.assign_coords(n_face=np.arange(uxds.n_face.size)) + uxds2 = uxds1.isel(n_face=range(0, 100, 5)) + uxds3 = uxds2 + 7 + # "check what the results look like on what were originally faces 20, 30, and 40" + uxds4 = uxds3.sel(n_face=[20,30,40]) + assert uxds4.sizes['n_face'] == uxds4.uxgrid.n_face == 3 # test corresponding to the workflow described in #1641, for UxDataArray uxarr = uxds['psi'] @@ -70,15 +70,14 @@ def test_sel_uses_grid_dim_labels(): def test_can_index_grid_dim_not_in_data(): """ensure isel() and sel() can both index a grid dim even if that dim is not present in the data itself; - for UxDataArrays and UxDatasets. TODO: fix #1713 then uncomment the UxDataset tests below. + for UxDataArrays and UxDatasets. The UxDataset checks serve as a regression test for #1713. """ ds = ux.tutorial.open_dataset("outCSne30-vortex") - # (uncomment the next few lines after fixing #1713) - # assert "n_face" in ds.dims - # result = ds.isel(n_edge=7) - # assert result.sizes["n_face"] == result.uxgrid.n_face == 2 - # result = ds.sel(n_edge=7) - # assert result.sizes["n_face"] == result.uxgrid.n_face == 2 + assert "n_face" in ds.dims + result = ds.isel(n_edge=7) + assert result.sizes["n_face"] == result.uxgrid.n_face == 2 + result = ds.sel(n_edge=7) + assert result.sizes["n_face"] == result.uxgrid.n_face == 2 arr = ds["psi"] result = arr.isel(n_edge=7) @@ -110,7 +109,7 @@ def test_sel_can_use_slice(): """ensure sel() can use slice() objects as indexers, and provides expected results, with expected sizes, for UxDataArrays and UxDatasets. Regression test inspired by reviewer comment in #1641, also related to #1639. - TODO: fix #1714 then uncomment the relevant UxDataset tests below + Also contains regression test for #1714. """ grid = ux.Grid.from_healpix(zoom=0) # 12 faces arr = ux.UxDataArray( @@ -129,10 +128,10 @@ def test_sel_can_use_slice(): uxds = ux.UxDataset({'data': arr.to_xarray()}, uxgrid=grid) result = uxds.sel(n_face=slice(0, 2)) assert result.n_face.size == result.uxgrid.n_face == 2 - # (uncomment the next few lines after fixing #1714) - # labeled_ds = uxds.assign_coords(n_face=np.arange(grid.n_face)) - # result = labeled_ds.sel(n_face=slice(0, 2)) - # assert result.n_face.size == result.uxgrid.n_face == 3 + # (the remaining lines also serve as a regression test for #1714) + labeled_ds = uxds.assign_coords(n_face=np.arange(grid.n_face)) + result = labeled_ds.sel(n_face=slice(0, 2)) + assert result.n_face.size == result.uxgrid.n_face == 3 def test_isel_can_use_bool(): """ensure isel() supports indexing by a boolean indexer array. @@ -253,7 +252,7 @@ def test_sel_crash_if_provided_selection_options_with_coordless_dims(): whenever any of the indexed dims have no associated coordinates. (Tests below also demonstrate that this behavior is consistent with xarray.) Regression test inspired by reviewer comment in #1641. - TODO: fix #1714 then uncomment the relevant UxDataset tests below + Also includes a regression test for #1714. """ kw_options = ({"method": "nearest"}, {"method": "nearest", "tolerance": 0.1}) @@ -264,13 +263,13 @@ def test_sel_crash_if_provided_selection_options_with_coordless_dims(): assert set(ds0.dims) == {'n_face'} ds0_labeled = ds0.assign_coords({'n_face': [0,10,20,30]}) - # (uncomment the next few lines after fixing #1714) - # ds0.sel(n_face=[0,1]) # (sanity check: no crash when no options provided) - # for kw in kw_options: - # with pytest.raises(ValueError, match=r"cannot supply selection options.+for dimension 'n_face'"): - # ds0.sel(n_face=[0,1], **kw) # provides method, tolerance, or both. - # # separately: checking to ensure that passing these options is fine in "labeled" case. - # ds0_labeled.sel(n_face=[0,10], **kw) + # (the next few lines also serve as a regression test for #1714) + ds0.sel(n_face=[0,1]) # (sanity check: no crash when no options provided) + for kw in kw_options: + with pytest.raises(ValueError, match=r"cannot supply selection options.+for dimension 'n_face'"): + ds0.sel(n_face=[0,1], **kw) # provides method, tolerance, or both. + # separately: checking to ensure that passing these options is fine in "labeled" case. + ds0_labeled.sel(n_face=[0,10], **kw) # ensure same behavior for xarray objects: ds0.to_xarray().sel(n_face=[0,1]) @@ -280,9 +279,9 @@ def test_sel_crash_if_provided_selection_options_with_coordless_dims(): ds0_labeled.to_xarray().sel(n_face=[0,10], **kw) # ensure supplying just tolerance raises a different error, if indexing is otherwise valid: - # (uncomment the next few lines after fixing #1714) - # with pytest.raises(ValueError, match=r"tolerance argument only valid if doing.+"): - # ds0_labeled.sel(n_face=[0,10], tolerance=0.1) + # (the following pytest.raises statement also serves as a regression test for #1714) + with pytest.raises(ValueError, match=r"tolerance argument only valid if doing.+"): + ds0_labeled.sel(n_face=[0,10], tolerance=0.1) with pytest.raises(ValueError, match=r"tolerance argument only valid if doing.+"): ds0_labeled.to_xarray().sel(n_face=[0,10], tolerance=0.1) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 1f397349b..8ec48e7dc 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -11,7 +11,6 @@ from xarray.core.options import OPTIONS from xarray.core.utils import UncachedAccessor -import uxarray from uxarray.constants import GRID_DIMS from uxarray.core.aggregation import _uxda_grid_aggregate from uxarray.core.gradient import ( @@ -596,8 +595,10 @@ def to_dataset( ------- uxds: UxDataSet """ + from uxarray.core.dataset import UxDataset + xrds = super().to_dataset(dim=dim, name=name, promote_attrs=promote_attrs) - uxds = uxarray.core.dataset.UxDataset(xrds, uxgrid=self._uxgrid) + uxds = UxDataset(xrds, uxgrid=self._uxgrid) return uxds @@ -2075,18 +2076,17 @@ def isel( **{grid_dim: grid_indexer}, inverse_indices=inverse_indices ) - da = self._slice_from_grid(sliced_grid) + result = self._slice_from_grid(sliced_grid) # if there are any remaining indexers, apply them if indexers: - xarr = super(UxDataArray, da).isel( + result = super(UxDataArray, result).isel( indexers=indexers, drop=drop, missing_dims=missing_dims ) # re‐wrap so the grid sticks around - return type(self)(xarr, uxgrid=sliced_grid) + result = type(self)(result, uxgrid=sliced_grid) - # no other dims, return the grid‐sliced da - return da + return result else: # len(grid_dims)>1; _validate_indexers should have crashed. raise AssertionError("internal implementation error if reached this line") @@ -2321,7 +2321,7 @@ def _slice_from_grid(self, sliced_grid): "Data variable must be either node, edge, or face centered." ) - return UxDataArray(da_sliced, uxgrid=sliced_grid) + return type(self)(da_sliced, uxgrid=sliced_grid) def get_dual(self): """Compute the dual mesh for a data array, returns a new data array @@ -2359,9 +2359,7 @@ def get_dual(self): dims = [dim_map.get(dim, dim) for dim in self.dims] # Construct the new data array - uxda = uxarray.UxDataArray( - uxgrid=dual, data=self.data, dims=dims, name=self.name - ) + uxda = type(self)(uxgrid=dual, data=self.data, dims=dims, name=self.name) return uxda diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index e81431707..ebd9a4d50 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -11,7 +11,7 @@ from xarray.core.options import OPTIONS from xarray.core.utils import UncachedAccessor -import uxarray +from uxarray.constants import GRID_DIMS from uxarray.core.dataarray import UxDataArray from uxarray.core.utils import ( _map_dims_to_ugrid, @@ -387,33 +387,22 @@ def from_healpix( return cls.from_xarray(ds, uxgrid, {face_dim: "n_face"}) - def _slice_dataset_from_grid(self, sliced_grid, grid_dim: str, grid_indexer): + def _slice_from_grid(self, sliced_grid): + """returns UxDataset based on slicing self according to sliced_grid. + sliced_grid should be a ``Grid`` which came directly from self.uxgrid.isel(...) + (or from slicing something equal to self.uxgrid), else behavior is undefined. + """ data_vars = {} for name, da in self.data_vars.items(): - if grid_dim in da.dims: - if hasattr(da, "_slice_from_grid"): - data_vars[name] = da._slice_from_grid(sliced_grid) - else: - data_vars[name] = da.isel({grid_dim: grid_indexer}) + if hasattr(da, "_slice_from_grid") and any( + dim in da.dims for dim in GRID_DIMS + ): + data_vars[name] = da._slice_from_grid(sliced_grid) else: data_vars[name] = da - coords = {} - for cname, cda in self.coords.items(): - if grid_dim in cda.dims: - # Prefer authoritative coords from the sliced grid if available - replacement = getattr(sliced_grid, cname, None) - coords[cname] = ( - replacement - if replacement is not None - else cda.isel({grid_dim: grid_indexer}) - ) - else: - coords[cname] = cda - - ds = xr.Dataset(data_vars=data_vars, coords=coords, attrs=self.attrs) - - return ds + ds_sliced = xr.Dataset(data_vars=data_vars, attrs=self.attrs) + return type(self)(ds_sliced, uxgrid=sliced_grid) def isel( self, @@ -500,23 +489,21 @@ def isel( indexers = indexers.copy() # don't modify the original dict grid_indexer = indexers.pop(grid_dim) - # slice the grid sliced_grid = self.uxgrid.isel( **{grid_dim: grid_indexer}, inverse_indices=inverse_indices ) - ds = self._slice_dataset_from_grid( - sliced_grid=sliced_grid, - grid_dim=grid_dim, - grid_indexer=grid_indexer, - ) + result = self._slice_from_grid(sliced_grid) + # if there are any remaining indexers, apply them if indexers: - ds = xr.Dataset.isel( - ds, indexers=indexers, drop=drop, missing_dims=missing_dims + result = super(UxDataset, result).isel( + indexers=indexers, drop=drop, missing_dims=missing_dims ) + # re‐wrap so the grid sticks around + result = type(self)(result, uxgrid=sliced_grid) - return type(self)(ds, uxgrid=sliced_grid) + return result else: # len(grid_dims)>1; _validate_indexers should have crashed. raise AssertionError("internal implementation error if reached this line") @@ -921,7 +908,7 @@ def get_dual(self): ) # Initialize new dataset - dataset = uxarray.UxDataset(uxgrid=dual) + dataset = type(self)(uxgrid=dual) # Dictionary to swap dimensions dim_map = {"n_face": "n_node", "n_node": "n_face"} @@ -932,9 +919,7 @@ def get_dual(self): dims = [dim_map.get(dim, dim) for dim in self[var].dims] # Construct the new data array - uxda = uxarray.UxDataArray( - uxgrid=dual, data=self[var].data, dims=dims, name=var - ) + uxda = UxDataArray(uxgrid=dual, data=self[var].data, dims=dims, name=var) # Add data array to dataset dataset[var] = uxda From 2f95f2ff030adec4808c08cb3da93bda539a031c Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Mon, 14 Sep 2026 17:51:05 -0400 Subject: [PATCH 02/14] fix: put isel grid dim indexer's coords in result (See issue #1712) --- test/core/test_indexing.py | 85 +++++++++++++++++++++++++++++++++----- uxarray/core/dataarray.py | 11 ++++- uxarray/core/dataset.py | 10 ++++- uxarray/utils/coords.py | 40 ++++++++++++++++++ 4 files changed, 134 insertions(+), 12 deletions(-) diff --git a/test/core/test_indexing.py b/test/core/test_indexing.py index 34bf0930c..27fadf8ce 100644 --- a/test/core/test_indexing.py +++ b/test/core/test_indexing.py @@ -153,51 +153,116 @@ def test_isel_can_use_bool(): def test_indexing_by_dataarray(): """ensure isel() and sel() with indexer=xr.DataArray(...) both work as expected. The dims/coords of the Grid object should never incorporate indexer's dims/coords. - The dims/coords of the data object (UxDataArray or UxDataset) should not incorporate + + The dims of the data object (UxDataArray or UxDataset) should not incorporate the indexer's dims when indexing along a grid dim (e.g. 'n_face') (this is already true), - but should probably incorporate its coords (this isn't true yet; see issue #1712). + (e.g. don't rename 'n_face' to match the indexer's dim name!) see issue #1712 for details. + + Though, the coords of the data object *should* incorporate the indexer's coords when possible: + - Always incorporate the indexer's scalar coords. + - For 1D coords, it is more complicated: + only incorporate 1D coords if grid dim is 'n_face' + (because 'n_edge' and 'n_node' indexing don't necessarily lead to same + size indexer as result) + and when data is located along 'n_face' + (because otherwise the inder's dim doesn't align with the result's grid dim). + I.e., only incorporate 1D coords when indexing face-centered data along 'n_face'. + See #1712 for more details. Regression test for bug in branch (fixed before merging to main) for PR 1729. + + TODO: update accordingly after fixing #1758. """ + # --- n_face indexing of n_face data --- # # ensure grid's dims/coords do not incorporate indexer's dims/coords: indexer0 = xr.DataArray(0, coords={"newcoord": 7}) indexer1 = xr.DataArray([1,2], dims="newdim", coords={"newdim": [7,8]}) + indexer2 = indexer1.assign_coords({"other1dcoord": ("newdim", [9,10]), "scalarcoord": 100}) ds = ux.tutorial.open_dataset("quad-hexagon") + assert "n_face" in ds.dims # behavior for #1712 depends on data location. result0_isel = ds.isel(n_face=indexer0) assert "newcoord" not in result0_isel.uxgrid._ds.coords - # assert "newcoord" in result0_isel.coords # uncomment after fixing #1712 + assert "newcoord" in result0_isel.coords # regression test for #1712 result0_sel = ds.sel(n_face=indexer0) assert "newcoord" not in result0_sel.uxgrid._ds.coords - # assert "newcoord" in result0_sel.coords # uncomment after fixing #1712 + assert "newcoord" in result0_sel.coords # regression test for #1712 result1_isel = ds.isel(n_face=indexer1) assert "newdim" not in result1_isel.uxgrid._ds.dims assert "newdim" not in result1_isel.uxgrid._ds.coords assert "newdim" not in result1_isel.dims and "n_face" in result1_isel.dims # didn't rename 'n_face'. - # assert "newdim" in result1_isel.coords # uncomment after fixing #1712 + assert "newdim" in result1_isel.coords # regression test for #1712 result1_sel = ds.sel(n_face=indexer1) assert "newdim" not in result1_sel.uxgrid._ds.dims assert "newdim" not in result1_sel.uxgrid._ds.coords assert "newdim" not in result1_sel.dims and "n_face" in result1_sel.dims - # assert "newdim" in result1_sel.coords # uncomment after fixing #1712 + assert "newdim" in result1_sel.coords # regression test for #1712 + result2_isel = ds.isel(n_face=indexer2) + assert np.all(result2_isel.coords["other1dcoord"] == [9,10]) + assert result2_isel.coords["scalarcoord"] == 100 + result2_sel = ds.sel(n_face=indexer2) + assert np.all(result2_sel.coords["other1dcoord"] == [9,10]) + assert result2_sel.coords["scalarcoord"] == 100 # repeat tests but with UxDataArray: arr = ds['t2m'] result0_isel = arr.isel(n_face=indexer0) assert "newcoord" not in result0_isel.uxgrid._ds.coords - # assert "newcoord" in result0_isel.coords + assert "newcoord" in result0_isel.coords result0_sel = arr.sel(n_face=indexer0) assert "newcoord" not in result0_sel.uxgrid._ds.coords - # assert "newcoord" in result0_sel.coords + assert "newcoord" in result0_sel.coords result1_isel = arr.isel(n_face=indexer1) assert "newdim" not in result1_isel.uxgrid._ds.dims assert "newdim" not in result1_isel.uxgrid._ds.coords assert "newdim" not in result1_isel.dims and "n_face" in result1_isel.dims - # assert "newdim" in result1_isel.coords + assert "newdim" in result1_isel.coords result1_sel = arr.sel(n_face=indexer1) assert "newdim" not in result1_sel.uxgrid._ds.dims assert "newdim" not in result1_sel.uxgrid._ds.coords assert "newdim" not in result1_sel.dims and "n_face" in result1_sel.dims - # assert "newdim" in result1_sel.coords + assert "newdim" in result1_sel.coords + result2_isel = arr.isel(n_face=indexer2) + assert np.all(result2_isel.coords["other1dcoord"] == [9,10]) + assert result2_isel.coords["scalarcoord"] == 100 + result2_sel = arr.sel(n_face=indexer2) + assert np.all(result2_sel.coords["other1dcoord"] == [9,10]) + assert result2_sel.coords["scalarcoord"] == 100 + + # --- non-n_face indexing and/or non-n_face data --- # + # using loops to avoid writing extremely long test; + # loops are slightly harder to debug but worthwhile to include in at least one test, + # to cover more combinations of cases (e.g., discovered #1758 while making this test). + ds_face = ds + ds_node = ux.tutorial.open_dataset("quad-hexagon-random-node") + #ds_edge = ux.tutorial.open_dataset("quad-hexagon-random-edge") # uncomment after fixing #1758 + assert "n_face" in ds_face.dims + assert "n_node" in ds_node.dims + #assert "n_edge" in ds_edge.dims # uncomment after fixing #1758 + counter = 0 # (count up during loop to make sure nothing is skipped unexpectedly) + for method in "isel", "sel": + for dataset in [ds_face, ds_node]: # include after fixing #1758 + for grid_dim in ("n_face", "n_edge", "n_node"): + for to_array in [False, True]: + if grid_dim == "n_face" and "n_face" in dataset.dims: + continue # already tested this above! + counter += 1 + obj = dataset[list(dataset.data_vars)[0]] if to_array else dataset + result0 = getattr(obj, method)({grid_dim: indexer0}) + assert "newcoord" not in result0.uxgrid._ds.coords + assert "newcoord" in result0.coords # 0D coord should always show up! + result1 = getattr(obj, method)({grid_dim: indexer1}) + assert "newdim" not in result1.uxgrid._ds.dims + assert "newdim" not in result1.uxgrid._ds.coords + assert "newdim" not in result1.dims + assert "newdim" not in result1.coords + result2 = getattr(obj, method)({grid_dim: indexer2}) + assert "other1dcoord" not in result2.uxgrid._ds.coords + assert "other1dcoord" not in result2.coords + assert result2.coords["scalarcoord"] == 100 + n_data_grid_dim_combos = 2 * 3 - 1 # after fixing #1758, update to: 3 * 3 - 1. + # the -1 accounts for skipping when both are "n_face" above. + assert counter == 2 * n_data_grid_dim_combos * 2 + def test_indexing_does_not_edit_indexers_dict(): """ensure isel() and sel() do not edit the provided indexers dict. diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 8ec48e7dc..6bbf048a1 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -44,7 +44,10 @@ from uxarray.plot.accessor import UxDataArrayPlotAccessor from uxarray.remap.accessor import RemapAccessor from uxarray.subset import DataArraySubsetAccessor -from uxarray.utils.coords import _preserve_valid_coords +from uxarray.utils.coords import ( + _assign_grid_dim_indexer_coords_if_appropriate, + _preserve_valid_coords, +) if TYPE_CHECKING: import cartopy.crs as ccrs @@ -2008,6 +2011,7 @@ def isel( the result would have 'n_face' with just those two faces. For data on 'n_edge', the result would have 'n_edge' with all edges located on either of those two faces. Grid dimension indexers cannot have more than 1 dimension (such as a 2D DataArray). + Grid dimensions are never renamed (even if indexed by 1D DataArray with different dim name). Parameters ---------- @@ -2078,6 +2082,10 @@ def isel( result = self._slice_from_grid(sliced_grid) + result = _assign_grid_dim_indexer_coords_if_appropriate( + result, grid_dim, grid_indexer + ) + # if there are any remaining indexers, apply them if indexers: result = super(UxDataArray, result).isel( @@ -2107,6 +2115,7 @@ def sel( the result would have 'n_face' with just those two faces. For data on 'n_edge', the result would have 'n_edge' with all edges located on either of those two faces. Grid dimension indexers cannot have more than 1 dimension (such as a 2D DataArray). + Grid dimensions are never renamed (even if indexed by 1D DataArray with different dim name). By default, grid dims do not have coordinates assigned. But, if they have been assigned, `.sel()` respects them in the intuitive way. For example, diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index ebd9a4d50..e1df92837 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -28,6 +28,7 @@ from uxarray.io._healpix import get_zoom_from_cells from uxarray.plot.accessor import UxDatasetPlotAccessor from uxarray.remap.accessor import RemapAccessor +from uxarray.utils.coords import _assign_grid_dim_indexer_coords_if_appropriate class UxDataset(xr.Dataset): @@ -425,6 +426,7 @@ def isel( the result would have 'n_face' with just those two faces. For data on 'n_edge', the result would have 'n_edge' with all edges located on either of those two faces. Grid dimension indexers cannot have more than 1 dimension (such as a 2D DataArray). + Grid dimensions are never renamed (even if indexed by 1D DataArray with different dim name). Parameters ---------- @@ -495,6 +497,10 @@ def isel( result = self._slice_from_grid(sliced_grid) + result = _assign_grid_dim_indexer_coords_if_appropriate( + result, grid_dim, grid_indexer + ) + # if there are any remaining indexers, apply them if indexers: result = super(UxDataset, result).isel( @@ -525,6 +531,7 @@ def sel( the result would have 'n_face' with just those two faces. For data on 'n_edge', the result would have 'n_edge' with all edges located on either of those two faces. Grid dimension indexers cannot have more than 1 dimension (such as a 2D DataArray). + Grid dimensions are never renamed (even if indexed by 1D DataArray with different dim name). By default, grid dims do not have coordinates assigned. But, if they have been assigned, `.sel()` respects them in the intuitive way. For example, @@ -553,7 +560,8 @@ def sel( multi-index, the indexer may also be a dict-like object with keys matching index level names. If DataArrays are passed as indexers, xarray-style indexing will be - carried out. See :ref:`indexing` for the details. + carried out (see :ref:`indexing` for the details), + with one exception: grid dimensions will never be renamed. One of indexers or indexers_kwargs must be provided. method : {None, "nearest", "pad", "ffill", "backfill", "bfill"}, optional Method to use for inexact matches: diff --git a/uxarray/utils/coords.py b/uxarray/utils/coords.py index caaccf525..6cc388454 100644 --- a/uxarray/utils/coords.py +++ b/uxarray/utils/coords.py @@ -8,6 +8,8 @@ import xarray as xr import xarray.core.utils as xr_core_utils +from uxarray.errors import DimensionError + def _preserve_valid_coords( obj: xr.DataArray | xr.Dataset, @@ -84,3 +86,41 @@ def _indices1d_from_indexing(xarray_obj, dim, indexer): if dim in xarray_obj.coords: xarray_obj = xarray_obj.drop_vars(dim) return xarray_obj[dim].isel({dim: indexer}).values + + +def _assign_grid_dim_indexer_coords_if_appropriate(uxarray_obj, grid_dim, indexer): + """returns uxarray_obj but with coords assigned from indexer if appropriate. + "appropriate" `indexer` is a 0D or 1D xr.DataArray and has any relevant coordinates to assign. + 2D+ xr.DataArray indexers are not supported for grid dimensions and cause DimensionError here. + All other indexers do not provide coordinate info, so uxarray_obj gets returned unchanged. + + For 0D indexer, just assign indexer.coords if nonempty (else, return uxarray_obj unchanged). + For 1D indexer, depends on grid_dim and uxarray_obj. + if grid_dim=="n_face" and "n_face" in uxarray_obj: + swap indexer's 1 dim to be "n_face" instead of its original name, + then assign indexer.coords. + in all other cases: + drop indexer's 1 dim, then assign indexer.coords if nonempty. + (For "n_edge" and "n_node" indexing, the result's shape won't necessarily + match the indexer's shape, so coords along the grid dim can't be assigned. + Meanwhile, if grid_dim not in uxarray_obj, it is impossible to assign coords + along that dim, so once again, coords along the grid dim can't be assigned.) + """ + if isinstance(indexer, xr.DataArray): + if indexer.ndim == 0: + coords = indexer.coords + elif indexer.ndim == 1: + the_dim = indexer.dims[0] + if grid_dim == "n_face" and "n_face" in uxarray_obj.dims: + coords = indexer.swap_dims({the_dim: "n_face"}).coords + else: + # remove any 1D coords (but keep scalar coords) + coords = indexer.isel({the_dim: 0}, drop=True).coords + else: + raise DimensionError( + f"2D+ indexers are not supported for grid dimensions. Got xr.DataArray " + f"indexer with ndim={indexer.ndim}, along grid_dim={grid_dim!r}." + ) + if coords: + return uxarray_obj.assign_coords(coords) + return uxarray_obj From 04f88f0da1a3f173b094687dcffdeb22824c81ea Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Mon, 14 Sep 2026 18:04:59 -0400 Subject: [PATCH 03/14] update sel/isel docstrings to clarify #1712 fix --- uxarray/core/dataarray.py | 7 ++++++- uxarray/core/dataset.py | 7 ++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 6bbf048a1..a1872db94 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -2047,6 +2047,10 @@ def isel( ------- UxDataArray A new UxDataArray indexed according to `indexers` and updated grid if applicable. + If indexer DataArrays have coordinates that do not conflict with + this object, then these coordinates will be attached, + except that 1D coordinates of indexers applied along a grid dimension will + only be included if it is 'n_face' and the data also has 'n_face' dimension. Raises ------ @@ -2176,7 +2180,8 @@ def sel( and the uxgrid indexed appropriately as well, if indexing any grid dim. If indexer DataArrays have coordinates that do not conflict with this object, then these coordinates will be attached, - except for indexers along a grid dimension (see issue #1712). + except that 1D coordinates of indexers applied along a grid dimension will + only be included if it is 'n_face' and the data also has 'n_face' dimension. In general, the result's data will be a view of the data in this array, unless indexing along a grid dimension or otherwise triggering vectorized indexing by using an array indexer, diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index e1df92837..0ee597053 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -462,6 +462,10 @@ def isel( ------- UxDataset A new UxDataset indexed according to `indexers` and updated grid if applicable. + If indexer DataArrays have coordinates that do not conflict with + this object, then these coordinates will be attached, + except that 1D coordinates of indexers applied along a grid dimension will + only be included if it is 'n_face' and the data also has 'n_face' dimension. Raises ------ @@ -594,7 +598,8 @@ def sel( and the uxgrid indexed appropriately as well, if indexing any grid dim. If indexer DataArrays have coordinates that do not conflict with this object, then these coordinates will be attached, - except for indexers along a grid dimension (see issue #1712). + except that 1D coordinates of indexers applied along a grid dimension will + only be included if it is 'n_face' and the data also has 'n_face' dimension. In general, each array's data will be a view of the array's data in this dataset, unless indexing along a grid dimension or otherwise triggering vectorized indexing by using an array indexer, From d79944dea2ad111712e22b790b8a8a0fb65f5f1f Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Fri, 18 Sep 2026 13:24:57 -0400 Subject: [PATCH 04/14] fix isel(boolean indexer array with coords) --- test/core/test_indexing.py | 44 ++++++++++++++++++++++++++++++++++++++ uxarray/utils/coords.py | 7 ++++++ 2 files changed, 51 insertions(+) diff --git a/test/core/test_indexing.py b/test/core/test_indexing.py index 27fadf8ce..de45e05eb 100644 --- a/test/core/test_indexing.py +++ b/test/core/test_indexing.py @@ -150,6 +150,50 @@ def test_isel_can_use_bool(): result = arr.isel(n_face=[False, False, False, False]) assert result.sizes['n_face'] == result.uxgrid.n_face == 0 + +def test_isel_can_use_bool_with_coords(): + """ensure isel() supports indexing by a boolean indexer array with coords. + Regression test for bug (1) discovered during review of PR #1759. + """ + ds = ux.tutorial.open_dataset("quad-hexagon") + arr = ds['t2m'] + + # simplest case: boolean xr.DataArray mask with coords. + tokeep0 = xr.DataArray([True, False, True, False], coords={'tokeep': [2,4,6,8]}) + # (also want to test UxDataArray boolean mask (ensure no infinite recursion)) + tokeep1 = arr * xr.DataArray([True, False, True, True], dims={'n_face'}) > 0 + assert isinstance(tokeep1, ux.UxDataArray) + # (also want to test UxDataArray boolean mask with coords + tokeep2 = tokeep1.assign_coords(tokeep=('n_face', [1,3,5,7])) + assert isinstance(tokeep2, ux.UxDataArray) + assert np.all(tokeep2.coords['tokeep'] == [1,3,5,7]) + + # test on UxDataArray + result0 = arr.isel(n_face=tokeep0) + assert result0.sizes['n_face'] == result0.uxgrid.n_face == 2 + assert np.all(result0.coords['tokeep'] == [2,6]) + result1 = arr.isel(n_face=tokeep1) + assert result1.sizes['n_face'] == result1.uxgrid.n_face == 3 + result2 = arr.isel(n_face=tokeep2) + assert result2.sizes['n_face'] == result2.uxgrid.n_face == 3 + assert np.all(result2.coords['tokeep'] == [1,5,7]) + # (reviewer found bug on arr.where(), so doing a spot check for that here too) + result_where2 = arr.where(tokeep2, drop=True) + assert result2.equals(result_where2) + + # repeat tests for UxDataset + result0 = ds.isel(n_face=tokeep0) + assert result0.sizes['n_face'] == result0.uxgrid.n_face == 2 + assert np.all(result0.coords['tokeep'] == [2,6]) + result1 = ds.isel(n_face=tokeep1) + assert result1.sizes['n_face'] == result1.uxgrid.n_face == 3 + result2 = ds.isel(n_face=tokeep2) + assert result2.sizes['n_face'] == result2.uxgrid.n_face == 3 + assert np.all(result2.coords['tokeep'] == [1,5,7]) + result_where2 = ds.where(tokeep2, drop=True) + assert result2.equals(result_where2) + + def test_indexing_by_dataarray(): """ensure isel() and sel() with indexer=xr.DataArray(...) both work as expected. The dims/coords of the Grid object should never incorporate indexer's dims/coords. diff --git a/uxarray/utils/coords.py b/uxarray/utils/coords.py index 6cc388454..ebbd4e050 100644 --- a/uxarray/utils/coords.py +++ b/uxarray/utils/coords.py @@ -105,6 +105,9 @@ def _assign_grid_dim_indexer_coords_if_appropriate(uxarray_obj, grid_dim, indexe match the indexer's shape, so coords along the grid dim can't be assigned. Meanwhile, if grid_dim not in uxarray_obj, it is impossible to assign coords along that dim, so once again, coords along the grid dim can't be assigned.) + Note: if indexer is booleans, instead use indexer.isel(indexer_dim=indexer).coords, + because the result will only keep values wherever indexer value is True. + (Also in this case, if indexer.to_xarray() exists, call it, to avoid recursion.) """ if isinstance(indexer, xr.DataArray): if indexer.ndim == 0: @@ -112,6 +115,10 @@ def _assign_grid_dim_indexer_coords_if_appropriate(uxarray_obj, grid_dim, indexe elif indexer.ndim == 1: the_dim = indexer.dims[0] if grid_dim == "n_face" and "n_face" in uxarray_obj.dims: + if indexer.dtype == bool: + if hasattr(indexer, "to_xarray"): + indexer = indexer.to_xarray() + indexer = indexer.isel({the_dim: indexer}) coords = indexer.swap_dims({the_dim: "n_face"}).coords else: # remove any 1D coords (but keep scalar coords) From 7083f8db5c2463049506ccf6bdb7c175a0a8605a Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Fri, 18 Sep 2026 13:47:14 -0400 Subject: [PATCH 05/14] add crash if indexer & arr coords incompatible --- test/core/test_indexing.py | 31 +++++++++++++++++++++++++++++++ uxarray/utils/coords.py | 3 +++ 2 files changed, 34 insertions(+) diff --git a/test/core/test_indexing.py b/test/core/test_indexing.py index de45e05eb..b0b401f2d 100644 --- a/test/core/test_indexing.py +++ b/test/core/test_indexing.py @@ -496,3 +496,34 @@ def test_sel_crash_if_provided_selection_options_with_coordless_dims(): arr1.to_xarray().sel(time=4, **kw) with pytest.raises(ValueError, match=r"cannot supply selection options"): arr1.to_xarray().sel(time=4, n_face=[3], **kw) + +def test_isel_crash_if_coordinates_conflict(): + """Ensure isel crashes if there is a coordinates conflict, + such as indexing an array with time dim by an array with a scalar time coord. + Regression test for bug (2) discovered during review of PR #1759. + """ + ds = ux.tutorial.open_dataset("outCSne30-timeseries") + arr = ds['psi'] + indexer0 = xr.DataArray(0, coords={'time': arr['time'][0].item()}) + indexer1 = arr.isel(time=0).argmax('n_face') + assert isinstance(indexer1, ux.UxDataArray) + assert 'time' in indexer1.coords and 'time' not in indexer1.dims + MATCH_ERRMSG = "dimension coordinate 'time' conflicts between indexed and indexing objects" + with pytest.raises(IndexError, match=MATCH_ERRMSG): + arr.to_xarray().isel(n_face=indexer0) # sanity check that xarray also crashes here. + with pytest.raises(IndexError, match=MATCH_ERRMSG): + arr.isel(n_face=indexer0) + with pytest.raises(IndexError, match=MATCH_ERRMSG): + arr.to_xarray().isel(n_face=indexer1) # sanity check that xarray also crashes here. + with pytest.raises(IndexError, match=MATCH_ERRMSG): + arr.isel(n_face=indexer1) + + # repeat tests for UxDataArray: + with pytest.raises(IndexError, match=MATCH_ERRMSG): + ds.to_xarray().isel(n_face=indexer0) + with pytest.raises(IndexError, match=MATCH_ERRMSG): + ds.isel(n_face=indexer0) + with pytest.raises(IndexError, match=MATCH_ERRMSG): + ds.to_xarray().isel(n_face=indexer1) + with pytest.raises(IndexError, match=MATCH_ERRMSG): + ds.isel(n_face=indexer1) diff --git a/uxarray/utils/coords.py b/uxarray/utils/coords.py index ebbd4e050..f5d419e69 100644 --- a/uxarray/utils/coords.py +++ b/uxarray/utils/coords.py @@ -6,6 +6,7 @@ import numpy as np import xarray as xr +import xarray.core.coordinates import xarray.core.utils as xr_core_utils from uxarray.errors import DimensionError @@ -110,6 +111,8 @@ def _assign_grid_dim_indexer_coords_if_appropriate(uxarray_obj, grid_dim, indexe (Also in this case, if indexer.to_xarray() exists, call it, to avoid recursion.) """ if isinstance(indexer, xr.DataArray): + xr.core.coordinates.assert_coordinate_consistent(uxarray_obj, indexer.coords) + # ^ e.g. if uxarray_obj has time dim but indexer has time scalar coord, crash! if indexer.ndim == 0: coords = indexer.coords elif indexer.ndim == 1: From 11312f05c4251aeb108b9bc69390cb7dbb180b96 Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Fri, 18 Sep 2026 14:27:09 -0400 Subject: [PATCH 06/14] account for bonus coords in UxDataset.isel --- test/core/test_indexing.py | 13 +++++++++++++ uxarray/core/dataset.py | 19 ++++++++++++++++++- 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/test/core/test_indexing.py b/test/core/test_indexing.py index b0b401f2d..86f976fcf 100644 --- a/test/core/test_indexing.py +++ b/test/core/test_indexing.py @@ -308,6 +308,19 @@ def test_indexing_by_dataarray(): assert counter == 2 * n_data_grid_dim_combos * 2 +def test_dataset_isel_keeps_bonus_coords(): + """ensure UxDataset.isel() keeps "bonus" coords, + i.e. coords in the dataset which do not actually appear in any data var. + Regression test for bug (3) discovered during review of PR #1759. + """ + ds = ux.tutorial.open_dataset('quad-hexagon') + ds = ds.assign_coords({'bonus_coord': xr.DataArray(['a', 'b'], dims=['bonus_dim'])}) + assert 'bonus_coord' in ds.coords + assert all('bonus_coord' not in arr for arr in ds.data_vars.values()) + result = ds.isel(n_face=0) + assert 'bonus_coord' in result.coords and 'bonus_dim' in result.dims + + def test_indexing_does_not_edit_indexers_dict(): """ensure isel() and sel() do not edit the provided indexers dict. Regression test for #1711. diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 0ee597053..03dcc1881 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -402,7 +402,24 @@ def _slice_from_grid(self, sliced_grid): else: data_vars[name] = da - ds_sliced = xr.Dataset(data_vars=data_vars, attrs=self.attrs) + # Also account for any coords which aren't attached to any data_var: + bonus_coords = {} + for coord in self.coords: + for data_var in data_vars.values(): + if coord in data_var.coords: + break + else: # didn't break + da = self.coords[coord] + if hasattr(da, "_slice_from_grid") and any( + dim in da.dims for dim in GRID_DIMS + ): + bonus_coords[coord] = da._slice_from_grid(sliced_grid) + else: + bonus_coords[coord] = da + + ds_sliced = xr.Dataset( + data_vars=data_vars, coords=bonus_coords, attrs=self.attrs + ) return type(self)(ds_sliced, uxgrid=sliced_grid) def isel( From 1fd7f61cef1e70018b78cb4d6afc8709ad41bf5f Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Tue, 22 Sep 2026 14:53:45 -0400 Subject: [PATCH 07/14] fix subtle isel bugs for 1D indexer with coords --- test/core/test_indexing.py | 49 ++++++++++++++++++++++++++++++++++++++ uxarray/core/dataarray.py | 7 ++++++ uxarray/core/dataset.py | 11 ++++++++- uxarray/utils/coords.py | 41 ++++++++++++++++++++++++++++++- 4 files changed, 106 insertions(+), 2 deletions(-) diff --git a/test/core/test_indexing.py b/test/core/test_indexing.py index 86f976fcf..e51e1498f 100644 --- a/test/core/test_indexing.py +++ b/test/core/test_indexing.py @@ -514,14 +514,17 @@ def test_isel_crash_if_coordinates_conflict(): """Ensure isel crashes if there is a coordinates conflict, such as indexing an array with time dim by an array with a scalar time coord. Regression test for bug (2) discovered during review of PR #1759. + (Also checks indexing an array with scalar time coord by an array with a time dim.) """ ds = ux.tutorial.open_dataset("outCSne30-timeseries") arr = ds['psi'] indexer0 = xr.DataArray(0, coords={'time': arr['time'][0].item()}) indexer1 = arr.isel(time=0).argmax('n_face') + indexer2 = xr.DataArray([0,1], dims='time', coords={'time': arr['time'][:2].values}) assert isinstance(indexer1, ux.UxDataArray) assert 'time' in indexer1.coords and 'time' not in indexer1.dims MATCH_ERRMSG = "dimension coordinate 'time' conflicts between indexed and indexing objects" + MATCH_ERRMSG_2 = "The indexer's dimension \('time'\) already exists as a scalar" with pytest.raises(IndexError, match=MATCH_ERRMSG): arr.to_xarray().isel(n_face=indexer0) # sanity check that xarray also crashes here. with pytest.raises(IndexError, match=MATCH_ERRMSG): @@ -530,6 +533,10 @@ def test_isel_crash_if_coordinates_conflict(): arr.to_xarray().isel(n_face=indexer1) # sanity check that xarray also crashes here. with pytest.raises(IndexError, match=MATCH_ERRMSG): arr.isel(n_face=indexer1) + arr_t0 = arr.isel(time=0) + assert 'time' in arr_t0.coords and 'time' not in arr_t0.dims + with pytest.raises(ux.errors.DimensionError, match=MATCH_ERRMSG_2): + arr_t0.isel(n_face=indexer2) # repeat tests for UxDataArray: with pytest.raises(IndexError, match=MATCH_ERRMSG): @@ -540,3 +547,45 @@ def test_isel_crash_if_coordinates_conflict(): ds.to_xarray().isel(n_face=indexer1) with pytest.raises(IndexError, match=MATCH_ERRMSG): ds.isel(n_face=indexer1) + ds_t0 = ds.isel(time=0) + assert 'time' in ds_t0.coords and 'time' not in ds_t0.dims + with pytest.raises(ux.errors.DimensionError, match=MATCH_ERRMSG_2): + ds_t0.isel(n_face=indexer2) + +def test_isel_when_indexer_extra_dims_match(): + """Ensure isel() crashes with NotImplementedError when the indexer's dim has + the same name as a dim or non-scalar coordinate in the uxarray object being indexed. + Regression test for follow-up to bug (2) discovered during review of PR #1759. + """ + ds = ux.tutorial.open_dataset('quad-hexagon').expand_dims(time=[100,200]) + arr = ds['t2m'] + imax = arr.argmax('n_face') + assert set(imax.dims) == {'time'} and set(imax.coords) == {'time'} + assert set(ds.dims) == {'time', 'n_face'} and set(ds.coords) == {'time'} + assert imax.coords['time'].equals(ds.coords['time']) + # pure xarray indexing here returns a result with only 'time' dimension; + # that's the main reason uxarray should raise NotImplementedError in this case. + xr_result = ds.to_xarray().isel(n_face=imax.to_xarray()) + assert set(xr_result.dims) == {'time'} + ERRMSG = ( + r"Indexing a {typestr} .+ using an xarray DataArray whose dimension \('time'\) " + r"is already present .+ is not yet supported" + ) + with pytest.raises(NotImplementedError, match=ERRMSG.format(typestr="UxDataset")): + ds.isel(n_face=imax) + ds1 = ds.swap_dims({'time': 'otherdim'}) + assert 'time' in ds1.coords and 'time' not in ds1.dims + with pytest.raises(NotImplementedError, match=ERRMSG.format(typestr="UxDataset")): + ds1.isel(n_face=imax) + + # repeat tests for UxDataArray: + assert set(arr.dims) == {'time', 'n_face'} and set(arr.coords) == {'time'} + assert imax.coords['time'].equals(arr.coords['time']) + xr_result = arr.to_xarray().isel(n_face=imax.to_xarray()) + assert set(xr_result.dims) == {'time'} + with pytest.raises(NotImplementedError, match=ERRMSG.format(typestr="UxDataArray")): + arr.isel(n_face=imax) + arr1 = arr.swap_dims({'time': 'otherdim'}) + assert 'time' in arr1.coords and 'time' not in arr1.dims + with pytest.raises(NotImplementedError, match=ERRMSG.format(typestr="UxDataArray")): + arr1.isel(n_face=imax) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index ae59de84f..7a19dadd4 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -46,6 +46,7 @@ from uxarray.subset import DataArraySubsetAccessor from uxarray.utils.coords import ( _assign_grid_dim_indexer_coords_if_appropriate, + _crash_if_1d_xarray_indexer_dim_in_uxarray_obj, _preserve_valid_coords, ) from uxarray.utils.imports import _raise_hint_if_optional_deps_missing @@ -2000,8 +2001,10 @@ def isel( using n_edge=7 selects just the two faces touching edge 7. For data on 'n_face', the result would have 'n_face' with just those two faces. For data on 'n_edge', the result would have 'n_edge' with all edges located on either of those two faces. + Grid dimension indexers cannot have more than 1 dimension (such as a 2D DataArray). Grid dimensions are never renamed (even if indexed by 1D DataArray with different dim name). + Grid dimension indexer cannot have a non-grid dimension which exists in the original UxDataArray. Parameters ---------- @@ -2070,6 +2073,8 @@ def isel( indexers = indexers.copy() # don't modify the original dict grid_indexer = indexers.pop(grid_dim) + _crash_if_1d_xarray_indexer_dim_in_uxarray_obj(self, grid_dim, grid_indexer) + sliced_grid = self.uxgrid.isel( **{grid_dim: grid_indexer}, inverse_indices=inverse_indices ) @@ -2108,8 +2113,10 @@ def sel( using n_edge=7 selects just the two faces touching edge 7. For data on 'n_face', the result would have 'n_face' with just those two faces. For data on 'n_edge', the result would have 'n_edge' with all edges located on either of those two faces. + Grid dimension indexers cannot have more than 1 dimension (such as a 2D DataArray). Grid dimensions are never renamed (even if indexed by 1D DataArray with different dim name). + Grid dimension indexer cannot have a non-grid dimension which exists in the original UxDataArray. By default, grid dims do not have coordinates assigned. But, if they have been assigned, `.sel()` respects them in the intuitive way. For example, diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 03dcc1881..31360ad7a 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -28,7 +28,10 @@ from uxarray.io._healpix import get_zoom_from_cells from uxarray.plot.accessor import UxDatasetPlotAccessor from uxarray.remap.accessor import RemapAccessor -from uxarray.utils.coords import _assign_grid_dim_indexer_coords_if_appropriate +from uxarray.utils.coords import ( + _assign_grid_dim_indexer_coords_if_appropriate, + _crash_if_1d_xarray_indexer_dim_in_uxarray_obj, +) class UxDataset(xr.Dataset): @@ -442,8 +445,10 @@ def isel( using n_edge=7 selects just the two faces touching edge 7. For data on 'n_face', the result would have 'n_face' with just those two faces. For data on 'n_edge', the result would have 'n_edge' with all edges located on either of those two faces. + Grid dimension indexers cannot have more than 1 dimension (such as a 2D DataArray). Grid dimensions are never renamed (even if indexed by 1D DataArray with different dim name). + Grid dimension indexer cannot have a non-grid dimension which exists in the original UxDataset. Parameters ---------- @@ -512,6 +517,8 @@ def isel( indexers = indexers.copy() # don't modify the original dict grid_indexer = indexers.pop(grid_dim) + _crash_if_1d_xarray_indexer_dim_in_uxarray_obj(self, grid_dim, grid_indexer) + sliced_grid = self.uxgrid.isel( **{grid_dim: grid_indexer}, inverse_indices=inverse_indices ) @@ -551,8 +558,10 @@ def sel( using n_edge=7 selects just the two faces touching edge 7. For data on 'n_face', the result would have 'n_face' with just those two faces. For data on 'n_edge', the result would have 'n_edge' with all edges located on either of those two faces. + Grid dimension indexers cannot have more than 1 dimension (such as a 2D DataArray). Grid dimensions are never renamed (even if indexed by 1D DataArray with different dim name). + Grid dimension indexer cannot have a non-grid dimension which exists in the original UxDataset. By default, grid dims do not have coordinates assigned. But, if they have been assigned, `.sel()` respects them in the intuitive way. For example, diff --git a/uxarray/utils/coords.py b/uxarray/utils/coords.py index f5d419e69..7ff37f7fb 100644 --- a/uxarray/utils/coords.py +++ b/uxarray/utils/coords.py @@ -89,6 +89,40 @@ def _indices1d_from_indexing(xarray_obj, dim, indexer): return xarray_obj[dim].isel({dim: indexer}).values +def _crash_if_1d_xarray_indexer_dim_in_uxarray_obj(uxarray_obj, grid_dim, indexer): + """if xarray indexer's dim is not grid_dim and is in uxarray object, raise NotImplementedError + rather than silently treating indexer as indexer.values, or silently giving a confusing result. + (Should only be applied where indexer is indexing uxarray_obj along grid_dim.) + + In this case, uxarray_obj.to_xarray().isel({grid_dim: indexer}) produces "pointwise" indexing, + but that is difficult to support reliably along a grid dimension. + + Example: uxarr.isel(n_face=uxarr.argmax('n_face')) for uxarr a UxDataArray with 'time' dimension, + should maybe produce "the maximum values of uxarr across all faces, at every face where there is + a maximum at any given time, for all times" or something like that? + (It very confusing, and unclear if any user is doing something like this.) + + In xarray, uxarr.to_xarray().isel(n_face=uxarr.argmax('n_face')) produces an array with only the + 'time' dimension, telling the maximum value of uxarray across all faces, at each time. + + (Only handles 1D xr.DataArray indexers; 0D, 2D+, and non-DataArray indexers are handled elsewhere.) + """ + if isinstance(indexer, xr.DataArray): + if indexer.ndim == 1: + the_dim = indexer.dims[0] + nonscalar_coords = [c for c in uxarray_obj.coords if len(uxarray_obj.coords[c].dims) > 0] + if the_dim != grid_dim and (the_dim in uxarray_obj.dims or the_dim in nonscalar_coords): + raise NotImplementedError( + f"Indexing a {type(uxarray_obj).__name__} along a grid dimension ({grid_dim!r}), using " + f"an xarray DataArray whose dimension ({the_dim!r}) is already present in the original " + f"{type(uxarray_obj).__name__}'s dims or non-scalar coords, is not yet supported. " + f"Consider using obj.to_xarray() for basic xarray indexing, " + f"using indexer.data to convert to indexer to a non-DataArray object, " + f"or using indexer.rename({{{the_dim!r}: 'any_unused_dim_name'}}) to avoid matching dims." + ) + # all other cases handled elsewhere. + + def _assign_grid_dim_indexer_coords_if_appropriate(uxarray_obj, grid_dim, indexer): """returns uxarray_obj but with coords assigned from indexer if appropriate. "appropriate" `indexer` is a 0D or 1D xr.DataArray and has any relevant coordinates to assign. @@ -111,7 +145,7 @@ def _assign_grid_dim_indexer_coords_if_appropriate(uxarray_obj, grid_dim, indexe (Also in this case, if indexer.to_xarray() exists, call it, to avoid recursion.) """ if isinstance(indexer, xr.DataArray): - xr.core.coordinates.assert_coordinate_consistent(uxarray_obj, indexer.coords) + xr.core.coordinates.assert_coordinate_consistent(uxarray_obj, indexer.coords.variables) # ^ e.g. if uxarray_obj has time dim but indexer has time scalar coord, crash! if indexer.ndim == 0: coords = indexer.coords @@ -122,6 +156,11 @@ def _assign_grid_dim_indexer_coords_if_appropriate(uxarray_obj, grid_dim, indexe if hasattr(indexer, "to_xarray"): indexer = indexer.to_xarray() indexer = indexer.isel({the_dim: indexer}) + if the_dim in uxarray_obj.coords and len(uxarray_obj.coords[the_dim].dims)==0: + raise DimensionError( + f"The indexer's dimension ({the_dim!r}) already exists as a scalar " + f"coordinate in the {type(uxarray_obj).__name__} object being indexed." + ) # coords = indexer.swap_dims({the_dim: "n_face"}).coords else: # remove any 1D coords (but keep scalar coords) From 59c25a18468476c96ddd3394c280748bb401d2e4 Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Tue, 22 Sep 2026 14:54:58 -0400 Subject: [PATCH 08/14] forgot pre-commit ruff formatting --- uxarray/utils/coords.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/uxarray/utils/coords.py b/uxarray/utils/coords.py index 7ff37f7fb..dab1f6f5d 100644 --- a/uxarray/utils/coords.py +++ b/uxarray/utils/coords.py @@ -110,8 +110,12 @@ def _crash_if_1d_xarray_indexer_dim_in_uxarray_obj(uxarray_obj, grid_dim, indexe if isinstance(indexer, xr.DataArray): if indexer.ndim == 1: the_dim = indexer.dims[0] - nonscalar_coords = [c for c in uxarray_obj.coords if len(uxarray_obj.coords[c].dims) > 0] - if the_dim != grid_dim and (the_dim in uxarray_obj.dims or the_dim in nonscalar_coords): + nonscalar_coords = [ + c for c in uxarray_obj.coords if len(uxarray_obj.coords[c].dims) > 0 + ] + if the_dim != grid_dim and ( + the_dim in uxarray_obj.dims or the_dim in nonscalar_coords + ): raise NotImplementedError( f"Indexing a {type(uxarray_obj).__name__} along a grid dimension ({grid_dim!r}), using " f"an xarray DataArray whose dimension ({the_dim!r}) is already present in the original " @@ -145,7 +149,9 @@ def _assign_grid_dim_indexer_coords_if_appropriate(uxarray_obj, grid_dim, indexe (Also in this case, if indexer.to_xarray() exists, call it, to avoid recursion.) """ if isinstance(indexer, xr.DataArray): - xr.core.coordinates.assert_coordinate_consistent(uxarray_obj, indexer.coords.variables) + xr.core.coordinates.assert_coordinate_consistent( + uxarray_obj, indexer.coords.variables + ) # ^ e.g. if uxarray_obj has time dim but indexer has time scalar coord, crash! if indexer.ndim == 0: coords = indexer.coords @@ -156,7 +162,10 @@ def _assign_grid_dim_indexer_coords_if_appropriate(uxarray_obj, grid_dim, indexe if hasattr(indexer, "to_xarray"): indexer = indexer.to_xarray() indexer = indexer.isel({the_dim: indexer}) - if the_dim in uxarray_obj.coords and len(uxarray_obj.coords[the_dim].dims)==0: + if ( + the_dim in uxarray_obj.coords + and len(uxarray_obj.coords[the_dim].dims) == 0 + ): raise DimensionError( f"The indexer's dimension ({the_dim!r}) already exists as a scalar " f"coordinate in the {type(uxarray_obj).__name__} object being indexed." From c83d6221d1a535bed4b6693065f414c940339f1e Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Tue, 22 Sep 2026 14:58:52 -0400 Subject: [PATCH 09/14] tiny test renaming for clarity --- test/core/test_indexing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/core/test_indexing.py b/test/core/test_indexing.py index e51e1498f..8d11aff6f 100644 --- a/test/core/test_indexing.py +++ b/test/core/test_indexing.py @@ -552,7 +552,7 @@ def test_isel_crash_if_coordinates_conflict(): with pytest.raises(ux.errors.DimensionError, match=MATCH_ERRMSG_2): ds_t0.isel(n_face=indexer2) -def test_isel_when_indexer_extra_dims_match(): +def test_isel_when_indexer_dim_in_uxarray_obj(): """Ensure isel() crashes with NotImplementedError when the indexer's dim has the same name as a dim or non-scalar coordinate in the uxarray object being indexed. Regression test for follow-up to bug (2) discovered during review of PR #1759. From 363dfc76168c8c9720a089aeb650cd0a873a0b1b Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Wed, 23 Sep 2026 08:57:33 -0400 Subject: [PATCH 10/14] add tests of 2nd reviewer's examples 1 & 2 --- test/core/test_indexing.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/test/core/test_indexing.py b/test/core/test_indexing.py index 8d11aff6f..19042f9b8 100644 --- a/test/core/test_indexing.py +++ b/test/core/test_indexing.py @@ -193,6 +193,13 @@ def test_isel_can_use_bool_with_coords(): result_where2 = ds.where(tokeep2, drop=True) assert result2.equals(result_where2) + # also test second reviewer's example (2) from PR #1759: + uxds = ux.tutorial.open_dataset("quad-hexagon") + uxds = uxds.assign_coords(node_id=("n_node", np.arange(uxds.uxgrid.n_node))) + mask = xr.DataArray([True, False, True, False], dims="n_face", + coords={"lab": ("n_face", [10, 20, 30, 40])}) + uxds.isel(n_face=mask) # (just ensuring it doesn't crash) + def test_indexing_by_dataarray(): """ensure isel() and sel() with indexer=xr.DataArray(...) both work as expected. @@ -320,6 +327,12 @@ def test_dataset_isel_keeps_bonus_coords(): result = ds.isel(n_face=0) assert 'bonus_coord' in result.coords and 'bonus_dim' in result.dims + # also test using second reviewer's example (1) from PR #1759: + uxds = ux.tutorial.open_dataset("quad-hexagon") + uxds = uxds.assign_coords(node_id=("n_node", np.arange(uxds.uxgrid.n_node))) + sub = uxds.isel(n_face=[0, 1]) + assert 'node_id' in sub.coords + def test_indexing_does_not_edit_indexers_dict(): """ensure isel() and sel() do not edit the provided indexers dict. From 7732e23310c5d9ab3648e99c2e29959a6376e7b9 Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Wed, 23 Sep 2026 09:26:42 -0400 Subject: [PATCH 11/14] fix isel() bug with size 0 xr.DataArray indexer --- test/core/test_indexing.py | 57 +++++++++++++++++++++++++++++++++++++- uxarray/utils/coords.py | 5 +++- 2 files changed, 60 insertions(+), 2 deletions(-) diff --git a/test/core/test_indexing.py b/test/core/test_indexing.py index 19042f9b8..180c29557 100644 --- a/test/core/test_indexing.py +++ b/test/core/test_indexing.py @@ -291,7 +291,7 @@ def test_indexing_by_dataarray(): #assert "n_edge" in ds_edge.dims # uncomment after fixing #1758 counter = 0 # (count up during loop to make sure nothing is skipped unexpectedly) for method in "isel", "sel": - for dataset in [ds_face, ds_node]: # include after fixing #1758 + for dataset in [ds_face, ds_node]: # include ds_edge after fixing #1758 for grid_dim in ("n_face", "n_edge", "n_node"): for to_array in [False, True]: if grid_dim == "n_face" and "n_face" in dataset.dims: @@ -334,6 +334,61 @@ def test_dataset_isel_keeps_bonus_coords(): assert 'node_id' in sub.coords +def test_indexing_by_size_0_array(): + """ensure indexing by size 0 indexer works as expected (size 0 result). + Regression test for second reviewer's example (3) from PR 1759. + """ + def _as_array(obj): # return UxDataArray from UxDataArray, or UxDataset's data_vars[0]. + return obj[list(obj.data_vars)[0]] if isinstance(obj, ux.UxDataset) else obj + + # example (3) from PR 1759 + uxds = ux.tutorial.open_dataset("quad-hexagon") + assert set(uxds.dims) == {'n_face'} + uxds = uxds.assign_coords(node_id=("n_node", np.arange(uxds.uxgrid.n_node))) + assert set(uxds.dims) == {'n_face', 'n_node'} + hits = np.where(uxds.uxgrid.edge_lon > 1e9)[0] # empty + assert hits.size == 0 + result = uxds.isel(n_edge=xr.DataArray(hits, dims="selected")) + assert _as_array(result).size == 0 + assert 'node_id' in result.coords + assert set(result.dims) == {'n_face', 'n_node'} + assert result.sizes['n_face'] == result.sizes['n_node'] == 0 + assert result.uxgrid.n_face == result.uxgrid.n_node == result.uxgrid.n_edge == 0 + + # simpler tests, but applied across isel, sel, UxDataArray, UxDataset, n_edge, n_node, and n_face: + ds_face = ux.tutorial.open_dataset("quad-hexagon-random-face") + ds_node = ux.tutorial.open_dataset("quad-hexagon-random-node") + #ds_edge = ux.tutorial.open_dataset("quad-hexagon-random-edge") # uncomment after fixing #1758 + assert "n_face" in ds_face.dims + assert "n_node" in ds_node.dims + #assert "n_edge" in ds_edge.dims # uncomment after fixing #1758 + counter = 0 # (count up during loop to make sure nothing is skipped unexpectedly) + for method in "isel", "sel": + for dataset in [ds_face, ds_node]: # include ds_edge after fixing #1758 + for grid_dim in ("n_face", "n_edge", "n_node"): + for to_array in [False, True]: + counter += 1 + obj = _as_array(dataset) if to_array else dataset + # index by empty list: + result = getattr(obj, method)({grid_dim: []}) + assert _as_array(result).size == 0 + assert set(result.dims) == set(obj.dims) + # index by empty numpy array: + result = getattr(obj, method)({grid_dim: np.array([])}) + assert _as_array(result).size == 0 + assert set(result.dims) == set(obj.dims) + # index by empty xr.DataArray: + result = getattr(obj, method)({grid_dim: xr.DataArray([])}) + assert _as_array(result).size == 0 + assert set(result.dims) == set(obj.dims) + # index by empty xr.DataArray with a scalar coord: + indexer = xr.DataArray([]).assign_coords({"scalar7": 7}) + result = getattr(obj, method)({grid_dim: indexer}) + assert _as_array(result).size == 0 + assert set(result.dims) == set(obj.dims) + assert result.coords["scalar7"] == 7 + + def test_indexing_does_not_edit_indexers_dict(): """ensure isel() and sel() do not edit the provided indexers dict. Regression test for #1711. diff --git a/uxarray/utils/coords.py b/uxarray/utils/coords.py index dab1f6f5d..9f837fb32 100644 --- a/uxarray/utils/coords.py +++ b/uxarray/utils/coords.py @@ -173,7 +173,10 @@ def _assign_grid_dim_indexer_coords_if_appropriate(uxarray_obj, grid_dim, indexe coords = indexer.swap_dims({the_dim: "n_face"}).coords else: # remove any 1D coords (but keep scalar coords) - coords = indexer.isel({the_dim: 0}, drop=True).coords + if indexer.size > 0: + coords = indexer.isel({the_dim: 0}, drop=True).coords + else: # there is nothing along the 1 dim, so there is nothing to remove! + coords = indexer.coords else: raise DimensionError( f"2D+ indexers are not supported for grid dimensions. Got xr.DataArray " From 9dcd23a0dd28d483195d69edda9d23946c51d330 Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Wed, 23 Sep 2026 10:27:37 -0400 Subject: [PATCH 12/14] fix sel() dropping indexer's coords (also ensure sel() does a consistency check with indexer's coords, if both contain coords along grid dim.) --- test/core/test_indexing.py | 52 ++++++++++++++++++++++++++++++++++++++ uxarray/core/dataarray.py | 5 ++++ uxarray/core/dataset.py | 5 ++++ uxarray/core/utils.py | 25 ++++++++++++++---- uxarray/utils/coords.py | 11 ++++++++ 5 files changed, 93 insertions(+), 5 deletions(-) diff --git a/test/core/test_indexing.py b/test/core/test_indexing.py index 180c29557..c4011bcf3 100644 --- a/test/core/test_indexing.py +++ b/test/core/test_indexing.py @@ -315,6 +315,58 @@ def test_indexing_by_dataarray(): assert counter == 2 * n_data_grid_dim_combos * 2 +def test_indexing_of_uxarray_obj_with_grid_dim_coords(): + """ensure obj.isel() and obj.sel() work properly when obj has coords along grid dim. + Regression test for second reviewer's example (4) from PR 1759. + """ + # example (4) from PR 1759 + uxds = ux.tutorial.open_dataset("quad-hexagon").assign_coords(n_face=[0, 10, 20, 30]) + pick_by_label = xr.DataArray([0, 20], dims="station", coords={"station": ["A", "B"]}) + pick_by_index = xr.DataArray([0, 2], dims="station", coords={"station": ["A", "B"]}) + + by_index = uxds.isel(n_face=pick_by_index) + assert set(by_index.coords) == {'n_face', 'station'} + assert np.all(by_index.coords['n_face'] == [0, 20]) + assert np.all(by_index.coords['station'] == ['A', 'B']) + by_label = uxds.sel(n_face=pick_by_label) + assert by_label.identical(by_index) + + # repeat above, for UxDataArray: + uxarr = uxds['t2m'] + by_index = uxarr.isel(n_face=pick_by_index) + assert set(by_index.coords) == {'n_face', 'station'} + assert np.all(by_index.coords['n_face'] == [0, 20]) + assert np.all(by_index.coords['station'] == ['A', 'B']) + by_label = uxarr.sel(n_face=pick_by_label) + assert by_label.identical(by_index) + + # sanity check / spot tests: indexer with grid dim and coords: + # (A) should be allowed if coords don't conflict with obj + # (B) should crash if coords conflict with obj + indexerA_sel = xr.DataArray([0, 20], dims="n_face", coords={"n_face": [0, 20]}) + indexerB_sel = xr.DataArray([0, 20], dims="n_face", coords={"n_face": [7, 99]}) + indexerA_isel = xr.DataArray([0, 2], dims="n_face", coords={"n_face": [0, 20]}) + indexerB_isel = xr.DataArray([0, 2], dims="n_face", coords={"n_face": [0, 2]}) + resultA_sel = uxds.sel(n_face=indexerA_sel) + assert np.all(resultA_sel.coords['n_face'] == [0, 20]) + with pytest.raises(IndexError, match="dimension coordinate 'n_face' conflicts"): + _resultB_sel = uxds.sel(n_face=indexerB_sel) + resultA_isel = uxds.isel(n_face=indexerA_isel) + assert np.all(resultA_isel.coords['n_face'] == [0, 20]) + with pytest.raises(IndexError, match="dimension coordinate 'n_face' conflicts"): + _resultA_sel = uxds.isel(n_face=indexerB_isel) + + # repeat sanity checks, for UxDataArray: + resultA_sel = uxarr.sel(n_face=indexerA_sel) + assert np.all(resultA_sel.coords['n_face'] == [0, 20]) + with pytest.raises(IndexError, match="dimension coordinate 'n_face' conflicts"): + _resultB_sel = uxarr.sel(n_face=indexerB_sel) + resultA_isel = uxarr.isel(n_face=indexerA_isel) + assert np.all(resultA_isel.coords['n_face'] == [0, 20]) + with pytest.raises(IndexError, match="dimension coordinate 'n_face' conflicts"): + _resultA_sel = uxarr.isel(n_face=indexerB_isel) + + def test_dataset_isel_keeps_bonus_coords(): """ensure UxDataset.isel() keeps "bonus" coords, i.e. coords in the dataset which do not actually appear in any data var. diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 7a19dadd4..96adfe4d9 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -45,6 +45,7 @@ from uxarray.remap.accessor import RemapAccessor from uxarray.subset import DataArraySubsetAccessor from uxarray.utils.coords import ( + _assert_grid_dim_coord_consistent_if_in_both, _assign_grid_dim_indexer_coords_if_appropriate, _crash_if_1d_xarray_indexer_dim_in_uxarray_obj, _preserve_valid_coords, @@ -2224,6 +2225,10 @@ def sel( # offload the grid-indexing work to isel(): result = self.isel({grid_dim: grid_indices}, drop=drop) + # special case: if grid_dim in indexer and result.coords, ensure consistency. + # (all other coords' consistency checks already occurred in isel().) + _assert_grid_dim_coord_consistent_if_in_both(result, grid_dim, grid_indexer) + # index by other dims if any remain: ds = result.to_xarray().sel( indexers=indexers, # (grid_dim indexer was popped) diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 31360ad7a..9de3be7bd 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -29,6 +29,7 @@ from uxarray.plot.accessor import UxDatasetPlotAccessor from uxarray.remap.accessor import RemapAccessor from uxarray.utils.coords import ( + _assert_grid_dim_coord_consistent_if_in_both, _assign_grid_dim_indexer_coords_if_appropriate, _crash_if_1d_xarray_indexer_dim_in_uxarray_obj, ) @@ -672,6 +673,10 @@ def sel( # offload the grid-indexing work to isel(): result = self.isel({grid_dim: grid_indices}, drop=drop) + # special case: if grid_dim in indexer and result.coords, ensure consistency. + # (all other coords' consistency checks already occurred in isel().) + _assert_grid_dim_coord_consistent_if_in_both(result, grid_dim, grid_indexer) + # index by other dims if any remain: ds = result.to_xarray().sel( indexers=indexers, # (grid_dim indexer was popped) diff --git a/uxarray/core/utils.py b/uxarray/core/utils.py index c82d69928..84c0a2dc6 100644 --- a/uxarray/core/utils.py +++ b/uxarray/core/utils.py @@ -1,3 +1,5 @@ +import warnings + import numpy as np import xarray as xr from xarray.core.utils import either_dict_or_kwargs @@ -186,10 +188,7 @@ def _resolve_coordinate_labels_to_indices( ): """returns indices which would be selected by coord_array.sel({dim: labels_to_sel}, ...) coord_array.isel({dim: result}) should be equivalent to coord_array.sel({dim: labels_to_sel}, ...). - - (Implementation here drops extra coordinates from any indexers, - but if it is being applied to grid dims for sel() then it will produce behavior - which is consistent with isel(), unless issue #1712 gets fixed.) + If labels_to_sel is an xr.DataArray, its coords/dims will also be attached to the result. dim: str dimension name to select along @@ -203,10 +202,26 @@ def _resolve_coordinate_labels_to_indices( # (Maybe a more efficient implementation exists, but this is simple and gives correct results.) indices = xr.DataArray(np.arange(coord_array.sizes[dim]), dims=dim) _indices_coord_name = f"__{dim}_indices__" # just needs to be any unused name. + if _indices_coord_name in coord_array.coords: + warnings.warn( + f"Coordinate {_indices_coord_name!r} already exists in coord_array.coords " + "and will be overwritten, which may cause errors or subtly incorrect results..." + ) if hasattr(coord_array, "to_xarray"): # convert to xarray to avoid recursive sel() coord_array = coord_array.to_xarray() coord_with_indices = coord_array.assign_coords({_indices_coord_name: indices}) selected = coord_with_indices.sel( {dim: labels_to_sel}, method=method, tolerance=tolerance ) - return selected[_indices_coord_name].values # (return as np.ndarray, not DataArray) + result = selected[_indices_coord_name] + if isinstance(labels_to_sel, xr.DataArray): + # handle coords appropriately + result = result.drop_vars((dim, _indices_coord_name)) + # (drop grid dim coords because the caller is expected to handle those directly; + # here the goal is just to properly propagate any other coords from labels_to_sel.) + result = result.rename(None) # no reason to keep the _indices_coord_name + # (and keeping it for longer could maybe cause confusing error later?) + else: + # drop all coords/name info which was added internally during this method. + result = result.values + return result diff --git a/uxarray/utils/coords.py b/uxarray/utils/coords.py index 9f837fb32..737f90006 100644 --- a/uxarray/utils/coords.py +++ b/uxarray/utils/coords.py @@ -185,3 +185,14 @@ def _assign_grid_dim_indexer_coords_if_appropriate(uxarray_obj, grid_dim, indexe if coords: return uxarray_obj.assign_coords(coords) return uxarray_obj + + +def _assert_grid_dim_coord_consistent_if_in_both(uxarray_obj, grid_dim, indexer): + """assert grid_dim's coordinate is consistent if in both uxarray_obj and indexer. + Otherwise, does nothing. + """ + if isinstance(indexer, xr.DataArray): + if grid_dim in uxarray_obj.coords and grid_dim in indexer.coords: + xr.core.coordinates.assert_coordinate_consistent( + uxarray_obj, {grid_dim: indexer.coords.variables[grid_dim]} + ) From 7f68e0958e8962734dce23db91499ce0bee98128 Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Wed, 23 Sep 2026 15:58:34 -0400 Subject: [PATCH 13/14] fix coord conflict in isel(1d bool along n_face) --- test/core/test_indexing.py | 11 +++++++++++ uxarray/utils/coords.py | 18 +++++++++++++++--- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/test/core/test_indexing.py b/test/core/test_indexing.py index c4011bcf3..c0dac15cf 100644 --- a/test/core/test_indexing.py +++ b/test/core/test_indexing.py @@ -200,6 +200,17 @@ def test_isel_can_use_bool_with_coords(): coords={"lab": ("n_face", [10, 20, 30, 40])}) uxds.isel(n_face=mask) # (just ensuring it doesn't crash) + # also test what happens if uxarray object has grid dim coords: + ds1 = ux.tutorial.open_dataset("quad-hexagon").assign_coords(n_face=[0, 10, 20, 30]) + arr1 = ds1['t2m'] + wherebig = arr1 > 297.6 # hard-coding just to make the example easily. + # (if quad-hexagon file changes, rework this example) + assert np.all(wherebig.values == [False, True, True, False]) + resultA = arr1.isel(n_face=wherebig) # shouldn't crash! + assert np.all(resultA.coords['n_face'] == [10, 20]) + resultB = ds1.isel(n_face=wherebig) # shouldn't crash! + assert np.all(resultB.coords['n_face'] == [10, 20]) + def test_indexing_by_dataarray(): """ensure isel() and sel() with indexer=xr.DataArray(...) both work as expected. diff --git a/uxarray/utils/coords.py b/uxarray/utils/coords.py index 737f90006..235851656 100644 --- a/uxarray/utils/coords.py +++ b/uxarray/utils/coords.py @@ -149,10 +149,19 @@ def _assign_grid_dim_indexer_coords_if_appropriate(uxarray_obj, grid_dim, indexe (Also in this case, if indexer.to_xarray() exists, call it, to avoid recursion.) """ if isinstance(indexer, xr.DataArray): - xr.core.coordinates.assert_coordinate_consistent( - uxarray_obj, indexer.coords.variables + indexing_1d_bool_along_n_face = ( + indexer.ndim == 1 + and indexer.dtype == bool + and grid_dim == "n_face" + and "n_face" in uxarray_obj.dims ) - # ^ e.g. if uxarray_obj has time dim but indexer has time scalar coord, crash! + if not indexing_1d_bool_along_n_face: + # make sure indexer.coords are consistent with uxarray_obj's coords. + # e.g., if uxarray_obj has time dim but indexer has time scalar coord, crash! + xr.core.coordinates.assert_coordinate_consistent( + uxarray_obj, indexer.coords.variables + ) + # else: handle that check below (using indexer.isel(...) instead.) if indexer.ndim == 0: coords = indexer.coords elif indexer.ndim == 1: @@ -162,6 +171,9 @@ def _assign_grid_dim_indexer_coords_if_appropriate(uxarray_obj, grid_dim, indexe if hasattr(indexer, "to_xarray"): indexer = indexer.to_xarray() indexer = indexer.isel({the_dim: indexer}) + xr.core.coordinates.assert_coordinate_consistent( + uxarray_obj, indexer.coords.variables + ) if ( the_dim in uxarray_obj.coords and len(uxarray_obj.coords[the_dim].dims) == 0 From 816e3f13dc12e8db91eaca99922fde537c71177e Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Wed, 23 Sep 2026 18:28:14 -0400 Subject: [PATCH 14/14] fix subtle issues with fancy 0d indexer --- test/core/test_indexing.py | 15 +++++++++++++++ uxarray/utils/coords.py | 12 +++++++----- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/test/core/test_indexing.py b/test/core/test_indexing.py index c0dac15cf..8c0faffd1 100644 --- a/test/core/test_indexing.py +++ b/test/core/test_indexing.py @@ -418,6 +418,21 @@ def _as_array(obj): # return UxDataArray from UxDataArray, or UxDataset's data_ assert result.sizes['n_face'] == result.sizes['n_node'] == 0 assert result.uxgrid.n_face == result.uxgrid.n_node == result.uxgrid.n_edge == 0 + # followup to example (3): fancy empty array (with dim that has coords (which are empty)) + uxds = ux.tutorial.open_dataset("quad-hexagon") + fancy_empty = xr.DataArray(np.array([]), dims="selected", + coords={"selected": np.array([], dtype=int)}) + for grid_dim in ("n_face", "n_edge", "n_node"): + resultA = uxds.isel({grid_dim: fancy_empty}) + assert resultA.sizes == {'n_face': 0} # and, must not pick up "selected" dim from indexer. + resultB = uxds.sel({grid_dim: fancy_empty}) + assert resultB.sizes == {'n_face': 0} + # repeat for UxDataArray: + resultC = uxds['t2m'].isel({grid_dim: fancy_empty}) + assert resultC.sizes == {'n_face': 0} + resultD = uxds['t2m'].sel({grid_dim: fancy_empty}) + assert resultD.sizes == {'n_face': 0} + # simpler tests, but applied across isel, sel, UxDataArray, UxDataset, n_edge, n_node, and n_face: ds_face = ux.tutorial.open_dataset("quad-hexagon-random-face") ds_node = ux.tutorial.open_dataset("quad-hexagon-random-node") diff --git a/uxarray/utils/coords.py b/uxarray/utils/coords.py index 235851656..15c0bcf2e 100644 --- a/uxarray/utils/coords.py +++ b/uxarray/utils/coords.py @@ -184,11 +184,13 @@ def _assign_grid_dim_indexer_coords_if_appropriate(uxarray_obj, grid_dim, indexe ) # coords = indexer.swap_dims({the_dim: "n_face"}).coords else: - # remove any 1D coords (but keep scalar coords) - if indexer.size > 0: - coords = indexer.isel({the_dim: 0}, drop=True).coords - else: # there is nothing along the 1 dim, so there is nothing to remove! - coords = indexer.coords + # remove any 1D coords (but keep scalar coords). + # Drop by name to properly handle any subtleties of a 0D indexer + # (which would crash .isel({the_dim: 0}, drop=True, + # and would subtly add a new dim to the result if the_dim isn't dropped.) + coords = indexer.drop_vars( + [c for c in indexer.coords if the_dim in indexer[c].dims] + ).coords else: raise DimensionError( f"2D+ indexers are not supported for grid dimensions. Got xr.DataArray "