基于前几行的数据帧计算速度非常慢

Question

我正在计算股票价格的RSI值，其中当前行的结果需要前一行。我目前正在通过循环完整的数据帧来完成这项工作，因为有很多条目需要花费很多时间（在我的电脑上执行时间约为15秒）

有什么方法可以改进代码吗

import pandas as pd
from pathlib import Path


filename = Path("Tesla.csv")
test = pd.read_csv(filename)
data = pd.DataFrame(test[["Date","Close"]])
data["Change"] = (data["Close"].shift(-1)-data["Close"]).shift(1)
data["Gain"] = 0.0
data["Loss"] = 0.0
data.loc[data["Change"] >= 0, "Gain"] = data["Change"]
data.loc[data["Change"] <= 0, "Loss"] = data["Change"]*-1
data.loc[:, "avgGain"] = 0.0
data.loc[:, "avgLoss"] = 0.0
data["avgGain"].iat[14] = data["Gain"][1:15].mean()
data["avgLoss"].iat[14] = data["Loss"][1:15].mean()


for index in data.iterrows():
    data.loc[15:, "avgGain"] = (data.loc[14:, "avgGain"].shift(1)*13 + data.loc[15:, "Gain"])/14
    data.loc[15:, "avgLoss"] = (data.loc[14:, "avgLoss"].shift(1)*13 + data.loc[15:, "Loss"])/14

使用的数据集可在此处下载： TSLA historic dataset from yahoo finance

目标是根据待计算的avgGain和avgLoss值计算RSI值。 第0:14行的avgGain值不存在。 第15行的avgGain值是增益列的第[1:14]行的平均值。 从第16行开始的avgGain值计算如下： （13*avgGain（前一行）+增益（当前行））/14

Answer 1

“itertuples”比“iterrows”更快，矢量化操作通常在time方面表现最好

在这里，您可以使用窗口大小为14的rolling方法计算14天内的平均收益和损失（滚动平均值）

%%timeit
data["avgGain"].iat[14] = data["Gain"][1:15].mean()
data["avgLoss"].iat[14] = data["Loss"][1:15].mean()

for index in data.iterrows():
    data.loc[15:, "avgGain"] = (data.loc[14:, "avgGain"].shift(1)*13 + data.loc[15:, "Gain"])/14
    data.loc[15:, "avgLoss"] = (data.loc[14:, "avgLoss"].shift(1)*13 + data.loc[15:, "Loss"])/14

1.12 s ± 3.73 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%%timeit
data['avgGain_alt'] = data['Gain'].rolling(window=14).mean().fillna(0)
data['avgLos_alt'] = data['Gain'].rolling(window=14).mean().fillna(0)

1.38 ms ± 2.31 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

data.head(15)

使用矢量化操作来计算移动平均值比使用循环计算快大约10倍

但是请注意，对于第一次之后的平均值，代码中也存在一些计算错误