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