Add corr, cov, std & var to .rolling_exp (#8307)

* Add `corr`, `cov`, `std` & `var` to `.rolling_exp`

From the new routines in numbagg.

Maybe needs better tests (though these are quite heavily tested in numbagg), docs, and potentially need to think about types (maybe existing binary ops can help here?)

Author: max-sixty

xarray/core/rolling_exp.py

@@ -4,6 +4,7 @@
 from typing import Any, Generic
 
 import numpy as np
+from packaging.version import Version
 
 from xarray.core.computation import apply_ufunc
 from xarray.core.options import _get_keep_attrs
@@ -14,9 +15,9 @@
     import numbagg
     from numbagg import move_exp_nanmean, move_exp_nansum
 
-    has_numbagg = numbagg.__version__
+    has_numbagg: Version | None = Version(numbagg.__version__)
 except ImportError:
-    has_numbagg = False
+    has_numbagg = None
 
 
 def _get_alpha(
@@ -99,17 +100,17 @@ def __init__(
         window_type: str = "span",
         min_weight: float = 0.0,
     ):
-        if has_numbagg is False:
+        if has_numbagg is None:
             raise ImportError(
                 "numbagg >= 0.2.1 is required for rolling_exp but currently numbagg is not installed"
             )
-        elif has_numbagg < "0.2.1":
+        elif has_numbagg < Version("0.2.1"):
             raise ImportError(
-                f"numbagg >= 0.2.1 is required for rolling_exp but currently version {has_numbagg} is installed"
+                f"numbagg >= 0.2.1 is required for `rolling_exp` but currently version {has_numbagg} is installed"
             )
-        elif has_numbagg < "0.3.1" and min_weight > 0:
+        elif has_numbagg < Version("0.3.1") and min_weight > 0:
             raise ImportError(
-                f"numbagg >= 0.3.1 is required for `min_weight > 0` but currently version {has_numbagg} is installed"
+                f"numbagg >= 0.3.1 is required for `min_weight > 0` within `.rolling_exp` but currently version {has_numbagg} is installed"
             )
 
         self.obj: T_DataWithCoords = obj
@@ -194,3 +195,133 @@ def sum(self, keep_attrs: bool | None = None) -> T_DataWithCoords:
             on_missing_core_dim="copy",
             dask="parallelized",
         ).transpose(*dim_order)
+
+    def std(self) -> T_DataWithCoords:
+        """
+        Exponentially weighted moving standard deviation.
+
+        `keep_attrs` is always True for this method. Drop attrs separately to remove attrs.
+
+        Examples
+        --------
+        >>> da = xr.DataArray([1, 1, 2, 2, 2], dims="x")
+        >>> da.rolling_exp(x=2, window_type="span").std()
+        <xarray.DataArray (x: 5)>
+        array([       nan, 0.        , 0.67936622, 0.42966892, 0.25389527])
+        Dimensions without coordinates: x
+        """
+
+        if has_numbagg is None or has_numbagg < Version("0.4.0"):
+            raise ImportError(
+                f"numbagg >= 0.4.0 is required for rolling_exp().std(), currently {has_numbagg} is installed"
+            )
+        dim_order = self.obj.dims
+
+        return apply_ufunc(
+            numbagg.move_exp_nanstd,
+            self.obj,
+            input_core_dims=[[self.dim]],
+            kwargs=self.kwargs,
+            output_core_dims=[[self.dim]],
+            keep_attrs=True,
+            on_missing_core_dim="copy",
+            dask="parallelized",
+        ).transpose(*dim_order)
+
+    def var(self) -> T_DataWithCoords:
+        """
+        Exponentially weighted moving variance.
+
+        `keep_attrs` is always True for this method. Drop attrs separately to remove attrs.
+
+        Examples
+        --------
+        >>> da = xr.DataArray([1, 1, 2, 2, 2], dims="x")
+        >>> da.rolling_exp(x=2, window_type="span").var()
+        <xarray.DataArray (x: 5)>
+        array([       nan, 0.        , 0.46153846, 0.18461538, 0.06446281])
+        Dimensions without coordinates: x
+        """
+
+        if has_numbagg is None or has_numbagg < Version("0.4.0"):
+            raise ImportError(
+                f"numbagg >= 0.4.0 is required for rolling_exp().var(), currently {has_numbagg} is installed"
+            )
+        dim_order = self.obj.dims
+
+        return apply_ufunc(
+            numbagg.move_exp_nanvar,
+            self.obj,
+            input_core_dims=[[self.dim]],
+            kwargs=self.kwargs,
+            output_core_dims=[[self.dim]],
+            keep_attrs=True,
+            on_missing_core_dim="copy",
+            dask="parallelized",
+        ).transpose(*dim_order)
+
+    def cov(self, other: T_DataWithCoords) -> T_DataWithCoords:
+        """
+        Exponentially weighted moving covariance.
+
+        `keep_attrs` is always True for this method. Drop attrs separately to remove attrs.
+
+        Examples
+        --------
+        >>> da = xr.DataArray([1, 1, 2, 2, 2], dims="x")
+        >>> da.rolling_exp(x=2, window_type="span").cov(da**2)
+        <xarray.DataArray (x: 5)>
+        array([       nan, 0.        , 1.38461538, 0.55384615, 0.19338843])
+        Dimensions without coordinates: x
+        """
+
+        if has_numbagg is None or has_numbagg < Version("0.4.0"):
+            raise ImportError(
+                f"numbagg >= 0.4.0 is required for rolling_exp().cov(), currently {has_numbagg} is installed"
+            )
+        dim_order = self.obj.dims
+
+        return apply_ufunc(
+            numbagg.move_exp_nancov,
+            self.obj,
+            other,
+            input_core_dims=[[self.dim], [self.dim]],
+            kwargs=self.kwargs,
+            output_core_dims=[[self.dim]],
+            keep_attrs=True,
+            on_missing_core_dim="copy",
+            dask="parallelized",
+        ).transpose(*dim_order)
+
+    def corr(self, other: T_DataWithCoords) -> T_DataWithCoords:
+        """
+        Exponentially weighted moving correlation.
+
+        `keep_attrs` is always True for this method. Drop attrs separately to remove attrs.
+
+        Examples
+        --------
+        >>> da = xr.DataArray([1, 1, 2, 2, 2], dims="x")
+        >>> da.rolling_exp(x=2, window_type="span").corr(da.shift(x=1))
+        <xarray.DataArray (x: 5)>
+        array([       nan,        nan,        nan, 0.4330127 , 0.48038446])
+        Dimensions without coordinates: x
+        """
+
+        if has_numbagg is None or has_numbagg < Version("0.4.0"):
+            raise ImportError(
+                f"numbagg >= 0.4.0 is required for rolling_exp().cov(), currently {has_numbagg} is installed"
+            )
+        dim_order = self.obj.dims
+
+        return apply_ufunc(
+            numbagg.move_exp_nancorr,
+            self.obj,
+            other,
+            input_core_dims=[[self.dim], [self.dim]],
+            kwargs=self.kwargs,
+            output_core_dims=[[self.dim]],
+            keep_attrs=True,
+            on_missing_core_dim="copy",
+            dask="parallelized",
+        ).transpose(*dim_order)

xarray/tests/__init__.py

@@ -75,7 +75,7 @@ def _importorskip(
 has_zarr, requires_zarr = _importorskip("zarr")
 has_fsspec, requires_fsspec = _importorskip("fsspec")
 has_iris, requires_iris = _importorskip("iris")
-has_numbagg, requires_numbagg = _importorskip("numbagg")
+has_numbagg, requires_numbagg = _importorskip("numbagg", "0.4.0")
 has_seaborn, requires_seaborn = _importorskip("seaborn")
 has_sparse, requires_sparse = _importorskip("sparse")
 has_cupy, requires_cupy = _importorskip("cupy")

xarray/tests/test_rolling.py

@@ -394,7 +394,7 @@ class TestDataArrayRollingExp:
         [["span", 5], ["alpha", 0.5], ["com", 0.5], ["halflife", 5]],
     )
     @pytest.mark.parametrize("backend", ["numpy"], indirect=True)
-    @pytest.mark.parametrize("func", ["mean", "sum"])
+    @pytest.mark.parametrize("func", ["mean", "sum", "var", "std"])
     def test_rolling_exp_runs(self, da, dim, window_type, window, func) -> None:
         da = da.where(da > 0.2)