如何使用Polars计算SMMA（平滑移动平均线）并获取Expression？

需求说明 #

• 目标：使用Polars计算SMMA（平滑移动平均线），优先通过单次select操作获取Expression而非直接返回DataFrame
• 可用参数：窗口大小frame_length、包含价格数据的DataFrame（示例价格列数据：[1.2, 5.4, 6.1, 4.3, ...]）

原生Python SMMA计算逻辑 #

原生Python中SMMA的计算逻辑可通过以下函数实现：

def get_smma(sum_in_frame: float, prev_smma: float, frame_length: int) -> float:
    return (sum_in_frame - prev_smma) / float(frame_length)

关于Polars的递归计算支持疑问 #

用户疑问：Polars是否支持滚动递归（rolling recursion）或访问前一个计算值（previous calculated value）？

基于迭代公式的Polars实现 #

用户最终采用的SMMA迭代公式如下：

N = 移动窗口大小
SMMA(0) = SUM(PRICE[0:N-1]) / N # 第一个SMMA值等于简单移动平均线
SMMA(i) = (SMMA(i-1) * (N-1) + PRICE[i]) / N

对应的Polars实现代码：

def get_smma_df(
        frame_len: int,
        price_candle_df: DataFrame,
        src_field='close_price'
) -> DataFrame:

    def get_smma_value(acc_series, row_series):
        price = row_series[0]
        rol_sum = row_series[1]
        prev_smma = acc_series[0]

        # 滚动和未生成时，沿用前一个SMMA值
        if rol_sum is None:
            return prev_smma

        # 初始SMMA值：用窗口内的总和除以窗口大小
        if prev_smma is None:
            smma = rol_sum / frame_len
        # 后续SMMA值：按迭代公式计算
        else:
            smma = (prev_smma * (frame_len - 1) + price) / frame_len
        return smma

    df = (
        price_candle_df
        # 选择价格列并计算窗口滚动和
        .select(
            price=pl.col(src_field),
            rol_sum=pl.col(src_field).rolling_sum(frame_len)
        )
        # 转置以适配cumfold的行处理逻辑
        .transpose()
        # 使用cumfold进行递归计算SMMA
        .select(
            smma=pl.cumfold(
                acc=None,
                exprs=pl.all(),
                function=get_smma_value
            )
        )
        # 展开SMMA结果
        .unnest('smma')
        # 转回原结构并命名列
        .transpose(column_names=["smma"])
    )

    return df

内容的提问来源于stack exchange，提问作者Joris Medeišis