Advice refactoring polars expr - python-polars

Question

I have a polars expr, and I cannot use a context 'cause my function has to return a polars expr.
I've implemented a RSI indicator in polars:
rsi_indicator = (100*pl.when(pl.col("close").pct_change() >= 0) \
.then(pl.col("close").pct_change()) \
.otherwise(0.0) \
.rolling_mean(window_size=window) \
/ (pl.when(pl.col("close").pct_change() >= 0) \
.then(pl.col("close").pct_change()) \
.otherwise(0.0).rolling_mean(window_size=window) + \
pl.when(pl.col("close").pct_change() < 0) \
.then(pl.col("close").pct_change()) \
.otherwise(0.0).abs().rolling_mean(window_size=window))).alias(f"rsi_{window}")
I would refactor this code isolating some quantities in order to easily maintaining code and readability.
For example I'd like to define a variable
U = pl.when(pl.col("close").pct_change() >= 0) \
.then(pl.col("close").pct_change()) \
.otherwise(0.0) \
.rolling_mean(window_size=window)
and its fiend V for negative returns in order to return simply 100*U/(U+V) but seems it doesn't work. Any advise?

Answer 1

First, let's use some real financial data, so that our results look reasonable. I'll also change the variable names so that your code will work, as-is.
import polars as pl
from yfinance import Ticker
ticker_data = Ticker("AAPL").history(period="3mo")
df = pl.from_pandas(ticker_data)
df = df.rename({col_nm: col_nm.lower() for col_nm in df.columns})
df.tail(10)
shape: (10, 7)
┌────────────┬────────────┬────────────┬────────────┬──────────┬───────────┬──────────────┐
│ open ┆ high ┆ low ┆ close ┆ volume ┆ dividends ┆ stock splits │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ f64 ┆ f64 ┆ f64 ┆ f64 ┆ i64 ┆ f64 ┆ i64 │
╞════════════╪════════════╪════════════╪════════════╪══════════╪═══════════╪══════════════╡
│ 136.820007 ┆ 138.589996 ┆ 135.630005 ┆ 138.270004 ┆ 72433800 ┆ 0.0 ┆ 0 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 139.899994 ┆ 141.910004 ┆ 139.770004 ┆ 141.660004 ┆ 89116800 ┆ 0.0 ┆ 0 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 142.699997 ┆ 143.490005 ┆ 140.970001 ┆ 141.660004 ┆ 70207900 ┆ 0.0 ┆ 0 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 142.130005 ┆ 143.419998 ┆ 137.320007 ┆ 137.440002 ┆ 67083400 ┆ 0.0 ┆ 0 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 137.460007 ┆ 140.669998 ┆ 136.669998 ┆ 139.229996 ┆ 66242400 ┆ 0.0 ┆ 0 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 137.25 ┆ 138.369995 ┆ 133.770004 ┆ 136.720001 ┆ 98964500 ┆ 0.0 ┆ 0 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 136.039993 ┆ 139.039993 ┆ 135.660004 ┆ 138.929993 ┆ 71007500 ┆ 0.0 ┆ 0 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 137.770004 ┆ 141.610001 ┆ 136.929993 ┆ 141.559998 ┆ 73353800 ┆ 0.0 ┆ 0 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 141.350006 ┆ 144.119995 ┆ 141.080002 ┆ 142.919998 ┆ 73972200 ┆ 0.0 ┆ 0 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 141.354996 ┆ 144.335007 ┆ 143.300003 ┆ 144.304993 ┆ 5503179 ┆ 0.0 ┆ 0 │
└────────────┴────────────┴────────────┴────────────┴──────────┴───────────┴──────────────┘
Next, I'll reformat your code, and construct U and V, per your question.
window = 10
rsi_indicator = (
100
* pl.when(pl.col("close").pct_change() >= 0)
.then(pl.col("close").pct_change())
.otherwise(0.0)
.rolling_mean(window_size=window)
/ (
pl.when(pl.col("close").pct_change() >= 0)
.then(pl.col("close").pct_change())
.otherwise(0.0)
.rolling_mean(window_size=window)
+ pl.when(pl.col("close").pct_change() < 0)
.then(pl.col("close").pct_change())
.otherwise(0.0)
.abs()
.rolling_mean(window_size=window)
)
).alias(f"rsi_{window}")
U = (
pl.when(pl.col("close").pct_change() >= 0)
.then(pl.col("close").pct_change())
.otherwise(0.0)
.rolling_mean(window_size=window)
)
V = (
pl.when(pl.col("close").pct_change() < 0)
.then(pl.col("close").pct_change())
.otherwise(0.0)
.abs()
.rolling_mean(window_size=window)
)
We can then express rsi_indicator in terms of U and V as follows:
df.select([
pl.col('close'),
rsi_indicator,
100 * (U / (U + V)).alias('rsi_UV')
]).tail(10)
shape: (10, 3)
┌────────────┬───────────┬───────────┐
│ close ┆ rsi_10 ┆ rsi_UV │
│ --- ┆ --- ┆ --- │
│ f64 ┆ f64 ┆ f64 │
╞════════════╪═══════════╪═══════════╡
│ 138.270004 ┆ 37.209581 ┆ 37.209581 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 141.660004 ┆ 49.321609 ┆ 49.321609 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 141.660004 ┆ 58.898881 ┆ 58.898881 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 137.440002 ┆ 61.526297 ┆ 61.526297 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 139.229996 ┆ 62.768001 ┆ 62.768001 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 136.720001 ┆ 53.110534 ┆ 53.110534 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 138.929993 ┆ 69.836919 ┆ 69.836919 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 141.559998 ┆ 71.086117 ┆ 71.086117 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 142.919998 ┆ 66.779857 ┆ 66.779857 │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 144.304993 ┆ 70.359592 ┆ 70.359592 │
└────────────┴───────────┴───────────┘