From 97e8f6d6e2cc2f5a5d336610d195c5142f3ea3ef Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Wed, 26 Aug 2026 14:23:58 -0400 Subject: [PATCH 1/2] fix: sel wasn't slicing uxgrid; add sel docstrings (sel had docstrings before but they were the xarray docstrings, so they didn't say anything about uxarray-specific behaviors.) --- uxarray/core/dataarray.py | 132 ++++++++++++++++++++++++++++++++- uxarray/core/dataset.py | 151 +++++++++++++++++++++++++++++++++----- uxarray/core/utils.py | 60 ++++++++++++++- 3 files changed, 317 insertions(+), 26 deletions(-) diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 75c90e4b3..bdb5938a3 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -2,7 +2,7 @@ import warnings from html import escape -from typing import TYPE_CHECKING, Any, Hashable, Literal, Mapping, Optional +from typing import TYPE_CHECKING, Any, Hashable, Iterable, Literal, Mapping, Optional from warnings import warn import numpy as np @@ -19,7 +19,11 @@ _calculate_edge_node_difference, _compute_gradient, ) -from uxarray.core.utils import _map_dims_to_ugrid +from uxarray.core.utils import ( + _map_dims_to_ugrid, + _resolve_coordinate_labels_to_indices, + _validate_indexers, +) from uxarray.core.zonal import ( _compute_conservative_zonal_mean_bands, _compute_non_conservative_zonal_mean, @@ -2026,8 +2030,6 @@ def isel( 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 - indexers, grid_dims = _validate_indexers( indexers, indexers_kwargs, "isel", ignore_grid ) @@ -2066,6 +2068,128 @@ def isel( else: # len(grid_dims)>1; _validate_indexers should have crashed. raise AssertionError("internal implementation error if reached this line") + def sel( + self, + indexers: Mapping[Any, Any] | None = None, + method: str | None = None, + tolerance: int | float | Iterable[int | float] | None = None, + drop: bool = False, + **indexers_kwargs: Any, + ): + """Returns a new array indexed by labels, instead of indices, along the specified dimension(s). + + Grid dimensions ('n_node', 'n_edge', 'n_face') are treated specially. + 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. + + By default, grid dims do not have coordinates assigned. But, if they have + been assigned, `.sel()` respects them in the intuitive way. For example, + using `.sel(n_face=30)` for data with `n_face` coordinates [0,10,20,30,40] + would be equivalent to using `.isel(n_face=3)`. Meanwhile, if the data + does not contain the specified grid dim (as in the n_edge=7 example above), + it also cannot contain coordinates along that grid dim, + so in that case `.sel()` performs index-based selection just like `.isel()`. + + Under the hood, this method is powered by using pandas's powerful Index + objects. This makes label based indexing essentially just as fast as + using integer indexing. + + It also means this method uses pandas's (well documented) logic for + indexing. This means you can use string shortcuts for datetime indexes + (e.g., '2000-01' to select all values in January 2000). It also means + that slices are treated as inclusive of both the start and stop values, + unlike normal Python indexing. + + Parameters + ---------- + indexers : dict, optional + A dict with keys matching dimensions and values given + by scalars, slices or arrays of tick labels. For dimensions with + 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. + One of indexers or indexers_kwargs must be provided. + method : {None, "nearest", "pad", "ffill", "backfill", "bfill"}, optional + Method to use for inexact matches: + + * None (default): only exact matches + * pad / ffill: propagate last valid index value forward + * backfill / bfill: propagate next valid index value backward + * nearest: use nearest valid index value + tolerance : optional + Maximum distance between original and new labels for inexact + matches. The values of the index at the matching locations must + satisfy the equation ``abs(index[indexer] - target) <= tolerance``. + drop : bool, optional + If ``drop=True``, drop coordinates variables in `indexers` instead + of making them scalar. + **indexers_kwargs : {dim: indexer, ...}, optional + The keyword arguments form of ``indexers``. + One of indexers or indexers_kwargs must be provided. + + Returns + ------- + obj : UxDataArray + A new UxDataArray with each dimension is indexed appropriately, + 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). + 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, + in which case the data will be a copy. + """ + indexers, grid_dims = _validate_indexers( + indexers, indexers_kwargs, "sel", ignore_grid=False + ) # (sel doesn't support ignore_grid=True option) + + if len(grid_dims) == 0: + # no grid dims --> just call xarray's sel + return type(self)( + self.to_xarray().sel( + indexers=indexers, + method=method, + tolerance=tolerance, + drop=drop, + ), + uxgrid=self.uxgrid, + ) + elif len(grid_dims) == 1: + # pop off the one grid‐dim indexer + grid_dim = list(grid_dims)[0] + indexers = indexers.copy() # don't modify the original dict + grid_indexer = indexers.pop(grid_dim) + if grid_dim in self.coords: # label-based indexing + grid_indices = _resolve_coordinate_labels_to_indices( + grid_dim, + grid_indexer, + self.coords[grid_dim], + method=method, + tolerance=tolerance, + ) + else: # index-based indexing + grid_indices = grid_indexer + + # offload the grid-indexing work to isel(): + result = self.isel({grid_dim: grid_indices}, drop=drop) + + # index by other dims if any remain: + ds = result.to_xarray().sel( + indexers=indexers, # (grid_dim indexer was popped) + method=method, + tolerance=tolerance, + drop=drop, + ) + + return type(self)(ds, uxgrid=result.uxgrid) + 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 87a3c45a2..6d9affc6e 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -3,7 +3,7 @@ import os import sys from html import escape -from typing import IO, Any, Hashable, Mapping +from typing import IO, Any, Hashable, Iterable, Mapping from warnings import warn import xarray as xr @@ -13,7 +13,12 @@ import uxarray from uxarray.core.dataarray import UxDataArray -from uxarray.core.utils import _map_dims_to_ugrid, _open_dataset_with_fallback +from uxarray.core.utils import ( + _map_dims_to_ugrid, + _open_dataset_with_fallback, + _resolve_coordinate_labels_to_indices, + _validate_indexers, +) from uxarray.errors import DimensionError, GridInvalidError from uxarray.formatting_html import dataset_repr from uxarray.grid import Grid @@ -473,8 +478,6 @@ def isel( 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 - indexers, grid_dims = _validate_indexers( indexers, indexers_kwargs, "isel", ignore_grid ) @@ -514,6 +517,130 @@ def isel( else: # len(grid_dims)>1; _validate_indexers should have crashed. raise AssertionError("internal implementation error if reached this line") + def sel( + self, + indexers: Mapping[Any, Any] | None = None, + method: str | None = None, + tolerance: int | float | Iterable[int | float] | None = None, + drop: bool = False, + **indexers_kwargs: Any, + ): + """Returns a new dataset with each array indexed by labels, instead of indices, + along the specified dimension(s). + + Grid dimensions ('n_node', 'n_edge', 'n_face') are treated specially. + 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. + + By default, grid dims do not have coordinates assigned. But, if they have + been assigned, `.sel()` respects them in the intuitive way. For example, + using `.sel(n_face=30)` for data with `n_face` coordinates [0,10,20,30,40] + would be equivalent to using `.isel(n_face=3)`. Meanwhile, if the data + does not contain the specified grid dim (as in the n_edge=7 example above), + it also cannot contain coordinates along that grid dim, + so in that case `.sel()` performs index-based selection just like `.isel()`. + + Under the hood, this method is powered by using pandas's powerful Index + objects. This makes label based indexing essentially just as fast as + using integer indexing. + + It also means this method uses pandas's (well documented) logic for + indexing. This means you can use string shortcuts for datetime indexes + (e.g., '2000-01' to select all values in January 2000). It also means + that slices are treated as inclusive of both the start and stop values, + unlike normal Python indexing. + + Parameters + ---------- + indexers : dict, optional + A dict with keys matching dimensions and values given + by scalars, slices or arrays of tick labels. For dimensions with + 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. + One of indexers or indexers_kwargs must be provided. + method : {None, "nearest", "pad", "ffill", "backfill", "bfill"}, optional + Method to use for inexact matches: + + * None (default): only exact matches + * pad / ffill: propagate last valid index value forward + * backfill / bfill: propagate next valid index value backward + * nearest: use nearest valid index value + tolerance : optional + Maximum distance between original and new labels for inexact + matches. The values of the index at the matching locations must + satisfy the equation ``abs(index[indexer] - target) <= tolerance``. + drop : bool, optional + If ``drop=True``, drop coordinates variables in `indexers` instead + of making them scalar. + **indexers_kwargs : {dim: indexer, ...}, optional + The keyword arguments form of ``indexers``. + One of indexers or indexers_kwargs must be provided. + + Returns + ------- + obj : UxDataset + A new UxDataset with the same contents as this dataset, except each + variable and dimension is indexed by the appropriate indexers, + 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). + 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, + in which case the data will be a copy. + """ + indexers, grid_dims = _validate_indexers( + indexers, indexers_kwargs, "sel", ignore_grid=False + ) # (sel doesn't support ignore_grid=True option) + + if len(grid_dims) == 0: + # no grid dims --> just call xarray's sel + return type(self)( + self.to_xarray().sel( + indexers=indexers, + method=method, + tolerance=tolerance, + drop=drop, + ), + uxgrid=self.uxgrid, + ) + elif len(grid_dims) == 1: + # pop off the one grid‐dim indexer + grid_dim = list(grid_dims)[0] + indexers = indexers.copy() # don't modify the original dict + grid_indexer = indexers.pop(grid_dim) + if grid_dim in self.coords: # label-based indexing + grid_indices = _resolve_coordinate_labels_to_indices( + grid_dim, + grid_indexer, + self.coords[grid_dim], + method=method, + tolerance=tolerance, + ) + else: # index-based indexing + grid_indices = grid_indexer + + # offload the grid-indexing work to isel(): + result = self.isel({grid_dim: grid_indices}, drop=drop) + + # index by other dims if any remain: + ds = result.to_xarray().sel( + indexers=indexers, # (grid_dim indexer was popped) + method=method, + tolerance=tolerance, + drop=drop, + ) + + return type(self)(ds, uxgrid=result.uxgrid) + 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.""" # Lazy import to avoid circular imports @@ -802,22 +929,6 @@ def where(self, cond: Any, other: Any = dtypes.NA, drop: bool = False): where.__doc__ = xr.Dataset.where.__doc__ - def sel( - self, indexers=None, method=None, tolerance=None, drop=False, **indexers_kwargs - ): - return UxDataset( - self.to_xarray().sel( - indexers=indexers, - method=method, - tolerance=tolerance, - drop=drop, - **indexers_kwargs, - ), - uxgrid=self.uxgrid, - ) - - sel.__doc__ = xr.Dataset.sel.__doc__ - def fillna(self, value: Any): return UxDataset(super().fillna(value), uxgrid=self._uxgrid) diff --git a/uxarray/core/utils.py b/uxarray/core/utils.py index bfa6c509b..3eb3ab836 100644 --- a/uxarray/core/utils.py +++ b/uxarray/core/utils.py @@ -1,5 +1,8 @@ +import numpy as np import xarray as xr +from xarray.core.utils import either_dict_or_kwargs +from uxarray.constants import GRID_DIMS from uxarray.errors import DimensionError from uxarray.io.utils import _get_source_dims_dict, _parse_grid_type @@ -123,9 +126,31 @@ def match_chunks_to_ugrid(grid_filename_or_obj, chunks): def _validate_indexers(indexers, indexers_kwargs, func_name, ignore_grid): - from xarray.core.utils import either_dict_or_kwargs + """returns (dict of indexers, set of grid_dim strs). - from uxarray.constants import GRID_DIMS + Parameters + ---------- + indexers: dict + indexers originally provided as dict. E.g., uxarr.isel({'n_face': 0}). + Provide indexers or indexers_kwargs but not both. + indexers_kwargs: dict + indexers originally provided as kwargs. E.g. uxarr.isel(n_face=0). + Provide indexers or indexers_kwargs but not both. + func_name: str + name of the function calling _validate_indexers. E.g. "isel". + Included in error message if provided both indexers and indexers_kwargs. + ignore_grid: bool + whether ignore_grid=True flag was set in the indexing operation. + If False, ensure len(grid_dims) <= 1 else raise DimensionError. + + Returns + ------- + indexers: dict + validated dict of indexers, including grid dims indexers if present. + grid_dims: set + set of grid dimension names (from ``GRID_DIMS``) present as keys in indexers; + values from {"n_face", "n_node", "n_edge"} (at most 1 value if ignore_grid=False). + """ # Used to filter out slices containing all Nones (causes subscription errors, i.e., var[0]) _is_full_none_slice = lambda v: ( @@ -147,3 +172,34 @@ def _validate_indexers(indexers, indexers_kwargs, func_name, ignore_grid): ) return indexers, grid_dims + + +def _resolve_coordinate_labels_to_indices( + dim, labels_to_sel, coord_array, *, method=None, tolerance=None +): + """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.) + + dim: str + dimension name to select along + labels_to_sel: any valid indexer which can be passed to .sel() + values to select along dim + coord_array: xr.DataArray or UxDataArray + coordinate array to select from. + method, tolerance: passed directly to .sel(). + """ + # just using xarray's .sel() on a simple np.arange(), to ensure exactly consistent behavior with sel(). + # (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 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) From 9afef65bcbb645fe29f2ec9988f33ea7785b65df Mon Sep 17 00:00:00 2001 From: Sam Evans <47793072+Sevans711@users.noreply.github.com> Date: Wed, 26 Aug 2026 17:36:44 -0400 Subject: [PATCH 2/2] adds test_indexing tests for sel() --- test/core/test_indexing.py | 85 ++++++++++++++++++++++++++++++++++++++ uxarray/core/dataset.py | 2 +- 2 files changed, 86 insertions(+), 1 deletion(-) create mode 100644 test/core/test_indexing.py diff --git a/test/core/test_indexing.py b/test/core/test_indexing.py new file mode 100644 index 000000000..cc9b148a4 --- /dev/null +++ b/test/core/test_indexing.py @@ -0,0 +1,85 @@ +""" +Purpose: tests related to indexing Grid, UxDataArray, and/or UxDataset, +e.g. UxDataArray's and UxDataset's .isel() and .sel() methods. + +(Some .isel() and sel() tests are in test_dataarray.py and test_dataset.py, +but could maybe be moved here? Having test_indexing.py as its own file helps +to ensure consistency between UxDataArray and UxDataset indexing.) +""" +import numpy as np +import uxarray as ux + +def test_sel_indexes_grid(): + """ensure obj.sel({grid_dim: ...}) actually indexes the result.uxgrid, too, + for UxDataArrays and UxDatasets. Regression test for #1641. + """ + # extremely simple case: + uxds = ux.tutorial.open_dataset("quad-hexagon") + result = uxds.sel(n_face=0) + assert result.sizes['n_face'] == result.uxgrid.n_face == 1 + # (similar check for UxDataArray) + uxarr = uxds['t2m'] + result = uxarr.sel(n_face=0) + assert result.sizes['n_face'] == result.uxgrid.n_face == 1 + + # more complicated case: + uxds = ux.tutorial.open_dataset("outCSne30-timeseries") + assert uxds.sizes['n_face'] > 5000 + result = uxds.sel(time=['2018-04-28T00', '2018-04-28T03'], n_face=range(0, 5000, 100)) + assert result.sizes == {'time': 2, 'n_face': 50} + assert result.uxgrid.n_face == 50 + # (similar check for UxDataArray) + uxarr = uxds['psi'] + result = uxarr.sel(time=['2018-04-28T00', '2018-04-28T03'], n_face=range(0, 5000, 100)) + assert result.sizes == {'time': 2, 'n_face': 50} + assert result.uxgrid.n_face == 50 + +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 + """ + # 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 + + # test corresponding to the workflow described in #1641, for UxDataArray + uxarr = uxds['psi'] + uxarr1 = uxarr.assign_coords(n_face=np.arange(uxarr.n_face.size)) + uxarr2 = uxarr1.isel(n_face=range(0, 100, 5)) + uxarr3 = uxarr2 + 7 + # "check what the results look like on what were originally faces 20, 30, and 40" + uxarr4 = uxarr3.sel(n_face=[20,30,40]) + assert uxarr4.sizes['n_face'] == uxarr4.uxgrid.n_face == 3 + + # test to check what happens if using coordinate labels not equal to indexes: + uxarr1 = uxarr.assign_coords(n_face=10*np.arange(uxarr.n_face.size)) + uxarr2 = uxarr1.isel(n_face=[5,6,7,8]) + assert np.all(uxarr2 == uxarr1.sel(n_face=[50,60,70,80])) + assert np.all(uxarr2['n_face'] == [50,60,70,80]) # (isel shouldn't drop coord labels) + uxarr3 = uxarr2.isel(n_face=2) + assert np.all(uxarr3 == uxarr2.sel(n_face=70)) + +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. + """ + 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 + + arr = ds["psi"] + result = arr.isel(n_edge=7) + assert result.sizes["n_face"] == result.uxgrid.n_face == 2 + result = arr.sel(n_edge=7) + assert result.sizes["n_face"] == result.uxgrid.n_face == 2 diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 6d9affc6e..4ceda31d1 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -424,7 +424,7 @@ def isel( inverse_indices: bool = False, **indexers_kwargs, ): - """Return a new UxDataset with indexed along the specified dimension(s). + """Return a new UxDataset with arrays indexed along the specified dimension(s). Each data array is indexed appropriately, along with the underlying grid when applicable.