python - 如何加快涉及 Pandas 前一行的计算？

Question

我正在尝试使用自己创建的列的移位值创建一个新的 Pandas DataFrame 列。

我能够做到这一点的唯一方法是遍历数据，这太慢了，会导致我的代码出现瓶颈。

import pandas as pd 

df = pd.DataFrame([1,6,2,8], columns=['a'])
df.at[0, 'b'] = 5

for i in range(1, len(df)):
    df.loc[i, ('b')] = (df.a[i-1] + df.b[i-1]) /2

我尝试使用 shift，但没有成功。它为第 1 行填充值，其余为 NaN。我假设此方法无法即时读取新创建的值。

df.loc[1:, ('b')] = (df.a.shift() + df.b.shift()) /2

更新

通过在迭代中使用 df.at 而不是 df.loc，我能够显着减少时间

def with_df_loc(df):
    for i in range(1, len(df)):
        df.loc[i, ('b')] = (df.a[i-1] + df.b[i-1]) /2
    return df

def with_df_at(df):
    for i in range(1, len(df)):
        df.at[i, 'b'] = (df.a[i-1] + df.b[i-1]) /2
    return df



%timeit with_df_loc(df)
183 ms ± 75.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit with_df_at(df)
19.4 ms ± 2.74 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

此计时基于 150 行的较大数据集。考虑到 df.rolling(20).mean() 大约需要 3 毫秒，我认为这可能是我能做的最好的了。

感谢您的回答，如果我需要进一步优化，我会查看 Asish M numba 的建议。

Answer 1

我们可以使用numba要加快计算速度，请参阅 Enhancing performance文档中的部分。

import numba 

@numba.njit
def func(a, b_0=5):
    n = len(a)
    b = np.full(n, b_0, dtype=np.float64)
    for i in range(1, n):
        b[i] = (b[i - 1] + a[i - 1]) / 2
    return b

df['b'] = func(df['a'].to_numpy())
df

   a     b
0  1  5.00
1  6  3.00
2  2  4.50
3  8  3.25

---

比较性能

Benchmarking code, for reference .

蓝线代表您当前方法的最快版本的性能(使用 .at)。橙色线代表 numba 的性能。