如何在Polars中用自定义累积表达式实现时间感知指数移动平均？

你可以利用Polars的cumulative_eval函数（Polars 0.19.0+版本支持）实现递归的时间感知EMA，完全保留惰性计算能力和原生性能，无需依赖numba。以下是具体实现步骤和代码示例：

核心思路 #

匹配你定义的递归公式：

• ( y_0 = x_0 )
• ( y_{i+1} = q_{i+1} \times y_i + (1 - q_{i+1}) \times x_{i+1} )，其中 ( q = exp(-\Delta t / \lambda) )，( \Delta t = t_{i+1} - t_i )

cumulative_eval支持定义带状态的累积计算，用前一个EMA值（状态）结合当前行的x和q值，迭代计算每一行的EMA结果，全程由Polars查询引擎执行，保持高效性。

完整代码示例 #

1. 基础单序列计算 #

import polars as pl

# 示例数据：时间列+数值列
df = pl.DataFrame({
    "t": pl.date_range(start="2023-01-01", end="2023-01-10", interval="1d"),
    "x": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]
})

λ = 3.0  # 平滑参数，单位需与Δt一致（此处为天）

# 步骤1：计算时间差Δt（转换为数值型，这里用天数）
df = df.with_columns(
    Δt=pl.col("t").diff().dt.total_days()
)

# 步骤2：计算q列，第一行填充1.0（不会被实际使用）
df = df.with_columns(
    q=pl.exp(-pl.col("Δt")/λ).fill_null(1.0)
)

# 步骤3：用cumulative_eval计算时间感知EMA
df = df.with_columns(
    y=pl.col("x").cumulative_eval(
        # 递归计算逻辑：前一个EMA值 * q + 当前x值 * (1-q)
        lambda acc, x_current, q_current: q_current * acc + (1 - q_current) * x_current,
        initial=pl.col("x").first(),  # 初始值y0=x0
        args=[pl.col("x"), pl.col("q")]  # 传入当前行的x和q值
    )
)

print(df)

2. 分组计算场景 #

如果需要按分组单独计算EMA，只需结合over子句：

import polars as pl

df = pl.DataFrame({
    "group": ["A", "A", "A", "B", "B", "B"],
    "t": pl.date_range(start="2023-01-01", end="2023-01-06", interval="1d"),
    "x": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
})

λ = 3.0

df = df.with_columns(
    # 按分组计算时间差
    Δt=pl.col("t").diff().dt.total_days().over("group")
).with_columns(
    # 按分组计算q列
    q=pl.exp(-pl.col("Δt")/λ).fill_null(1.0).over("group")
).with_columns(
    # 按分组执行累积计算
    y=pl.col("x").cumulative_eval(
        lambda acc, x_current, q_current: q_current * acc + (1 - q_current) * x_current,
        initial=pl.col("x").first().over("group"),
        args=[pl.col("x"), pl.col("q")]
    ).over("group")
)

print(df)

关键说明 #

• 惰性计算支持：如果使用pl.scan_csv等惰性读取方式，上述逻辑可完全链式执行，不会触发提前计算，保留Polars的内存效率优势。
• 时间差适配：根据你的时间单位（秒/分/天），调整dt.total_seconds()/dt.total_minutes()/dt.total_days()转换Δt的数值类型，确保与λ的单位一致。
• 空值处理：第一行的Δt为null，对应的q值填充为1.0，但由于cumulative_eval从第二行开始计算，初始值直接取第一个x值，因此该填充值不会影响结果。

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