From 24ac32f4f103b6a886d56c149a9b2c68252c9f07 Mon Sep 17 00:00:00 2001 From: Matt Roeschke Date: Thu, 1 Oct 2020 19:32:08 -0700 Subject: [PATCH 1/2] CLN: Use more pytest idioms in test_momemts_ewm.py --- .../tests/window/moments/test_moments_ewm.py | 195 +++++++++++++----- 1 file changed, 144 insertions(+), 51 deletions(-) diff --git a/pandas/tests/window/moments/test_moments_ewm.py b/pandas/tests/window/moments/test_moments_ewm.py index a83bfabc4a048..797cf122ce990 100644 --- a/pandas/tests/window/moments/test_moments_ewm.py +++ b/pandas/tests/window/moments/test_moments_ewm.py @@ -7,21 +7,15 @@ import pandas._testing as tm -def check_ew(name=None, preserve_nan=False, series=None, frame=None, nan_locs=None): - series_result = getattr(series.ewm(com=10), name)() +def test_ewma(series, frame): + series_result = series.ewm(com=10).mean() assert isinstance(series_result, Series) - frame_result = getattr(frame.ewm(com=10), name)() + frame_result = frame.ewm(com=10).mean() assert type(frame_result) == DataFrame - result = getattr(series.ewm(com=10), name)() - if preserve_nan: - assert result[nan_locs].isna().all() - - -def test_ewma(series, frame, nan_locs): - check_ew(name="mean", frame=frame, series=series, nan_locs=nan_locs) +def test_ewma_adjust(): vals = pd.Series(np.zeros(1000)) vals[5] = 1 result = vals.ewm(span=100, adjust=False).mean().sum() @@ -53,63 +47,162 @@ def test_ewma_nan_handling(): result = s.ewm(com=5).mean() tm.assert_series_equal(result, Series([np.nan] * 2 + [1.0] * 4)) - # GH 7603 - s0 = Series([np.nan, 1.0, 101.0]) - s1 = Series([1.0, np.nan, 101.0]) - s2 = Series([np.nan, 1.0, np.nan, np.nan, 101.0, np.nan]) - s3 = Series([1.0, np.nan, 101.0, 50.0]) - com = 2.0 - alpha = 1.0 / (1.0 + com) - - def simple_wma(s, w): - return (s.multiply(w).cumsum() / w.cumsum()).fillna(method="ffill") - - for (s, adjust, ignore_na, w) in [ - (s0, True, False, [np.nan, (1.0 - alpha), 1.0]), - (s0, True, True, [np.nan, (1.0 - alpha), 1.0]), - (s0, False, False, [np.nan, (1.0 - alpha), alpha]), - (s0, False, True, [np.nan, (1.0 - alpha), alpha]), - (s1, True, False, [(1.0 - alpha) ** 2, np.nan, 1.0]), - (s1, True, True, [(1.0 - alpha), np.nan, 1.0]), - (s1, False, False, [(1.0 - alpha) ** 2, np.nan, alpha]), - (s1, False, True, [(1.0 - alpha), np.nan, alpha]), - (s2, True, False, [np.nan, (1.0 - alpha) ** 3, np.nan, np.nan, 1.0, np.nan]), - (s2, True, True, [np.nan, (1.0 - alpha), np.nan, np.nan, 1.0, np.nan]), + +@pytest.mark.parametrize( + "s, adjust, ignore_na, w", + [ + ( + Series([np.nan, 1.0, 101.0]), + True, + False, + [np.nan, (1.0 - (1.0 / (1.0 + 2.0))), 1.0], + ), + ( + Series([np.nan, 1.0, 101.0]), + True, + True, + [np.nan, (1.0 - (1.0 / (1.0 + 2.0))), 1.0], + ), + ( + Series([np.nan, 1.0, 101.0]), + False, + False, + [np.nan, (1.0 - (1.0 / (1.0 + 2.0))), (1.0 / (1.0 + 2.0))], + ), + ( + Series([np.nan, 1.0, 101.0]), + False, + True, + [np.nan, (1.0 - (1.0 / (1.0 + 2.0))), (1.0 / (1.0 + 2.0))], + ), + ( + Series([1.0, np.nan, 101.0]), + True, + False, + [(1.0 - (1.0 / (1.0 + 2.0))) ** 2, np.nan, 1.0], + ), ( - s2, + Series([1.0, np.nan, 101.0]), + True, + True, + [(1.0 - (1.0 / (1.0 + 2.0))), np.nan, 1.0], + ), + ( + Series([1.0, np.nan, 101.0]), False, False, - [np.nan, (1.0 - alpha) ** 3, np.nan, np.nan, alpha, np.nan], + [(1.0 - (1.0 / (1.0 + 2.0))) ** 2, np.nan, (1.0 / (1.0 + 2.0))], ), - (s2, False, True, [np.nan, (1.0 - alpha), np.nan, np.nan, alpha, np.nan]), - (s3, True, False, [(1.0 - alpha) ** 3, np.nan, (1.0 - alpha), 1.0]), - (s3, True, True, [(1.0 - alpha) ** 2, np.nan, (1.0 - alpha), 1.0]), ( - s3, + Series([1.0, np.nan, 101.0]), + False, + True, + [(1.0 - (1.0 / (1.0 + 2.0))), np.nan, (1.0 / (1.0 + 2.0))], + ), + ( + Series([np.nan, 1.0, np.nan, np.nan, 101.0, np.nan]), + True, + False, + [np.nan, (1.0 - (1.0 / (1.0 + 2.0))) ** 3, np.nan, np.nan, 1.0, np.nan], + ), + ( + Series([np.nan, 1.0, np.nan, np.nan, 101.0, np.nan]), + True, + True, + [np.nan, (1.0 - (1.0 / (1.0 + 2.0))), np.nan, np.nan, 1.0, np.nan], + ), + ( + Series([np.nan, 1.0, np.nan, np.nan, 101.0, np.nan]), False, False, [ - (1.0 - alpha) ** 3, np.nan, - (1.0 - alpha) * alpha, - alpha * ((1.0 - alpha) ** 2 + alpha), + (1.0 - (1.0 / (1.0 + 2.0))) ** 3, + np.nan, + np.nan, + (1.0 / (1.0 + 2.0)), + np.nan, ], ), - (s3, False, True, [(1.0 - alpha) ** 2, np.nan, (1.0 - alpha) * alpha, alpha]), - ]: - expected = simple_wma(s, Series(w)) - result = s.ewm(com=com, adjust=adjust, ignore_na=ignore_na).mean() + ( + Series([np.nan, 1.0, np.nan, np.nan, 101.0, np.nan]), + False, + True, + [ + np.nan, + (1.0 - (1.0 / (1.0 + 2.0))), + np.nan, + np.nan, + (1.0 / (1.0 + 2.0)), + np.nan, + ], + ), + ( + Series([1.0, np.nan, 101.0, 50.0]), + True, + False, + [ + (1.0 - (1.0 / (1.0 + 2.0))) ** 3, + np.nan, + (1.0 - (1.0 / (1.0 + 2.0))), + 1.0, + ], + ), + ( + Series([1.0, np.nan, 101.0, 50.0]), + True, + True, + [ + (1.0 - (1.0 / (1.0 + 2.0))) ** 2, + np.nan, + (1.0 - (1.0 / (1.0 + 2.0))), + 1.0, + ], + ), + ( + Series([1.0, np.nan, 101.0, 50.0]), + False, + False, + [ + (1.0 - (1.0 / (1.0 + 2.0))) ** 3, + np.nan, + (1.0 - (1.0 / (1.0 + 2.0))) * (1.0 / (1.0 + 2.0)), + (1.0 / (1.0 + 2.0)) + * ((1.0 - (1.0 / (1.0 + 2.0))) ** 2 + (1.0 / (1.0 + 2.0))), + ], + ), + ( + Series([1.0, np.nan, 101.0, 50.0]), + False, + True, + [ + (1.0 - (1.0 / (1.0 + 2.0))) ** 2, + np.nan, + (1.0 - (1.0 / (1.0 + 2.0))) * (1.0 / (1.0 + 2.0)), + (1.0 / (1.0 + 2.0)), + ], + ), + ], +) +def test_ewma_nan_handling_cases(s, adjust, ignore_na, w): + # GH 7603 + expected = (s.multiply(w).cumsum() / Series(w).cumsum()).fillna(method="ffill") + result = s.ewm(com=2.0, adjust=adjust, ignore_na=ignore_na).mean() + tm.assert_series_equal(result, expected) + if ignore_na is False: + # check that ignore_na defaults to False + result = s.ewm(com=2.0, adjust=adjust).mean() tm.assert_series_equal(result, expected) - if ignore_na is False: - # check that ignore_na defaults to False - result = s.ewm(com=com, adjust=adjust).mean() - tm.assert_series_equal(result, expected) @pytest.mark.parametrize("name", ["var", "vol"]) -def test_ewmvar_ewmvol(series, frame, nan_locs, name): - check_ew(name=name, frame=frame, series=series, nan_locs=nan_locs) +def test_ewmvar_ewmvol(series, frame, name): + series_result = getattr(series.ewm(com=10), name)() + assert isinstance(series_result, Series) + + frame_result = getattr(frame.ewm(com=10), name)() + assert type(frame_result) == DataFrame def test_ewma_span_com_args(series): From 1bab68e1bed654c794d7952dd54669296746982a Mon Sep 17 00:00:00 2001 From: Matt Roeschke Date: Fri, 2 Oct 2020 10:50:50 -0700 Subject: [PATCH 2/2] Address comments --- .../tests/window/moments/test_moments_ewm.py | 21 +++++++------------ 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/pandas/tests/window/moments/test_moments_ewm.py b/pandas/tests/window/moments/test_moments_ewm.py index 797cf122ce990..287cd7ebba536 100644 --- a/pandas/tests/window/moments/test_moments_ewm.py +++ b/pandas/tests/window/moments/test_moments_ewm.py @@ -7,12 +7,16 @@ import pandas._testing as tm -def test_ewma(series, frame): - series_result = series.ewm(com=10).mean() +@pytest.mark.parametrize("name", ["var", "vol", "mean"]) +def test_ewma_series(series, name): + series_result = getattr(series.ewm(com=10), name)() assert isinstance(series_result, Series) - frame_result = frame.ewm(com=10).mean() - assert type(frame_result) == DataFrame + +@pytest.mark.parametrize("name", ["var", "vol", "mean"]) +def test_ewma_frame(frame, name): + frame_result = getattr(frame.ewm(com=10), name)() + assert isinstance(frame_result, DataFrame) def test_ewma_adjust(): @@ -196,15 +200,6 @@ def test_ewma_nan_handling_cases(s, adjust, ignore_na, w): tm.assert_series_equal(result, expected) -@pytest.mark.parametrize("name", ["var", "vol"]) -def test_ewmvar_ewmvol(series, frame, name): - series_result = getattr(series.ewm(com=10), name)() - assert isinstance(series_result, Series) - - frame_result = getattr(frame.ewm(com=10), name)() - assert type(frame_result) == DataFrame - - def test_ewma_span_com_args(series): A = series.ewm(com=9.5).mean() B = series.ewm(span=20).mean()