pandas - 使用 cumprod 进行快速 pandas groupby 计算

Question

此问题链接至Speedup of pandas groupby 。它是为了加速肮脏的 cumproduct 计算。 DataFrame 是 2D 的，具有由 3 个整数组成的多重索引。

可以在此处找到数据帧的 HDF5 文件:http://filebin.ca/2Csy0E2QuF2w/phi.h5

我正在执行的实际计算与此类似:

   >>> phi = pd.read_hdf('phi.h5', 'phi')
   >>> %timeit phi.groupby(level='atomic_number').cumprod()
   100 loops, best of 3: 5.45 ms per loop

另一个可能的加速是我使用相同的索引结构但使用不同的数字执行此计算大约 100 次。我想知道它是否可以以某种方式缓存索引。

任何帮助将不胜感激。

Answer 1

Numba 似乎在这里工作得很好。事实上，这些结果似乎好得令人难以置信，下面的 numba 函数比原始方法快约 4,000 倍，比普通方法快 5 倍 cumprod没有groupby 。希望这些是正确的，如果有错误请告诉我。

np.random.seed(1234)
df=pd.DataFrame({ 'x':np.repeat(range(200),4), 'y':np.random.randn(800) })
df = df.sort('x')
df['cp_groupby'] = df.groupby('x').cumprod()

from numba import jit

@jit
def group_cumprod(x,y):
    z = np.ones(len(x))
    for i in range(len(x)):
        if x[i] == x[i-1]:
            z[i] = y[i] * z[i-1]
        else:
            z[i] = y[i]
    return z

df['cp_numba'] = group_cumprod(df.x.values,df.y.values)

df['dif'] = df.cp_groupby - df.cp_numba

测试两种方法是否给出相同的答案:

all(df.cp_groupby==df.cp_numba)
Out[1447]: True

时间安排:

%timeit df.groupby('x').cumprod()
10 loops, best of 3: 102 ms per loop

%timeit df['y'].cumprod()
10000 loops, best of 3: 133 µs per loop

%timeit group_cumprod(df.x.values,df.y.values)
10000 loops, best of 3: 24.4 µs per loop