Polars EMA 计算：掌握空值处理与高效实现

import polars as pl
import numpy as np # 尽管不用于空值，但可能在其他场景有用

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

    参数:
        close (pl.Series): 输入的收盘价或其他数值序列。
        length (int): EMA 的周期长度。
        adjust (bool): 是否使用调整因子。当 adjust=True 时，权重会考虑初始值的存在。
                       Polars 的 ewm_mean 默认 adjust=True，与 pandas 行为一致。
        sma_initial (bool): 是否在 EMA 计算前，用前 length 个值的简单移动平均 (SMA)
                            填充第 length-1 个位置，并用 None 填充前 length-1 个位置。
                            这是一种常见的 EMA 初始化方法。

    返回:
        pl.Series: 计算出的 EMA 序列。
    """
    if not isinstance(close, pl.Series):
        raise TypeError("close 必须是一个 Polars Series。")
    if close.is_empty():
        return pl.Series(dtype=pl.Float64) # 如果输入为空，返回一个空的 Float64 Series

    length = int(length) if length and length > 0 else 10
    if length <= 0:
        raise ValueError("EMA 周期长度 (length) 必须大于 0。")

    # 复制 Series 以避免修改原始数据
    processed_close = close.clone()

    if sma_initial:
        if len(processed_close) < length:
            # 如果数据长度不足 EMA 周期，则无法计算初始 SMA
            # 返回一个包含 None 的 Series，长度与输入相同
            return pl.Series([None] * len(processed_close), dtype=pl.Float64)

        # 计算前 length 个值的简单移动平均
        initial_sma = processed_close.slice(0, length).mean()

        # 创建一个包含 (length - 1) 个 None 值的 Series
        # 明确指定 dtype 为 pl.Float64 以确保类型兼容性
        nones_series = pl.Series([None] * (length - 1), dtype=pl.Float64)

        # 创建一个包含初始 SMA 值的 Series
        sma_value_series = pl.Series([initial_sma], dtype=pl.Float64) # 确保类型一致

        # 拼接 Series: (length-1)个 None + 1个 SMA值 + 剩余数据
        # 确保所有 Series 的数据类型兼容
        processed_close = nones_series.append(sma_value_series).append(processed_close.slice(length))

    # 使用 Polars 的 ewm_mean 函数计算 EMA
    # ignore_nulls=False: 关键！确保 None 被视为占位符而非跳过。
    #                     这意味着 Polars 会在计算权重时考虑这些 None 的位置，
    #                     但它们本身不贡献数值，从而使得第一个有效 EMA 值
    #                     （即 initial_sma）能够正确地启动指数加权。
    # min_periods=0: 允许在数据点不足 span 长度时就开始计算，只要有数据就计算。
    #                这与 sma_initial 逻辑结合，确保第一个有效 EMA 值出现在正确位置。
    ema_series = processed_close.ewm_mean(
        span=length,
        adjust=adjust,
        ignore_nulls=False,
        min_periods=0
    )
    return ema_series

# 示例数据
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_result = polars_ema(sample_data, length=10)
print("计算出的 EMA 结果:")
print(ema_result)