initial implementation of min_weight

Author: max-sixty

pandas/core/generic.py

@@ -5175,30 +5175,34 @@ def _add_series_or_dataframe_operations(cls):
         from pandas.core import window as rwindow
 
         @Appender(rwindow.rolling.__doc__)
-        def rolling(self, window, min_periods=None, freq=None, center=False,
-                    win_type=None, axis=0):
+        def rolling(self, window, min_periods=None, min_weight=None,
+                    freq=None, center=False, win_type=None, axis=0):
             axis = self._get_axis_number(axis)
-            return rwindow.rolling(self, window=window,
-                                   min_periods=min_periods, freq=freq,
-                                   center=center, win_type=win_type, axis=axis)
+            return rwindow.rolling(
+                self, window=window, min_periods=min_periods,
+                min_weight=min_weight, freq=freq, center=center,
+                win_type=win_type, axis=axis)
 
         cls.rolling = rolling
 
         @Appender(rwindow.expanding.__doc__)
-        def expanding(self, min_periods=1, freq=None, center=False, axis=0):
+        def expanding(self, min_periods=1, min_weight=None, freq=None,
+                      center=False, axis=0):
             axis = self._get_axis_number(axis)
-            return rwindow.expanding(self, min_periods=min_periods, freq=freq,
+            return rwindow.expanding(self, min_periods=min_periods,
+                                     min_weight=min_weight, freq=freq,
                                      center=center, axis=axis)
 
         cls.expanding = expanding
 
         @Appender(rwindow.ewm.__doc__)
         def ewm(self, com=None, span=None, halflife=None, alpha=None,
-                min_periods=0, freq=None, adjust=True, ignore_na=False,
-                axis=0):
+                min_periods=0, min_weight=None, freq=None, adjust=True,
+                ignore_na=False, axis=0):
             axis = self._get_axis_number(axis)
             return rwindow.ewm(self, com=com, span=span, halflife=halflife,
-                               alpha=alpha, min_periods=min_periods, freq=freq,
+                               alpha=alpha, min_periods=min_periods,
+                               min_weight=min_weight, freq=freq,
                                adjust=adjust, ignore_na=ignore_na, axis=axis)
 
         cls.ewm = ewm

pandas/core/window.py

@@ -36,12 +36,12 @@
 
 
 class _Window(PandasObject, SelectionMixin):
-    _attributes = ['window', 'min_periods', 'freq', 'center', 'win_type',
-                   'axis']
+    _attributes = ['window', 'min_periods', 'min_weight', 'freq', 'center',
+                   'win_type', 'axis']
     exclusions = set()
 
-    def __init__(self, obj, window=None, min_periods=None, freq=None,
-                 center=False, win_type=None, axis=0, **kwargs):
+    def __init__(self, obj, window=None, min_periods=None, min_weight=None,
+                 freq=None, center=False, win_type=None, axis=0, **kwargs):
 
         if freq is not None:
             warnings.warn("The freq kw is deprecated and will be removed in a "
@@ -52,6 +52,7 @@ def __init__(self, obj, window=None, min_periods=None, freq=None,
         self.obj = obj
         self.window = window
         self.min_periods = min_periods
+        self.min_weight = min_weight
         self.freq = freq
         self.center = center
         self.win_type = win_type
@@ -564,7 +565,12 @@ def calc(x):
 
             results.append(result)
 
-        return self._wrap_results(results, blocks, obj)
+        result = self._wrap_results(results, blocks, obj)
+
+        if self.min_weight:
+            result = result.where(_min_weight_mask(self, self.min_weight))
+
+        return result
 
 
 class _Rolling_and_Expanding(_Rolling):
@@ -804,7 +810,7 @@ def _get_corr(a, b):
 
 class Rolling(_Rolling_and_Expanding):
     """
-    Provides rolling window calculcations.
+    Provides rolling window calculations.
 
     .. versionadded:: 0.18.0
 
@@ -981,14 +987,14 @@ class Expanding(_Rolling_and_Expanding):
     of :meth:`~pandas.Series.resample` (i.e. using the `mean`).
     """
 
-    _attributes = ['min_periods', 'freq', 'center', 'axis']
+    _attributes = ['min_periods', 'min_weight', 'freq', 'center', 'axis']
 
-    def __init__(self, obj, min_periods=1, freq=None, center=False, axis=0,
-                 **kwargs):
-        return super(Expanding, self).__init__(obj=obj,
-                                               min_periods=min_periods,
-                                               freq=freq, center=center,
-                                               axis=axis)
+    def __init__(self, obj, min_periods=1, min_weight=None, freq=None,
+                 center=False, axis=0, **kwargs):
+        super(Expanding, self).__init__(
+            obj=obj, min_periods=min_periods,
+            min_weight=min_weight, freq=freq, center=center,
+            axis=axis)
 
     @property
     def _constructor(self):
@@ -1204,14 +1210,16 @@ class EWM(_Rolling):
     More details can be found at
     http://pandas.pydata.org/pandas-docs/stable/computation.html#exponentially-weighted-windows
     """
-    _attributes = ['com', 'min_periods', 'freq', 'adjust', 'ignore_na', 'axis']
+    _attributes = ['com', 'min_periods', 'min_weight', 'freq', 'adjust',
+                   'ignore_na', 'axis']
 
     def __init__(self, obj, com=None, span=None, halflife=None, alpha=None,
-                 min_periods=0, freq=None, adjust=True, ignore_na=False,
-                 axis=0):
+                 min_periods=0, min_weight=None, freq=None, adjust=True,
+                 ignore_na=False, axis=0):
         self.obj = obj
         self.com = _get_center_of_mass(com, span, halflife, alpha)
         self.min_periods = min_periods
+        self.min_weight = min_weight
         self.freq = freq
         self.adjust = adjust
         self.ignore_na = ignore_na
@@ -1271,7 +1279,12 @@ def func(arg):
 
             results.append(np.apply_along_axis(func, self.axis, values))
 
-        return self._wrap_results(results, blocks, obj)
+        result = self._wrap_results(results, blocks, obj)
+
+        if self.min_weight:
+            result = result.where(_min_weight_mask(self, self.min_weight))
+
+        return result
 
     @Substitution(name='ewm')
     @Appender(_doc_template)
@@ -1477,6 +1490,19 @@ def _check_func(minp, window):
     return _check_func
 
 
+def _min_weight_mask(rolling, min_weight):
+    data = rolling.obj
+    valid_data = data.notnull()
+    # fill the object with 1s
+    potential_valid = rolling.obj.copy()
+    potential_valid[:] = 1
+
+    valid_proportion = rolling._shallow_copy(obj=valid_data, min_periods=0,
+                                             min_weight=None).mean()
+
+    return valid_proportion >= min_weight
+
+
 def _use_window(minp, window):
     if minp is None:
         return window

pandas/tests/test_window.py

@@ -19,6 +19,7 @@
 import pandas.core.datetools as datetools
 import pandas.stats.moments as mom
 import pandas.core.window as rwindow
+from pandas.core.window import _min_weight_mask
 from pandas.core.base import SpecificationError
 import pandas.util.testing as tm
 from pandas.compat import range, zip, PY3
@@ -365,7 +366,7 @@ def test_deprecations(self):
 # make sure rolling functions works for different dtypes
 #
 # NOTE that these are yielded tests and so _create_data is
-# explicity called, nor do these inherit from unittest.TestCase
+# explicitly called, nor do these inherit from unittest.TestCase
 #
 # further note that we are only checking rolling for fully dtype
 # compliance (though both expanding and ewm inherit)
@@ -1482,6 +1483,58 @@ def _check_ew_structures(self, func, name):
         frame_result = getattr(self.frame.ewm(com=10), name)()
         self.assertEqual(type(frame_result), DataFrame)
 
+    def test_min_weight_mask_series(self):
+
+        rolling = Series([pd.np.NaN, -8, 3, 10, pd.np.NaN, 5]).rolling(3)
+
+        # 30%
+        result = _min_weight_mask(rolling, 0.3)
+        expected = Series([False, True, True, True, True, True])
+        assert_series_equal(result, expected)
+
+        # 50%
+        result = _min_weight_mask(rolling, 0.6)
+        expected = Series([False, False, True, True, True, True])
+        assert_series_equal(result, expected)
+
+        # 70%
+        result = _min_weight_mask(rolling, 0.7)
+        expected = Series([False, False, False, True, False, False])
+        assert_series_equal(result, expected)
+
+    def test_min_weight_rolling(self):
+
+        series = Series([pd.np.NaN, -8, 3, 10, pd.np.NaN, 5])
+        rolling = series.rolling(3, min_periods=1, min_weight=0.6)
+
+        result = rolling.sum()
+        expected = Series([pd.np.NaN, pd.np.NaN, -5, 5, 13, 15])
+
+        assert_series_equal(result, expected)
+
+    def test_min_weight_expanding(self):
+
+        series = Series([pd.np.NaN, -8, 3, pd.np.NaN, 10, 5])
+        rolling = series.expanding(min_periods=1, min_weight=0.51)
+
+        result = rolling.sum()
+        expected = Series([pd.np.NaN, pd.np.NaN, -5, pd.np.NaN, 5, 10])
+
+        assert_series_equal(result, expected)
+
+    def test_min_weight_ewm(self):
+
+        from itertools import chain
+
+        # create a series with a big gap in the middle
+        series = Series(list(chain(range(9), [pd.np.NaN] * 80, range(9, 0))))
+        rolling = series.ewm(span=10, min_weight=0.5)
+
+        result = rolling.mean()
+
+        # check that all points between 25 and 90 are NaN
+        self.assertTrue(result.iloc[24:89].isnull().all())
+
 
 # create the data only once as we are not setting it
 def _create_consistency_data():