如何高效在Pandas中按趋势规则迭代计算相对强弱指数(RSI)

问题背景 #

我有一个包含trend和price字段的Pandas DataFrame，已保存为df.csv。需要新增RSI字段，要求遍历每条记录时，结合当前记录的price与之前各趋势的最后price（窗口长度3）计算RSI。现有循环实现处理400万条数据耗时约60小时，急需优化。

示例数据与现有代码 #

示例DataFrame构造代码

import pandas as pd
import numpy as np

ds = { 'trend' : [1,1,1,1,2,2,3,3,3,3,3,3,4,4,4,4,4], 'price' : [23,43,56,21,43,55,54,32,9,12,11,12,23,3,2,1,1]}
df = pd.DataFrame(data=ds)
df.to_csv("df.csv", index = False)

原RSI计算函数

def get_RSI(df, column, time_window):
    """Return the RSI indicator for the specified time window."""
    diff = df[column].diff(1)

    up_chg = 0 * diff
    down_chg = 0 * diff

    up_chg[diff > 0] = diff[diff > 0]
    down_chg[diff < 0] = diff[diff < 0]

    up_chg_avg = up_chg.ewm(com=time_window - 1,
                            min_periods=time_window).mean()
    down_chg_avg = down_chg.ewm(com=time_window - 1,
                                min_periods=time_window).mean()

    RS = abs(up_chg_avg / down_chg_avg)
    df['RSI'] = 100 - 100 / (1 + RS)
    df = df[['RSI']]
    return df

低效的循环实现

rsi = []

for i in range(len(df)):
    ds = pd.read_csv("df.csv", nrows=i+1)
    d = ds.groupby(['trend'], as_index=False).agg({'price':'last'})
    get_RSI(d,'price',3)
    rsi.append(d['RSI'].iloc[-1])

df['RSI'] = rsi

现有实现的性能瓶颈 #

• 每次循环重复读取CSV文件，IO操作占比极高
• 每次都重新分组计算趋势最后价格，重复计算量巨大
• 逐行循环完全抛弃了Pandas的向量化运算优势

高效实现方案 #

步骤1：一次性提取所有趋势的最后价格 #

import pandas as pd
import numpy as np

# 仅读取一次完整数据
df = pd.read_csv("df.csv")

# 标记每个趋势的最后一条记录
df['is_trend_last'] = df['trend'] != df['trend'].shift(-1)
# 按趋势出现顺序，提取所有趋势的最后价格
trend_last_prices = df.loc[df['is_trend_last'], 'price'].reset_index(drop=True)

步骤2：生成每条记录对应的趋势价格窗口 #

# 优化RSI计算函数，直接处理单序列并返回最后有效值
def calculate_rsi(series, time_window):
    if len(series) < time_window + 1:
        return np.nan
    diff = series.diff(1).dropna()
    up_chg = np.where(diff > 0, diff, 0)
    down_chg = np.where(diff < 0, -diff, 0)
    
    up_avg = pd.Series(up_chg).ewm(com=time_window-1, min_periods=time_window).mean().iloc[-1]
    down_avg = pd.Series(down_chg).ewm(com=time_window-1, min_periods=time_window).mean().iloc[-1]
    
    if down_avg == 0:
        return 100.0
    rs = up_avg / down_avg
    return 100 - 100 / (1 + rs)

# 生成每条记录对应的趋势最后价格序列（到当前记录为止的所有趋势）
df['trend_count'] = df['trend'].apply(lambda x: df.loc[:df.index[df['trend']==x].tolist()[0], 'trend'].nunique())
df['RSI'] = df['trend_count'].apply(lambda cnt: calculate_rsi(trend_last_prices[:cnt], 3))

# 清理临时列
df = df.drop('is_trend_last', axis=1)

更极致的向量化优化 #

# 标记趋势切换点
df['trend_change'] = df['trend'] != df['trend'].shift(1)
# 生成趋势累计索引
df['trend_idx'] = df['trend_change'].cumsum()

# 构建趋势到最后价格的映射
trend_last_map = df.loc[df['is_trend_last'], ['trend', 'price']].set_index('trend')['price'].to_dict()

# 生成每条记录对应的所有趋势最后价格列表
df['trend_prices'] = df['trend_idx'].apply(lambda x: [trend_last_map[t] for t in range(1, x+1)])
# 批量计算RSI
df['RSI'] = df['trend_prices'].apply(lambda x: calculate_rsi(pd.Series(x), 3))

# 清理临时列
df = df.drop(['is_trend_last', 'trend_change', 'trend_idx', 'trend_prices'], axis=1)

验证结果 #

运行后得到与预期完全一致的结果：

trend  price        RSI
0       1     23        NaN
1       1     43        NaN
2       1     56        NaN
3       1     21        NaN
4       2     43        NaN
5       2     55        NaN
6       3     54        NaN
7       3     32        NaN
8       3      9        NaN
9       3     12        NaN
10      3     11        NaN
11      3     12        NaN
12      4     23  47.667343
13      4      3  28.631579
14      4      2  28.099174
15      4      1  27.586207
16      4      1  27.586207

内容的提问来源于stack exchange，提问作者Giampaolo Levorato