From 49b67c38d5e776697eb9e28ff3ac1bb5de73213a Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sat, 7 Oct 2023 18:05:30 -0700 Subject: [PATCH 1/4] Enable `.rolling_exp` to work on dask arrays Another benefit of the move to `.apply_ufunc`... --- xarray/core/rolling_exp.py | 4 ++-- xarray/tests/test_rolling.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/xarray/core/rolling_exp.py b/xarray/core/rolling_exp.py index c56bf6a384e..cb77358869c 100644 --- a/xarray/core/rolling_exp.py +++ b/xarray/core/rolling_exp.py @@ -147,9 +147,9 @@ def mean(self, keep_attrs: bool | None = None) -> T_DataWithCoords: input_core_dims=[[self.dim]], kwargs=dict(alpha=self.alpha, axis=-1), output_core_dims=[[self.dim]], - exclude_dims={self.dim}, keep_attrs=keep_attrs, on_missing_core_dim="copy", + dask="parallelized", ).transpose(*dim_order) def sum(self, keep_attrs: bool | None = None) -> T_DataWithCoords: @@ -183,7 +183,7 @@ def sum(self, keep_attrs: bool | None = None) -> T_DataWithCoords: input_core_dims=[[self.dim]], kwargs=dict(alpha=self.alpha, axis=-1), output_core_dims=[[self.dim]], - exclude_dims={self.dim}, keep_attrs=keep_attrs, on_missing_core_dim="copy", + dask="parallelized", ).transpose(*dim_order) diff --git a/xarray/tests/test_rolling.py b/xarray/tests/test_rolling.py index 2dc8ae24438..70302f69af2 100644 --- a/xarray/tests/test_rolling.py +++ b/xarray/tests/test_rolling.py @@ -788,7 +788,7 @@ def test_raise_no_warning_dask_rolling_assert_close(self, ds, name) -> None: @requires_numbagg class TestDatasetRollingExp: - @pytest.mark.parametrize("backend", ["numpy"], indirect=True) + @pytest.mark.parametrize("backend", ["numpy", "dask"], indirect=True) def test_rolling_exp(self, ds) -> None: result = ds.rolling_exp(time=10, window_type="span").mean() assert isinstance(result, Dataset) From 756baf5131018be8a1869de1a58eb84a28c41f23 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sat, 7 Oct 2023 18:07:35 -0700 Subject: [PATCH 2/4] . --- doc/whats-new.rst | 3 +++ 1 file changed, 3 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 92acc3f90c0..8f576f486dc 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -29,6 +29,9 @@ New Features - :py:meth:`DataArray.sortby` & :py:meth:`Dataset.sortby` accept a callable for the ``variables`` parameter, passing the object as the only argument. By `Maximilian Roos `_. +- ``.rolling_exp`` functions can now operate on dask-backed arrays, assuming the + core dim has exactly one chunk. (:pull:`8284`). + By `Maximilian Roos `_. Breaking changes ~~~~~~~~~~~~~~~~ From 0f5265d80a920287bb8029795b37a703ebeebcab Mon Sep 17 00:00:00 2001 From: Maximilian Roos <5635139+max-sixty@users.noreply.github.com> Date: Mon, 9 Oct 2023 08:48:18 -0700 Subject: [PATCH 3/4] Update test_rolling.py Co-authored-by: Deepak Cherian --- xarray/tests/test_rolling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/tests/test_rolling.py b/xarray/tests/test_rolling.py index 70302f69af2..6fc8e2eacbe 100644 --- a/xarray/tests/test_rolling.py +++ b/xarray/tests/test_rolling.py @@ -788,7 +788,7 @@ def test_raise_no_warning_dask_rolling_assert_close(self, ds, name) -> None: @requires_numbagg class TestDatasetRollingExp: - @pytest.mark.parametrize("backend", ["numpy", "dask"], indirect=True) + @pytest.mark.parametrize("backend", ["numpy", pytest.param("dask", marks=requires_dask )], indirect=True) def test_rolling_exp(self, ds) -> None: result = ds.rolling_exp(time=10, window_type="span").mean() assert isinstance(result, Dataset) From be94faaa805105a9d83929be19cb37b828166c49 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 9 Oct 2023 15:49:28 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/tests/test_rolling.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/xarray/tests/test_rolling.py b/xarray/tests/test_rolling.py index 6fc8e2eacbe..da834b76124 100644 --- a/xarray/tests/test_rolling.py +++ b/xarray/tests/test_rolling.py @@ -788,7 +788,9 @@ def test_raise_no_warning_dask_rolling_assert_close(self, ds, name) -> None: @requires_numbagg class TestDatasetRollingExp: - @pytest.mark.parametrize("backend", ["numpy", pytest.param("dask", marks=requires_dask )], indirect=True) + @pytest.mark.parametrize( + "backend", ["numpy", pytest.param("dask", marks=requires_dask)], indirect=True + ) def test_rolling_exp(self, ds) -> None: result = ds.rolling_exp(time=10, window_type="span").mean() assert isinstance(result, Dataset)