diff --git a/xarray/tests/conftest.py b/xarray/tests/conftest.py index 6a8cf008f9f..f153c2f4dc0 100644 --- a/xarray/tests/conftest.py +++ b/xarray/tests/conftest.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import numpy as np import pandas as pd import pytest @@ -77,3 +79,44 @@ def da(request, backend): return da else: raise ValueError + + +@pytest.fixture(params=[Dataset, DataArray]) +def type(request): + return request.param + + +@pytest.fixture(params=[1]) +def d(request, backend, type) -> DataArray | Dataset: + """ + For tests which can test either a DataArray or a Dataset. + """ + result: DataArray | Dataset + if request.param == 1: + ds = Dataset( + dict( + a=(["x", "z"], np.arange(24).reshape(2, 12)), + b=(["y", "z"], np.arange(100, 136).reshape(3, 12).astype(np.float64)), + ), + dict( + x=("x", np.linspace(0, 1.0, 2)), + y=range(3), + z=("z", pd.date_range("2000-01-01", periods=12)), + w=("x", ["a", "b"]), + ), + ) + if type == DataArray: + result = ds["a"].assign_coords(w=ds.coords["w"]) + elif type == Dataset: + result = ds + else: + raise ValueError + else: + raise ValueError + + if backend == "dask": + return result.chunk() + elif backend == "numpy": + return result + else: + raise ValueError diff --git a/xarray/tests/test_rolling.py b/xarray/tests/test_rolling.py index 645ec1f85e6..7cb2cd70d29 100644 --- a/xarray/tests/test_rolling.py +++ b/xarray/tests/test_rolling.py @@ -36,6 +36,31 @@ def compute_backend(request): yield request.param +@pytest.mark.parametrize("func", ["mean", "sum"]) +@pytest.mark.parametrize("min_periods", [1, 10]) +def test_cumulative(d, func, min_periods) -> None: + # One dim + result = getattr(d.cumulative("z", min_periods=min_periods), func)() + expected = getattr(d.rolling(z=d["z"].size, min_periods=min_periods), func)() + assert_identical(result, expected) + + # Multiple dim + result = getattr(d.cumulative(["z", "x"], min_periods=min_periods), func)() + expected = getattr( + d.rolling(z=d["z"].size, x=d["x"].size, min_periods=min_periods), + func, + )() + assert_identical(result, expected) + + +def test_cumulative_vs_cum(d) -> None: + result = d.cumulative("z").sum() + expected = d.cumsum("z") + # cumsum drops the coord of the dimension; cumulative doesn't + expected = expected.assign_coords(z=result["z"]) + assert_identical(result, expected) + + class TestDataArrayRolling: @pytest.mark.parametrize("da", (1, 2), indirect=True) @pytest.mark.parametrize("center", [True, False]) @@ -485,29 +510,6 @@ def test_rolling_exp_keep_attrs(self, da, func) -> None: ): da.rolling_exp(time=10, keep_attrs=True) - @pytest.mark.parametrize("func", ["mean", "sum"]) - @pytest.mark.parametrize("min_periods", [1, 20]) - def test_cumulative(self, da, func, min_periods) -> None: - # One dim - result = getattr(da.cumulative("time", min_periods=min_periods), func)() - expected = getattr( - da.rolling(time=da.time.size, min_periods=min_periods), func - )() - assert_identical(result, expected) - - # Multiple dim - result = getattr(da.cumulative(["time", "a"], min_periods=min_periods), func)() - expected = getattr( - da.rolling(time=da.time.size, a=da.a.size, min_periods=min_periods), - func, - )() - assert_identical(result, expected) - - def test_cumulative_vs_cum(self, da) -> None: - result = da.cumulative("time").sum() - expected = da.cumsum("time") - assert_identical(result, expected) - class TestDatasetRolling: @pytest.mark.parametrize( @@ -832,25 +834,6 @@ def test_raise_no_warning_dask_rolling_assert_close(self, ds, name) -> None: expected = getattr(getattr(ds.rolling(time=4), name)().rolling(x=3), name)() assert_allclose(actual, expected) - @pytest.mark.parametrize("func", ["mean", "sum"]) - @pytest.mark.parametrize("ds", (2,), indirect=True) - @pytest.mark.parametrize("min_periods", [1, 10]) - def test_cumulative(self, ds, func, min_periods) -> None: - # One dim - result = getattr(ds.cumulative("time", min_periods=min_periods), func)() - expected = getattr( - ds.rolling(time=ds.time.size, min_periods=min_periods), func - )() - assert_identical(result, expected) - - # Multiple dim - result = getattr(ds.cumulative(["time", "x"], min_periods=min_periods), func)() - expected = getattr( - ds.rolling(time=ds.time.size, x=ds.x.size, min_periods=min_periods), - func, - )() - assert_identical(result, expected) - @requires_numbagg class TestDatasetRollingExp: