如何将Pine Script二阶极点SSF代码转Python并匹配结果？

问题背景 #

需将实现二阶极点SSF（二阶超级平滑滤波器）的Pine Script代码转换为Python，但自行编写的代码运行结果与原Pine Script输出完全不符，需修正实现逻辑以匹配原结果。

原Pine Script代码 #

get2PoleSSF(src, length) =>
PI = 2 * asin(1)
arg = sqrt(2) * PI / length
a1 = exp(-arg)
b1 = 2 * a1 * cos(arg)
c2 = b1
c3 = -pow(a1, 2)
c1 = 1 - c2 - c3

ssf = 0.0
ssf := c1 * src + c2 * nz(ssf[1]) + c3 * nz(ssf[2])

自行编写的Python代码 #

def get2PoleSSD(src,length):
    PI = 2* np.arcsin(1)
    arg = np.sqrt(2)* PI / length
    a1 = np.exp(-arg)
    b1 = 2 * a1 *np. cos(arg)
    c2 = b1
    c3 = -pow(a1, 2)
    c1 = 1 - c2 - c3
    df['ssf']  = 0.0
    df['ssf'] = c1* src
    df['ssf_1p'] = df['ssf'].shift(1)
    df['ssf_2p'] = df['ssf'].shift(2)
    df['ssf_n'] = c1 * src + c2 * df['ssf_1p'] + c3 * df['ssf_2p']

错误核心原因 #

原Pine Script中的ssf :=是递归迭代赋值，每一步的计算结果都依赖前两次的ssf输出；而你编写的Python代码仅用初始的c1*src做移位计算，完全没有实现递归迭代的核心逻辑，这是结果不符的关键。此外还存在全局依赖df、未处理NaN值等问题。

正确Python实现 #

import numpy as np

def get2PoleSSF(src, length):
    PI = 2 * np.arcsin(1)
    arg = np.sqrt(2) * PI / length
    a1 = np.exp(-arg)
    b1 = 2 * a1 * np.cos(arg)
    c2 = b1
    c3 = -a1 ** 2
    c1 = 1 - c2 - c3

    # 初始化结果数组，兼容NumPy数组和Pandas Series
    src_arr = np.asarray(src)
    ssf = np.zeros_like(src_arr)
    
    # 处理前两个元素，对应Pine Script中nz(ssf[1])=0、nz(ssf[2])=0的初始规则
    ssf[0] = c1 * src_arr[0]
    if len(src_arr) >= 2:
        ssf[1] = c1 * src_arr[1] + c2 * ssf[0]
    
    # 从第三个元素开始递归计算，处理NaN值对应nz函数
    for i in range(2, len(src_arr)):
        prev1 = ssf[i-1] if not np.isnan(ssf[i-1]) else 0
        prev2 = ssf[i-2] if not np.isnan(ssf[i-2]) else 0
        ssf[i] = c1 * src_arr[i] + c2 * prev1 + c3 * prev2
    
    # 保持输入类型一致，若原输入是Series则返回Series
    if isinstance(src, pd.Series):
        return pd.Series(ssf, index=src.index)
    return ssf

使用示例 #

import pandas as pd

# 示例数据
df = pd.DataFrame({'close': [10,11,12,13,14,15,16,17,18,19]})
# 计算SSF并赋值到DataFrame
df['ssf'] = get2PoleSSF(df['close'], length=5)

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