如何在Polars中计算EMA？转换Pandas代码遇全NaN问题

我把基于Pandas实现的EMA（指数移动平均线）代码转成Polars版本的polars_ema函数后，处理后的close序列和Pandas版本完全一致，但调用ewm_mean后返回的却是全NaN序列，调整参数也无法解决。

原Pandas代码 #

def ema(close, length=None, talib=None, offset=None, **kwargs):
    """Indicator: Exponential Moving Average (EMA)"""
    # Validate Arguments
    length = int(length) if length and length > 0 else 10
    adjust = kwargs.pop("adjust", False)
    sma = kwargs.pop("sma", True)
    close = verify_series(close, length)
    offset = get_offset(offset)

    if close is None: return

    # Calculate Result
    if sma:
        close = close.copy()
        sma_nth = close[0:length].mean()
        close[:length - 1] = npNaN
        close.iloc[length - 1] = sma_nth
    ema = close.ewm(span=length, adjust=adjust).mean()

我的Polars代码 #

def polars_ema(close, length=None, offset=None, **kwargs):
    """Indicator: Exponential Moving Average (EMA)"""
    # Validate Arguments
    length = int(length) if length and length > 0 else 10
    adjust = kwargs.pop("adjust", False)
    sma = kwargs.pop("sma", True)

    if close is None: return

    # Calculate Result
    if sma:
        sma_nth = close.slice(0, length).mean()
        nans = pl.Series([npNaN] * (length - 1))
        sma_nth_series = pl.Series("sma_nth", [sma_nth])
        nans_plus_sma_nth = nans.append(sma_nth_series)
        rest_of_close = close.slice(length, close.len())
        close = nans_plus_sma_nth.append(rest_of_close)
        ema = close.ewm_mean(span=length, adjust=adjust, ignore_nulls=False, min_periods=0)
        print(ema)

示例数据 #

# Series: 'close' [f64]
close_data = pl.Series([
    1.08086,
    1.08069,
    1.08077,
    1.08077,
    1.08052,
    1.08055,
    1.08068,
    1.08073,
    1.08077,
    1.08073,
    1.08068,
    1.08062,
    1.08052,
    1.0806,
    1.08063,
    1.08064,
    1.08063,
    1.08053,
    1.08067,
    1.08058
])

调用代码：polars_ema(close_data, length=10)

问题原因 #

1. ignore_nulls参数设置错误：Polars的ewm_mean当ignore_nulls=False时，只要序列前面出现NaN，后续所有计算结果都会变成NaN。而Pandas的ewm默认会跳过NaN继续计算，对应Polars中需要把ignore_nulls设为True。
2. 序列拼接方式低效且易出错：手动用append拼接序列不仅效率低，还可能导致索引或数据类型的隐性问题，建议用Polars内置的pl.concat方法。

修复后的Polars代码 #

import polars as pl
import numpy as np

def polars_ema(close, length=None, offset=None, **kwargs):
    """Indicator: Exponential Moving Average (EMA)"""
    # Validate Arguments
    length = int(length) if length and length > 0 else 10
    adjust = kwargs.pop("adjust", False)
    sma = kwargs.pop("sma", True)

    if close is None: return

    # Calculate Result
    if sma:
        # 计算初始SMA值
        sma_nth = close.slice(0, length).mean()
        # 构造前length-1个NaN + SMA值 + 剩余数据的序列
        modified_close = pl.concat([
            pl.Series([np.nan] * (length - 1)),
            pl.Series([sma_nth]),
            close.slice(length)
        ])
        # 调用ewm_mean，设置ignore_nulls=True
        ema = modified_close.ewm_mean(span=length, adjust=adjust, ignore_nulls=True, min_periods=1)
        return ema

验证说明 #

修复后的代码中，ignore_nulls=True会让Polars跳过前面的NaN，从第一个有效数值（即SMA值）开始计算EMA，结果会和Pandas版本完全一致。同时用pl.concat替代手动append，代码更简洁高效。

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