0

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?

Sigi
  • 53
  • 8

1 Answers1

2

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 │
└────────────┴───────────┴───────────┘
  • Thanks. Actually you have used a select environment. I need only expressions, this is why I have a function that takes a string and returns an Expr. Another function collects the set of strings and returns a list of Expr that can be used inside the context to create the polars data frame. – Sigi Jul 07 '22 at 13:55
  • 1
    No worries, the result is an expression, and needs no context. `type( (100 * (U / (U + V)).alias('rsi_UV')) ) ` returns `` –  Jul 07 '22 at 14:14