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/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])}, diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index b62b7f1c9..050cfeef7 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -1971,30 +1971,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). @@ -2010,6 +2020,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 @@ -2017,57 +2030,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. + 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 129466589..e78a52bda 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -418,13 +418,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 ---------- @@ -440,26 +443,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 @@ -467,43 +478,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. + raise AssertionError("internal implementation error if reached this line") def __getattribute__(self, name): """Intercept accessor method calls to return Ux-aware accessors."""