[Python] Problem converting Pandas-based ATR function to be Polars-based

I have just started trying to convert some stock trading python code from being Pandas-based to use Polars. So, I am a complete newbie to Polars, not much better at Python, but I get by. I am migrating to Polars because of the increased performance of Polars. The following is my Pandas-based ATR function:

def ATR(df: pa.DataFrame, window_size: int = 14) -> pa.DataFrame:
high, low, prev_close = df['high'], df['low'], df['close'].shift()
tr_all = [high - low, high - prev_close, low - prev_close]
tr_all = [tr.abs() for tr in tr_all]

tr = pa.concat(tr_all, axis = 1).max(axis = 1)
df['ATR'] = tr.ewm(alpha = 1/window_size, min_periods = window_size, adjust = False, ignore_na = True).mean()
return df


I call this function like this:

raw_dd_df = ATR(raw_dd_df, window_size = slowline_period)


and produces results similar to the following:

date open high low ... volume vwap mid ATR
300 2024-01-01 1.27300 1.27330 1.26936 ... 0 1.27189 1.27300 NaN
299 2024-01-02 1.27291 1.27597 1.26105 ... 242445 1.26779 1.26707 NaN
298 2024-01-03 1.26123 1.26765 1.26123 ... 296035 1.26414 1.26384 NaN
297 2024-01-04 1.26644 1.27295 1.26565 ... 270883 1.26830 1.26730 NaN
296 2024-01-05 1.26816 1.27710 1.26113 ... 333038 1.26949 1.26987 NaN
.. ... ... ... ... ... ... ... ... ...
4 2024-12-26 1.25296 1.25474 1.25005 ... 245898 1.25241 1.25242 0.009308
3 2024-12-27 1.25187 1.25925 1.25046 ... 234639 1.25464 1.25443 0.009282
2 2024-12-29 1.25756 1.25756 1.25756 ... 0 1.25756 1.25756 0.008846
1 2024-12-30 1.25726 1.26070 1.25059 ... 243089 1.25587 1.25610 0.008910
0 2024-12-31 1.25493 1.25688 1.25048 ... 227765 1.25336 1.25303 0.008784


The following is my attempt to rewite the function to use Polars:

def pl_ATR(df: pl.DataFrame, window_size: int = 14) -> pl.DataFrame:
high, low, prev_close = df['High'], df['Low'], df['Close'].shift()
tr_all = [high - low, high - prev_close, low - prev_close]
tr_all = [tr.abs() for tr in tr_all]

tr = pl.concat(tr_all, rechunk = True).max()
df['ATR'] = pl.Expr.ewm_mean(tr, alpha = 1/window_size, min_samples = window_size, adjust = False, ignore_nulls = True)
return df


I am not sure if i have got tr = pl.concat(tr_all, rechunk = True).max() right, but I get errors at the next line (when trying to do pl.Expr.ewm_mean):

Traceback (most recent call last):
File "/home/stuart/Projects/Python/Trading/Scratches/Polars/demo_polars.py", line 48, in <module>
df = pl_ATR(df, window_size = 20)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/stuart/Projects/Python/Trading/Scratches/Polars/demo_polars.py", line 37, in pl_ATR
df['ATR'] = pl.Expr.ewm_mean(tr, alpha = 1/window_size, min_samples = window_size, adjust = False, ignore_nulls = True)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/stuart/Projects/Python/Env/lib/python3.12/site-packages/polars/_utils/deprecation.py", line 119, in wrapper
return function(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/stuart/Projects/Python/Env/lib/python3.12/site-packages/polars/expr/expr.py", line 9488, in ewm_mean
return self._from_pyexpr(
^^^^^^^^^^^^^^^^^
AttributeError: 'float' object has no attribute '_from_pyexpr'


I have spent hours trying to get it working, but to no avail. Can anybody help please?

Regards, Stuart

Continue reading...