TST/ENH: finish implementing centered moving window functions from scikits.timeseries

Author: changhiskhan

pandas/src/moments.pyx

@@ -952,3 +952,70 @@ def roll_generic(ndarray[float64_t, cast=True] input, int win,
     bufarr.data = <char*> oldbuf
 
     return output
+
+def roll_window(ndarray[float64_t, ndim=1, cast=True] input,
+                ndarray[float64_t, ndim=1, cast=True] weights,
+                int minp, bint avg=True, bint avg_wgt=False):
+    """
+    Assume len(weights) << len(input)
+    """
+    cdef:
+        ndarray[double_t] output, tot_wgt
+        ndarray[int64_t] counts
+        Py_ssize_t in_i, win_i, win_n, win_k, in_n, in_k
+        float64_t val_in, val_win, c, w
+
+    in_n = len(input)
+    win_n = len(weights)
+    output = np.zeros(in_n, dtype=float)
+    counts = np.zeros(in_n, dtype=int)
+    if avg:
+        tot_wgt = np.zeros(in_n, dtype=float)
+
+    minp = _check_minp(len(weights), minp, in_n)
+
+    if avg_wgt:
+        for win_i from 0 <= win_i < win_n:
+            val_win = weights[win_i]
+            if val_win != val_win:
+                continue
+
+            for in_i from 0 <= in_i < in_n - (win_n - win_i) + 1:
+                val_in = input[in_i]
+                if val_in == val_in:
+                    output[in_i + (win_n - win_i) - 1] += val_in * val_win
+                    counts[in_i + (win_n - win_i) - 1] += 1
+                    tot_wgt[in_i + (win_n - win_i) - 1] += val_win
+
+        for in_i from 0 <= in_i < in_n:
+            c = counts[in_i]
+            if c < minp:
+                output[in_i] = NaN
+            else:
+                w = tot_wgt[in_i]
+                if w == 0:
+                    output[in_i] = NaN
+                else:
+                    output[in_i] /= tot_wgt[in_i]
+
+    else:
+        for win_i from 0 <= win_i < win_n:
+            val_win = weights[win_i]
+            if val_win != val_win:
+                continue
+
+            for in_i from 0 <= in_i < in_n - (win_n - win_i) + 1:
+                val_in = input[in_i]
+
+                if val_in == val_in:
+                    output[in_i + (win_n - win_i) - 1] += val_in * val_win
+                    counts[in_i + (win_n - win_i) - 1] += 1
+
+        for in_i from 0 <= in_i < in_n:
+            c = counts[in_i]
+            if c < minp:
+                output[in_i] = NaN
+            elif avg:
+                output[in_i] /= c
+
+    return output

pandas/stats/moments.py

@@ -11,6 +11,7 @@
 
 from pandas.core.api import DataFrame, Series, notnull, Panel
 import pandas.lib as lib
+import pandas.core.common as com
 
 from pandas.util.decorators import Substitution, Appender
 
@@ -572,26 +573,29 @@ 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):
+def rolling_window(arg, window=None, win_type=None, min_periods=None,
+                   freq=None, center=False, mean=True, time_rule=None,
+                   axis=0, **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'
+    arg : Series, DataFrame
+    window : int or ndarray
+        Filtering window specification. If the window is an integer, then it is
+        treated as the window length and win_type is required
+    win_type : str, default None
         Window type (see Notes)
-
     min_periods : int
-        Minimum number of observations in window required to have a value
+        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
+    mean : boolean, default True
+        If True computes weighted mean, else weighted sum
 
     Returns
     -------
@@ -616,26 +620,52 @@ def rolling_window(arg, window, window_type='boxcar', min_periods=None,
     * ``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]
+    if isinstance(window, (list, tuple, np.ndarray)):
+        if win_type is not None:
+            raise ValueError(('Do not specify window type if using custom '
+                              'weights'))
+        window = com._asarray_tuplesafe(window).astype(float)
+    elif com.is_integer(window): #window size
+        if win_type is None:
+            raise ValueError('Must specify window type')
+        import scipy.signal as sig
+        win_type = _validate_win_type(win_type, kwargs) # may pop from kwargs
+        window = sig.get_window(win_type, window).astype(float)
     else:
-        raise ValueError, "Data should be at most 2D"
-    data._mask[:k] = data._mask[-k:] = True
-    return data
+        raise ValueError('Invalid window %s' % str(window))
+
+    minp = _use_window(min_periods, len(window))
+
+    arg = _conv_timerule(arg, freq, time_rule)
+    return_hook, values = _process_data_structure(arg)
+
+    f = lambda x: lib.roll_window(x, window, minp, avg=mean)
+    result = np.apply_along_axis(f, axis, values)
+
+    rs = return_hook(result)
+    if center:
+        rs = _center_window(rs, len(window), axis)
+    return rs
+
+def _validate_win_type(win_type, kwargs):
+    # may pop from kwargs
+    arg_map = {'kaiser' : ['beta'],
+               'gaussian' : ['std'],
+               'general_gaussian' : ['power', 'width'],
+               'slepian' : ['width']}
+    if win_type in arg_map:
+        return tuple([win_type] +
+                     _pop_args(win_type, arg_map[win_type], kwargs))
+    return win_type
+
+def _pop_args(win_type, arg_names, kwargs):
+    msg = '%s window requires %%s' % win_type
+    all_args = []
+    for n in arg_names:
+        if n not in kwargs:
+            raise ValueError(msg % n)
+        all_args.append(kwargs.pop(n))
+    return all_args
 
 
 def _expanding_func(func, desc, check_minp=_use_window):

pandas/stats/tests/test_moments.py

@@ -74,15 +74,66 @@ def test_cmov_window(self):
         vals = np.random.randn(10)
         xp = cmov_window(vals, 5, 'boxcar')
 
-        rs = mom.rolling_window(vals, 5, center=True)
+        rs = mom.rolling_window(vals, 5, 'boxcar', 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)
+        rs = mom.rolling_window(Series(vals), 5, 'boxcar', center=True)
         assert_series_equal(xp, rs)
 
-    def test_cmov_window_types(self):
+    def test_cmov_window_corner(self):
+        try:
+            from scikits.timeseries.lib import cmov_window
+        except ImportError:
+            raise nose.SkipTest
+
+        # all nan
+        vals = np.empty(10, dtype=float)
+        vals.fill(np.nan)
+        rs = mom.rolling_window(vals, 5, 'boxcar', center=True)
+        self.assert_(np.isnan(rs).all())
+
+        # empty
+        vals = np.array([])
+        rs = mom.rolling_window(vals, 5, 'boxcar', center=True)
+        self.assert_(len(rs) == 0)
+
+        # shorter than window
+        vals = np.random.randn(5)
+        rs = mom.rolling_window(vals, 10, 'boxcar')
+        self.assert_(np.isnan(rs).all())
+        self.assert_(len(rs) == 5)
+
+    def test_cmov_window_frame(self):
+        try:
+            from scikits.timeseries.lib import cmov_window
+        except ImportError:
+            raise nose.SkipTest
+
+        # DataFrame
+        vals = np.random.randn(10, 2)
+        xp = cmov_window(vals, 5, 'boxcar')
+        rs = mom.rolling_window(DataFrame(vals), 5, 'boxcar', center=True)
+        assert_frame_equal(DataFrame(xp), rs)
+
+    def test_cmov_window_na_min_periods(self):
+        try:
+            from scikits.timeseries.lib import cmov_window
+        except ImportError:
+            raise nose.SkipTest
+
+        # min_periods
+        vals = Series(np.random.randn(10))
+        vals[4] = np.nan
+        vals[8] = np.nan
+
+        xp = mom.rolling_mean(vals, 5, min_periods=4, center=True)
+        rs = mom.rolling_window(vals, 5, 'boxcar', min_periods=4, center=True)
+
+        assert_series_equal(xp, rs)
+
+    def test_cmov_window_regular(self):
         try:
             from scikits.timeseries.lib import cmov_window
         except ImportError:
@@ -94,28 +145,26 @@ def test_cmov_window_types(self):
             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))
+            rs = mom.rolling_window(Series(vals), 5, wt, center=True)
+            assert_series_equal(Series(xp), rs)
 
-    def test_cmov_special(self):
+    def test_cmov_window_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.}]
+                {'width' : 0.5}]
 
         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,
+            rs = mom.rolling_window(Series(vals), 5, wt, center=True,
                                     **k)
-            assert_almost_equal(xp.compressed(), rs[2:-2])
-            assert_almost_equal(xp.mask, np.isnan(rs))
+            assert_series_equal(Series(xp), rs)
 
     def test_rolling_median(self):
         self._check_moment_func(mom.rolling_median, np.median)