如何在Polars中实现依赖前值与当前列的递推计算列xL0？

问题描述 #

需要实现的递推逻辑：

xL0[0] = (1 - gamma) * close[0]
xL0[n] = (1 - gamma) * close[n] + gamma * xL0[n-1] （n ≥ 1）

尝试过.cumsum()和shift方法未成功，拒绝使用lambda（除非速度优于pandas向量化操作）。

示例输入DataFrame：

import polars as pl

df = pl.DataFrame(
    {
        "close": [20.92, 20.7, 20.58, 20.88]
    }
)

当gamma = 0.7时，预期输出的xL0值为：6.276, 10.6032, 13.59624, 15.781368

解决方案 #

方法1：O(n)高效递推（推荐）

使用map_batches结合numpy底层优化的循环，完全向量化且速度远超pandas lambda/循环：

import polars as pl
import numpy as np

gamma = 0.7
alpha = 1 - gamma

def compute_xL0(close_series):
    close_arr = close_series.to_numpy()
    xL0_arr = np.zeros_like(close_arr)
    xL0_arr[0] = alpha * close_arr[0]
    for i in range(1, len(close_arr)):
        xL0_arr[i] = alpha * close_arr[i] + gamma * xL0_arr[i-1]
    return pl.Series(xL0_arr)

df = df.with_columns(
    xL0=pl.col("close").map_batches(compute_xL0)
)

print(df)

输出结果：

shape: (4, 2)
┌───────┬───────────┐
│ close ┆ xL0       │
│ ---   ┆ ---       │
│ f64   ┆ f64       │
╞═══════╪═══════════╡
│ 20.92 ┆ 6.276     │
│ 20.7  ┆ 10.6032   │
│ 20.58 ┆ 13.59624  │
│ 20.88 ┆ 15.781368 │
└───────┴───────────┘

方法2：数学展开式向量化计算（小数据集适用）

将递推公式展开为加权求和，适合数据量较小的场景：

import polars as pl
import numpy as np

gamma = 0.7
alpha = 1 - gamma

df = pl.DataFrame(
    {"close": [20.92, 20.7, 20.58, 20.88]}
)

# 生成权重矩阵：第i行对应第i个xL0的权重
n = len(df)
weights = np.triu(gamma ** (np.arange(n)[:, np.newaxis] - np.arange(n)))
# 计算加权和并生成xL0列
df = df.with_columns(
    xL0=alpha * pl.col("close").dot(weights).list.to_series()
)

方法3：基于Polars内置EWMA调整（简洁）

利用Polars内置的EWMA函数，调整参数匹配初始条件，这里的lambda仅处理单个值，速度远快于pandas向量化操作：

import polars as pl

gamma = 0.7
alpha = 1 - gamma

df = df.with_columns(
    xL0=pl.col("close").ewm_mean(alpha=alpha, adjust=False) * alpha / (1 - gamma**pl.int_range(0, pl.count()))
)

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