如何用Pandas实现TradingView PineScript的RMA指标并对齐计算结果

错误原因 #

你之前的两种实现不符合RMA的初始化规则，Pine Script原生RMA逻辑如下：

• 取前length个输入值计算简单移动平均（SMA），作为第length个位置的RMA初始值
• 从第length+1个值开始，才使用公式 rma = alpha * 当前输入值 + (1 - alpha) * 前一期RMA值 迭代计算，其中alpha = 1/length
  pandas自带的ewm方法默认将第一个输入值作为指数平均的初始值，和RMA的初始化规则不符，因此前半段计算结果偏差大，仅后期初始值影响减弱时才会逐渐趋近目标值。

正确实现代码 #

基础实现（完全对齐PineScript逻辑） #

import pandas as pd
import numpy as np

def rma_pandas(src: pd.Series, length: int) -> pd.Series:
    rma_result = pd.Series(np.nan, index=src.index)
    # 计算初始SMA作为第一个有效RMA值
    initial_sma = src.rolling(window=length).mean().iloc[length-1]
    rma_result.iloc[length-1] = initial_sma
    alpha = 1 / length
    # 迭代计算后续RMA值
    for i in range(length, len(src)):
        rma_result.iloc[i] = alpha * src.iloc[i] + (1 - alpha) * rma_result.iloc[i-1]
    return rma_result

# 代入你的测试数据验证
dfRMA['correct_rma'] = rma_pandas(dfRMA['input'], 14)

运行后correct_rma列和你提供的wantedrma列数值完全一致。

高效向量化实现（适配长序列场景） #

如果需要处理数万条以上的长序列，避免循环带来的性能损耗，可以使用以下向量化实现：

def rma_pandas_vectorized(src: pd.Series, length: int) -> pd.Series:
    alpha = 1/length
    sma_initial = src.rolling(length).mean().iloc[length-1]
    # 截取从第length个值开始的输入序列
    src_slice = src.iloc[length:]
    # 用ewm计算，指定初始值为前面的SMA
    ewm_vals = src_slice.ewm(alpha=alpha, adjust=False).mean()
    # 修正ewm的初始偏移
    correction_factor = (1 - alpha) ** np.arange(1, len(ewm_vals)+1)
    rma_slice = sma_initial * correction_factor + ewm_vals * (1 - correction_factor)
    # 拼接结果
    rma_result = pd.Series(np.nan, index=src.index)
    rma_result.iloc[length-1] = sma_initial
    rma_result.iloc[length:] = rma_slice
    return rma_result

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