ENH: centered moving window functions #1270

Author: changhiskhan

pandas/stats/moments.py

@@ -9,7 +9,7 @@
 from numpy import NaN
 import numpy as np
 
-from pandas.core.api import DataFrame, Series, notnull
+from pandas.core.api import DataFrame, Series, notnull, Panel
 import pandas.lib as lib
 
 from pandas.util.decorators import Substitution, Appender
@@ -124,7 +124,7 @@
 """
 
 
-def rolling_count(arg, window, freq=None, time_rule=None):
+def rolling_count(arg, window, freq=None, center=False, time_rule=None):
     """
     Rolling count of number of non-NaN observations inside provided window.
 
@@ -134,6 +134,8 @@ def rolling_count(arg, window, freq=None, time_rule=None):
     window : Number of observations used for calculating statistic
     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
     -------
@@ -146,7 +148,7 @@ def rolling_count(arg, window, freq=None, time_rule=None):
 
     converted = np.isfinite(values).astype(float)
     result = rolling_sum(converted, window, min_periods=1,
-                         time_rule=time_rule)
+                         center=center) # already converted
 
     # putmask here?
     result[np.isnan(result)] = 0
@@ -156,22 +158,37 @@ def rolling_count(arg, window, freq=None, time_rule=None):
 
 @Substitution("Unbiased moving covariance", _binary_arg_flex, _flex_retval)
 @Appender(_doc_template)
-def rolling_cov(arg1, arg2, window, min_periods=None, time_rule=None):
+def rolling_cov(arg1, arg2, window, min_periods=None, freq=None,
+                center=False, time_rule=None):
+    arg1 = _conv_timerule(arg1, freq, time_rule)
+    arg2 = _conv_timerule(arg2, freq, time_rule)
+    window = min(window, len(arg1), len(arg2))
     def _get_cov(X, Y):
-        mean = lambda x: rolling_mean(x, window, min_periods, time_rule)
-        count = rolling_count(X + Y, window, time_rule)
+        mean = lambda x: rolling_mean(x, window, min_periods)
+        count = rolling_count(X + Y, window)
         bias_adj = count / (count - 1)
         return (mean(X * Y) - mean(X) * mean(Y)) * bias_adj
-    return _flex_binary_moment(arg1, arg2, _get_cov)
-
+    rs = _flex_binary_moment(arg1, arg2, _get_cov)
+    if center:
+        if isinstance(rs, (Series, DataFrame, Panel)):
+            rs = rs.shift(-int((window + 1) / 2.))
+        else:
+            offset = int((window + 1) / 2.)
+            rs[:-offset] = rs[offset:]
+            rs[-offset:] = np.nan
+    return rs
 
 @Substitution("Moving sample correlation", _binary_arg_flex, _flex_retval)
 @Appender(_doc_template)
-def rolling_corr(arg1, arg2, window, min_periods=None, time_rule=None):
+def rolling_corr(arg1, arg2, window, min_periods=None, freq=None,
+                 center=False, time_rule=None):
     def _get_corr(a, b):
-        num = rolling_cov(a, b, window, min_periods, time_rule)
-        den = (rolling_std(a, window, min_periods, time_rule) *
-                rolling_std(b, window, min_periods, time_rule))
+        num = rolling_cov(a, b, window, min_periods, freq=freq,
+                          center=center, time_rule=time_rule)
+        den = (rolling_std(a, window, min_periods, freq=freq,
+                           center=center, time_rule=time_rule) *
+                rolling_std(b, window, min_periods, freq=freq,
+                            center=center, time_rule=time_rule))
         return num / den
     return _flex_binary_moment(arg1, arg2, _get_corr)
 
@@ -234,7 +251,7 @@ def rolling_corr_pairwise(df, window, min_periods=None):
 
 
 def _rolling_moment(arg, window, func, minp, axis=0, freq=None,
-                    time_rule=None, **kwargs):
+                    center=False, time_rule=None, **kwargs):
     """
     Rolling statistical measure using supplied function. Designed to be
     used with passed-in Cython array-based functions.
@@ -249,6 +266,8 @@ def _rolling_moment(arg, window, func, minp, axis=0, freq=None,
     axis : int, default 0
     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
     -------
@@ -260,8 +279,16 @@ def _rolling_moment(arg, window, func, minp, axis=0, freq=None,
     # actually calculate the moment. Faster way to do this?
     result = np.apply_along_axis(calc, axis, values)
 
-    return return_hook(result)
+    rs = return_hook(result)
+    if center:
+        if isinstance(rs, (Series, DataFrame, Panel)):
+            rs = rs.shift(-int((window + 1) / 2.), axis=axis)
+        else:
+            offset = int((window + 1)/ 2.)
+            rs[:-offset] = rs[offset:]
+            rs[-offset:] = np.nan
 
+    return rs
 
 def _process_data_structure(arg, kill_inf=True):
     if isinstance(arg, DataFrame):
@@ -450,12 +477,14 @@ def _rolling_func(func, desc, check_minp=_use_window):
     @Substitution(desc, _unary_arg, _type_of_input)
     @Appender(_doc_template)
     @wraps(func)
-    def f(arg, window, min_periods=None, freq=None, time_rule=None, **kwargs):
+    def f(arg, window, min_periods=None, freq=None, center=False,
+          time_rule=None, **kwargs):
         def call_cython(arg, window, minp, **kwds):
             minp = check_minp(minp, window)
             return func(arg, window, minp, **kwds)
         return _rolling_moment(arg, window, call_cython, min_periods,
-                               freq=freq, time_rule=time_rule, **kwargs)
+                               freq=freq, center=center,
+                               time_rule=time_rule, **kwargs)
 
     return f
 
@@ -477,7 +506,7 @@ def call_cython(arg, window, minp, **kwds):
 
 
 def rolling_quantile(arg, window, quantile, min_periods=None, freq=None,
-                     time_rule=None):
+                     center=False, time_rule=None):
     """Moving quantile
 
     Parameters
@@ -489,6 +518,8 @@ def rolling_quantile(arg, window, quantile, min_periods=None, freq=None,
         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
     -------
@@ -499,11 +530,11 @@ def call_cython(arg, window, minp):
         minp = _use_window(minp, window)
         return lib.roll_quantile(arg, window, minp, quantile)
     return _rolling_moment(arg, window, call_cython, min_periods,
-                           freq=freq, time_rule=time_rule)
+                           freq=freq, center=center, time_rule=time_rule)
 
 
 def rolling_apply(arg, window, func, min_periods=None, freq=None,
-                  time_rule=None):
+                  center=False, time_rule=None):
     """Generic moving function application
 
     Parameters
@@ -516,6 +547,8 @@ def rolling_apply(arg, window, func, min_periods=None, freq=None,
         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
     -------
@@ -525,21 +558,23 @@ def call_cython(arg, window, minp):
         minp = _use_window(minp, window)
         return lib.roll_generic(arg, window, minp, func)
     return _rolling_moment(arg, window, call_cython, min_periods,
-                           freq=freq, time_rule=time_rule)
+                           freq=freq, center=center, time_rule=time_rule)
 
 
 def _expanding_func(func, desc, check_minp=_use_window):
     @Substitution(desc, _unary_arg, _type_of_input)
     @Appender(_expanding_doc)
     @wraps(func)
-    def f(arg, min_periods=1, freq=None, time_rule=None, **kwargs):
+    def f(arg, min_periods=1, freq=None, center=False, time_rule=None,
+          **kwargs):
         window = len(arg)
 
         def call_cython(arg, window, minp, **kwds):
             minp = check_minp(minp, window)
             return func(arg, window, minp, **kwds)
         return _rolling_moment(arg, window, call_cython, min_periods,
-                               freq=freq, time_rule=time_rule, **kwargs)
+                               freq=freq, center=center,
+                               time_rule=time_rule, **kwargs)
 
     return f
 
@@ -560,7 +595,7 @@ def call_cython(arg, window, minp, **kwds):
                              check_minp=_require_min_periods(4))
 
 
-def expanding_count(arg, freq=None, time_rule=None):
+def expanding_count(arg, freq=None, center=False, time_rule=None):
     """
     Expanding count of number of non-NaN observations.
 
@@ -569,16 +604,19 @@ def expanding_count(arg, freq=None, time_rule=None):
     arg :  DataFrame or numpy ndarray-like
     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
     -------
     expanding_count : type of caller
     """
-    return rolling_count(arg, len(arg), freq=freq, time_rule=time_rule)
+    return rolling_count(arg, len(arg), freq=freq, center=center,
+                         time_rule=time_rule)
 
 
 def expanding_quantile(arg, quantile, min_periods=1, freq=None,
-                     time_rule=None):
+                       center=False, time_rule=None):
     """Expanding quantile
 
     Parameters
@@ -589,29 +627,35 @@ def expanding_quantile(arg, quantile, min_periods=1, freq=None,
         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
     """
     return rolling_quantile(arg, len(arg), quantile, min_periods=min_periods,
-                            freq=freq, time_rule=time_rule)
+                            freq=freq, center=center, time_rule=time_rule)
 
 
 @Substitution("Unbiased expanding covariance", _binary_arg_flex, _flex_retval)
 @Appender(_expanding_doc)
-def expanding_cov(arg1, arg2, min_periods=1, time_rule=None):
+def expanding_cov(arg1, arg2, min_periods=1, freq=None, center=False,
+                  time_rule=None):
     window = max(len(arg1), len(arg2))
     return rolling_cov(arg1, arg2, window,
-                       min_periods=min_periods, time_rule=time_rule)
+                       min_periods=min_periods, freq=freq,
+                       center=center, time_rule=time_rule)
 
 
 @Substitution("Expanding sample correlation", _binary_arg_flex, _flex_retval)
 @Appender(_expanding_doc)
-def expanding_corr(arg1, arg2, min_periods=1, time_rule=None):
+def expanding_corr(arg1, arg2, min_periods=1, freq=None, center=False,
+                   time_rule=None):
     window = max(len(arg1), len(arg2))
     return rolling_corr(arg1, arg2, window,
-                        min_periods=min_periods, time_rule=time_rule)
+                        min_periods=min_periods,
+                        freq=freq, center=center, time_rule=time_rule)
 
 
 def expanding_corr_pairwise(df, min_periods=1):
@@ -634,7 +678,8 @@ def expanding_corr_pairwise(df, min_periods=1):
     return rolling_corr_pairwise(df, window, min_periods=min_periods)
 
 
-def expanding_apply(arg, func, min_periods=1, freq=None, time_rule=None):
+def expanding_apply(arg, func, min_periods=1, freq=None, center=False,
+                    time_rule=None):
     """Generic expanding function application
 
     Parameters
@@ -646,11 +691,13 @@ def expanding_apply(arg, func, min_periods=1, freq=None, time_rule=None):
         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
     """
     window = len(arg)
     return rolling_apply(arg, window, func, min_periods=min_periods, freq=freq,
-                         time_rule=time_rule)
+                         center=center, time_rule=time_rule)