在 Polars 中高效计算指数移动平均线 (EMA) 并避免常见陷阱

import polars as pl
import numpy as np # 仅用于上下文说明，实际解决方案中避免使用 np.NaN

def polars_ema(close: pl.Series, length: int = 10, adjust: bool = False, sma: bool = True) -> pl.Series:
    """
    计算 Polars Series 的指数移动平均线 (EMA)。
    该函数支持在 EMA 启动前，使用初始简单移动平均线 (SMA) 进行预热。

    参数:
        close (pl.Series): 输入的收盘价或待计算 EMA 的序列。
        length (int): EMA 的周期长度。默认为 10。
        adjust (bool): 是否根据权重调整除数。默认为 False (与 pandas 默认行为不同)。
                       当 adjust=False 时，递归公式为：
                       EMA_t = (value_t + (1 - alpha) * EMA_{t-1}) / (1 + (1 - alpha))
                       当 adjust=True 时，EMA_t = alpha * value_t + (1 - alpha) * EMA_{t-1}
                       其中 alpha = 2 / (span + 1)，span 等同于 length。
        sma (bool): 是否使用初始 SMA 来预热 EMA 序列。如果为 True，
                    前 length-1 个值将为 None，第 length 个值为前 length 个数据的 SMA。
                    默认为 True。

    返回:
        pl.Series: 包含 EMA 值的 Polars Series。
    """
    if close is None or not isinstance(close, pl.Series):
        raise ValueError("输入 'close' 必须是一个 Polars Series。")
    if length <= 0:
        raise ValueError("长度 'length' 必须大于 0。")

    # 如果需要初始 SMA 预热
    if sma:
        # 1. 计算前 'length' 个数据点的简单移动平均 (SMA)
        # 使用 .head(length) 或 .slice(0, length) 获取前 'length' 个元素
        initial_sma = close.head(length).mean()

        # 2. 创建一个包含 'length - 1' 个 None 值的 Polars Series
        # 关键修复: 使用 None 而不是 np.NaN，并明确指定 dtype 为 pl.Float64
        nones_series = pl.Series([None] * (length - 1), dtype=pl.Float64)

        # 3. 创建一个包含初始 SMA 值的 Series
        initial_sma_series = pl.Series([initial_sma], dtype=pl.Float64)

        # 4. 拼接 None 值和初始 SMA 值，形成序列的起始部分
        initial_segment = nones_series.append(initial_sma_series)

        # 5. 获取原始 'close' 序列中剩余的部分
        rest_of_close = close.slice(length, close.len())

        # 6. 将起始部分与剩余部分拼接，形成最终用于 EMA 计算的序列
        # 这个序列的结构是：[None, ..., None, initial_sma, close_value_at_length, ...]
        processed_close = initial_segment.append(rest_of_close)
    else:
        # 如果不使用 SMA 预热，直接使用原始序列
        processed_close = close

    # 7. 计算指数移动平均线 (EMA)
    # span: 等同于 length，用于计算指数加权因子 alpha = 2 / (span + 1)
    # adjust: 控制加权方式，False 更接近 pandas 的默认行为（但仍有细微差异，需注意）
    # ignore_nulls=False: 表示在计算时考虑 None 值。
    # min_periods=0: 允许在少于 span 个非 None 值的情况下也进行计算。
    #                这对于我们序列开头有 None 的情况至关重要，确保从第一个有效值开始计算。
    ema = processed_close.ewm_mean(
        span=length,
        adjust=adjust,
        ignore_nulls=False,
        min_periods=0
    )
    return ema

import polars as pl

# 示例数据
sample_data = pl.Series('close', [
    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
], dtype=pl.Float64)

# 调用函数计算 EMA，周期为 10
ema_result = polars_ema(sample_data, length=10)

print("计算得到的 EMA 序列:")
print(ema_result)

Shape: (20,)
Series: 'close' [f64]
[
    null
    null
    null
    null
    null
    null
    null
    null
    null
    1.080697    // 前10个值的SMA
    1.0806954545454546
    1.0806817363636363
    1.0806605115684364
    1.0806489640096297
    1.0806436978260244
    1.0806429345851109
    1.0806414919332726
    1.080628493399945
    1.0806387673272277
    1.0806298914500954
]