Fix: .sel() now actually slices the uxgrid too, for UxDataArray and UxDataset

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

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):
         """