Add DataArrayCoarsen.reduce and DatasetCoarsen.reduce methods

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,14 +873,48 @@ 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)
+        <xarray.DataArray (a: 2, b: 2)>
+        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__ = ()
 
     _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)

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)

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",
     [