From 5b30f5ca03b46d7074b7ffd18bdf81abcf0c2c89 Mon Sep 17 00:00:00 2001 From: Spencer Clark Date: Sun, 21 Feb 2021 13:09:57 -0500 Subject: [PATCH 01/11] Add DataArrayCoarsen.reduce and DatasetCoarsen.reduce methods --- xarray/core/rolling.py | 62 ++++++++++++++++++++++++++++++++-- xarray/tests/test_dataarray.py | 18 ++++++++++ xarray/tests/test_dataset.py | 21 ++++++++++++ 3 files changed, 99 insertions(+), 2 deletions(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index 6087fd4c806..52f28ba48a3 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -836,7 +836,9 @@ class DataArrayCoarsen(Coarsen): _reduce_extra_args_docstring = """""" @classmethod - def _reduce_method(cls, func: Callable, include_skipna: bool, numeric_only: bool): + def _reduce_method( + cls, func: Callable, include_skipna: bool = False, numeric_only: bool = False + ): """ Return a wrapped function for injecting reduction methods. see ops.inject_reduce_methods @@ -871,6 +873,38 @@ def wrapped_func(self, **kwargs): return wrapped_func + def reduce(self, func: Callable, **kwargs): + """Reduce the items in this group by applying `func` along some + dimension(s). + + Parameters + ---------- + func : callable + Function which can be called in the form + `func(x, axis, **kwargs)` to return the result of collapsing an + np.ndarray over the coarsening dimensions. It must be possible to + provide the `axis` argument with a tuple of integers. + **kwargs : dict + Additional keyword arguments passed on to `func`. + + Returns + ------- + reduced : DataArray + Array with summarized data. + + Examples + -------- + >>> da = xr.DataArray(np.arange(8).reshape(2, 4), dims=("a", "b")) + >>> coarsen = da.coarsen(b=2) + >>> coarsen.reduce(np.sum) + + array([[ 1, 5], + [ 9, 13]]) + Dimensions without coordinates: a, b + """ + wrapped_func = self._reduce_method(func) + return wrapped_func(self, **kwargs) + class DatasetCoarsen(Coarsen): __slots__ = () @@ -878,7 +912,9 @@ class DatasetCoarsen(Coarsen): _reduce_extra_args_docstring = """""" @classmethod - def _reduce_method(cls, func: Callable, include_skipna: bool, numeric_only: bool): + def _reduce_method( + cls, func: Callable, include_skipna: bool = False, numeric_only: bool = False + ): """ Return a wrapped function for injecting reduction methods. see ops.inject_reduce_methods @@ -917,6 +953,28 @@ def wrapped_func(self, **kwargs): return wrapped_func + def reduce(self, func: Callable, **kwargs): + """Reduce the items in this group by applying `func` along some + dimension(s). + + Parameters + ---------- + func : callable + Function which can be called in the form + `func(x, axis, **kwargs)` to return the result of collapsing an + np.ndarray over the coarsening dimensions. It must be possible to + provide the `axis` argument with a tuple of integers. + **kwargs : dict + Additional keyword arguments passed on to `func`. + + Returns + ------- + reduced : Dataset + Arrays with summarized data. + """ + wrapped_func = self._reduce_method(func) + return wrapped_func(self, **kwargs) + inject_reduce_methods(DataArrayCoarsen) inject_reduce_methods(DatasetCoarsen) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 68e31cd123a..6e025194b7b 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -6382,6 +6382,24 @@ def test_coarsen_keep_attrs(): xr.testing.assert_identical(da, da2) +@pytest.mark.parametrize("da", (1, 2), indirect=True) +@pytest.mark.parametrize("window", (1, 2, 3, 4)) +@pytest.mark.parametrize("name", ("sum", "mean", "std", "max")) +def test_coarsen_reduce(da, window, name): + if da.isnull().sum() > 1 and window == 1: + # this causes all nan slices + window = 2 + + # Use boundary="trim" to accomodate all window sizes used in tests + coarsen_obj = da.coarsen(time=window, boundary="trim") + + # add nan prefix to numpy methods to get similar # behavior as bottleneck + actual = coarsen_obj.reduce(getattr(np, "nan%s" % name)) + expected = getattr(coarsen_obj, name)() + assert_allclose(actual, expected) + assert actual.dims == expected.dims + + @pytest.mark.parametrize("da", (1, 2), indirect=True) def test_rolling_iter(da): rolling_obj = da.rolling(time=7) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index c4161093837..af554e02e99 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -6055,6 +6055,27 @@ def test_coarsen_keep_attrs(): xr.testing.assert_identical(ds, ds2) +@pytest.mark.slow +@pytest.mark.parametrize("ds", (1, 2), indirect=True) +@pytest.mark.parametrize("window", (1, 2, 3, 4)) +@pytest.mark.parametrize("name", ("sum", "mean", "std", "var", "min", "max", "median")) +def test_coarsen_reduce(ds, window, name): + # Use boundary="trim" to accomodate all window sizes used in tests + coarsen_obj = ds.coarsen(time=window, boundary="trim") + + # add nan prefix to numpy methods to get similar behavior as bottleneck + actual = coarsen_obj.reduce(getattr(np, "nan%s" % name)) + expected = getattr(coarsen_obj, name)() + assert_allclose(actual, expected) + assert actual.dims == expected.dims + # make sure the order of data_var are not changed. + assert list(ds.data_vars.keys()) == list(actual.data_vars.keys()) + + # Make sure the dimension order is restored + for key, src_var in ds.data_vars.items(): + assert src_var.dims == actual[key].dims + + @pytest.mark.parametrize( "funcname, argument", [ From a734b4204fece25b022dce850418e147ba91beb8 Mon Sep 17 00:00:00 2001 From: Spencer Clark Date: Tue, 23 Feb 2021 08:39:33 -0500 Subject: [PATCH 02/11] Address review comments --- doc/whats-new.rst | 4 ++++ xarray/core/rolling.py | 28 +++++++++++++--------------- xarray/tests/test_dataarray.py | 6 ++---- xarray/tests/test_dataset.py | 4 ++-- 4 files changed, 21 insertions(+), 21 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 05066f883bf..8a1b28b3261 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -91,6 +91,10 @@ New Features (including globs for the latter) for ``engine="zarr"``, and so allow reading from many remote and other file systems (:pull:`4461`) By `Martin Durant `_ +- :py:class:`DataArrayCoarsen` and :py:class:`DatasetCoarsen` now implement a + ``reduce`` method, enabling coarsening operations with custom reduction + functions (:issue:`3741`, :pull:`4939`). By `Spencer Clark + `. Bug fixes ~~~~~~~~~ diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index 52f28ba48a3..ccc8dd2fe6f 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -874,18 +874,18 @@ def wrapped_func(self, **kwargs): return wrapped_func def reduce(self, func: Callable, **kwargs): - """Reduce the items in this group by applying `func` along some + """Reduce the items in this group by applying ``func`` along some dimension(s). Parameters ---------- func : callable - Function which can be called in the form - `func(x, axis, **kwargs)` to return the result of collapsing an - np.ndarray over the coarsening dimensions. It must be possible to - provide the `axis` argument with a tuple of integers. + Function which can be called in the form ``func(x, axis, **kwargs)`` + to return the result of collapsing an np.ndarray over the coarsening + dimensions. It must be possible to provide the ``axis`` argument + with a tuple of integers. **kwargs : dict - Additional keyword arguments passed on to `func`. + Additional keyword arguments passed on to ``func``. Returns ------- @@ -954,23 +954,21 @@ def wrapped_func(self, **kwargs): return wrapped_func def reduce(self, func: Callable, **kwargs): - """Reduce the items in this group by applying `func` along some + """Reduce the items in this group by applying ``func`` along some dimension(s). Parameters ---------- - func : callable - Function which can be called in the form - `func(x, axis, **kwargs)` to return the result of collapsing an - np.ndarray over the coarsening dimensions. It must be possible to - provide the `axis` argument with a tuple of integers. + func : callable Function which can be called in the form ``func(x, axis, + **kwargs)`` to return the result of collapsing an np.ndarray over + the coarsening dimensions. It must be possible to provide the + ``axis`` argument with a tuple of integers. **kwargs : dict - Additional keyword arguments passed on to `func`. + Additional keyword arguments passed on to ``func``. Returns ------- - reduced : Dataset - Arrays with summarized data. + reduced : Dataset Arrays with summarized data. """ wrapped_func = self._reduce_method(func) return wrapped_func(self, **kwargs) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 6e025194b7b..d1d36dd93b3 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -6387,17 +6387,15 @@ def test_coarsen_keep_attrs(): @pytest.mark.parametrize("name", ("sum", "mean", "std", "max")) def test_coarsen_reduce(da, window, name): if da.isnull().sum() > 1 and window == 1: - # this causes all nan slices - window = 2 + pytest.skip("These parameters lead to all-NaN slices") # Use boundary="trim" to accomodate all window sizes used in tests coarsen_obj = da.coarsen(time=window, boundary="trim") # add nan prefix to numpy methods to get similar # behavior as bottleneck - actual = coarsen_obj.reduce(getattr(np, "nan%s" % name)) + actual = coarsen_obj.reduce(getattr(np, f"nan{name}")) expected = getattr(coarsen_obj, name)() assert_allclose(actual, expected) - assert actual.dims == expected.dims @pytest.mark.parametrize("da", (1, 2), indirect=True) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index af554e02e99..2118bc8b780 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -6064,10 +6064,10 @@ def test_coarsen_reduce(ds, window, name): coarsen_obj = ds.coarsen(time=window, boundary="trim") # add nan prefix to numpy methods to get similar behavior as bottleneck - actual = coarsen_obj.reduce(getattr(np, "nan%s" % name)) + actual = coarsen_obj.reduce(getattr(np, f"nan{name}")) expected = getattr(coarsen_obj, name)() assert_allclose(actual, expected) - assert actual.dims == expected.dims + # make sure the order of data_var are not changed. assert list(ds.data_vars.keys()) == list(actual.data_vars.keys()) From 5bdcb468c236e211bd973d243e9735cfeed5f399 Mon Sep 17 00:00:00 2001 From: Spencer Clark Date: Tue, 23 Feb 2021 08:46:43 -0500 Subject: [PATCH 03/11] Add new methods to api-hidden.rst --- doc/api-hidden.rst | 2 ++ xarray/core/rolling.py | 18 +++++++++--------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/doc/api-hidden.rst b/doc/api-hidden.rst index e5492ec73a4..14d79039a3a 100644 --- a/doc/api-hidden.rst +++ b/doc/api-hidden.rst @@ -47,6 +47,7 @@ core.rolling.DatasetCoarsen.median core.rolling.DatasetCoarsen.min core.rolling.DatasetCoarsen.prod + core.rolling.DatasetCoarsen.reduce core.rolling.DatasetCoarsen.std core.rolling.DatasetCoarsen.sum core.rolling.DatasetCoarsen.var @@ -190,6 +191,7 @@ core.rolling.DataArrayCoarsen.median core.rolling.DataArrayCoarsen.min core.rolling.DataArrayCoarsen.prod + core.rolling.DataArrayCoarsen.reduce core.rolling.DataArrayCoarsen.std core.rolling.DataArrayCoarsen.sum core.rolling.DataArrayCoarsen.var diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index ccc8dd2fe6f..9a09b49162c 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -874,18 +874,18 @@ def wrapped_func(self, **kwargs): return wrapped_func def reduce(self, func: Callable, **kwargs): - """Reduce the items in this group by applying ``func`` along some + """Reduce the items in this group by applying `func` along some dimension(s). Parameters ---------- func : callable - Function which can be called in the form ``func(x, axis, **kwargs)`` + Function which can be called in the form `func(x, axis, **kwargs)` to return the result of collapsing an np.ndarray over the coarsening - dimensions. It must be possible to provide the ``axis`` argument + dimensions. It must be possible to provide the `axis` argument with a tuple of integers. **kwargs : dict - Additional keyword arguments passed on to ``func``. + Additional keyword arguments passed on to `func`. Returns ------- @@ -954,17 +954,17 @@ def wrapped_func(self, **kwargs): return wrapped_func def reduce(self, func: Callable, **kwargs): - """Reduce the items in this group by applying ``func`` along some + """Reduce the items in this group by applying `func` along some dimension(s). Parameters ---------- - func : callable Function which can be called in the form ``func(x, axis, - **kwargs)`` to return the result of collapsing an np.ndarray over + func : callable Function which can be called in the form `func(x, axis, + **kwargs)` to return the result of collapsing an np.ndarray over the coarsening dimensions. It must be possible to provide the - ``axis`` argument with a tuple of integers. + `axis` argument with a tuple of integers. **kwargs : dict - Additional keyword arguments passed on to ``func``. + Additional keyword arguments passed on to `func. Returns ------- From 2e029ca6441f6224e47ca81405d4a9fd56f88ef9 Mon Sep 17 00:00:00 2001 From: Spencer Clark Date: Tue, 23 Feb 2021 09:00:58 -0500 Subject: [PATCH 04/11] Fix docstring --- xarray/core/rolling.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index 9a09b49162c..e8e492f60c7 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -959,10 +959,11 @@ def reduce(self, func: Callable, **kwargs): Parameters ---------- - func : callable Function which can be called in the form `func(x, axis, - **kwargs)` to return the result of collapsing an np.ndarray over - the coarsening dimensions. It must be possible to provide the - `axis` argument with a tuple of integers. + func : callable + Function which can be called in the form `func(x, axis, **kwargs)` + to return the result of collapsing an np.ndarray over the coarsening + dimensions. It must be possible to provide the `axis` argument with + a tuple of integers. **kwargs : dict Additional keyword arguments passed on to `func. From bd8050b2e92e7e25131d25564ea5c62fc53b5114 Mon Sep 17 00:00:00 2001 From: Spencer Clark Date: Tue, 23 Feb 2021 09:01:40 -0500 Subject: [PATCH 05/11] Fix docstring --- xarray/core/rolling.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index e8e492f60c7..e7427de9f66 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -960,10 +960,10 @@ def reduce(self, func: Callable, **kwargs): Parameters ---------- func : callable - Function which can be called in the form `func(x, axis, **kwargs)` - to return the result of collapsing an np.ndarray over the coarsening - dimensions. It must be possible to provide the `axis` argument with - a tuple of integers. + Function which can be called in the form `func(x, axis, **kwargs)` + to return the result of collapsing an np.ndarray over the coarsening + dimensions. It must be possible to provide the `axis` argument with + a tuple of integers. **kwargs : dict Additional keyword arguments passed on to `func. From 798f67330f71330b5a84dda314b726e4967fd33e Mon Sep 17 00:00:00 2001 From: Spencer Clark Date: Tue, 23 Feb 2021 09:02:28 -0500 Subject: [PATCH 06/11] lint --- xarray/core/rolling.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index e7427de9f66..16a7ff297fe 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -959,10 +959,10 @@ def reduce(self, func: Callable, **kwargs): Parameters ---------- - func : callable - Function which can be called in the form `func(x, axis, **kwargs)` - to return the result of collapsing an np.ndarray over the coarsening - dimensions. It must be possible to provide the `axis` argument with + func : callable + Function which can be called in the form `func(x, axis, **kwargs)` + to return the result of collapsing an np.ndarray over the coarsening + dimensions. It must be possible to provide the `axis` argument with a tuple of integers. **kwargs : dict Additional keyword arguments passed on to `func. From e9414d26516c580dda600f62254feb82d714a9d8 Mon Sep 17 00:00:00 2001 From: Spencer Clark Date: Tue, 23 Feb 2021 09:09:22 -0500 Subject: [PATCH 07/11] More docstring formatting --- xarray/core/rolling.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index 16a7ff297fe..e4e8b736f0a 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -969,7 +969,8 @@ def reduce(self, func: Callable, **kwargs): Returns ------- - reduced : Dataset Arrays with summarized data. + reduced : Dataset + Arrays with summarized data. """ wrapped_func = self._reduce_method(func) return wrapped_func(self, **kwargs) From b977701d7d9f182bee0abb0b2c31aa416920640c Mon Sep 17 00:00:00 2001 From: Spencer Clark Date: Tue, 23 Feb 2021 09:36:48 -0500 Subject: [PATCH 08/11] Update xarray/core/rolling.py Co-authored-by: Mathias Hauser --- xarray/core/rolling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index e4e8b736f0a..dbb2e794140 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -965,7 +965,7 @@ def reduce(self, func: Callable, **kwargs): dimensions. It must be possible to provide the `axis` argument with a tuple of integers. **kwargs : dict - Additional keyword arguments passed on to `func. + Additional keyword arguments passed on to `func`. Returns ------- From 70155247df57a5af07f9b613ee3b7d0cba221522 Mon Sep 17 00:00:00 2001 From: Spencer Clark Date: Tue, 23 Feb 2021 10:06:17 -0500 Subject: [PATCH 09/11] Fix what's new link --- 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 8a1b28b3261..ea9c5f89bb5 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -94,7 +94,7 @@ New Features - :py:class:`DataArrayCoarsen` and :py:class:`DatasetCoarsen` now implement a ``reduce`` method, enabling coarsening operations with custom reduction functions (:issue:`3741`, :pull:`4939`). By `Spencer Clark - `. + `_. Bug fixes ~~~~~~~~~ From 030674699524feaae9230a2f0886555683cc8917 Mon Sep 17 00:00:00 2001 From: Spencer Clark Date: Tue, 23 Feb 2021 10:07:58 -0500 Subject: [PATCH 10/11] Update xarray/core/rolling.py Co-authored-by: Mathias Hauser --- xarray/core/rolling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index dbb2e794140..9ec75f0b821 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -960,7 +960,7 @@ def reduce(self, func: Callable, **kwargs): Parameters ---------- func : callable - Function which can be called in the form `func(x, axis, **kwargs)` + Function which can be called in the form `func(x, axis, **kwargs)` to return the result of collapsing an np.ndarray over the coarsening dimensions. It must be possible to provide the `axis` argument with a tuple of integers. From 8416383a1b427804f46a7b3c076e4b7503c3bafb Mon Sep 17 00:00:00 2001 From: Spencer Clark Date: Tue, 23 Feb 2021 10:08:04 -0500 Subject: [PATCH 11/11] Update xarray/core/rolling.py Co-authored-by: Mathias Hauser --- xarray/core/rolling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py index 9ec75f0b821..dbdd9595069 100644 --- a/xarray/core/rolling.py +++ b/xarray/core/rolling.py @@ -962,7 +962,7 @@ def reduce(self, func: Callable, **kwargs): func : callable Function which can be called in the form `func(x, axis, **kwargs)` to return the result of collapsing an np.ndarray over the coarsening - dimensions. It must be possible to provide the `axis` argument with + dimensions. It must be possible to provide the `axis` argument with a tuple of integers. **kwargs : dict Additional keyword arguments passed on to `func`.