From 56e0dd720544ca4f6f010329cb217112c907bdfb Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sat, 14 Oct 2023 00:25:16 -0700 Subject: [PATCH] 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?) --- xarray/core/rolling_exp.py | 130 +++++++++++++++++++++++++++++++++++ xarray/tests/test_rolling.py | 2 +- 2 files changed, 131 insertions(+), 1 deletion(-) 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)