tweaks

Author: changhiskhan

pandas/stats/moments.py

@@ -18,7 +18,7 @@
            'rolling_sum', 'rolling_mean', 'rolling_std', 'rolling_cov',
            'rolling_corr', 'rolling_var', 'rolling_skew', 'rolling_kurt',
            'rolling_quantile', 'rolling_median', 'rolling_apply',
-           'rolling_corr_pairwise',
+           'rolling_corr_pairwise', 'rolling_window',
            'ewma', 'ewmvar', 'ewmstd', 'ewmvol', 'ewmcorr', 'ewmcov',
            'expanding_count', 'expanding_max', 'expanding_min',
            'expanding_sum', 'expanding_mean', 'expanding_std',
@@ -285,11 +285,10 @@ def _rolling_moment(arg, window, func, minp, axis=0, freq=None,
     return rs
 
 def _center_window(rs, window, axis):
+    offset = int((window - 1) / 2.)
     if isinstance(rs, (Series, DataFrame, Panel)):
-        rs = rs.shift(-int((window + 1) / 2.), axis=axis)
+        rs = rs.shift(-offset, axis=axis)
     else:
-        offset = int((window + 1) / 2.)
-
         rs_indexer = [slice(None)] * rs.ndim
         rs_indexer[axis] = slice(None, -offset)
 
@@ -573,6 +572,71 @@ def call_cython(arg, window, minp):
     return _rolling_moment(arg, window, call_cython, min_periods,
                            freq=freq, center=center, time_rule=time_rule)
 
+def rolling_window(arg, window, window_type='boxcar', min_periods=None,
+                   freq=None, center=False, time_rule=None, **kwargs):
+    """
+    Applies a centered moving window of type ``window_type`` and size ``window``
+    on the data.
+
+    Parameters
+    ----------
+    data : Series, DataFrame
+    window : int
+        Size of the filtering window.
+    window_type : str, default 'boxcar'
+        Window type (see Notes)
+
+    min_periods : int
+        Minimum number of observations in window required to have a value
+    freq : None or string alias / date offset object, default=None
+        Frequency to conform to before computing statistic
+    center : boolean, default False
+        Whether the label should correspond with center of window
+
+    Returns
+    -------
+    y : type of input argument
+
+    Notes
+    -----
+    The recognized window types are:
+
+    * ``boxcar``
+    * ``triang``
+    * ``blackman``
+    * ``hamming``
+    * ``bartlett``
+    * ``parzen``
+    * ``bohman``
+    * ``blackmanharris``
+    * ``nuttall``
+    * ``barthann``
+    * ``kaiser`` (needs beta)
+    * ``gaussian`` (needs std)
+    * ``general_gaussian`` (needs power, width)
+    * ``slepian`` (needs width).
+    """
+    from scipy.signal import convolve, get_window
+
+    data = marray(data, copy=True, subok=True)
+    if data._mask is nomask:
+        data._mask = np.zeros(data.shape, bool_)
+    window = get_window(window_type, span, fftbins=False)
+    (n, k) = (len(data), span//2)
+    #
+    if data.ndim == 1:
+        data._data.flat = convolve(data._data, window)[k:n+k] / float(span)
+        data._mask[:] = ((convolve(getmaskarray(data), window) > 0)[k:n+k])
+    elif data.ndim == 2:
+        for i in range(data.shape[-1]):
+            _data = data._data[:,i]
+            _data.flat = convolve(_data, window)[k:n+k] / float(span)
+            data._mask[:,i] = (convolve(data._mask[:,i], window) > 0)[k:n+k]
+    else:
+        raise ValueError, "Data should be at most 2D"
+    data._mask[:k] = data._mask[-k:] = True
+    return data
+
 
 def _expanding_func(func, desc, check_minp=_use_window):
     @Substitution(desc, _unary_arg, _type_of_input)

pandas/stats/tests/test_moments.py

@@ -48,6 +48,75 @@ def test_rolling_count(self):
     def test_rolling_mean(self):
         self._check_moment_func(mom.rolling_mean, np.mean)
 
+    def test_cmov_mean(self):
+        try:
+            from scikits.timeseries.lib import cmov_mean
+        except ImportError:
+            raise nose.SkipTest
+
+        vals = np.random.randn(10)
+        xp = cmov_mean(vals, 5)
+
+        rs = mom.rolling_mean(vals, 5, center=True)
+        assert_almost_equal(xp.compressed(), rs[2:-2])
+        assert_almost_equal(xp.mask, np.isnan(rs))
+
+        xp = Series(rs)
+        rs = mom.rolling_mean(Series(vals), 5, center=True)
+        assert_series_equal(xp, rs)
+
+    def test_cmov_window(self):
+        try:
+            from scikits.timeseries.lib import cmov_window
+        except ImportError:
+            raise nose.SkipTest
+
+        vals = np.random.randn(10)
+        xp = cmov_window(vals, 5, 'boxcar')
+
+        rs = mom.rolling_window(vals, 5, center=True)
+        assert_almost_equal(xp.compressed(), rs[2:-2])
+        assert_almost_equal(xp.mask, np.isnan(rs))
+
+        xp = Series(rs)
+        rs = mom.rolling_window(Series(vals), 5, center=True)
+        assert_series_equal(xp, rs)
+
+    def test_cmov_window_types(self):
+        try:
+            from scikits.timeseries.lib import cmov_window
+        except ImportError:
+            raise nose.SkipTest
+
+        win_types = ['triang', 'blackman', 'hamming', 'bartlett', 'bohman',
+                     'blackmanharris', 'nuttall', 'barthann']
+        for wt in win_types:
+            vals = np.random.randn(10)
+            xp = cmov_window(vals, 5, wt)
+
+            rs = mom.rolling_window(vals, 5, window_type=wt, center=True)
+            assert_almost_equal(xp.compressed(), rs[2:-2])
+            assert_almost_equal(xp.mask, np.isnan(rs))
+
+    def test_cmov_special(self):
+        try:
+            from scikits.timeseries.lib import cmov_window
+        except ImportError:
+            raise nose.SkipTest
+
+        win_types = ['kaiser', 'gaussian', 'general_gaussian', 'slepian']
+        kwds = [{'beta' : 1.}, {'std' : 1.}, {'power' : 2., 'width' : 2.},
+                {'width' : 2.}]
+
+        for wt, k in zip(win_types, kwds):
+            vals = np.random.randn(10)
+            xp = cmov_window(vals, 5, (wt,) + tuple(k.values()))
+
+            rs = mom.rolling_window(vals, 5, window_type=wt, center=True,
+                                    **k)
+            assert_almost_equal(xp.compressed(), rs[2:-2])
+            assert_almost_equal(xp.mask, np.isnan(rs))
+
     def test_rolling_median(self):
         self._check_moment_func(mom.rolling_median, np.median)
 
@@ -259,16 +328,16 @@ def _check_ndarray(self, func, static_comp, window=50,
                 result = func(arr, 20, center=True)
                 expected = func(arr, 20)
 
-            assert_almost_equal(result[0], expected[10])
+            assert_almost_equal(result[1], expected[10])
             if fill_value is None:
-                self.assert_(np.isnan(result[-10:]).all())
+                self.assert_(np.isnan(result[-9:]).all())
             else:
-                self.assert_((result[-10:] == 0).all())
+                self.assert_((result[-9:] == 0).all())
             if has_min_periods:
                 self.assert_(np.isnan(expected[23]))
-                self.assert_(np.isnan(result[13]))
+                self.assert_(np.isnan(result[14]))
                 self.assert_(np.isnan(expected[-5]))
-                self.assert_(np.isnan(result[-15]))
+                self.assert_(np.isnan(result[-14]))
 
     def _check_structures(self, func, static_comp,
                           has_min_periods=True, has_time_rule=True,
@@ -309,17 +378,17 @@ def _check_structures(self, func, static_comp,
         if has_center:
             if has_min_periods:
                 minp = 10
-                series_xp = func(self.series, 25, min_periods=minp).shift(-13)
-                frame_xp = func(self.frame, 25, min_periods=minp).shift(-13)
+                series_xp = func(self.series, 25, min_periods=minp).shift(-12)
+                frame_xp = func(self.frame, 25, min_periods=minp).shift(-12)
 
                 series_rs = func(self.series, 25, min_periods=minp,
                                  center=True)
                 frame_rs = func(self.frame, 25, min_periods=minp,
                                 center=True)
 
             else:
-                series_xp = func(self.series, 25).shift(-13)
-                frame_xp = func(self.frame, 25).shift(-13)
+                series_xp = func(self.series, 25).shift(-12)
+                frame_xp = func(self.frame, 25).shift(-12)
 
                 series_rs = func(self.series, 25, center=True)
                 frame_rs = func(self.frame, 25, center=True)