I have the following code which works.

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()

I like the ability to compose different steps using polars, because it keeps my code readable and easy to maintain (as opposed to method chaining). Note that ultimately calculations are more complex.

However, now I want to calculate the above column but grouped by "company". I tried adding .over("company") where relevant. However, this doesn't work.

# 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()

Questions

1.) What is the best way to fix this "window expression not allowed in aggregation" error while keeping the above code approach?

2.) Related question: why is window expression not allowed in aggregation in the first place. What is the problem with this from a technical perspective? Can someone explain to me in laymans terms?

Thanks!