Polars Technical Indicators and Strategy

243 Views Asked by At

As pandas is very slow trying to use Polars. I was able to calculate the SMA and MACD values to the data frame( Not sure the values are correct for now). How can I convert the said code to use with Polars.

def implement_macd_strategy(df):    
buy_price = []
sell_price = []
macd_signal = []
signal = 0

for i in range(len(df)):
    if df['MACD_8_21_5'][i] > df['MACDs_8_21_5'][i]:
        if signal != 1:
            buy_price.append(df['Close'][i])
            sell_price.append(np.nan)
            signal = 1
            macd_signal.append(signal)
        else:
            buy_price.append(np.nan)
            sell_price.append(np.nan)
            macd_signal.append(0)
    elif df['MACD_8_21_5'][i] < df['MACDs_8_21_5'][i]:
        if signal != -1:
            buy_price.append(np.nan)
            sell_price.append(df['Close'][i])
            signal = -1
            macd_signal.append(signal)
        else:
            buy_price.append(np.nan)
            sell_price.append(np.nan)
            macd_signal.append(0)
    else:
        buy_price.append(np.nan)
        sell_price.append(np.nan)
        macd_signal.append(0)
        
df['buy_price MACD'] = buy_price
df['sell_price MACD'] = sell_price
df['signal MACD'] = signal

        
return df
2

There are 2 best solutions below

2
Iskander14yo On

it is better in such cases to give some input and desired output for testing purposes. Can you try this and say if it is ok for you:

import polars as pl

def implement_macd_strategy(df):
    buy_price = pl.col("Close").apply(lambda x: pl.when(pl.col("MACD_8_21_5") > pl.col("MACDs_8_21_5")).then(x).otherwise(pl.nan()))
    sell_price = pl.col("Close").apply(lambda x: pl.when(pl.col("MACD_8_21_5") < pl.col("MACDs_8_21_5")).then(x).otherwise(pl.nan()))
    macd_signal = pl.col("MACD_8_21_5").apply(lambda x: pl.when(x > pl.col("MACDs_8_21_5")).then(1).when(x < pl.col("MACDs_8_21_5")).then(-1).otherwise(0))

    df = df.with_column(buy_price.alias("buy_price MACD"))
    df = df.with_column(sell_price.alias("sell_price MACD"))
    df = df.with_column(macd_signal.alias("signal MACD"))

    return df
0
Wayoshi On

The iterative approach as in the original post isn't best for Polars. To fully take advantage of Polars, let's use the patterns that emerge when constructing the signal column.

We can calculate the signal column in this way, all in parallelizable Polars expressions:

  • Generate the sign of the difference of the two columns with (pl.col('MACD_8_21_5') - pl.col('MACDs_8_21_5')).sign(). This will generate a column of -1 and 1s (potentially 0 but it looks like you're not worried about equal values here, assuming it's a floating point column).
  • Note that signal is 0 on any consecutive -1s and 1s after the first one.
  • rle_id is a fairly recent addition to polars that maps a column's values to a unique integer ID for each "run" of the same value in a column.
  • Further note that is_first_distinct gives a mask that will be True once, then all Falses over each run.
  • We are now ready to create the signal column with a when expression (Polars' equivalent of if/else branches), zeroing out the non-first rows of each run.

Once we have the correct signal column, filling in the buy/sell price columns is one more simple when-then expression.

Example with dummy values (random seed set for reproduciblity):

random.seed(1)

df = pl.from_dict(
    {
        'MACD_8_21_5': [random.uniform(-1000, 1000) for _ in range(20)],
        'MACDs_8_21_5': [random.uniform(-1000, 1000) for _ in range(20)],
        'Close' : [random.uniform(10000, 100000) for _ in range(20)]
    }
)

raw_signal_expr = (pl.col('MACD_8_21_5') - pl.col('MACDs_8_21_5')).sign()
signal_expr = (
    pl.when(raw_signal_expr.is_first_distinct().over(raw_signal_expr.rle_id()))
    .then(raw_signal_expr)
    .otherwise(0)
)

df.with_columns(
    buy_price=pl.when(signal_expr > 0).then(pl.col('Close')).otherwise(np.nan),
    sell_price=pl.when(signal_expr < 0).then(pl.col('Close')).otherwise(np.nan),
    signal=signal_expr,
)
shape: (20, 6)
┌─────────────┬──────────────┬──────────────┬──────────────┬──────────────┬────────┐
│ MACD_8_21_5 ┆ MACDs_8_21_5 ┆ Close        ┆ buy_price    ┆ sell_price   ┆ signal │
│ ---         ┆ ---          ┆ ---          ┆ ---          ┆ ---          ┆ ---    │
│ f64         ┆ f64          ┆ f64          ┆ f64          ┆ f64          ┆ i64    │
╞═════════════╪══════════════╪══════════════╪══════════════╪══════════════╪════════╡
│ -731.271512 ┆ -949.108278  ┆ 99328.907096 ┆ 99328.907096 ┆ NaN          ┆ 1      │
│ 694.867474  ┆ 82.824946    ┆ 87395.187592 ┆ NaN          ┆ NaN          ┆ 0      │
│ 527.549238  ┆ 878.298326   ┆ 20880.096383 ┆ NaN          ┆ 20880.096383 ┆ -1     │
│ -489.861949 ┆ -237.591525  ┆ 39942.566682 ┆ NaN          ┆ NaN          ┆ 0      │
│ -9.129826   ┆ -566.801206  ┆ 74933.596682 ┆ 74933.596682 ┆ NaN          ┆ 1      │
│ -101.01787  ┆ -155.766849  ┆ 74007.259273 ┆ NaN          ┆ NaN          ┆ 0      │
│ 303.185945  ┆ -941.918425  ┆ 94279.652812 ┆ NaN          ┆ NaN          ┆ 0      │
│ 577.446702  ┆ -556.616667  ┆ 47989.629997 ┆ NaN          ┆ NaN          ┆ 0      │
│ -812.280826 ┆ -124.224813  ┆ 84703.212395 ┆ NaN          ┆ 84703.212395 ┆ -1     │
│ -943.305047 ┆ -8.375517    ┆ 70327.500977 ┆ NaN          ┆ NaN          ┆ 0      │
│ 671.530208  ┆ -533.831099  ┆ 37303.165984 ┆ 37303.165984 ┆ NaN          ┆ 1      │
│ -134.465864 ┆ -538.266917  ┆ 62882.254553 ┆ NaN          ┆ NaN          ┆ 0      │
│ 524.560165  ┆ -562.437925  ┆ 89423.110075 ┆ NaN          ┆ NaN          ┆ 0      │
│ -995.787893 ┆ -80.793069   ┆ 86157.767659 ┆ NaN          ┆ 86157.767659 ┆ -1     │
│ -109.225612 ┆ -420.436771  ┆ 55475.543852 ┆ 55475.543852 ┆ NaN          ┆ 1      │
│ 443.080065  ┆ -957.020589  ┆ 63010.203218 ┆ NaN          ┆ NaN          ┆ 0      │
│ -542.475557 ┆ 675.155951   ┆ 13107.324714 ┆ NaN          ┆ 13107.324714 ┆ -1     │
│ 890.541391  ┆ 112.908645   ┆ 31846.597619 ┆ 31846.597619 ┆ NaN          ┆ 1      │
│ 802.854915  ┆ 284.588726   ┆ 81766.38228  ┆ NaN          ┆ NaN          ┆ 0      │
│ -938.820034 ┆ -628.187468  ┆ 47288.259937 ┆ NaN          ┆ 47288.259937 ┆ -1     │
└─────────────┴──────────────┴──────────────┴──────────────┴──────────────┴────────┘

Putting raw_signal_expr and signal_expr in variables to reuse for readability is fine - Polars will recognize the reused expressions in the query plan and auto-cache it.