From 70fd5c704a9564867d1a2dde64112233f8d9a505 Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:12:25 -0400 Subject: [PATCH 1/4] fix/cleanup/improve isel() see #1683 for details. Still need to add regression test for the UxDataArray.isel(..., ignore_grid=True) case --- test/core/test_dataarray.py | 6 +- uxarray/core/dataarray.py | 121 +++++++++++++++++------------------- uxarray/core/dataset.py | 102 ++++++++++++++++-------------- 3 files changed, 116 insertions(+), 113 deletions(-) diff --git a/test/core/test_dataarray.py b/test/core/test_dataarray.py index 5932809dc..be0226604 100644 --- a/test/core/test_dataarray.py +++ b/test/core/test_dataarray.py @@ -125,14 +125,14 @@ def test_isel_invalid_dim(gridpath, datasetpath): uxda = UxDataArray(data, dims=["time", "n_face"], uxgrid=uxds.uxgrid) with pytest.raises( - DimensionError, - match=r"Dimensions \{'invalid_dim'\} do not exist\..*Available dimensions: \('time', 'n_face'\)", + ValueError, + match=r"Dimensions \{'invalid_dim'\} do not exist\. Expected one or more of \('time', 'n_face'\)", ): uxda.isel(invalid_dim=0) with pytest.raises( ValueError, - match=r"Dimensions \{'level'\} do not exist\..*Available dimensions: \('time', 'n_face'\)", + match=r"Dimensions \{'level'\} do not exist\. Expected one or more of \('time', 'n_face'\)", ): uxda.isel(level=0) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index f265ad53b..ef3ea1cc1 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -1966,30 +1966,40 @@ def isel( inverse_indices: bool = False, **indexers_kwargs, ): - """ - Return a new DataArray whose data is given by selecting indexes along the specified dimension(s). + """Return a new UxDataArray indexed along the specified dimension(s). + The data is indexed, as well as the underlying grid when applicable. - Performs xarray-style integer-location indexing along specified dimensions. - If a single grid dimension ('n_node', 'n_edge', or 'n_face') is provided - and `ignore_grid=False`, the underlying grid is sliced accordingly, - and remaining indexers are applied to the resulting DataArray. + Grid dimensions ('n_node', 'n_edge', 'n_face') are treated specially + when `ignore_grid=False`. Providing one of them will slice to the specified + nodes, edges, or faces, regardless of data location. If the data does not + contain the specified dimension, the result will have the minimal grid + region containing everything specified. For example, using n_edge=7 for data + on 'n_face' makes a result with 'n_face' with just the two faces on edge 7. Parameters ---------- indexers : Mapping[Any, Any], optional - A mapping of dimension names to indexers. Each indexer may be an integer, - slice, array-like, or DataArray. Mutually exclusive with indexing via kwargs. + A dict with keys matching dimensions and values given + by integers, slice objects or arrays. + indexer can be a integer, slice, array-like or DataArray. + If DataArrays are passed as indexers, xarray-style indexing will be + carried out. See :ref:`indexing` for the details. + One of indexers or indexers_kwargs must be provided. drop : bool, default=False - If True, drop any coordinate variables indexed by integers instead of - retaining them as length-1 dimensions. + If ``drop=True``, drop coordinates variables indexed by integers + instead of making them scalar. missing_dims : {'raise', 'warn', 'ignore'}, default='raise' - Behavior when indexers reference dimensions not present in the array. - - 'raise': raise an error - - 'warn': emit a warning and ignore missing dimensions - - 'ignore': ignore missing dimensions silently + What to do if dimensions that should be selected from are not present in the + UxDataArray: + - "raise": raise an exception + - "warn": raise a warning, and ignore the missing dimensions + - "ignore": ignore the missing dimensions ignore_grid : bool, default=False - If False (default), allow slicing on one grid dimension to automatically - update the associated UXarray grid. If True, fall back to pure xarray behavior. + If False (default), slice the underlying UXarray grid appropriately too, + ensuring the resulting data actually lies on the result's underlying grid. + If True, slice the data only; attach self.uxgrid to the result, unchanged. + CAUTION: using ignore_grid=True will cause the result's data to be + inconsistent with its underlying grid, if any grid dimensions were sliced. inverse_indices : bool, default=False For grid-based slicing, pass this flag to `Grid.isel` to invert indices when selecting (useful for staggering or reversing order). @@ -2005,6 +2015,9 @@ def isel( ------ DimensionError (subclass of ValueError) If more than one grid dimension is selected and `ignore_grid=False`. + ValueError + If parameters are invalid for xarray's .isel(), such as if + slicing by a nonexistent dimension, or using invalid indexers. """ from uxarray.core.utils import _validate_indexers @@ -2012,57 +2025,39 @@ def isel( indexers, indexers_kwargs, "isel", ignore_grid ) - try: - # Grid Branch - if not ignore_grid: - if len(grid_dims) == 1: - # pop off the one grid‐dim indexer - grid_dim = grid_dims.pop() - grid_indexer = indexers.pop(grid_dim) - - sliced_grid = self.uxgrid.isel( - **{grid_dim: grid_indexer}, inverse_indices=inverse_indices - ) + if ignore_grid or len(grid_dims)==0: + # no grid dims, or ignore_grid=True --> just call xarray's isel + return type(self)( + super().isel( + indexers=indexers or None, + drop=drop, + missing_dims=missing_dims, + ), + uxgrid=self.uxgrid, + ) + elif len(grid_dims) == 1: + # pop off the one grid‐dim indexer + grid_dim = grid_dims.pop() + grid_indexer = indexers.pop(grid_dim) - da = self._slice_from_grid(sliced_grid) + sliced_grid = self.uxgrid.isel( + **{grid_dim: grid_indexer}, inverse_indices=inverse_indices + ) - # if there are any remaining indexers, apply them - if indexers: - xarr = super(UxDataArray, da).isel( - indexers=indexers, drop=drop, missing_dims=missing_dims - ) - # re‐wrap so the grid sticks around - return type(self)(xarr, uxgrid=sliced_grid) + da = self._slice_from_grid(sliced_grid) - # no other dims, return the grid‐sliced da - return da - else: - return type(self)( - super().isel( - indexers=indexers or None, - drop=drop, - missing_dims=missing_dims, - ), - uxgrid=self.uxgrid, - ) + # if there are any remaining indexers, apply them + if indexers: + xarr = super(UxDataArray, da).isel( + indexers=indexers, drop=drop, missing_dims=missing_dims + ) + # re‐wrap so the grid sticks around + return type(self)(xarr, uxgrid=sliced_grid) - return super().isel( - indexers=indexers or None, - drop=drop, - missing_dims=missing_dims, - ) - except ValueError as e: - if "Dimensions" in str(e) and "do not exist" in str(e): - # The error message from xarray is quite good, but we can add to it. - # e.g. "Dimensions {'level'} do not exist. Expected one of ('n_face', 'time', 'lev')" - # Let's just append the available dimensions. - original_error_msg = str(e) - raise DimensionError( - f"{original_error_msg}. Available dimensions: {self.dims}" - ) from e - else: - # re-raise other ValueErrors - raise e + # no other dims, return the grid‐sliced da + return da + else: # len(grid_dims)>1; _validate_indexers should have crashed. + assert False, 'internal implementation error if reached this line' @classmethod def from_xarray(cls, da: xr.DataArray, uxgrid: Grid, ugrid_dims: dict = None): diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 41194dfe2..c8724f675 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -419,13 +419,16 @@ def isel( inverse_indices: bool = False, **indexers_kwargs, ): - """Returns a new dataset with each array indexed along the specified - dimension(s). + """Return a new UxDataset with indexed along the specified dimension(s). + Each data array is indexed appropriately, + along with the underlying grid when applicable. - Performs xarray-style integer-location indexing along specified dimensions. - If a single grid dimension ('n_node', 'n_edge', or 'n_face') is provided - and `ignore_grid=False`, the underlying grid is sliced accordingly, - and remaining indexers are applied to the resulting Dataset. + Grid dimensions ('n_node', 'n_edge', 'n_face') are treated specially + when `ignore_grid=False`. Providing one of them will slice to the specified + nodes, edges, or faces, regardless of data location. If the data does not + contain the specified dimension, the result will have the minimal grid + region containing everything specified. For example, using n_edge=7 for data + on 'n_face' makes a result with 'n_face' with just the two faces on edge 7. Parameters ---------- @@ -441,26 +444,34 @@ def isel( instead of making them scalar. missing_dims : {"raise", "warn", "ignore"}, default: "raise" What to do if dimensions that should be selected from are not present in the - Dataset: + UxDataset: - "raise": raise an exception - "warn": raise a warning, and ignore the missing dimensions - "ignore": ignore the missing dimensions ignore_grid : bool, default=False - If False (default), allow slicing on one grid dimension to automatically - update the associated UXarray grid. If True, fall back to pure xarray behavior. + If False (default), slice the underlying UXarray grid appropriately too, + ensuring the resulting data actually lies on the result's underlying grid. + If True, slice the data only; attach self.uxgrid to the result, unchanged. + CAUTION: using ignore_grid=True will cause the result's data to be + inconsistent with its underlying grid, if any grid dimensions were sliced. inverse_indices : bool, default=False For grid-based slicing, pass this flag to `Grid.isel` to invert indices when selecting (useful for staggering or reversing order). **indexers_kwargs : dimension=indexer pairs, optional + The keyword arguments form of `indexers`. - **indexers_kwargs : {dim: indexer, ...}, optional - The keyword arguments form of ``indexers``. - One of indexers or indexers_kwargs must be provided. - - Returns + Returns ------- UxDataset A new UxDataset indexed according to `indexers` and updated grid if applicable. + + Raises + ------ + DimensionError (subclass of ValueError) + If more than one grid dimension is selected and `ignore_grid=False`. + ValueError + If parameters are invalid for xarray's .isel(), such as if + slicing by a nonexistent dimension, or using invalid indexers. """ from uxarray.core.utils import _validate_indexers @@ -468,43 +479,40 @@ def isel( indexers, indexers_kwargs, "isel", ignore_grid ) - if not ignore_grid: - if len(grid_dims) == 1: - grid_dim = grid_dims.pop() - 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, - ) + if ignore_grid or len(grid_dims)==0: + # no grid dims, or ignore_grid=True --> just call xarray's isel + return type(self)( + super().isel( + indexers=indexers or None, + drop=drop, + missing_dims=missing_dims, + ), + uxgrid=self.uxgrid, + ) + elif len(grid_dims) == 1: + # pop off the one grid‐dim indexer + grid_dim = grid_dims.pop() + grid_indexer = indexers.pop(grid_dim) + + # slice the grid + sliced_grid = self.uxgrid.isel( + **{grid_dim: grid_indexer}, inverse_indices=inverse_indices + ) - if indexers: - ds = xr.Dataset.isel( - ds, indexers=indexers, drop=drop, missing_dims=missing_dims - ) + ds = self._slice_dataset_from_grid( + sliced_grid=sliced_grid, + grid_dim=grid_dim, + grid_indexer=grid_indexer, + ) - return type(self)(ds, uxgrid=sliced_grid) - else: - return type(self)( - super().isel( - indexers=indexers or None, - drop=drop, - missing_dims=missing_dims, - ), - uxgrid=self.uxgrid, + if indexers: + ds = xr.Dataset.isel( + ds, indexers=indexers, drop=drop, missing_dims=missing_dims ) - return super().isel( - indexers=indexers or None, - drop=drop, - missing_dims=missing_dims, - ) + return type(self)(ds, uxgrid=sliced_grid) + else: # len(grid_dims)>1; _validate_indexers should have crashed. + assert False, 'internal implementation error if reached this line' def __getattribute__(self, name): """Intercept accessor method calls to return Ux-aware accessors.""" From 46020ef9dd5285eaca3f7b70296adcb15f344602 Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Tue, 18 Aug 2026 14:30:05 -0400 Subject: [PATCH 2/4] test: UxDataset.isel(..., ignore_grid=True) --- test/core/test_dataset.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/test/core/test_dataset.py b/test/core/test_dataset.py index 9b26d84d3..2a503af6a 100644 --- a/test/core/test_dataset.py +++ b/test/core/test_dataset.py @@ -129,6 +129,20 @@ def test_sel_method_forwarded(gridpath, datasetpath): np.array(uxds["time"].values[2], dtype="datetime64[ns]"), ) +def test_isel_ignore_grid(): + """ensure UxDataset.isel(..., ignore_grid=True) still attaches result.uxgrid. + Regression test for issue #1683. + """ + uxds = ux.tutorial.open_dataset("outCSne30-timeseries") + result = uxds.isel(time=0, ignore_grid=True) + result.uxgrid # (will cause crash if uxgrid not properly attached to result) + assert result.uxgrid == uxds.uxgrid + + result = uxds.isel(n_face=0, ignore_grid=True) + result.uxgrid # (will cause crash if uxgrid not properly attached to result) + assert result.uxgrid == uxds.uxgrid # ignore_grid means grid never gets sliced here + + def test_uxdataset_init_from_xarray_dataset(): ds = xr.Dataset( data_vars={"a": ("x", [1, 2])}, From 9787af526b4aa873f1b42f81d2f4b579422f4ea0 Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Tue, 18 Aug 2026 15:05:10 -0400 Subject: [PATCH 3/4] forgot pre-commit ruff formatting --- uxarray/core/dataarray.py | 4 ++-- uxarray/core/dataset.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index d741470c0..5a1ee7e62 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -2030,7 +2030,7 @@ def isel( indexers, indexers_kwargs, "isel", ignore_grid ) - if ignore_grid or len(grid_dims)==0: + if ignore_grid or len(grid_dims) == 0: # no grid dims, or ignore_grid=True --> just call xarray's isel return type(self)( super().isel( @@ -2062,7 +2062,7 @@ def isel( # no other dims, return the grid‐sliced da return da else: # len(grid_dims)>1; _validate_indexers should have crashed. - assert False, 'internal implementation error if reached this line' + assert False, "internal implementation error if reached this line" @classmethod def from_xarray(cls, da: xr.DataArray, uxgrid: Grid, ugrid_dims: dict = None): diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 53befd96e..1f511db33 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -478,7 +478,7 @@ def isel( indexers, indexers_kwargs, "isel", ignore_grid ) - if ignore_grid or len(grid_dims)==0: + if ignore_grid or len(grid_dims) == 0: # no grid dims, or ignore_grid=True --> just call xarray's isel return type(self)( super().isel( @@ -511,7 +511,7 @@ def isel( return type(self)(ds, uxgrid=sliced_grid) else: # len(grid_dims)>1; _validate_indexers should have crashed. - assert False, 'internal implementation error if reached this line' + assert False, "internal implementation error if reached this line" def __getattribute__(self, name): """Intercept accessor method calls to return Ux-aware accessors.""" From 631f98c5d0c1b2bb9f602a9416354b610f5db4ca Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Wed, 19 Aug 2026 13:09:15 -0400 Subject: [PATCH 4/4] raise AssertionError instead of assert False --- uxarray/core/dataarray.py | 2 +- uxarray/core/dataset.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 5a1ee7e62..050cfeef7 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -2062,7 +2062,7 @@ def isel( # no other dims, return the grid‐sliced da return da else: # len(grid_dims)>1; _validate_indexers should have crashed. - assert False, "internal implementation error if reached this line" + raise AssertionError("internal implementation error if reached this line") @classmethod def from_xarray(cls, da: xr.DataArray, uxgrid: Grid, ugrid_dims: dict = None): diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 1f511db33..e78a52bda 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -511,7 +511,7 @@ def isel( return type(self)(ds, uxgrid=sliced_grid) else: # len(grid_dims)>1; _validate_indexers should have crashed. - assert False, "internal implementation error if reached this line" + raise AssertionError("internal implementation error if reached this line") def __getattribute__(self, name): """Intercept accessor method calls to return Ux-aware accessors."""