如何在 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 但没有成功。它会填充第一行的值，而其他行则为 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。

- JDavda

---

pandas并没有处理递归问题的简单方法 - 推荐使用numba来加速循环。 - Asish M.

作为另一种选择 - scipy.signal.filter 在这里似乎很有用 - 请参阅 https://dev59.com/Y18d5IYBdhLWcg3w8F-- - Asish M.

你还在寻找比上面循环更高效的解决方案吗？我有点困惑，因为问题和答案不匹配。 - cs95

@cs95 我稍微编辑了一下问题，因为其中一个答案有效，但仅因为df ['a']列值按数字顺序排列，如果该列包含随机值，则应该有效。我通过在迭代中使用df.at而不是df.loc来加快速度，但如果您知道其他方法，我很乐意听取。始终寻找改进代码的方法。 - JDavda

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

---

性能比较

参考的基准测试代码.

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

- cs95

---

谢谢 @cs95 - 我能够使用你的代码并且相对于使用df.at所用的19.4ms，使用另一个更大的数据框架将时间减少了9ms，这太神奇了。我真的很高兴。还有一个问题，在你的基准代码中，我没有看到func函数或numba导入 - 我错过了什么？ - JDavda

@cs95 大佬，回答得非常好，继续保持 :) - Shubham Sharma

@cs95 我已经标记它为接受的答案 - 再次感谢，使用numba一直在我的清单上，感谢您指引的方向！ - JDavda

Answer 2

你可以尝试使用shift和cumsum，从5开始并使用fillna进行填充：

import pandas as pd


df = pd.DataFrame([1,2,3,4], columns=['a'])
df['b'] = df['a'].shift().fillna(5).cumsum()

print(df)

输出

   a     b
0  1   5.0
1  2   6.0
2  3   8.0
3  4  11.0

- Dani Mesejo

---

谢谢，我可能可以使用fillna()而不是我的当前方法df.at，但计算需要基于a+b的先前值。在这种情况下，您的代码有效，因为'a'值为1,2,3,4，但我为此示例简化了我的实际代码 - 在我的原始代码中，'a'的值是随机的。我将编辑问题以澄清这一点。 - JDavda

Answer 3

我很可能会误解你的问题，但如果你想创建一个移位列，请尝试这个：

df = pd.DataFrame([1,2,3,4], columns=['a'])
df["b"] = df.a.shift()

- Tony Ng