如何加速Pandas DataFrame的rolling_apply?

Question

在改进这个问题的基础上，之前提供了一个聪明的解决方案，可以在DataFrame的多个列上应用一个函数。我在想，这个解决方案是否还能进一步优化，以提高速度。

环境：Python 2.7.8，Pandas 14.1，Numpy 1.8。

下面是一个示例设置：

import pandas as pd
import numpy as np
import random

def meanmax(ii,df):
    xdf = df.iloc[map(int,ii)]
    n = max(xdf['A']) + max(xdf['B'])
    return n / 2.0

df  = pd.DataFrame(np.random.randn(2500,2)/10000, 
                    index=pd.date_range('2001-01-01',periods=2500),
                    columns=['A','B'])              
df['ii'] = range(len(df))      

res = pd.rolling_apply(df.ii, 26, lambda x: meanmax(x, df))

需要注意的是，meanmax函数并不是成对的，所以像rolling_mean(df['A'] + df['B'],26)这样的写法是行不通的。

不过，我可以这样做：

res2 = (pd.rolling_max(df['A'],26) + pd.rolling_max(df['B'],26)) / 2

这样做的速度大约快了3000倍：

%timeit res = pd.rolling_apply(df.ii, 26, lambda x: meanmax(x, df))
1 loops, best of 3: 1 s per loop

%timeit res2 = (pd.rolling_max(df['A'],26) + pd.rolling_max(df['B'],26)) / 2
1000 loops, best of 3: 325 µs per loop

有没有比上面第二种方法更好或等效的方案，考虑到示例函数并使用rolling_apply？虽然第二种方法更快，但它没有使用rolling_apply，而这个方法可以应用于更广泛的问题。

编辑：性能计时修正

Answer 1

你可能无法达到rolling_max的速度，但通过使用.values转到numpy，你通常可以提高速度，大约能快一个数量级：

def meanmax_np(ii, df):
    ii = ii.astype(int)
    n = df["A"].values[ii].max() + df["B"].values[ii].max()
    return n/2.0

这让我得到了

>>> %timeit res = pd.rolling_apply(df.ii, 26, lambda x: meanmax(x, df))
1 loops, best of 3: 701 ms per loop
>>> %timeit res_np = pd.rolling_apply(df.ii, 26, lambda x: meanmax_np(x, df))
10 loops, best of 3: 31.2 ms per loop
>>> %timeit res2 = (pd.rolling_max(df['A'],26) + pd.rolling_max(df['B'],26)) / 2
1000 loops, best of 3: 247 µs per loop

虽然这个速度还是比优化后的情况慢100倍，但比原来的快多了。有时候，我只需要速度快十倍，就足够不让它成为主要的时间消耗了。

Answer 2

在一个大小为 n 的数组上计算一个通用的滚动函数，窗口大小为 m，大概需要 O(n*m) 的时间。不过，内置的 rollin_xxx 方法使用了一些非常聪明的算法，能让运行时间远低于这个，通常可以保证在 O(n) 的时间内完成，这其实是相当了不起的。

特别是 rolling_min 和 rolling_max 的实现是借鉴自 bottleneck，而这个库提到的算法来源是 Richard Harter，不过我发现了一篇我认为是更早描述同样算法的 论文。

所以，经过这段历史介绍后，你很可能会发现，想要两全其美并不容易。虽然 rolling_apply 非常方便，但几乎总是会牺牲一些性能来适应特定的算法。根据我的经验，使用 Python 科学计算库时，最有趣的部分之一就是想出高效的计算方法，利用提供的快速原语进行创造性的组合。你自己调用 rolling_max 两次的做法就是一个很好的例子。所以放轻松，享受这个过程吧，知道如果你或者 Stack Overflow 的朋友们想不出更聪明的解决方案，你总是可以依赖 rolling_apply。