基于Polars实现MACD交易策略的代码转换咨询

Hey there! I totally get why you're switching to Polars—Pandas loops can crawl on large datasets, and Polars' vectorized approach is way faster. Let's convert your MACD strategy code, and I'll also fix a small bug I noticed in your original Pandas code along the way.

First, let's recap what your code is supposed to do:

• When the MACD line crosses above the signal line (and we weren't already holding a long position), we log the current close as a buy price and mark a buy signal
• When the MACD line crosses below the signal line (and we weren't already in a short position), we log the current close as a sell price and mark a sell signal
• For all other rows, we leave buy/sell prices as null and the signal as 0

Quick note on your original code: It looks like there's a typo at the end—you assigned the single signal variable to df['signal MACD'] instead of the macd_signal list. I'll make sure our Polars code fixes that so the signal_MACD column correctly shows 1/-1 only when a signal triggers, 0 otherwise.

Here's the efficient Polars implementation, no loops required:

import polars as pl

def implement_macd_strategy_pl(df: pl.DataFrame) -> pl.DataFrame:
    return df.with_columns(
        # Step 1: Detect golden crosses (MACD crosses above signal) and death crosses (MACD crosses below signal)
        golden_cross = pl.when(
            (pl.col("MACD_8_21_5") > pl.col("MACDs_8_21_5")) & 
            (pl.col("MACD_8_21_5").shift(1) <= pl.col("MACDs_8_21_5").shift(1))
        ).then(1).otherwise(0),
        death_cross = pl.when(
            (pl.col("MACD_8_21_5") < pl.col("MACDs_8_21_5")) & 
            (pl.col("MACD_8_21_5").shift(1) >= pl.col("MACDs_8_21_5").shift(1))
        ).then(1).otherwise(0),
    ).with_columns(
        # Step 2: Track our current position state (1 = long, -1 = short, 0 = flat)
        # We only update state when a cross happens, otherwise carry over the previous state
        position_state = pl.when(pl.col("golden_cross") == 1)
                   .then(1)
                   .when(pl.col("death_cross") == 1)
                   .then(-1)
                   .otherwise(pl.col("position_state").shift(1))
                   .fill_null(0),
    ).with_columns(
        # Step 3: Generate buy/sell prices only when our position changes
        buy_price_MACD = pl.when(
            (pl.col("position_state") == 1) & (pl.col("position_state").shift(1) != 1)
        ).then(pl.col("Close")).otherwise(None),
        sell_price_MACD = pl.when(
            (pl.col("position_state") == -1) & (pl.col("position_state").shift(1) != -1)
        ).then(pl.col("Close")).otherwise(None),
        # Step 4: Create the signal column (1/-1 on new signals, 0 otherwise)
        signal_MACD = pl.when(
            (pl.col("position_state") == 1) & (pl.col("position_state").shift(1) != 1)
        ).then(1).when(
            (pl.col("position_state") == -1) & (pl.col("position_state").shift(1) != -1)
        ).then(-1).otherwise(0),
    ).drop("golden_cross", "death_cross", "position_state")  # Clean up temporary helper columns

How this works: #

• Cross Detection: We use shift(1) to compare the current MACD/signal line values with the previous row, so we only flag actual cross events (not just when MACD is above/below the signal line).
• Position Tracking: Instead of a loop variable, we use a column that carries over the previous position state unless a cross triggers a change. This mimics the signal variable in your original loop but in a vectorized way.
• Signal & Price Generation: We check when the position state changes to a new value (from non-long to long, or non-short to short) to log buy/sell prices and signal values. All other rows get null (Polars' equivalent of np.nan) for prices and 0 for signals.

This code will run drastically faster than your original Pandas loop, especially on large datasets, and it's more maintainable since it follows Polars' idiomatic patterns. Just make sure your input Polars DataFrame already has the MACD_8_21_5, MACDs_8_21_5, and Close columns you mentioned you've calculated.

备注：内容来源于stack exchange，提问作者Imtiaz Sattar