Polars中ewm_mean计算EMA遇版本差异问题及优化需求

测试数据 #

values = [143.15,143.1,143.06,143.01,143.03,143.09,143.14,143.18,143.2,143.2,143.2,143.31,143.38,143.35,143.34,143.25,143.33,143.3,143.33,143.36]

参考正确实现（Pandas+TA-Lib） #

使用Pandas结合TA-Lib可得到预期的EMA结果：

import pandas as pd
import talib as ta

df_pan = pd.DataFrame(
    {
        'value': values
    }
)

df_pan['ema_9'] = ta.EMA(df_pan['value'], timeperiod=9)

旧版Polars（0.19.14）的有效代码 #

在Polars 0.19.14版本中，以下代码计算出的EMA结果与TA-Lib完全一致：

import polars as pl

df_pol = (
    pl.DataFrame(
        {
            'value': values
        }
    )
    .with_columns(
        pl.when(pl.col('value').cum_count() < 9)
        .then(pl.col('value').head(9).mean())
        .otherwise(pl.col('value'))
        .ewm_mean(span=9, min_periods=9, ignore_nulls=True, adjust=False)
        .alias('ema_9')
    )
)

新版Polars（0.20.21）的问题 #

升级到Polars 0.20.21后，上述代码的计算结果与预期不符。尝试调整ewm_mean的参数，但由于不了解其底层数学原理，未能解决问题。

自定义低效实现 #

目前已实现一个基于循环的自定义函数，能得到正确结果，但效率极低：

def add_ema_col(df: pl.DataFrame, column: str, periods: int, smoothing:int = 2) -> pl.DataFrame:
    """为OHLCV数据表添加*EMA*列。"""

    values = df.select(column).to_series().to_list()
    ema = [sum(values[:periods]) / periods]

    nones = [None] * (periods - 1)
    
    for price in values[periods:]:
        ema.append((price * (smoothing / (1 + periods))) + ema[-1] * (1 - (smoothing / (1 + periods))))
    
    final_values = pl.Series(nones + ema)

    return (
        df
        .with_columns(
            final_values
        )
    )

df = (
    pl.DataFrame({'Close': values})
    .pipe(lambda df: add_ema_col(df, 'Close', 9))
)

需求 #

寻求以下两个方向的解决方案：

• 调整ewm_mean的参数配置，使其计算结果与TA-Lib的EMA一致
• 将上述自定义循环函数改写为纯Polars表达式，提升计算效率

已提交相关缺陷报告，暂不考虑转回Pandas+TA-Lib的方案。

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