Как применитьroll_map к полярам и создать два столбца? ⇐ Python

[Anonymous]

Смотрите код ниже, это замена того, что мне нужно. ChatGPT и недоумение здесь не помогут. На практике я могу вычислять столбцы col1 и col2 по одному, но это неоправданно медленно.

Код: Выделить всё

import numpy as np
import polars as pl

def _compute_coef1(series):
y = series[1:]
x = series[:-1]

mean_x = x.mean()
mean_y = y.mean()

x_centered = x - mean_x
y_centered = y - mean_y

den = (x_centered * x_centered).sum() or np.nan
coef1 = (x_centered * y_centered).sum() / den

return coef1

def _compute_coef2(series):
y = series[1:]
x = series[:-1]

mean_x = x.mean()
mean_y = y.mean()

x_centered = x - mean_x
y_centered = y - mean_y

den = (x_centered * x_centered).sum() or np.nan
coef1 = (x_centered * y_centered).sum() / den
coef2 = mean_y - coef1 * mean_x

return coef2

# Apply the rolling computation of AR(1) coefficients and coef2
expr1 = (
pl.col("input_col")
.rolling_map(_compute_coef1, window_size=10, min_periods=3)
.alias("coef1")
)

expr2 = (
pl.col("input_col")
.rolling_map(_compute_coef2, window_size=10, min_periods=3)
.alias("coef2")
)

# Testing:
df = pl.DataFrame({"input_col": [1, 1, 2, 2]})

df = df.with_columns(expr1).with_columns(expr2)

print(df)

Хотя приведенный выше код выполняется и работает, он работает неоправданно медленно. Настоящая версия, которую я использую, невыносима и требует слишком много памяти.

Подробнее здесь: https://stackoverflow.com/questions/786 ... wo-columns