diff --git a/xarray/core/rolling_exp.py b/xarray/core/rolling_exp.py index 4935f3c8172..f3a9aed42c6 100644 --- a/xarray/core/rolling_exp.py +++ b/xarray/core/rolling_exp.py @@ -194,3 +194,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() + + array([ nan, 0. , 0.67936622, 0.42966892, 0.25389527]) + Dimensions without coordinates: x + """ + + if has_numbagg is False or has_numbagg < "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() + + array([ nan, 0. , 0.46153846, 0.18461538, 0.06446281]) + Dimensions without coordinates: x + """ + + if has_numbagg is False or has_numbagg < "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) + + array([ nan, 0. , 1.38461538, 0.55384615, 0.19338843]) + Dimensions without coordinates: x + """ + + if has_numbagg is False or has_numbagg < "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)) + + array([ nan, nan, nan, 0.4330127 , 0.48038446]) + Dimensions without coordinates: x + """ + + if has_numbagg is False or has_numbagg < "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) diff --git a/xarray/tests/test_rolling.py b/xarray/tests/test_rolling.py index da834b76124..3b213db0b88 100644 --- a/xarray/tests/test_rolling.py +++ b/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)