Add corr, cov, std & var to .rolling_exp

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)