如何替代for循环，在Pandas中用shift/偏移处理RSI多序列计算？

问题说明 #

我是Python新手，编写了计算相对强弱指数（RSI，股市常用指标）的函数。需要更新存储平均值的Pandas序列，规则是结合另一序列的当前值与序列自身的前一个更新后的值。目前用for循环遍历索引完成计算，想改用更简洁的方法替代循环，但尝试的向量化代码未得到预期结果——因为代码使用了序列的旧值而非更新后的值：

avg_ups[period+1::] = up[period+1::]*alpha+(1-alpha)*avg_ups[period+1::].shift(1)
avg_downs[period+1::] = down[period+1::]*alpha+(1-alpha)*avg_downs[period+1::].shift(1)

原函数代码 #

import pandas as pd
from datetime import date
import yfinance as yf 

def yahoo_rsi(data:pd.DataFrame, period:int=14, column:str="Close"):

       delta = data[column].diff()
       delta = delta[1:]
       up = delta.clip(lower=0)
       down = delta.clip(upper=0)
       up[ up < 0 ] = 0
       down[down > 0] = 0

       down = down.abs()

       avg_ups:pd.Series = up.ewm(span=period).mean()
       avg_downs:pd.Series = down.abs().ewm(span=period).mean()

       alpha:float = 1/period

       # 希望简化的for循环
       for x in range(period+1, len(avg_ups)):
           avg_ups[x] = up[x]*alpha+(1-alpha)*avg_ups[x-1]
           avg_downs[x] = down[x]*alpha+(1-alpha)*avg_downs[x-1]

       rsis:pd.Series = pd.Series(100 - (100/(1+avg_ups/avg_downs)), avg_ups.index)
       return rsis


def run():
   ticker = yf.Ticker("AAPL")
   df = ticker.history("1d","1d", date(2022, 2, 2), date(2022, 8, 2), 
   prepost=True,auto_adjust=False, back_adjust=False)
   yahoo_rsi(df)

run() 

解决方案 #

你的循环逻辑是实现Wilder式RSI的递推平滑计算（初始用简单平均，后续用avg = 当前值*alpha + 前一次平均值*(1-alpha)更新），这种依赖前序结果的计算无法直接用向量化shift实现，因为shift取的是原始序列的前值，而非迭代更新后的值。可以用itertools.accumulate替代循环，实现简洁的递推计算：

import pandas as pd
from datetime import date
import yfinance as yf
from itertools import accumulate

def yahoo_rsi(data:pd.DataFrame, period:int=14, column:str="Close"):
    # 计算价格变动，丢弃空值
    delta = data[column].diff().dropna()
    # 分离上涨和下跌幅度
    up = delta.clip(lower=0)
    down = delta.clip(upper=0).abs()

    # 计算初始平均值：前period个值的简单平均
    initial_avg_up = up.iloc[:period].mean()
    initial_avg_down = down.iloc[:period].mean()
    alpha = 1 / period

    # 处理周期后的数据递推
    if len(up) > period:
        # 提取周期后的上涨/下跌数据
        up_rest = up.iloc[period:]
        down_rest = down.iloc[period:]

        # 用accumulate实现递推计算，每个值依赖前一次结果
        avg_ups_rest = list(accumulate(
            up_rest,
            lambda prev_avg, curr: curr * alpha + prev_avg * (1 - alpha),
            initial=initial_avg_up
        ))[1:]  # 去掉初始值，只保留递推结果

        avg_downs_rest = list(accumulate(
            down_rest,
            lambda prev_avg, curr: curr * alpha + prev_avg * (1 - alpha),
            initial=initial_avg_down
        ))[1:]

        # 拼接完整的平均值序列
        avg_ups = pd.Series(
            [initial_avg_up] * period + avg_ups_rest,
            index=up.index
        )
        avg_downs = pd.Series(
            [initial_avg_down] * period + avg_downs_rest,
            index=down.index
        )
    else:
        # 数据长度不足周期时，直接用初始平均值填充
        avg_ups = pd.Series([initial_avg_up] * len(up), index=up.index)
        avg_downs = pd.Series([initial_avg_down] * len(down), index=down.index)

    # 计算RSI
    rsis = 100 - (100 / (1 + avg_ups / avg_downs))
    return rsis

def run():
    ticker = yf.Ticker("AAPL")
    df = ticker.history(
        period="1d", interval="1d", start=date(2022, 2, 2), end=date(2022, 8, 2),
        prepost=True, auto_adjust=False, back_adjust=False
    )
    print(yahoo_rsi(df))

run()

关键说明 #

1. 移除了原代码中冗余的up[up<0]=0和down[down>0]=0，因为clip方法已经完成了相同的过滤。
2. itertools.accumulate会按顺序迭代计算，每一步的结果作为下一次计算的输入，完美匹配原循环的递推逻辑。
3. 初始平均值采用前period个值的简单平均，符合Wilder RSI的标准计算方式，修正了原代码用ewm初始化的偏差。

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