diff --git a/pandas/core/generic.py b/pandas/core/generic.py index 3678168890444..a1be2cc470a8d 100644 --- a/pandas/core/generic.py +++ b/pandas/core/generic.py @@ -5568,12 +5568,12 @@ def _add_series_or_dataframe_operations(cls): @Appender(rwindow.rolling.__doc__) def rolling(self, window, min_periods=None, freq=None, center=False, - win_type=None, on=None, axis=0): + win_type=None, on=None, axis=0, closed='right'): axis = self._get_axis_number(axis) return rwindow.rolling(self, window=window, min_periods=min_periods, freq=freq, center=center, win_type=win_type, - on=on, axis=axis) + on=on, axis=axis, closed=closed) cls.rolling = rolling diff --git a/pandas/core/window.py b/pandas/core/window.py index b7276aed506de..99d79e0c45934 100644 --- a/pandas/core/window.py +++ b/pandas/core/window.py @@ -54,11 +54,12 @@ class _Window(PandasObject, SelectionMixin): _attributes = ['window', 'min_periods', 'freq', 'center', 'win_type', - 'axis', 'on'] + 'axis', 'on', 'closed'] exclusions = set() def __init__(self, obj, window=None, min_periods=None, freq=None, - center=False, win_type=None, axis=0, on=None, **kwargs): + center=False, win_type=None, axis=0, on=None, closed='right', + **kwargs): if freq is not None: warnings.warn("The freq kw is deprecated and will be removed in a " @@ -69,6 +70,7 @@ def __init__(self, obj, window=None, min_periods=None, freq=None, self.blocks = [] self.obj = obj self.on = on + self.closed = closed self.window = window self.min_periods = min_periods self.freq = freq @@ -99,6 +101,8 @@ def validate(self): if self.min_periods is not None and not \ is_integer(self.min_periods): raise ValueError("min_periods must be an integer") + if self.closed not in ['right', 'both']: + raise ValueError("closed must be right or both") def _convert_freq(self, how=None): """ resample according to the how, return a new object """ @@ -377,6 +381,12 @@ class Window(_Window): axis : int or string, default 0 + .. versionadded:: 0.19.0 + + closed: 'right' or 'both', default 'right' + For offset-based windows, make the interval closed only on the right + or on both endpoints. + Returns ------- a Window or Rolling sub-classed for the particular operation @@ -455,6 +465,17 @@ class Window(_Window): 2013-01-01 09:00:05 NaN 2013-01-01 09:00:06 4.0 + For time-based windows, it is possible to make the resulting window + contain its left edge by setting closed='both'. + + >>> df.rolling('2s', closed='both').sum() + B + 2013-01-01 09:00:00 0.0 + 2013-01-01 09:00:02 1.0 + 2013-01-01 09:00:03 3.0 + 2013-01-01 09:00:05 2.0 + 2013-01-01 09:00:06 4.0 + Notes ----- By default, the result is set to the right edge of the window. This can be @@ -1047,16 +1068,21 @@ def validate(self): # this will raise ValueError on non-fixed freqs self.window = freq.nanos + if self.closed == 'both': + self.window += 1 self.win_type = 'freq' # min_periods must be an integer if self.min_periods is None: self.min_periods = 1 - - elif not is_integer(self.window): - raise ValueError("window must be an integer") - elif self.window < 0: - raise ValueError("window must be non-negative") + else: + if self.closed == 'both': + raise ValueError("closed=both only valid for datetimelike " + "and offset based windows") + elif not is_integer(self.window): + raise ValueError("window must be an integer") + elif self.window < 0: + raise ValueError("window must be non-negative") @Substitution(name='rolling') @Appender(SelectionMixin._see_also_template) diff --git a/pandas/tests/test_window.py b/pandas/tests/test_window.py index 929ff43bfaaad..a9edb6e2bf12b 100644 --- a/pandas/tests/test_window.py +++ b/pandas/tests/test_window.py @@ -3288,6 +3288,65 @@ def test_min_periods(self): result = df.rolling('2s', min_periods=1).sum() tm.assert_frame_equal(result, expected) + def test_closed(self): + + # closed=both only valid for datetimelike + with self.assertRaises(ValueError): + self.regular.rolling(window=3, closed='both') + + # closed must be 'right' or 'both' + with self.assertRaises(ValueError): + self.regular.rolling(window='1min', closed="'tis wrong") + + df = DataFrame({'B': [1, 1, 2, np.nan, 4]}, + index=[Timestamp('20130101 09:00:00'), + Timestamp('20130101 09:00:02'), + Timestamp('20130101 09:00:03'), + Timestamp('20130101 09:00:05'), + Timestamp('20130101 09:00:06')]) + expected = df.rolling('4s').count() + result = df.rolling('3s', closed='both').count() + tm.assert_frame_equal(result, expected) + expected = df.rolling('3s').count() + result = df.rolling('3s', closed='right').count() + tm.assert_frame_equal(result, expected) + + df.index = df.index - Timestamp('20130101 09:00:00') + expected = df.rolling('4s').count() + result = df.rolling('3s', closed='both').count() + tm.assert_frame_equal(result, expected) + + df = DataFrame({'B': [1, 1, 2, np.nan, 4], + 'timestamp': [Timestamp('20130101 09:00:00'), + Timestamp('20130101 09:00:02'), + Timestamp('20130101 09:00:03'), + Timestamp('20130101 09:00:05'), + Timestamp('20130101 09:00:06')]}) + expected = df.rolling('4s', on='timestamp').count() + result = df.rolling('3s', on='timestamp', closed='both').count() + tm.assert_frame_equal(result, expected) + + df = DataFrame({'B': [1, 1, 2, np.nan, 4]}, + index=[Timestamp('20130101'), + Timestamp('20130103'), + Timestamp('20130104'), + Timestamp('20130106'), + Timestamp('20130107')]) + expected = df.rolling('4d').count() + result = df.rolling('3d', closed='both').count() + tm.assert_frame_equal(result, expected) + + df = DataFrame({'B': [1] * 3}, + index=[Timestamp('20130101 09:00:30'), + Timestamp('20130101 09:01:00'), + Timestamp('20130101 09:02:00')]) + expected = DataFrame({'B': [1., 2., 2.]}, + index=[Timestamp('20130101 09:00:30'), + Timestamp('20130101 09:01:00'), + Timestamp('20130101 09:02:00')]) + result = df.rolling('1min', closed='both').count() + tm.assert_frame_equal(result, expected) + def test_ragged_sum(self): df = self.ragged