Polars是否提供指数加权移动和(EWMS)函数？求更优实现方案

背景说明 #

• Polars内置的ewm_mean采用的更新公式为：

  y = (1 - alpha) * y + alpha * x
  

• 我实际需要的是指数加权移动和（EWMS），其更新公式为：

  y = (1 - alpha) * y + x
  

• 补充：当时间间隔相等时，EWMS的值仅为EWMA的1/α倍；但在时间间隔不等的场景下，两者的差异会更复杂。

当前实现方式 #

我目前通过.map_batches()实现了EWMS的计算，但觉得这个方案的效率不够理想，想寻求更优的替代实现方式。以下是我的代码：

import numpy as np
import polars as pl
from typing import Union

def ems_iter(values: np.array, intervals: Union[np.array, float], half_lifes: Union[np.array, float]):
    """
    >>> ems_iter([1, 0, 0, 1, 0, 0], 1, 1)
    [1.0, 0.5, 0.25, 1.125, 0.5625, 0.28125]
    """
    res = []
    agg = 0.
    values = np.array(values)
    decays = (1/2) ** (np.array(intervals) / half_lifes)
    decays = np.ones(values.shape) * decays
    for value, decay in zip(values, decays):
        agg = agg * decay + value
        res.append(agg)
    return res

def ems(val_col: str, interval_col: str, hl: int) -> pl.Expr:
    return (
        pl.struct([val_col, interval_col])
        .map_batches(
            lambda x: pl.Series(
                ems_iter(
                    x.struct.field(val_col),
                    x.struct.field(interval_col).dt.total_seconds(),
                    hl
                )
            )
        )
    )

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