Rolling exp sum (#5178)

* Small simplification to RollingExp

* Add rolling_exp().sum()

* lint

* Update xarray/core/rolling_exp.py

Co-authored-by: Mathias Hauser <mathause@users.noreply.github.com>

* Check for 0.2.1 version

* Add tests

* whatsnew

* Skip sum tests on older version of numbagg

* .

Co-authored-by: Mathias Hauser <mathause@users.noreply.github.com>

Author: max-sixty

doc/whats-new.rst

@@ -25,6 +25,11 @@ New Features
   indexing (:issue:`3015`, :pull:`5362`). By `Matthias Göbel <https://github.com/matzegoebel>`_.
 - Attempting to reduce a weighted object over missing dimensions now raises an error (:pull:`5362`).
   By `Mattia Almansi <https://github.com/malmans2>`_.
+- Add ``.sum`` to :py:meth:`~xarray.DataArray.rolling_exp` and
+  :py:meth:`~xarray.Dataset.rolling_exp` for exponentially weighted rolling
+  sums. These require numbagg 0.2.1;
+  (:pull:`5178`).
+  By `Maximilian Roos <https://github.com/max-sixty>`_.
 - :py:func:`xarray.cov` and :py:func:`xarray.corr` now lazily check for missing
   values if inputs are dask arrays (:issue:`4804`, :pull:`5284`).
   By `Andrew Williams <https://github.com/AndrewWilliams3142>`_.

xarray/core/rolling_exp.py

@@ -1,4 +1,5 @@
-from typing import TYPE_CHECKING, Generic, Hashable, Mapping, Optional, TypeVar, Union
+from distutils.version import LooseVersion
+from typing import TYPE_CHECKING, Generic, Hashable, Mapping, TypeVar, Union
 
 import numpy as np
 
@@ -26,12 +27,25 @@ def move_exp_nanmean(array, *, axis, alpha):
         raise TypeError("rolling_exp is not currently support for dask-like arrays")
     import numbagg
 
+    # No longer needed in numbag > 0.2.0; remove in time
     if axis == ():
         return array.astype(np.float64)
     else:
         return numbagg.move_exp_nanmean(array, axis=axis, alpha=alpha)
 
 
+def move_exp_nansum(array, *, axis, alpha):
+    if is_duck_dask_array(array):
+        raise TypeError("rolling_exp is not currently supported for dask-like arrays")
+    import numbagg
+
+    # numbagg <= 0.2.0 did not have a __version__ attribute
+    if LooseVersion(getattr(numbagg, "__version__", "0.1.0")) < LooseVersion("0.2.0"):
+        raise ValueError("`rolling_exp(...).sum() requires numbagg>=0.2.1.")
+
+    return numbagg.move_exp_nansum(array, axis=axis, alpha=alpha)
+
+
 def _get_center_of_mass(comass, span, halflife, alpha):
     """
     Vendored from pandas.core.window.common._get_center_of_mass
@@ -98,9 +112,9 @@ def __init__(
         self.dim = dim
         self.alpha = _get_alpha(**{window_type: window})
 
-    def mean(self, keep_attrs: Optional[bool] = None) -> T_DSorDA:
+    def mean(self, keep_attrs: bool = None) -> T_DSorDA:
         """
-        Exponentially weighted moving average
+        Exponentially weighted moving average.
 
         Parameters
         ----------
@@ -124,3 +138,30 @@ def mean(self, keep_attrs: Optional[bool] = None) -> T_DSorDA:
         return self.obj.reduce(
             move_exp_nanmean, dim=self.dim, alpha=self.alpha, keep_attrs=keep_attrs
         )
+
+    def sum(self, keep_attrs: bool = None) -> T_DSorDA:
+        """
+        Exponentially weighted moving sum.
+
+        Parameters
+        ----------
+        keep_attrs : bool, default: None
+            If True, the attributes (``attrs``) will be copied from the original
+            object to the new one. If False, the new object will be returned
+            without attributes. If None uses the global default.
+
+        Examples
+        --------
+        >>> da = xr.DataArray([1, 1, 2, 2, 2], dims="x")
+        >>> da.rolling_exp(x=2, window_type="span").sum()
+        <xarray.DataArray (x: 5)>
+        array([1.        , 1.33333333, 2.44444444, 2.81481481, 2.9382716 ])
+        Dimensions without coordinates: x
+        """
+
+        if keep_attrs is None:
+            keep_attrs = _get_keep_attrs(default=True)
+
+        return self.obj.reduce(
+            move_exp_nansum, dim=self.dim, alpha=self.alpha, keep_attrs=keep_attrs
+        )

xarray/tests/test_dataarray.py

@@ -7268,10 +7268,32 @@ def test_fallback_to_iris_AuxCoord(self, coord_values):
     "window_type, window", [["span", 5], ["alpha", 0.5], ["com", 0.5], ["halflife", 5]]
 )
 @pytest.mark.parametrize("backend", ["numpy"], indirect=True)
-def test_rolling_exp(da, dim, window_type, window):
-    da = da.isel(a=0)
+@pytest.mark.parametrize("func", ["mean", "sum"])
+def test_rolling_exp_runs(da, dim, window_type, window, func):
+    import numbagg
+
+    if (
+        LooseVersion(getattr(numbagg, "__version__", "0.1.0")) < "0.2.1"
+        and func == "sum"
+    ):
+        pytest.skip("rolling_exp.sum requires numbagg 0.2.1")
+
     da = da.where(da > 0.2)
 
+    rolling_exp = da.rolling_exp(window_type=window_type, **{dim: window})
+    result = getattr(rolling_exp, func)()
+    assert isinstance(result, DataArray)
+
+
+@requires_numbagg
+@pytest.mark.parametrize("dim", ["time", "x"])
+@pytest.mark.parametrize(
+    "window_type, window", [["span", 5], ["alpha", 0.5], ["com", 0.5], ["halflife", 5]]
+)
+@pytest.mark.parametrize("backend", ["numpy"], indirect=True)
+def test_rolling_exp_mean_pandas(da, dim, window_type, window):
+    da = da.isel(a=0).where(lambda x: x > 0.2)
+
     result = da.rolling_exp(window_type=window_type, **{dim: window}).mean()
     assert isinstance(result, DataArray)
 
@@ -7288,30 +7310,42 @@ def test_rolling_exp(da, dim, window_type, window):
 
 @requires_numbagg
 @pytest.mark.parametrize("backend", ["numpy"], indirect=True)
-def test_rolling_exp_keep_attrs(da):
+@pytest.mark.parametrize("func", ["mean", "sum"])
+def test_rolling_exp_keep_attrs(da, func):
+    import numbagg
+
+    if (
+        LooseVersion(getattr(numbagg, "__version__", "0.1.0")) < "0.2.1"
+        and func == "sum"
+    ):
+        pytest.skip("rolling_exp.sum requires numbagg 0.2.1")
+
     attrs = {"attrs": "da"}
     da.attrs = attrs
 
+    # Equivalent of `da.rolling_exp(time=10).mean`
+    rolling_exp_func = getattr(da.rolling_exp(time=10), func)
+
     # attrs are kept per default
-    result = da.rolling_exp(time=10).mean()
+    result = rolling_exp_func()
     assert result.attrs == attrs
 
     # discard attrs
-    result = da.rolling_exp(time=10).mean(keep_attrs=False)
+    result = rolling_exp_func(keep_attrs=False)
     assert result.attrs == {}
 
     # test discard attrs using global option
     with set_options(keep_attrs=False):
-        result = da.rolling_exp(time=10).mean()
+        result = rolling_exp_func()
     assert result.attrs == {}
 
     # keyword takes precedence over global option
     with set_options(keep_attrs=False):
-        result = da.rolling_exp(time=10).mean(keep_attrs=True)
+        result = rolling_exp_func(keep_attrs=True)
     assert result.attrs == attrs
 
     with set_options(keep_attrs=True):
-        result = da.rolling_exp(time=10).mean(keep_attrs=False)
+        result = rolling_exp_func(keep_attrs=False)
     assert result.attrs == {}
 
     with pytest.warns(