From d40e7c4395738f3c103ffcb4e8ebbb5550f2b21b Mon Sep 17 00:00:00 2001 From: Keewis Date: Sat, 26 Mar 2022 11:59:23 +0100 Subject: [PATCH 1/7] use the DaskIndexingAdapter for all duck-dask-arrays --- xarray/core/indexing.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index c797e6652de..c8851788c29 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -25,13 +25,7 @@ from . import duck_array_ops, nputils, utils from .npcompat import DTypeLike -from .pycompat import ( - dask_array_type, - dask_version, - integer_types, - is_duck_dask_array, - sparse_array_type, -) +from .pycompat import dask_version, integer_types, is_duck_dask_array, sparse_array_type from .types import T_Xarray from .utils import either_dict_or_kwargs, get_valid_numpy_dtype @@ -682,7 +676,7 @@ def as_indexable(array): return NumpyIndexingAdapter(array) if isinstance(array, pd.Index): return PandasIndexingAdapter(array) - if isinstance(array, dask_array_type): + if is_duck_dask_array(array): return DaskIndexingAdapter(array) if hasattr(array, "__array_function__"): return NdArrayLikeIndexingAdapter(array) From c8bb1e68dc286fa4ae51a8a97f4c4e256afcd33b Mon Sep 17 00:00:00 2001 From: Keewis Date: Sat, 26 Mar 2022 12:27:09 +0100 Subject: [PATCH 2/7] add tests covering the change --- xarray/tests/test_units.py | 35 +++++++++++++++++++++++++++-------- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py index bc3cc367c0e..ed36617dfe1 100644 --- a/xarray/tests/test_units.py +++ b/xarray/tests/test_units.py @@ -1838,20 +1838,39 @@ def test_broadcast_equals(self, unit, dtype): assert expected == actual @pytest.mark.parametrize( - "indices", + ["variable", "indexers"], ( - pytest.param(4, id="single index"), - pytest.param([5, 2, 9, 1], id="multiple indices"), + pytest.param( + xr.Variable("x", np.linspace(0, 5, 10)), + {"x": 4}, + id="single value–single indexer", + ), + pytest.param( + xr.Variable("x", np.linspace(0, 5, 10)), + {"x": [5, 2, 9, 1]}, + id="multiple values–single indexer", + ), + pytest.param( + xr.Variable(("x", "y"), np.linspace(0, 5, 20).reshape(4, 5)), + {"x": 1, "y": 4}, + id="single value—multiple indexers", + ), + pytest.param( + xr.Variable(("x", "y"), np.linspace(0, 5, 20).reshape(4, 5)), + {"x": [0, 1, 2], "y": [0, 2, 4]}, + id="multiple values—multiple indexers", + ), ), ) - def test_isel(self, indices, dtype): - array = np.linspace(0, 5, 10).astype(dtype) * unit_registry.s - variable = xr.Variable("x", array) + def test_isel(self, variable, indexers, dtype): + quantified = xr.Variable( + variable.dims, variable.data.astype(dtype) * unit_registry.s + ) expected = attach_units( - strip_units(variable).isel(x=indices), extract_units(variable) + strip_units(quantified).isel(indexers), extract_units(quantified) ) - actual = variable.isel(x=indices) + actual = quantified.isel(indexers) assert_units_equal(expected, actual) assert_identical(expected, actual) From 27c60da68d4fe44f014c857e573c7e1b14eb2bb5 Mon Sep 17 00:00:00 2001 From: Keewis Date: Sat, 26 Mar 2022 12:38:55 +0100 Subject: [PATCH 3/7] use force_ndarray_like to allow duck arrays in quantities --- xarray/tests/test_units.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py index ed36617dfe1..9aa282384a2 100644 --- a/xarray/tests/test_units.py +++ b/xarray/tests/test_units.py @@ -31,7 +31,7 @@ # make sure scalars are converted to 0d arrays so quantities can # always be treated like ndarrays -unit_registry = pint.UnitRegistry(force_ndarray=True) +unit_registry = pint.UnitRegistry(force_ndarray_like=True) Quantity = unit_registry.Quantity From 1a05d83e16741d6a44adb5375df72cca75caa456 Mon Sep 17 00:00:00 2001 From: Keewis Date: Sat, 26 Mar 2022 12:39:34 +0100 Subject: [PATCH 4/7] fix the test ids --- xarray/tests/test_units.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py index 9aa282384a2..1d9a3f2ebf7 100644 --- a/xarray/tests/test_units.py +++ b/xarray/tests/test_units.py @@ -1843,22 +1843,22 @@ def test_broadcast_equals(self, unit, dtype): pytest.param( xr.Variable("x", np.linspace(0, 5, 10)), {"x": 4}, - id="single value–single indexer", + id="single value-single indexer", ), pytest.param( xr.Variable("x", np.linspace(0, 5, 10)), {"x": [5, 2, 9, 1]}, - id="multiple values–single indexer", + id="multiple values-single indexer", ), pytest.param( xr.Variable(("x", "y"), np.linspace(0, 5, 20).reshape(4, 5)), {"x": 1, "y": 4}, - id="single value—multiple indexers", + id="single value-multiple indexers", ), pytest.param( xr.Variable(("x", "y"), np.linspace(0, 5, 20).reshape(4, 5)), {"x": [0, 1, 2], "y": [0, 2, 4]}, - id="multiple values—multiple indexers", + id="multiple values-multiple indexers", ), ), ) From dc4ea763c066aa10cb5394942121e43a18949416 Mon Sep 17 00:00:00 2001 From: Keewis Date: Sat, 26 Mar 2022 12:40:06 +0100 Subject: [PATCH 5/7] also test all combinations with dask --- xarray/tests/test_units.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py index 1d9a3f2ebf7..c18b7d18c04 100644 --- a/xarray/tests/test_units.py +++ b/xarray/tests/test_units.py @@ -1837,6 +1837,7 @@ def test_broadcast_equals(self, unit, dtype): assert expected == actual + @pytest.mark.parametrize("dask", [False, pytest.param(True, marks=[requires_dask])]) @pytest.mark.parametrize( ["variable", "indexers"], ( @@ -1862,7 +1863,9 @@ def test_broadcast_equals(self, unit, dtype): ), ), ) - def test_isel(self, variable, indexers, dtype): + def test_isel(self, variable, indexers, dask, dtype): + if dask: + variable = variable.chunk({dim: 2 for dim in variable.dims}) quantified = xr.Variable( variable.dims, variable.data.astype(dtype) * unit_registry.s ) From c580c7439e28be760c051f6d13f0afee8709ae55 Mon Sep 17 00:00:00 2001 From: Keewis Date: Sat, 26 Mar 2022 13:04:12 +0100 Subject: [PATCH 6/7] fix whats-new.rst --- doc/whats-new.rst | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 405a96806be..999e678d05f 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -24,7 +24,8 @@ New Features - Add a ``create_index=True`` parameter to :py:meth:`Dataset.stack` and :py:meth:`DataArray.stack` so that the creation of multi-indexes is optional - (:pull:`5692`). By `Benoît Bovy `_. + (:pull:`5692`). + By `Benoît Bovy `_. - Multi-index levels are now accessible through their own, regular coordinates instead of virtual coordinates (:pull:`5692`). By `Benoît Bovy `_. @@ -33,7 +34,8 @@ Breaking changes ~~~~~~~~~~~~~~~~ - The Dataset and DataArray ``rename*`` methods do not implicitly add or drop - indexes. (:pull:`5692`). By `Benoît Bovy `_. + indexes. (:pull:`5692`). + By `Benoît Bovy `_. - Many arguments like ``keep_attrs``, ``axis``, and ``skipna`` are now keyword only for all reduction operations like ``.mean``. By `Deepak Cherian `_, `Jimmy Westling `_. @@ -47,12 +49,14 @@ Bug fixes - Set ``skipna=None`` for all ``quantile`` methods (e.g. :py:meth:`Dataset.quantile`) and ensure it skips missing values for float dtypes (consistent with other methods). This should - not change the behavior (:pull:`6303`). By `Mathias Hauser `_. + not change the behavior (:pull:`6303`). + By `Mathias Hauser `_. - Many bugs fixed by the explicit indexes refactor, mainly related to multi-index (virtual) coordinates. See the corresponding pull-request on GitHub for more details. (:pull:`5692`). By `Benoît Bovy `_. - Fixed "unhashable type" error trying to read NetCDF file with variable having its 'units' - attribute not ``str`` (e.g. ``numpy.ndarray``) (:issue:`6368`). By `Oleh Khoma `_. + attribute not ``str`` (e.g. ``numpy.ndarray``) (:issue:`6368`). + By `Oleh Khoma `_. Documentation ~~~~~~~~~~~~~ @@ -65,7 +69,6 @@ Internal Changes corresponding pull-request on GitHub for more details. (:pull:`5692`). By `Benoît Bovy `_. -.. _whats-new.2022.02.0: .. _whats-new.2022.03.0: v2022.03.0 (2 March 2022) From 7f21f676937e17b6b389fff93772ff5fad75cceb Mon Sep 17 00:00:00 2001 From: Keewis Date: Sat, 26 Mar 2022 13:06:16 +0100 Subject: [PATCH 7/7] whats-new.rst entry --- doc/whats-new.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 999e678d05f..0704f1bf291 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -57,6 +57,8 @@ Bug fixes - Fixed "unhashable type" error trying to read NetCDF file with variable having its 'units' attribute not ``str`` (e.g. ``numpy.ndarray``) (:issue:`6368`). By `Oleh Khoma `_. +- Allow fancy indexing of duck dask arrays along multiple dimensions. (:pull:`6414`) + By `Justus Magin `_. Documentation ~~~~~~~~~~~~~