From 0c1d04c7b1ec5cc8dcb0b03c2aed722f0587d847 Mon Sep 17 00:00:00 2001 From: Phobos Date: Tue, 9 Jun 2020 11:13:02 -0400 Subject: [PATCH 1/3] Correct dask handling for 1D idxmax/min on ND data --- xarray/core/computation.py | 2 +- xarray/tests/test_dataarray.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/xarray/core/computation.py b/xarray/core/computation.py index cecd4fd8e70..4f4fd475c82 100644 --- a/xarray/core/computation.py +++ b/xarray/core/computation.py @@ -1563,7 +1563,7 @@ def _calc_idxminmax( chunks = dict(zip(array.dims, array.chunks)) dask_coord = dask.array.from_array(array[dim].data, chunks=chunks[dim]) - res = indx.copy(data=dask_coord[(indx.data,)]) + res = indx.copy(data=dask_coord[indx.data.ravel()].reshape(indx.shape)) # we need to attach back the dim name res.name = dim else: diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 95f0ad9f612..c6999807690 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -5257,6 +5257,25 @@ def test_idxmax(self, x, minindex, maxindex, nanindex, use_dask): assert_identical(result7, expected7) +class TestReduceND(TestReduce): + @pytest.mark.parametrize("op", ['idxmin', 'idxmax']) + @pytest.mark.parametrize("ndim", [3, 5]) + def test_idxminmax_dask(self, op, ndim): + if not has_dask: + pytest.skip("requires dask") + + ar0_raw = xr.DataArray( + np.random.random_sample(size=[10] * ndim), + dims=[i for i in 'abcdefghij'[:ndim - 1]] + ["x"], + coords={"x": np.arange(10)}, + attrs=self.attrs + ) + + ar0_dsk = ar0_raw.chunk({}) + # Assert idx is the same with dask and without + assert_equal(getattr(ar0_dsk, op)(dim='x'), getattr(ar0_raw, op)(dim='x')) + + @pytest.fixture(params=[1]) def da(request): if request.param == 1: From 8ee652509d3f52f30317c8777312615b7e83e213 Mon Sep 17 00:00:00 2001 From: Phobos Date: Tue, 9 Jun 2020 11:35:31 -0400 Subject: [PATCH 2/3] Passing black and others --- doc/whats-new.rst | 4 ++-- xarray/tests/test_dataarray.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 85e73e1b7e8..ac5d6348150 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -70,8 +70,8 @@ New Features :py:meth:`Dataset.idxmax`, :py:meth:`Dataset.idxmin`. (:issue:`60`, :pull:`3871`) By `Todd Jennings `_ - Support dask handling for :py:meth:`DataArray.idxmax`, :py:meth:`DataArray.idxmin`, - :py:meth:`Dataset.idxmax`, :py:meth:`Dataset.idxmin`. (:pull:`3922`) - By `Kai Mühlbauer `_. + :py:meth:`Dataset.idxmax`, :py:meth:`Dataset.idxmin`. (:pull:`3922`, :pull:`xxxx`) + By `Kai Mühlbauer `_ and `Pascal Bourgault `_. - More support for unit aware arrays with pint (:pull:`3643`) By `Justus Magin `_. - Support overriding existing variables in ``to_zarr()`` with ``mode='a'`` even diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index c6999807690..86554e84045 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -5258,7 +5258,7 @@ def test_idxmax(self, x, minindex, maxindex, nanindex, use_dask): class TestReduceND(TestReduce): - @pytest.mark.parametrize("op", ['idxmin', 'idxmax']) + @pytest.mark.parametrize("op", ["idxmin", "idxmax"]) @pytest.mark.parametrize("ndim", [3, 5]) def test_idxminmax_dask(self, op, ndim): if not has_dask: @@ -5266,14 +5266,14 @@ def test_idxminmax_dask(self, op, ndim): ar0_raw = xr.DataArray( np.random.random_sample(size=[10] * ndim), - dims=[i for i in 'abcdefghij'[:ndim - 1]] + ["x"], + dims=[i for i in "abcdefghij"[: ndim - 1]] + ["x"], coords={"x": np.arange(10)}, - attrs=self.attrs + attrs=self.attrs, ) ar0_dsk = ar0_raw.chunk({}) # Assert idx is the same with dask and without - assert_equal(getattr(ar0_dsk, op)(dim='x'), getattr(ar0_raw, op)(dim='x')) + assert_equal(getattr(ar0_dsk, op)(dim="x"), getattr(ar0_raw, op)(dim="x")) @pytest.fixture(params=[1]) From 9c00bb934f3914d0a913d6e90aa96cdef6a86ebb Mon Sep 17 00:00:00 2001 From: Phobos Date: Tue, 9 Jun 2020 11:36:38 -0400 Subject: [PATCH 3/3] Edit Whats New --- doc/whats-new.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index ac5d6348150..736e440e409 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -70,7 +70,7 @@ New Features :py:meth:`Dataset.idxmax`, :py:meth:`Dataset.idxmin`. (:issue:`60`, :pull:`3871`) By `Todd Jennings `_ - Support dask handling for :py:meth:`DataArray.idxmax`, :py:meth:`DataArray.idxmin`, - :py:meth:`Dataset.idxmax`, :py:meth:`Dataset.idxmin`. (:pull:`3922`, :pull:`xxxx`) + :py:meth:`Dataset.idxmax`, :py:meth:`Dataset.idxmin`. (:pull:`3922`, :pull:`4135`) By `Kai Mühlbauer `_ and `Pascal Bourgault `_. - More support for unit aware arrays with pint (:pull:`3643`) By `Justus Magin `_.