diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 21f3671188a..76db46bcc75 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -165,6 +165,10 @@ Deprecations Bug Fixes ~~~~~~~~~ +- Fix CuPy indexers with pandas-backed coordinates, basic and label-based + assignment to CuPy arrays, and scalar truth conversion while preserving + pandas extension-array behavior (:issue:`11739`, :pull:`11797`). + By `Huashuo Lei `_. - Results of the string accessor methods ``strip``, ``pad``, ``center``, ``ljust``, ``rjust``, ``zfill`` and ``replace`` on dask arrays are no longer truncated to a single character, and ``repeat`` no longer fails on dask arrays diff --git a/xarray/core/common.py b/xarray/core/common.py index fbd3ebc57a1..466075c8f0c 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -166,7 +166,10 @@ class AbstractArray(Generic[DimType_co]): # noqa: UP046 __slots__ = () def __bool__(self: Any) -> bool: - return bool(self.values) + data = self.data + if isinstance(data, pd.api.extensions.ExtensionArray): + return bool(self.values) + return bool(data) def __float__(self: Any) -> float: return float(self.values) diff --git a/xarray/core/indexes.py b/xarray/core/indexes.py index 2a513c5c938..3fd1fbdfd56 100644 --- a/xarray/core/indexes.py +++ b/xarray/core/indexes.py @@ -26,12 +26,15 @@ get_valid_numpy_dtype, is_allowed_extension_array_dtype, is_dict_like, + is_duck_array, is_scalar, ) +from xarray.namedarray.pycompat import to_numpy if TYPE_CHECKING: from xarray.core.types import ErrorOptions, JoinOptions, Self from xarray.core.variable import Variable + from xarray.namedarray._typing import duckarray IndexVars = dict[Any, "Variable"] @@ -836,6 +839,9 @@ def isel( if isinstance(indxr, slice) and indxr == slice(None): return self + if is_duck_array(indxr) and not isinstance(indxr, np.ndarray): + indxr = to_numpy(cast("duckarray[Any, Any]", indxr)) + return self._replace(self.index[indxr]) # type: ignore[index,unused-ignore] @override diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 2de92f9bcdb..113d0850678 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -2090,16 +2090,31 @@ def _convert_scalar(self, item) -> np.ndarray: def _index_get( self, indexer: ExplicitIndexer, func_name: str ) -> PandasIndexingAdapter | np.ndarray: - key = indexer.tuple + key: Any = indexer.tuple if len(key) == 1: # unpack key so it can index a pandas.Index object (pandas.Index # objects don't like tuples) (key,) = key + cupy_array_type = array_type("cupy") + if isinstance(key, cupy_array_type): + key = to_numpy(key) + elif isinstance(key, tuple): + key = tuple( + to_numpy(item) if isinstance(item, cupy_array_type) else item + for item in key + ) + # if multidimensional key, convert the index to numpy array and index the latter if getattr(key, "ndim", 0) > 1: indexable = NumpyIndexingAdapter(np.asarray(self)) + indexer = type(indexer)( + tuple( + to_numpy(item) if isinstance(item, cupy_array_type) else item + for item in indexer.tuple + ) + ) return getattr(indexable, func_name)(indexer) # otherwise index the pandas index then re-wrap or convert the result diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 02ee047f1e4..d7c0b23747e 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -1129,6 +1129,7 @@ def __setitem__(self, key, value): value = Variable(dims[-value.ndim :], value) # broadcast to become assignable value = value.set_dims(dims).data + value = to_like_array(value, self._data) if new_order: value = duck_array_ops.asarray(value) diff --git a/xarray/tests/test_cupy.py b/xarray/tests/test_cupy.py index cb8945e8ae5..c75315e993c 100644 --- a/xarray/tests/test_cupy.py +++ b/xarray/tests/test_cupy.py @@ -90,3 +90,86 @@ def test_isel_0d_cupy_indexer( assert actual.dims == expected_dims assert isinstance(actual.data, cp.ndarray) np.testing.assert_array_equal(actual.data.get(), expected) + + +def test_isel_cupy_vector_indexer_with_pandas_index() -> None: + da = xr.DataArray( + cp.arange(24).reshape(4, 6), + dims=("x", "y"), + coords={"x": [10, 20, 30, 40]}, + ) + + actual = da.isel(x=cp.array([1, 2])) + expected_data = cp.arange(24).reshape(4, 6)[cp.array([1, 2])] + + assert actual.get_index("x").equals(pd.Index([20, 30])) + np.testing.assert_array_equal(actual.data.get(), expected_data.get()) + + +def test_isel_cupy_multidimensional_vector_indexer_with_pandas_index() -> None: + da = xr.DataArray( + cp.arange(4), + dims="x", + coords={"x": [10, 20, 30, 40]}, + ) + indexer = xr.DataArray(cp.array([[0, 2], [3, 1]]), dims=("row", "col")) + + actual = da.isel(x=indexer) + + assert actual.dims == ("row", "col") + np.testing.assert_array_equal(actual.data.get(), [[0, 2], [3, 1]]) + np.testing.assert_array_equal(actual.coords["x"].data, [[10, 30], [40, 20]]) + + +def test_isel_cupy_argmax_indexer_with_pandas_index() -> None: + da = xr.DataArray( + cp.arange(24).reshape(4, 6), + dims=("x", "y"), + coords={"x": [10, 20, 30, 40]}, + ) + indexer = da.argmax("x") + + actual = da.isel(x=indexer) + expected_indexer = xr.DataArray(indexer.data.get(), dims=indexer.dims) + expected = da.isel(x=expected_indexer) + + assert actual.dims == expected.dims + np.testing.assert_array_equal(actual.data.get(), expected.data.get()) + np.testing.assert_array_equal(actual.coords["x"].data, expected.coords["x"].data) + + +@pytest.mark.parametrize("indexer", [2, slice(1, 3)]) +def test_cupy_basic_indexer_assignment(indexer: int | slice) -> None: + data = cp.arange(12).reshape(4, 3) + da = xr.DataArray(data.copy(), dims=("x", "y")) + expected = data.copy() + expected[indexer] = -1 + + da[{"x": indexer}] = -1 + + np.testing.assert_array_equal(da.data.get(), expected.get()) + + +def test_cupy_loc_assignment() -> None: + data = cp.arange(12).reshape(4, 3) + da = xr.DataArray(data.copy(), dims=("x", "y"), coords={"x": [0, 1, 2, 3]}) + expected = data.copy() + expected[2] = -1 + + da.loc[{"x": 2}] = -1 + + np.testing.assert_array_equal(da.data.get(), expected.get()) + + +@pytest.mark.parametrize("value", [True, False]) +def test_bool_zero_dim_cupy_dataarray(value: bool) -> None: + assert bool(xr.DataArray(cp.array(value))) is value + + +@pytest.mark.parametrize(("method", "expected"), [("idxmax", 1), ("idxmin", 0)]) +def test_cupy_idxminmax_skipna(method: str, expected: int) -> None: + da = xr.DataArray(cp.array([1, 3, 2]), dims="x", coords={"x": [0, 1, 2]}) + + result = getattr(da, method)(skipna=True) + + assert result.item() == expected diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 407cd08c6cd..c384a231fa4 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -86,6 +86,16 @@ def setup(self): np.uint64 ) + def test_bool_pandas_extension_array(self) -> None: + false_value = DataArray(pd.array([False], dtype="boolean"), dims="x") + true_value = DataArray(pd.array([True], dtype="boolean"), dims="x") + missing_value = DataArray(pd.array([pd.NA], dtype="boolean"), dims="x") + + assert not bool(false_value) + assert bool(true_value) + with pytest.raises(TypeError, match="boolean value of NA is ambiguous"): + bool(missing_value) + def test_repr(self) -> None: v = Variable(["time", "x"], [[1, 2, 3], [4, 5, 6]], {"foo": "bar"}) v = v.astype(np.uint64)