1import numpy as np
2import pytest
3
4from pandas import DataFrame, Series, concat, isna, notna
5import pandas._testing as tm
6
7import pandas.tseries.offsets as offsets
8
9
10@pytest.mark.parametrize(
11    "compare_func, roll_func, kwargs",
12    [
13        [np.mean, "mean", {}],
14        [np.nansum, "sum", {}],
15        pytest.param(
16            lambda x: np.isfinite(x).astype(float).sum(),
17            "count",
18            {},
19            marks=pytest.mark.filterwarnings("ignore:min_periods:FutureWarning"),
20        ),
21        [np.median, "median", {}],
22        [np.min, "min", {}],
23        [np.max, "max", {}],
24        [lambda x: np.std(x, ddof=1), "std", {}],
25        [lambda x: np.std(x, ddof=0), "std", {"ddof": 0}],
26        [lambda x: np.var(x, ddof=1), "var", {}],
27        [lambda x: np.var(x, ddof=0), "var", {"ddof": 0}],
28    ],
29)
30def test_series(series, compare_func, roll_func, kwargs):
31    result = getattr(series.rolling(50), roll_func)(**kwargs)
32    assert isinstance(result, Series)
33    tm.assert_almost_equal(result.iloc[-1], compare_func(series[-50:]))
34
35
36@pytest.mark.parametrize(
37    "compare_func, roll_func, kwargs",
38    [
39        [np.mean, "mean", {}],
40        [np.nansum, "sum", {}],
41        pytest.param(
42            lambda x: np.isfinite(x).astype(float).sum(),
43            "count",
44            {},
45            marks=pytest.mark.filterwarnings("ignore:min_periods:FutureWarning"),
46        ),
47        [np.median, "median", {}],
48        [np.min, "min", {}],
49        [np.max, "max", {}],
50        [lambda x: np.std(x, ddof=1), "std", {}],
51        [lambda x: np.std(x, ddof=0), "std", {"ddof": 0}],
52        [lambda x: np.var(x, ddof=1), "var", {}],
53        [lambda x: np.var(x, ddof=0), "var", {"ddof": 0}],
54    ],
55)
56def test_frame(raw, frame, compare_func, roll_func, kwargs):
57    result = getattr(frame.rolling(50), roll_func)(**kwargs)
58    assert isinstance(result, DataFrame)
59    tm.assert_series_equal(
60        result.iloc[-1, :],
61        frame.iloc[-50:, :].apply(compare_func, axis=0, raw=raw),
62        check_names=False,
63    )
64
65
66@pytest.mark.parametrize(
67    "compare_func, roll_func, kwargs, minp",
68    [
69        [np.mean, "mean", {}, 10],
70        [np.nansum, "sum", {}, 10],
71        [lambda x: np.isfinite(x).astype(float).sum(), "count", {}, 0],
72        [np.median, "median", {}, 10],
73        [np.min, "min", {}, 10],
74        [np.max, "max", {}, 10],
75        [lambda x: np.std(x, ddof=1), "std", {}, 10],
76        [lambda x: np.std(x, ddof=0), "std", {"ddof": 0}, 10],
77        [lambda x: np.var(x, ddof=1), "var", {}, 10],
78        [lambda x: np.var(x, ddof=0), "var", {"ddof": 0}, 10],
79    ],
80)
81def test_time_rule_series(series, compare_func, roll_func, kwargs, minp):
82    win = 25
83    ser = series[::2].resample("B").mean()
84    series_result = getattr(ser.rolling(window=win, min_periods=minp), roll_func)(
85        **kwargs
86    )
87    last_date = series_result.index[-1]
88    prev_date = last_date - 24 * offsets.BDay()
89
90    trunc_series = series[::2].truncate(prev_date, last_date)
91    tm.assert_almost_equal(series_result[-1], compare_func(trunc_series))
92
93
94@pytest.mark.parametrize(
95    "compare_func, roll_func, kwargs, minp",
96    [
97        [np.mean, "mean", {}, 10],
98        [np.nansum, "sum", {}, 10],
99        [lambda x: np.isfinite(x).astype(float).sum(), "count", {}, 0],
100        [np.median, "median", {}, 10],
101        [np.min, "min", {}, 10],
102        [np.max, "max", {}, 10],
103        [lambda x: np.std(x, ddof=1), "std", {}, 10],
104        [lambda x: np.std(x, ddof=0), "std", {"ddof": 0}, 10],
105        [lambda x: np.var(x, ddof=1), "var", {}, 10],
106        [lambda x: np.var(x, ddof=0), "var", {"ddof": 0}, 10],
107    ],
108)
109def test_time_rule_frame(raw, frame, compare_func, roll_func, kwargs, minp):
110    win = 25
111    frm = frame[::2].resample("B").mean()
112    frame_result = getattr(frm.rolling(window=win, min_periods=minp), roll_func)(
113        **kwargs
114    )
115    last_date = frame_result.index[-1]
116    prev_date = last_date - 24 * offsets.BDay()
117
118    trunc_frame = frame[::2].truncate(prev_date, last_date)
119    tm.assert_series_equal(
120        frame_result.xs(last_date),
121        trunc_frame.apply(compare_func, raw=raw),
122        check_names=False,
123    )
124
125
126@pytest.mark.parametrize(
127    "compare_func, roll_func, kwargs",
128    [
129        [np.mean, "mean", {}],
130        [np.nansum, "sum", {}],
131        [np.median, "median", {}],
132        [np.min, "min", {}],
133        [np.max, "max", {}],
134        [lambda x: np.std(x, ddof=1), "std", {}],
135        [lambda x: np.std(x, ddof=0), "std", {"ddof": 0}],
136        [lambda x: np.var(x, ddof=1), "var", {}],
137        [lambda x: np.var(x, ddof=0), "var", {"ddof": 0}],
138    ],
139)
140def test_nans(compare_func, roll_func, kwargs):
141    obj = Series(np.random.randn(50))
142    obj[:10] = np.NaN
143    obj[-10:] = np.NaN
144
145    result = getattr(obj.rolling(50, min_periods=30), roll_func)(**kwargs)
146    tm.assert_almost_equal(result.iloc[-1], compare_func(obj[10:-10]))
147
148    # min_periods is working correctly
149    result = getattr(obj.rolling(20, min_periods=15), roll_func)(**kwargs)
150    assert isna(result.iloc[23])
151    assert not isna(result.iloc[24])
152
153    assert not isna(result.iloc[-6])
154    assert isna(result.iloc[-5])
155
156    obj2 = Series(np.random.randn(20))
157    result = getattr(obj2.rolling(10, min_periods=5), roll_func)(**kwargs)
158    assert isna(result.iloc[3])
159    assert notna(result.iloc[4])
160
161    if roll_func != "sum":
162        result0 = getattr(obj.rolling(20, min_periods=0), roll_func)(**kwargs)
163        result1 = getattr(obj.rolling(20, min_periods=1), roll_func)(**kwargs)
164        tm.assert_almost_equal(result0, result1)
165
166
167def test_nans_count():
168    obj = Series(np.random.randn(50))
169    obj[:10] = np.NaN
170    obj[-10:] = np.NaN
171    result = obj.rolling(50, min_periods=30).count()
172    tm.assert_almost_equal(
173        result.iloc[-1], np.isfinite(obj[10:-10]).astype(float).sum()
174    )
175
176
177@pytest.mark.parametrize(
178    "roll_func, kwargs",
179    [
180        ["mean", {}],
181        ["sum", {}],
182        ["median", {}],
183        ["min", {}],
184        ["max", {}],
185        ["std", {}],
186        ["std", {"ddof": 0}],
187        ["var", {}],
188        ["var", {"ddof": 0}],
189    ],
190)
191@pytest.mark.parametrize("minp", [0, 99, 100])
192def test_min_periods(series, minp, roll_func, kwargs):
193    result = getattr(series.rolling(len(series) + 1, min_periods=minp), roll_func)(
194        **kwargs
195    )
196    expected = getattr(series.rolling(len(series), min_periods=minp), roll_func)(
197        **kwargs
198    )
199    nan_mask = isna(result)
200    tm.assert_series_equal(nan_mask, isna(expected))
201
202    nan_mask = ~nan_mask
203    tm.assert_almost_equal(result[nan_mask], expected[nan_mask])
204
205
206def test_min_periods_count(series):
207    result = series.rolling(len(series) + 1, min_periods=0).count()
208    expected = series.rolling(len(series), min_periods=0).count()
209    nan_mask = isna(result)
210    tm.assert_series_equal(nan_mask, isna(expected))
211
212    nan_mask = ~nan_mask
213    tm.assert_almost_equal(result[nan_mask], expected[nan_mask])
214
215
216@pytest.mark.parametrize(
217    "roll_func, kwargs, minp",
218    [
219        ["mean", {}, 15],
220        ["sum", {}, 15],
221        ["count", {}, 0],
222        ["median", {}, 15],
223        ["min", {}, 15],
224        ["max", {}, 15],
225        ["std", {}, 15],
226        ["std", {"ddof": 0}, 15],
227        ["var", {}, 15],
228        ["var", {"ddof": 0}, 15],
229    ],
230)
231def test_center(roll_func, kwargs, minp):
232    obj = Series(np.random.randn(50))
233    obj[:10] = np.NaN
234    obj[-10:] = np.NaN
235
236    result = getattr(obj.rolling(20, min_periods=minp, center=True), roll_func)(
237        **kwargs
238    )
239    expected = getattr(
240        concat([obj, Series([np.NaN] * 9)]).rolling(20, min_periods=minp), roll_func
241    )(**kwargs)[9:].reset_index(drop=True)
242    tm.assert_series_equal(result, expected)
243
244
245@pytest.mark.parametrize(
246    "roll_func, kwargs, minp, fill_value",
247    [
248        ["mean", {}, 10, None],
249        ["sum", {}, 10, None],
250        ["count", {}, 0, 0],
251        ["median", {}, 10, None],
252        ["min", {}, 10, None],
253        ["max", {}, 10, None],
254        ["std", {}, 10, None],
255        ["std", {"ddof": 0}, 10, None],
256        ["var", {}, 10, None],
257        ["var", {"ddof": 0}, 10, None],
258    ],
259)
260def test_center_reindex_series(series, roll_func, kwargs, minp, fill_value):
261    # shifter index
262    s = [f"x{x:d}" for x in range(12)]
263
264    series_xp = (
265        getattr(
266            series.reindex(list(series.index) + s).rolling(window=25, min_periods=minp),
267            roll_func,
268        )(**kwargs)
269        .shift(-12)
270        .reindex(series.index)
271    )
272    series_rs = getattr(
273        series.rolling(window=25, min_periods=minp, center=True), roll_func
274    )(**kwargs)
275    if fill_value is not None:
276        series_xp = series_xp.fillna(fill_value)
277    tm.assert_series_equal(series_xp, series_rs)
278
279
280@pytest.mark.parametrize(
281    "roll_func, kwargs, minp, fill_value",
282    [
283        ["mean", {}, 10, None],
284        ["sum", {}, 10, None],
285        ["count", {}, 0, 0],
286        ["median", {}, 10, None],
287        ["min", {}, 10, None],
288        ["max", {}, 10, None],
289        ["std", {}, 10, None],
290        ["std", {"ddof": 0}, 10, None],
291        ["var", {}, 10, None],
292        ["var", {"ddof": 0}, 10, None],
293    ],
294)
295def test_center_reindex_frame(frame, roll_func, kwargs, minp, fill_value):
296    # shifter index
297    s = [f"x{x:d}" for x in range(12)]
298
299    frame_xp = (
300        getattr(
301            frame.reindex(list(frame.index) + s).rolling(window=25, min_periods=minp),
302            roll_func,
303        )(**kwargs)
304        .shift(-12)
305        .reindex(frame.index)
306    )
307    frame_rs = getattr(
308        frame.rolling(window=25, min_periods=minp, center=True), roll_func
309    )(**kwargs)
310    if fill_value is not None:
311        frame_xp = frame_xp.fillna(fill_value)
312    tm.assert_frame_equal(frame_xp, frame_rs)
313