Polars: умный способ избежать «выражения окна, не разрешенного при агрегации» ⇐ Python

[Anonymous]

У меня есть следующий код, который работает.

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

import numpy as np
import polars as pl

data = {
"date": ["2021-01-01", "2021-01-02", "2021-01-03", "2021-01-04", "2021-01-05", "2021-01-06", "2021-01-07", "2021-01-08", "2021-01-09", "2021-01-10", "2021-01-11", "2021-01-12", "2021-01-13", "2021-01-14", "2021-01-15", "2021-01-16", "2021-01-17", "2021-01-18", "2021-01-19", "2021-01-20"],
"close": np.random.randint(100, 110, 10).tolist() + np.random.randint(200, 210, 10).tolist(),
"company": ["A", "A", "A", "A", "A", "A", "A", "A", "A", "A", "B", "B", "B", "B", "B", "B", "B", "B", "B", "B"]
}
df = pl.DataFrame(data).with_columns(date = pl.col("date").cast(pl.Date))

# Calculate Returns
R = pl.col("close").pct_change()

# Calculate Gains and Losses
G = pl.when(R > 0).then(R).otherwise(0).alias("gain")
L = pl.when(R < 0).then(R).otherwise(0).alias("loss")

# Calculate Moving Averages for Gains and Losses
window = 3
MA_G = G.rolling_mean(window).alias("MA_gain")
MA_L = L.rolling_mean(window).alias("MA_loss")

# Calculate Relative Strength Index based on Moving Averages
RSI = (100 - (100 / (1 + MA_G / MA_L))).alias("RSI")

df = df.with_columns(R, G, L, MA_G, MA_L, RSI)

df.head()

Мне нравится возможность составлять различные шаги с использованием поляров, потому что это делает мой код читабельным и простым в обслуживании (в отличие от цепочки методов). Обратите внимание, что в конечном итоге расчеты становятся более сложными.
Однако теперь я хочу вычислить указанный выше столбец, но сгруппированный по «компании». Я попробовал добавить .over("company") там, где это необходимо. Однако это не работает.

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

# Calculate Returns
R = pl.col("close").pct_change().over("company")

# Calculate Gains and Losses
G = pl.when(R > 0).then(R).otherwise(0).alias("gain")
L = pl.when(R < 0).then(R).otherwise(0).alias("loss")

# Calculate Moving Averages for Gains and Losses
window = 3
MA_G = G.rolling_mean(window).alias("MA_gain").over("company")
MA_L = L.rolling_mean(window).alias("MA_loss").over("company")

# Calculate Relative Strength Index based on Moving Averages
RSI = (100 - (100 / (1 + MA_G / MA_L))).over("company").alias("RSI")

df = df.with_columns(R, G, L, MA_G, MA_L, RSI)

df.head()

Вопросы
1.) Как лучше всего исправить ошибку «оконное выражение не разрешено при агрегации», сохраняя при этом описанный выше подход к коду?
2.) Связанный вопрос: почему оконное выражение вообще не разрешено при агрегации. В чем проблема с технической точки зрения? Может кто-нибудь объяснить мне простым языком?
Спасибо!

Подробнее здесь: https://stackoverflow.com/questions/779 ... ggregation