在Pandas中使用rolling_apply与需要两个参数的函数

Question

我正在尝试使用rollapply这个函数，它需要两个参数。根据我的了解，计算kendall tau相关性（包括标准的平局修正）唯一的方法就是：

>>> import scipy
>>> x = [5.05, 6.75, 3.21, 2.66]
>>> y = [1.65, 26.5, -5.93, 7.96]
>>> z = [1.65, 2.64, 2.64, 6.95]
>>> print scipy.stats.stats.kendalltau(x, y)[0]
0.333333333333

我也知道rollapply在处理两个参数时会遇到问题，这在这里有详细说明：

• 相关问题 1
• Github 问题
• 相关问题 2

不过，我还是在努力寻找一种方法，在一个包含多个列的数据框上进行滚动的kendalltau计算。

我的数据框大概是这样的：

A = pd.DataFrame([[1, 5, 1], [2, 4, 1], [3, 3, 1], [4, 2, 1], [5, 1, 1]], 
                 columns=['A', 'B', 'C'], index = [1, 2, 3, 4, 5])

我想创建一个可以做到这一点的函数：

In [1]:function(A, 3)  # A is df, 3 is the rolling window
Out[2]:
   A  B  C     AB     AC     BC  
1  1  5  2    NaN    NaN    NaN
2  2  4  4    NaN    NaN    NaN
3  3  3  1  -0.99  -0.33   0.33
4  4  2  2  -0.99  -0.33   0.33
5  5  1  4  -0.99   0.99  -0.99

在一个非常初步的尝试中，我考虑将函数定义成这样：

def tau1(x):
    y = np.array(A['A']) #  keep one column fix and run it in the other two
    tau, p_value = sp.stats.kendalltau(x, y)
    return tau

 A['AB'] = pd.rolling_apply(A['B'], 3, lambda x: tau1(x))

当然，这并没有成功。我得到了：

ValueError: all keys need to be the same shape

我明白这不是一个简单的问题。非常感谢任何建议。

Answer 1

从Pandas 0.14版本开始，rolling_apply这个函数只会把NumPy数组传给你自己写的函数。如果你想要解决这个问题，可以把np.arange(len(A))作为第一个参数传给rolling_apply，这样你自己写的tau函数就能接收到你想用的行的索引。然后在tau函数里面，

B = A[[col1, col2]].iloc[idx]

返回一个包含所有需要行的DataFrame。

---

import numpy as np
import pandas as pd
import scipy.stats as stats
import itertools as IT

A = pd.DataFrame([[1, 5, 2], [2, 4, 4], [3, 3, 1], [4, 2, 2], [5, 1, 4]], 
                 columns=['A', 'B', 'C'], index = [1, 2, 3, 4, 5])

for col1, col2 in IT.combinations(A.columns, 2):
    def tau(idx):
        B = A[[col1, col2]].iloc[idx]
        return stats.kendalltau(B[col1], B[col2])[0]
    A[col1+col2] = pd.rolling_apply(np.arange(len(A)), 3, tau)

print(A)    

会得到

   A  B  C  AB        AC        BC
1  1  5  2 NaN       NaN       NaN
2  2  4  4 NaN       NaN       NaN
3  3  3  1  -1 -0.333333  0.333333
4  4  2  2  -1 -0.333333  0.333333
5  5  1  4  -1  1.000000 -1.000000