ENH: Folded the *_pairwise rolling moment functions into the base function API.

Author: Tobias Brandt

pandas/stats/moments.py

@@ -5,30 +5,28 @@
 from __future__ import division
 
 from functools import wraps
+from collections import defaultdict
 
 from numpy import NaN
 import numpy as np
 
 from pandas.core.api import DataFrame, Series, Panel, notnull
 import pandas.algos as algos
-import pandas.core.common as com
-from pandas.core.common import _values_from_object
+import pandas.core.common as pdcom
 
 from pandas.util.decorators import Substitution, Appender
 
 __all__ = ['rolling_count', 'rolling_max', 'rolling_min',
            'rolling_sum', 'rolling_mean', 'rolling_std', 'rolling_cov',
            'rolling_corr', 'rolling_var', 'rolling_skew', 'rolling_kurt',
            'rolling_quantile', 'rolling_median', 'rolling_apply',
-           'rolling_cov_pairwise', 'rolling_corr_pairwise', 'rolling_window',
+           'rolling_corr_pairwise', 'rolling_window',
            'ewma', 'ewmvar', 'ewmstd', 'ewmvol', 'ewmcorr', 'ewmcov',
-           'ewmcorr_pairwise', 'ewmcov_pairwise',
            'expanding_count', 'expanding_max', 'expanding_min',
            'expanding_sum', 'expanding_mean', 'expanding_std',
            'expanding_cov', 'expanding_corr', 'expanding_var',
            'expanding_skew', 'expanding_kurt', 'expanding_quantile',
-           'expanding_median', 'expanding_apply',
-           'expanding_cov_pairwise', 'expanding_corr_pairwise']
+           'expanding_median', 'expanding_apply', 'expanding_corr_pairwise']
 
 #------------------------------------------------------------------------------
 # Docs
@@ -203,25 +201,43 @@ def rolling_count(arg, window, freq=None, center=False, time_rule=None):
 
 @Substitution("Unbiased moving covariance.", _binary_arg_flex, _flex_retval)
 @Appender(_doc_template)
-def rolling_cov(arg1, arg2, window, min_periods=None, freq=None,
-                center=False, time_rule=None):
+def rolling_cov(arg1, arg2=None, window=None, min_periods=None, freq=None,
+                center=False, time_rule=None, pairwise=None):
+    if window is None and isinstance(arg2, (int, float)):
+        window = arg2
+        arg2 = arg1
+        pairwise = True if pairwise is None else pairwise  # only default unset
+    elif arg2 is None:
+        arg2 = arg1
+        pairwise = True if pairwise is None else pairwise  # only default unset
     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,center=center)
-        count = rolling_count(X + Y, window,center=center)
+        mean = lambda x: rolling_mean(x, window, min_periods, center=center)
+        count = rolling_count(X + Y, window, center=center)
         bias_adj = count / (count - 1)
         return (mean(X * Y) - mean(X) * mean(Y)) * bias_adj
-    rs = _flex_binary_moment(arg1, arg2, _get_cov)
+    rs = _flex_binary_moment(arg1, arg2, _get_cov, pairwise=bool(pairwise))
     return rs
 
 
 @Substitution("Moving sample correlation.", _binary_arg_flex, _flex_retval)
 @Appender(_doc_template)
-def rolling_corr(arg1, arg2, window, min_periods=None, freq=None,
-                 center=False, time_rule=None):
+def rolling_corr(arg1, arg2=None, window=None, min_periods=None, freq=None,
+                 center=False, time_rule=None, pairwise=None):
+    if window is None and isinstance(arg2, (int, float)):
+        window = arg2
+        arg2 = arg1
+        pairwise = True if pairwise is None else pairwise  # only default unset
+    elif arg2 is None:
+        arg2 = arg1
+        pairwise = True if pairwise is None else pairwise  # only default unset
+    arg1 = _conv_timerule(arg1, freq, time_rule)
+    arg2 = _conv_timerule(arg2, freq, time_rule)
+    window = min(window, len(arg1), len(arg2))
+
     def _get_corr(a, b):
         num = rolling_cov(a, b, window, min_periods, freq=freq,
                           center=center, time_rule=time_rule)
@@ -230,10 +246,10 @@ def _get_corr(a, b):
                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)
+    return _flex_binary_moment(arg1, arg2, _get_corr, pairwise=bool(pairwise))
 
 
-def _flex_binary_moment(arg1, arg2, f):
+def _flex_binary_moment(arg1, arg2, f, pairwise=False):
     if not (isinstance(arg1,(np.ndarray, Series, DataFrame)) and
             isinstance(arg2,(np.ndarray, Series, DataFrame))):
         raise TypeError("arguments to moment function must be of type "
@@ -249,10 +265,23 @@ def _flex_binary_moment(arg1, arg2, f):
             X, Y = arg1.align(arg2, join='outer')
             X = X + 0 * Y
             Y = Y + 0 * X
-            res_columns = arg1.columns.union(arg2.columns)
-            for col in res_columns:
-                if col in X and col in Y:
-                    results[col] = f(X[col], Y[col])
+            if pairwise is False:
+                res_columns = arg1.columns.union(arg2.columns)
+                for col in res_columns:
+                    if col in X and col in Y:
+                        results[col] = f(X[col], Y[col])
+            elif pairwise is True:
+                results = defaultdict(dict)
+                for i, k1 in enumerate(arg1.columns):
+                    for j, k2 in enumerate(arg2.columns):
+                        if j<i and arg2 is arg1:
+                            # Symmetric case
+                            results[k1][k2] = results[k2][k1]
+                        else:
+                            results[k1][k2] = f(arg1[k1], arg2[k2])
+                return Panel.from_dict(results).swapaxes('items', 'major')
+            else:
+                raise ValueError("'pairwise' is not True/False")
         else:
             res_columns = arg1.columns
             X, Y = arg1.align(arg2, axis=0, join='outer')
@@ -266,53 +295,14 @@ def _flex_binary_moment(arg1, arg2, f):
         return _flex_binary_moment(arg2, arg1, f)
 
 
-def _flex_pairwise_moment(moment_func, df1, df2, **kwargs):
-    from collections import defaultdict
-
-    # Detect symmetry
-    if df2 is df1:
-        symmetric = True
-    else:
-        symmetric = False
-
-    all_results = defaultdict(dict)
-
-    for i, k1 in enumerate(df1.columns):
-        for j, k2 in enumerate(df2.columns):
-            if j<i and symmetric:
-                all_results[k1][k2] = all_results[k2][k1]
-            else:
-                all_results[k1][k2] = moment_func(df1[k1], df2[k2], **kwargs)
-
-    return Panel.from_dict(all_results).swapaxes('items', 'major')
-
-
-@Substitution("Pairwise unbiased moving covariance", _pairwise_arg,
-              _pairwise_retval)
-@Appender(_doc_template)
-def rolling_cov_pairwise(df1, df2, window=None, min_periods=None, freq=None,
-                center=False, time_rule=None):
-    # Try to preserve the previous API
-    if window is None and isinstance(df2, (int, float)):
-        window = df2
-        df2 = df1
-    return _flex_pairwise_moment(rolling_cov, df1, df2, window=window,
-                                 min_periods=min_periods, freq=freq,
-                                 center=center, time_rule=time_rule)
-
-
 @Substitution("Pairwise moving sample correlation", _pairwise_arg,
               _pairwise_retval)
 @Appender(_doc_template)
-def rolling_corr_pairwise(df1, df2, window=None, min_periods=None, freq=None,
-                 center=False, time_rule=None):
-    # Try to preserve the previous API
-    if window is None and isinstance(df2, (int, float)):
-        window = df2
-        df2 = df1
-    return _flex_pairwise_moment(rolling_corr, df1, df2, window=window,
-                                 min_periods=min_periods, freq=freq,
-                                 center=center, time_rule=time_rule)
+def rolling_corr_pairwise(df1, df2=None, window=None, min_periods=None,
+                          freq=None, center=False, time_rule=None):
+    return rolling_corr(df1, df2, window=window, min_periods=min_periods,
+                        freq=freq, center=center, time_rule=time_rule,
+                        pairwise=True)
 
 
 def _rolling_moment(arg, window, func, minp, axis=0, freq=None, center=False,
@@ -361,7 +351,8 @@ def _rolling_moment(arg, window, func, minp, axis=0, freq=None, center=False,
 
 def _center_window(rs, window, axis):
     if axis > rs.ndim-1:
-        raise ValueError("Requested axis is larger then no. of argument dimensions")
+        raise ValueError("Requested axis is larger then no. of argument "
+                         "dimensions")
 
     offset = int((window - 1) / 2.)
     if isinstance(rs, (Series, DataFrame, Panel)):
@@ -480,61 +471,55 @@ def ewmstd(arg, com=None, span=None, halflife=None, min_periods=0, bias=False,
 @Substitution("Exponentially-weighted moving covariance", _binary_arg, "",
               _type_of_input)
 @Appender(_ewm_doc)
-def ewmcov(arg1, arg2, com=None, span=None, halflife=None, min_periods=0, bias=False,
-           freq=None, time_rule=None):
+def ewmcov(arg1, arg2=None, com=None, span=None, halflife=None, min_periods=0, bias=False,
+           freq=None, time_rule=None, pairwise=None):
+    if arg2 is None:
+        arg2 = arg1
+        pairwise = True if pairwise is None else pairwise
+    elif isinstance(arg2, (int, float)) and com is None:
+        com = arg2
+        arg2 = arg1
+        pairwise = True if pairwise is None else pairwise
     arg1 = _conv_timerule(arg1, freq, time_rule)
     arg2 = _conv_timerule(arg2, freq, time_rule)
 
     def _get_ewmcov(X, Y):
         mean = lambda x: ewma(x, com=com, span=span, halflife=halflife, min_periods=min_periods)
         return (mean(X * Y) - mean(X) * mean(Y))
-    result = _flex_binary_moment(arg1, arg2, _get_ewmcov)
+    result = _flex_binary_moment(arg1, arg2, _get_ewmcov,
+                                 pairwise=bool(pairwise))
     if not bias:
         com = _get_center_of_mass(com, span, halflife)
         result *= (1.0 + 2.0 * com) / (2.0 * com)
 
     return result
 
 
-@Substitution("Pairwise exponentially-weighted moving covariance",
-              _pairwise_arg, "", _pairwise_retval)
-@Appender(_ewm_doc)
-def ewmcov_pairwise(df1, df2=None, com=None, span=None, min_periods=0,
-                    bias=False, freq=None, time_rule=None):
-    if df2 is None:
-        df2 = df1
-    return _flex_pairwise_moment(ewmcov, df1, df2, com=com, span=span,
-            min_periods=min_periods, bias=bias, freq=freq, time_rule=time_rule)
-
-
 @Substitution("Exponentially-weighted moving correlation", _binary_arg, "",
               _type_of_input)
 @Appender(_ewm_doc)
-def ewmcorr(arg1, arg2, com=None, span=None, halflife=None, min_periods=0,
-            freq=None, time_rule=None):
+def ewmcorr(arg1, arg2=None, com=None, span=None, halflife=None, min_periods=0,
+            freq=None, time_rule=None, pairwise=None):
+    if arg2 is None:
+        arg2 = arg1
+        pairwise = True if pairwise is None else pairwise
+    elif isinstance(arg2, (int, float)) and com is None:
+        com = arg2
+        arg2 = arg1
+        pairwise = True if pairwise is None else pairwise
     arg1 = _conv_timerule(arg1, freq, time_rule)
     arg2 = _conv_timerule(arg2, freq, time_rule)
 
     def _get_ewmcorr(X, Y):
         mean = lambda x: ewma(x, com=com, span=span, halflife=halflife, min_periods=min_periods)
         var = lambda x: ewmvar(x, com=com, span=span, halflife=halflife, min_periods=min_periods,
-                            bias=True)
+                               bias=True)
         return (mean(X * Y) - mean(X) * mean(Y)) / _zsqrt(var(X) * var(Y))
-    result = _flex_binary_moment(arg1, arg2, _get_ewmcorr)
+    result = _flex_binary_moment(arg1, arg2, _get_ewmcorr,
+                                 pairwise=bool(pairwise))
     return result
 
 
-@Substitution("Pairwise exponentially-weighted moving correlation",
-              _pairwise_arg, "", _pairwise_retval)
-@Appender(_ewm_doc)
-def ewmcorr_pairwise(df1, df2=None, com=None, span=None, min_periods=0,
-                     freq=None, time_rule=None):
-    if df2 is None:
-        df2 = df1
-    return _flex_pairwise_moment(ewmcorr, df1, df2, com=com, span=span,
-            min_periods=min_periods, freq=freq, time_rule=time_rule)
-
-
 def _zsqrt(x):
     result = np.sqrt(x)
     mask = x < 0
@@ -779,8 +764,8 @@ def rolling_window(arg, window=None, win_type=None, min_periods=None,
         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
+        window = pdcom._asarray_tuplesafe(window).astype(float)
+    elif pdcom.is_integer(window):  # window size
         if win_type is None:
             raise ValueError('Must specify window type')
         try:
@@ -928,46 +913,47 @@ def expanding_quantile(arg, quantile, min_periods=1, freq=None,
 
 @Substitution("Unbiased expanding covariance.", _binary_arg_flex, _flex_retval)
 @Appender(_expanding_doc)
-def expanding_cov(arg1, arg2, min_periods=1, freq=None, center=False,
-                  time_rule=None):
+def expanding_cov(arg1, arg2=None, min_periods=1, freq=None, center=False,
+                  time_rule=None, pairwise=None):
+    if arg2 is None:
+        arg2 = arg1
+        pairwise = True if pairwise is None else pairwise
+    elif isinstance(arg2, (int, float)) and min_periods is None:
+        min_periods = arg2
+        arg2 = arg1
+        pairwise = True if pairwise is None else pairwise
     window = max(len(arg1), len(arg2))
     return rolling_cov(arg1, arg2, window,
                        min_periods=min_periods, freq=freq,
-                       center=center, time_rule=time_rule)
-
-
-@Substitution("Pairwise unbiased expanding covariance", _pairwise_arg,
-              _pairwise_retval)
-@Appender(_expanding_doc)
-def expanding_cov_pairwise(df1, df2=None, min_periods=1, freq=None,
-                            center=False, time_rule=None):
-    if df2 is None:
-        df2 = df1
-    return _flex_pairwise_moment(expanding_cov, df1, df2,
-                                 min_periods=min_periods, freq=freq,
-                                 center=center, time_rule=time_rule)
+                       center=center, time_rule=time_rule, pairwise=pairwise)
 
 
 @Substitution("Expanding sample correlation.", _binary_arg_flex, _flex_retval)
 @Appender(_expanding_doc)
-def expanding_corr(arg1, arg2, min_periods=1, freq=None, center=False,
-                   time_rule=None):
+def expanding_corr(arg1, arg2=None, min_periods=1, freq=None, center=False,
+                   time_rule=None, pairwise=None):
+    if arg2 is None:
+        arg2 = arg1
+        pairwise = True if pairwise is None else pairwise
+    elif isinstance(arg2, (int, float)) and min_periods is None:
+        min_periods = arg2
+        arg2 = arg1
+        pairwise = True if pairwise is None else pairwise
     window = max(len(arg1), len(arg2))
     return rolling_corr(arg1, arg2, window,
                         min_periods=min_periods,
-                        freq=freq, center=center, time_rule=time_rule)
+                        freq=freq, center=center, time_rule=time_rule,
+                        pairwise=pairwise)
 
 
 @Substitution("Pairwise expanding sample correlation", _pairwise_arg,
               _pairwise_retval)
 @Appender(_expanding_doc)
 def expanding_corr_pairwise(df1, df2=None, min_periods=1, freq=None,
                             center=False, time_rule=None):
-    if df2 is None:
-        df2 = df1
-    return _flex_pairwise_moment(expanding_corr, df1, df2,
-                                 min_periods=min_periods, freq=freq,
-                                 center=center, time_rule=time_rule)
+    return expanding_corr(df1, df2, min_periods=min_periods,
+                          freq=freq, center=center, time_rule=time_rule,
+                          pairwise=True)
 
 
 def expanding_apply(arg, func, min_periods=1, freq=None, center=False,