Polars中指数移动平均线(EMA)的实现与空值处理详解

import polars as pl
import numpy as np # 虽然问题出在np.NaN，但np仍然是常用库，这里保留以示对比

def polars_ema(close: pl.Series, length: int = 10, adjust: bool = False, sma_init: bool = True) -> pl.Series:
    """
    在Polars中计算指数移动平均线 (EMA)。

    参数:
        close (pl.Series): 输入的收盘价或其他数值序列。
        length (int): EMA的计算周期。
        adjust (bool): 是否使用调整因子。当为True时，权重会根据实际观察到的数据点数量进行调整。
                       当为False时，所有权重都基于完整的span。
        sma_init (bool): 是否使用前length个元素的简单移动平均作为EMA的初始值。
                         如果为True，序列前length-1个元素将被填充为None，第length个元素为SMA。

    返回:
        pl.Series: 包含EMA值的Polars Series。
    """
    if close is None:
        return pl.Series([], dtype=pl.Float64) # 返回空Series

    # 验证并设置周期长度
    length = int(length) if length and length > 0 else 10

    # 如果需要SMA初始化
    if sma_init:
        # 计算前length个元素的简单移动平均
        sma_nth = close.slice(0, length).mean()

        # 创建一个包含 length-1 个 None 值的 Series，指定数据类型为 Float64
        # 这是关键：使用 None 而非 np.NaN 来表示 Polars 中的缺失值
        nones_series = pl.Series([None] * (length - 1), dtype=pl.Float64)

        # 将SMA值转换为一个Series
        sma_nth_series = pl.Series("sma_nth", [sma_nth])

        # 将 None 系列与 SMA 值系列拼接
        initial_part = nones_series.append(sma_nth_series)

        # 获取原始序列中从 length 索引开始的剩余部分
        rest_of_close = close.slice(length, close.len() - length)

        # 将初始部分与剩余部分拼接，形成新的 close 序列
        close = initial_part.append(rest_of_close)

    # 使用 Polars 的 ewm_mean 方法计算 EMA
    # ignore_nulls=False 确保 ewm_mean 在遇到 None 值时不会将其跳过，
    # 而是根据 adjust 参数和权重进行计算。
    # min_periods=0 允许在数据点不足 span 时也能计算 EMA（如果 adjust=True）。
    ema = close.ewm_mean(span=length, adjust=adjust, ignore_nulls=False, min_periods=0)

    return ema

# 错误的做法 (可能导致 ewm_mean 返回全 NaN)
# nans_series = pl.Series([np.NaN] * (length - 1)) 

# 正确的做法
nones_series = pl.Series([None] * (length - 1), dtype=pl.Float64)

# 示例数据
sample_data = [
    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.08060, 1.08063,
    1.08064, 1.08063, 1.08053, 1.08067, 1.08058
]
close_series = pl.Series("close", sample_data, dtype=pl.Float64)

# 计算周期为10的EMA
ema_result = polars_ema(close_series, length=10)

print("原始收盘价系列:")
print(close_series)
print("\n计算得到的EMA系列 (周期=10):")
print(ema_result)

原始收盘价系列:
shape: (20,)
Series: 'close' [f64]
[
    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
]

计算得到的EMA系列 (周期=10):
shape: (20,)
Series: 'ema' [f64]
[
    null
    null
    null
    null
    null
    null
    null
    null
    null
    1.080697
    1.0806957272727273
    1.0806792314049586
    1.080642999494966
    1.0806351813136994
    1.080634241074845
    1.0806343808794186
    1.0806336388904334
    1.0806193409012637
    1.080633990828309
    1.0806205387459074
]