在Polars中高效计算指数移动平均线(EMA)及其初始化策略

import polars as pl
import numpy as np

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

    参数:
        close (pl.Series): 输入的数值序列。
        length (int): EMA的周期长度。
        adjust (bool): 是否进行调整。在Polars中，通常建议根据需求明确设置。
        sma_init (bool): 如果为True，则前length个周期的EMA将以第一个length周期的SMA初始化。

    返回:
        pl.Series: 计算出的EMA序列。
    """
    # 参数校验
    if not isinstance(length, int) or length <= 0:
        length = 10
    if close is None:
        return pl.Series(dtype=pl.Float64)
    if not isinstance(close, pl.Series):
        raise TypeError("Input 'close' must be a Polars Series.")

    # 确保序列为浮点类型，以便处理空值和计算
    if close.dtype not in [pl.Float32, pl.Float64]:
        close = close.cast(pl.Float64)

    # 如果不需要SMA初始化，则直接计算标准EMA
    if not sma_init:
        return close.ewm_mean(span=length, adjust=adjust, ignore_nulls=False, min_periods=length)

    # 2. 计算初始SMA
    # 获取前length个数据并计算其均值
    initial_sma = close.head(length).mean()

    # 3. 构建带空值的序列
    # 创建 length-1 个 None 值序列。
    # 关键点：使用 Python 的 None 而非 np.NaN，并明确指定 dtype 为 Float64。
    # Polars的ewm_mean在处理None时表现更一致。
    nones_series = pl.Series([None] * (length - 1), dtype=pl.Float64)

    # 将计算出的SMA值封装为单元素Series
    sma_value_series = pl.Series("initial_sma", [initial_sma], dtype=pl.Float64)

    # 将 None 序列、SMA值序列和原始序列的剩余部分拼接起来
    # 这样构造的序列在第 length-1 索引处是SMA值，之前是None
    processed_close = nones_series.append(sma_value_series).append(close.tail(close.len() - length))

    # 4. 应用指数移动平均
    # 使用 ewm_mean 计算EMA。
    # ignore_nulls=False: 确保空值在计算中被考虑（即在遇到空值时，除非min_periods允许，否则不进行计算）。
    # min_periods=0: 允许从第一个非空值（即我们插入的SMA值）开始计算EMA。
    ema = processed_close.ewm_mean(span=length, adjust=adjust, ignore_nulls=False, min_periods=0)

    return ema

# 示例数据
sample_data = [
    1.08086, 1.08069, 1.08077, 1.08077, 1.08052,
    1.08055, 1.08068, 1.08073, 1.08077, 1.08073, # 第10个值
    1.08068, 1.08062, 1.08052, 1.0806,  1.08063,
    1.08064, 1.08063, 1.08053, 1.08067, 1.08058
]
close_series = pl.Series("close", sample_data)

# 调用函数计算EMA
length = 10
calculated_ema = polars_ema(close_series, length=length)

print("原始序列:")
print(close_series)
print(f"\n计算的EMA (周期={length}):")
print(calculated_ema)

# 验证第一个非空EMA值是否为前length个数据的SMA
first_n_sma = close_series.head(length).mean()
print(f"\n前 {length} 个数据的SMA: {first_n_sma}")
# 找到第一个非None的EMA值
first_ema_value = calculated_ema.drop_nulls().head(1).item()
print(f"计算出的EMA序列中第一个非空值: {first_ema_value}")
# 预期：first_ema_value 应该约等于 first_n_sma