import numpy as np import pytest from pandas import ( DataFrame, MultiIndex, Series, concat, date_range, ) import pandas._testing as tm from pandas.api.indexers import ( BaseIndexer, FixedForwardWindowIndexer, ) from pandas.core.indexers.objects import ( ExpandingIndexer, FixedWindowIndexer, VariableOffsetWindowIndexer, ) from pandas.tseries.offsets import BusinessDay def test_bad_get_window_bounds_signature(): class BadIndexer(BaseIndexer): def get_window_bounds(self): return None indexer = BadIndexer() with pytest.raises(ValueError, match="BadIndexer does not implement"): Series(range(5)).rolling(indexer) def test_expanding_indexer(): s = Series(range(10)) indexer = ExpandingIndexer() result = s.rolling(indexer).mean() expected = s.expanding().mean() tm.assert_series_equal(result, expected) def test_indexer_constructor_arg(): # Example found in computation.rst use_expanding = [True, False, True, False, True] df = DataFrame({"values": range(5)}) class CustomIndexer(BaseIndexer): def get_window_bounds(self, num_values, min_periods, center, closed, step): start = np.empty(num_values, dtype=np.int64) end = np.empty(num_values, dtype=np.int64) for i in range(num_values): if self.use_expanding[i]: start[i] = 0 end[i] = i + 1 else: start[i] = i end[i] = i + self.window_size return start, end indexer = CustomIndexer(window_size=1, use_expanding=use_expanding) result = df.rolling(indexer).sum() expected = DataFrame({"values": [0.0, 1.0, 3.0, 3.0, 10.0]}) tm.assert_frame_equal(result, expected) def test_indexer_accepts_rolling_args(): df = DataFrame({"values": range(5)}) class CustomIndexer(BaseIndexer): def get_window_bounds(self, num_values, min_periods, center, closed, step): start = np.empty(num_values, dtype=np.int64) end = np.empty(num_values, dtype=np.int64) for i in range(num_values): if ( center and min_periods == 1 and closed == "both" and step == 1 and i == 2 ): start[i] = 0 end[i] = num_values else: start[i] = i end[i] = i + self.window_size return start, end indexer = CustomIndexer(window_size=1) result = df.rolling( indexer, center=True, min_periods=1, closed="both", step=1 ).sum() expected = DataFrame({"values": [0.0, 1.0, 10.0, 3.0, 4.0]}) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize("constructor", [Series, DataFrame]) @pytest.mark.parametrize( "func,np_func,expected,np_kwargs", [ ("count", len, [3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 2.0, np.nan], {}), ("min", np.min, [0.0, 1.0, 2.0, 3.0, 4.0, 6.0, 6.0, 7.0, 8.0, np.nan], {}), ( "max", np.max, [2.0, 3.0, 4.0, 100.0, 100.0, 100.0, 8.0, 9.0, 9.0, np.nan], {}, ), ( "std", np.std, [ 1.0, 1.0, 1.0, 55.71654452, 54.85739087, 53.9845657, 1.0, 1.0, 0.70710678, np.nan, ], {"ddof": 1}, ), ( "var", np.var, [ 1.0, 1.0, 1.0, 3104.333333, 3009.333333, 2914.333333, 1.0, 1.0, 0.500000, np.nan, ], {"ddof": 1}, ), ( "median", np.median, [1.0, 2.0, 3.0, 4.0, 6.0, 7.0, 7.0, 8.0, 8.5, np.nan], {}, ), ], ) @pytest.mark.filterwarnings("ignore:min_periods:FutureWarning") def test_rolling_forward_window(constructor, func, np_func, expected, np_kwargs, step): # GH 32865 values = np.arange(10.0) values[5] = 100.0 indexer = FixedForwardWindowIndexer(window_size=3) match = "Forward-looking windows can't have center=True" with pytest.raises(ValueError, match=match): rolling = constructor(values).rolling(window=indexer, center=True) getattr(rolling, func)() match = "Forward-looking windows don't support setting the closed argument" with pytest.raises(ValueError, match=match): rolling = constructor(values).rolling(window=indexer, closed="right") getattr(rolling, func)() rolling = constructor(values).rolling(window=indexer, min_periods=2, step=step) result = getattr(rolling, func)() # Check that the function output matches the explicitly provided array expected = constructor(expected)[::step] tm.assert_equal(result, expected) # Check that the rolling function output matches applying an alternative # function to the rolling window object expected2 = constructor(rolling.apply(lambda x: np_func(x, **np_kwargs))) tm.assert_equal(result, expected2) # Check that the function output matches applying an alternative function # if min_periods isn't specified # GH 39604: After count-min_periods deprecation, apply(lambda x: len(x)) # is equivalent to count after setting min_periods=0 min_periods = 0 if func == "count" else None rolling3 = constructor(values).rolling(window=indexer, min_periods=min_periods) result3 = getattr(rolling3, func)() expected3 = constructor(rolling3.apply(lambda x: np_func(x, **np_kwargs))) tm.assert_equal(result3, expected3) @pytest.mark.parametrize("constructor", [Series, DataFrame]) def test_rolling_forward_skewness(constructor, step): values = np.arange(10.0) values[5] = 100.0 indexer = FixedForwardWindowIndexer(window_size=5) rolling = constructor(values).rolling(window=indexer, min_periods=3, step=step) result = rolling.skew() expected = constructor( [ 0.0, 2.232396, 2.229508, 2.228340, 2.229091, 2.231989, 0.0, 0.0, np.nan, np.nan, ] )[::step] tm.assert_equal(result, expected) @pytest.mark.parametrize( "func,expected", [ ("cov", [2.0, 2.0, 2.0, 97.0, 2.0, -93.0, 2.0, 2.0, np.nan, np.nan]), ( "corr", [ 1.0, 1.0, 1.0, 0.8704775290207161, 0.018229084250926637, -0.861357304646493, 1.0, 1.0, np.nan, np.nan, ], ), ], ) def test_rolling_forward_cov_corr(func, expected): values1 = np.arange(10).reshape(-1, 1) values2 = values1 * 2 values1[5, 0] = 100 values = np.concatenate([values1, values2], axis=1) indexer = FixedForwardWindowIndexer(window_size=3) rolling = DataFrame(values).rolling(window=indexer, min_periods=3) # We are interested in checking only pairwise covariance / correlation result = getattr(rolling, func)().loc[(slice(None), 1), 0] result = result.reset_index(drop=True) expected = Series(expected).reset_index(drop=True) expected.name = result.name tm.assert_equal(result, expected) @pytest.mark.parametrize( "closed,expected_data", [ ["right", [0.0, 1.0, 2.0, 3.0, 7.0, 12.0, 6.0, 7.0, 8.0, 9.0]], ["left", [0.0, 0.0, 1.0, 2.0, 5.0, 9.0, 5.0, 6.0, 7.0, 8.0]], ], ) def test_non_fixed_variable_window_indexer(closed, expected_data): index = date_range("2020", periods=10) df = DataFrame(range(10), index=index) offset = BusinessDay(1) indexer = VariableOffsetWindowIndexer(index=index, offset=offset) result = df.rolling(indexer, closed=closed).sum() expected = DataFrame(expected_data, index=index) tm.assert_frame_equal(result, expected) def test_fixed_forward_indexer_count(step): # GH: 35579 df = DataFrame({"b": [None, None, None, 7]}) indexer = FixedForwardWindowIndexer(window_size=2) result = df.rolling(window=indexer, min_periods=0, step=step).count() expected = DataFrame({"b": [0.0, 0.0, 1.0, 1.0]})[::step] tm.assert_frame_equal(result, expected) @pytest.mark.parametrize( ("end_value", "values"), [(1, [0.0, 1, 1, 3, 2]), (-1, [0.0, 1, 0, 3, 1])] ) @pytest.mark.parametrize(("func", "args"), [("median", []), ("quantile", [0.5])]) def test_indexer_quantile_sum(end_value, values, func, args): # GH 37153 class CustomIndexer(BaseIndexer): def get_window_bounds(self, num_values, min_periods, center, closed, step): start = np.empty(num_values, dtype=np.int64) end = np.empty(num_values, dtype=np.int64) for i in range(num_values): if self.use_expanding[i]: start[i] = 0 end[i] = max(i + end_value, 1) else: start[i] = i end[i] = i + self.window_size return start, end use_expanding = [True, False, True, False, True] df = DataFrame({"values": range(5)}) indexer = CustomIndexer(window_size=1, use_expanding=use_expanding) result = getattr(df.rolling(indexer), func)(*args) expected = DataFrame({"values": values}) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize( "indexer_class", [FixedWindowIndexer, FixedForwardWindowIndexer, ExpandingIndexer] ) @pytest.mark.parametrize("window_size", [1, 2, 12]) @pytest.mark.parametrize( "df_data", [ {"a": [1, 1], "b": [0, 1]}, {"a": [1, 2], "b": [0, 1]}, {"a": [1] * 16, "b": [np.nan, 1, 2, np.nan] + list(range(4, 16))}, ], ) def test_indexers_are_reusable_after_groupby_rolling( indexer_class, window_size, df_data ): # GH 43267 df = DataFrame(df_data) num_trials = 3 indexer = indexer_class(window_size=window_size) original_window_size = indexer.window_size for i in range(num_trials): df.groupby("a")["b"].rolling(window=indexer, min_periods=1).mean() assert indexer.window_size == original_window_size @pytest.mark.parametrize( "window_size, num_values, expected_start, expected_end", [ (1, 1, [0], [1]), (1, 2, [0, 1], [1, 2]), (2, 1, [0], [1]), (2, 2, [0, 1], [2, 2]), (5, 12, range(12), list(range(5, 12)) + [12] * 5), (12, 5, range(5), [5] * 5), (0, 0, np.array([]), np.array([])), (1, 0, np.array([]), np.array([])), (0, 1, [0], [0]), ], ) def test_fixed_forward_indexer_bounds( window_size, num_values, expected_start, expected_end, step ): # GH 43267 indexer = FixedForwardWindowIndexer(window_size=window_size) start, end = indexer.get_window_bounds(num_values=num_values, step=step) tm.assert_numpy_array_equal( start, np.array(expected_start[::step]), check_dtype=False ) tm.assert_numpy_array_equal(end, np.array(expected_end[::step]), check_dtype=False) assert len(start) == len(end) @pytest.mark.parametrize( "df, window_size, expected", [ ( DataFrame({"b": [0, 1, 2], "a": [1, 2, 2]}), 2, Series( [0, 1.5, 2.0], index=MultiIndex.from_arrays([[1, 2, 2], range(3)], names=["a", None]), name="b", dtype=np.float64, ), ), ( DataFrame( { "b": [np.nan, 1, 2, np.nan] + list(range(4, 18)), "a": [1] * 7 + [2] * 11, "c": range(18), } ), 12, Series( [ 3.6, 3.6, 4.25, 5.0, 5.0, 5.5, 6.0, 12.0, 12.5, 13.0, 13.5, 14.0, 14.5, 15.0, 15.5, 16.0, 16.5, 17.0, ], index=MultiIndex.from_arrays( [[1] * 7 + [2] * 11, range(18)], names=["a", None] ), name="b", dtype=np.float64, ), ), ], ) def test_rolling_groupby_with_fixed_forward_specific(df, window_size, expected): # GH 43267 indexer = FixedForwardWindowIndexer(window_size=window_size) result = df.groupby("a")["b"].rolling(window=indexer, min_periods=1).mean() tm.assert_series_equal(result, expected) @pytest.mark.parametrize( "group_keys", [ (1,), (1, 2), (2, 1), (1, 1, 2), (1, 2, 1), (1, 1, 2, 2), (1, 2, 3, 2, 3), (1, 1, 2) * 4, (1, 2, 3) * 5, ], ) @pytest.mark.parametrize("window_size", [1, 2, 3, 4, 5, 8, 20]) def test_rolling_groupby_with_fixed_forward_many(group_keys, window_size): # GH 43267 df = DataFrame( { "a": np.array(list(group_keys)), "b": np.arange(len(group_keys), dtype=np.float64) + 17, "c": np.arange(len(group_keys), dtype=np.int64), } ) indexer = FixedForwardWindowIndexer(window_size=window_size) result = df.groupby("a")["b"].rolling(window=indexer, min_periods=1).sum() result.index.names = ["a", "c"] groups = df.groupby("a")[["a", "b", "c"]] manual = concat( [ g.assign( b=[ g["b"].iloc[i : i + window_size].sum(min_count=1) for i in range(len(g)) ] ) for _, g in groups ] ) manual = manual.set_index(["a", "c"])["b"] tm.assert_series_equal(result, manual) def test_unequal_start_end_bounds(): class CustomIndexer(BaseIndexer): def get_window_bounds(self, num_values, min_periods, center, closed, step): return np.array([1]), np.array([1, 2]) indexer = CustomIndexer() roll = Series(1).rolling(indexer) match = "start" with pytest.raises(ValueError, match=match): roll.mean() with pytest.raises(ValueError, match=match): next(iter(roll)) with pytest.raises(ValueError, match=match): roll.corr(pairwise=True) with pytest.raises(ValueError, match=match): roll.cov(pairwise=True) def test_unequal_bounds_to_object(): # GH 44470 class CustomIndexer(BaseIndexer): def get_window_bounds(self, num_values, min_periods, center, closed, step): return np.array([1]), np.array([2]) indexer = CustomIndexer() roll = Series([1, 1]).rolling(indexer) match = "start and end" with pytest.raises(ValueError, match=match): roll.mean() with pytest.raises(ValueError, match=match): next(iter(roll)) with pytest.raises(ValueError, match=match): roll.corr(pairwise=True) with pytest.raises(ValueError, match=match): roll.cov(pairwise=True)