Выполнение прокручивающегося окна на основе целых чисел с использованием Python Polars ⇐ Python

[Anonymous]

У меня есть функция на основе внешнего/внутреннего цикла, которую я пытаюсь векторизовать с помощью Python Polars DataFrames. Эта функция представляет собой тип скользящего среднего и будет использоваться для фильтрации финансовых данных временных рядов. Вот функция:

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

def ma_j(df_src: pl.DataFrame, depth: float):

jrc04 = 0.0
jrc05 = 0.0
jrc06 = 0.0
jrc08 = 0.0

series = df_src['close']

for x in range(0, len(series)):
if x >= x - depth*2:
for k in np.arange(start=math.ceil(depth), stop=0, step=-1):
jrc04 = jrc04 + abs(series[x-k] - series[x-(k+1)])
jrc05 = jrc05 + (depth + k) * abs(series[x-k] - series[x-(k+1)])
jrc06 = jrc06 + series[x-(k+1)]
else:
jrc03 = abs(series - (series[1]))
jrc13 = abs(series[x-depth] - series[x - (depth+1)])
jrc04 = jrc04 - jrc13 + jrc03
jrc05 = jrc05 - jrc04 + jrc03 * depth
jrc06 = jrc06 - series[x - (depth+1)] + series[x-1]
jrc08 = abs(depth * series[x] - jrc06)

if jrc05 == 0.0:
ma = 0.0
else:
ma = jrc08/jrc05

return ma

Сложность для меня заключается в множественных повторах внутреннего цикла (для k in...). Я просмотрел несколько примеров, которые используют groupby_dynamic для данных временных рядов. Например, здесь. Я также видел пример для groupby_rolling, но там, похоже, все еще используется точка.
Однако я бы хотел убрать временные ряды и просто используйте исходную серию. Означает ли это, что мне нужно группировать по целочисленному диапазону?
Используя этот пример данных:

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

import polars as pl
import numpy as np

i, t, v = np.arange(0, 50, 1), np.arange(0, 100, 2), np.random.randint(1,101,50)
df = pl.DataFrame({"i": i, "t": t, "rand": v})
df = df.with_column((pl.datetime(2022,10,30) + pl.duration(seconds=df["t"])).alias("datetime")).drop("t")
cols = ["i", "datetime", "rand"]
df = df.select(cols)

DataFrame выглядит следующим образом:

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

shape: (50, 3)
┌─────┬─────────────────────┬──────┐
│ i   ┆ datetime            ┆ rand │
│ --- ┆ ---                 ┆ ---  │
│ i32 ┆ datetime[μs]        ┆ i32  │
╞═════╪═════════════════════╪══════╡
│ 0   ┆ 2022-10-30 00:00:00 ┆ 27   │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 1   ┆ 2022-10-30 00:00:02 ┆ 82   │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 2   ┆ 2022-10-30 00:00:04 ┆ 22   │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 3   ┆ 2022-10-30 00:00:06 ┆ 58   │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┤
│ ... ┆ ...                 ┆ ...  │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 46  ┆ 2022-10-30 00:01:32 ┆ 39   │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 47  ┆ 2022-10-30 00:01:34 ┆ 48   │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 48  ┆ 2022-10-30 00:01:36 ┆ 26   │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┤
│ 49  ┆ 2022-10-30 00:01:38 ┆ 53   │
└─────┴─────────────────────┴──────┘

...Я могу выполнить группировку по дате и времени вот так":

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

df.groupby_dynamic("datetime", every="10s").agg([
pl.col("v0").mean().alias('rolling mean')
])

что дает следующее:


Но здесь есть три проблемы:

• Я не хочу группировать дату и время... Я хочу группировать в каждой строке (может быть, i?) в ячейках размером [x].
• Мне нужны значения для каждой строки.
• Я бы хотелось бы определить функцию агрегирования, как в различных случаях в функции выше

Есть какие-нибудь советы о том, как я могу решить эту проблему с помощью Polars? Спасибо.
---------- Редактировать 1
Следуя замечательному совету @ritchie46 (спасибо, приятель!), вот группа:

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

result_grp = (
df
.groupby_rolling(index_column="i", period="10i")
.agg(
[
pl.count().alias("rolling_slots"),
pl.col("rand").mean().alias("roll_mean")
]
)
)

df2 = df.select(
[
pl.all(),
result_grp.get_column("rolling_slots"),
result_grp.get_column("roll_mean"),
]
)

Теперь это дает:

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

shape: (50, 5)
┌─────┬─────────────────────┬──────┬───────────────┬───────────┐
│ i   ┆ datetime            ┆ rand ┆ rolling_slots ┆ roll_mean │
│ --- ┆ ---                 ┆ ---  ┆ ---           ┆ ---       │
│ i32 ┆ datetime[μs]        ┆ i32  ┆ u32           ┆ f64       │
╞═════╪═════════════════════╪══════╪═══════════════╪═══════════╡
│ 0   ┆ 2022-10-30 00:00:00 ┆ 55   ┆ 1             ┆ 55.0      │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 1   ┆ 2022-10-30 00:00:02 ┆ 52   ┆ 2             ┆ 53.5      │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2022-10-30 00:00:04 ┆ 30   ┆ 3             ┆ 45.666667 │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 3   ┆ 2022-10-30 00:00:06 ┆ 63   ┆ 4             ┆ 50.0      │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ ... ┆ ...                 ┆ ...  ┆ ...           ┆ ...        │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 46  ┆ 2022-10-30 00:01:32 ┆ 51   ┆ 10            ┆ 68.3      │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 47  ┆ 2022-10-30 00:01:34 ┆ 94   ┆ 10            ┆ 69.6      │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 48  ┆ 2022-10-30 00:01:36 ┆ 26   ┆ 10            ┆ 68.6      │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 49  ┆ 2022-10-30 00:01:38 ┆ 56   ┆ 10            ┆ 64.8      │
└─────┴─────────────────────┴──────┴───────────────┴───────────┘

Это великолепно; теперь вместо mean(), как мне применить пользовательскую функцию к сгруппированным значениям, например:

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

f_jparams(depth_array, jrc04, jrc05, jrc06, jrc08):

_depth = len(depth_array)

if len(depth_array) > 3:
for x in np.arange(start=1, stop=len(depth_array), step=1):
jrc04 = jrc04 + abs(depth_array[x] - depth_array[x-1])
jrc05 = jrc05 + (_depth+x) * abs(depth_array[x] - depth_array[x-1])
jrc06 = jrc06 + depth_array[x-1]
else:
jrc03 = abs(depth_array[_depth-1] - depth_array[_depth-2])
jrc13 = abs(depth_array[0] - depth_array[1])
jrc04 = jrc04 - jrc13 + jrc03
jrc05 = jrc05 - jrc04 + jrc03*_depth
jrc06 = jrc06 - depth_array[1] + depth_array[_depth-2]

jrc08 = abs(_depth * depth_array[0] - jrc06)

if jrc05 == 0.0:
ma = 0.0
else:
ma = jrc08/jrc05

return ma, jrc04, jrc05, jrc06, jrc08

Спасибо!
---- Редактировать 2:
Благодаря этому сообщению я могу собирать объединить элементы в прокручиваемой группе rand в список для каждой строки:

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

depth = 10

result_grp = (
df
.groupby_rolling(
index_column="i",
period=str(depth) + "i",
# offset="0i",
# closed="left"
)
.agg(
[
pl.count().alias("rolling_slots"),
pl.col("rand").mean().alias("roll_mean"),
pl.col("rand").suffix('_val_list'),
]
)
)

df2 = df.select(
[
pl.all(),
result_grp.get_column("rolling_slots"),
result_grp.get_column("roll_mean"),
result_grp.get_column("rand_val_list"),
]
)

Также из этого поста я увидел способ сделать период скользящего окна переменной; здорово!
Есть ли способ использовать get_columns и исключать вместе, чтобы мне не приходилось перечислять все нужные столбцы?
Теперь кадр данных выглядит так:

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

shape: (50, 6)
┌─────┬─────────────────────┬──────┬───────────────┬───────────┬──────────────────┐
│ i   ┆ datetime            ┆ rand ┆ rolling_slots ┆ roll_mean ┆ rand_val_list    │
│ --- ┆ ---                 ┆ ---  ┆ ---           ┆ ---       ┆ ---              │
│ i32 ┆ datetime[μs]        ┆ i32  ┆ u32           ┆ f64       ┆ list[i32]        │
╞═════╪═════════════════════╪══════╪═══════════════╪═══════════╪══════════════════╡
│ 0   ┆ 2022-10-30 00:00:00 ┆ 64   ┆ 1             ┆ 64.0      ┆ [64]             │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1   ┆ 2022-10-30 00:00:02 ┆ 80   ┆ 2             ┆ 72.0      ┆ [64, 80]         │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ 2022-10-30 00:00:04 ┆ 23   ┆ 3             ┆ 55.666667 ┆ [64, 80, 23]     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 3   ┆ 2022-10-30 00:00:06 ┆ 30   ┆ 4             ┆ 49.25     ┆ [64, 80, ... 30] │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ ... ┆ ...                 ┆ ...  ┆ ...           ┆ ...       ┆ ...              │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 46  ┆ 2022-10-30 00:01:32 ┆ 25   ┆ 8             ┆ 22.625    ┆ [38, 32, ... 25] │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 47  ┆ 2022-10-30 00:01:34 ┆ 69   ┆ 8             ┆ 26.5      ┆ [32, 12, ... 69] │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 48  ┆ 2022-10-30 00:01:36 ┆ 72   ┆ 8             ┆ 31.5      ┆ [12, 3, ...  72]  │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 49  ┆ 2022-10-30 00:01:38 ┆ 100  ┆ 8             ┆ 42.5      ┆ [3, 26, ... 100] │
└─────┴─────────────────────┴──────┴───────────────┴───────────┴──────────────────┘

Должен ли я сейчас вернуться к циклическому обходу столбца rand_val_list и отправить каждый список сгруппированных значений в мою функцию? Или есть лучший полярный способ?
Еще раз спасибо!

Подробнее здесь: https://stackoverflow.com/questions/749 ... hon-polars