基于先前行的 Dataframe 计算速度非常慢

Question

我正在计算当前行的结果需要前一行的股票价格的 RSI 值。我目前正在通过 for 循环完整数据帧多次执行此操作，这些条目需要花费大量时间（在我的电脑上执行时间大约为 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 值是 Gain 列的第 [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 秒 ± 3.73 毫秒（平均值 ± 标准偏差。7 次运行，每个循环 1 个）

%%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（7 次运行的平均值 ± 标准偏差，每次 1000 个循环）

data.head(15)

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

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