Reduce input for helper functions of plot_signals.py so that only necessary dateframe columns are considered.

Author: niklasniedermeier

src/backtest_bay/plot/plot_signals.py

@@ -4,15 +4,17 @@
 
 def plot_signals(df, title):
     """Create Plotly figure to plot trading signals."""
-    df = _map_signals_for_plotting(df)
+    df["signal_plot"] = _map_signals_for_plotting(df["signal"])
 
     fig = make_subplots(
         rows=2, cols=1, shared_xaxes=True, row_heights=[0.7, 0.3], vertical_spacing=0.1
     )
 
-    fig.add_trace(_create_candlestick_trace(df), row=1, col=1)
+    fig.add_trace(
+        _create_candlestick_trace(df[["Open", "High", "Low", "Close"]]), row=1, col=1
+    )
 
-    buy_trace, sell_trace = _create_signal_traces(df)
+    buy_trace, sell_trace = _create_signal_traces(df["signal_plot"])
     fig.add_trace(buy_trace, row=2, col=1)
     fig.add_trace(sell_trace, row=2, col=1)
 
@@ -35,11 +37,10 @@ def plot_signals(df, title):
     return fig
 
 
-def _map_signals_for_plotting(df):
+def _map_signals_for_plotting(signal):
     """Map signals to -1, 0, 1 for plotting."""
     signal_mapping = {0: 0, 1: -1, 2: 1}
-    df["signal_plot"] = df["signal"].map(signal_mapping)
-    return df
+    return signal.map(signal_mapping)
 
 
 def _create_candlestick_trace(df):
@@ -54,20 +55,20 @@ def _create_candlestick_trace(df):
     )
 
 
-def _create_signal_traces(df):
+def _create_signal_traces(signal_plot):
     """Create signal bar traces for buy and sell signals."""
-    buy_signals = df["signal_plot"] > 0
+    buy_signals = signal_plot > 0
     buy_trace = go.Bar(
-        x=df.index[buy_signals],
-        y=df.loc[buy_signals, "signal_plot"],
+        x=signal_plot.index[buy_signals],
+        y=signal_plot[buy_signals],
         marker_color="green",
         name="Buy Signal",
     )
 
-    sell_signals = df["signal_plot"] < 0
+    sell_signals = signal_plot < 0
     sell_trace = go.Bar(
-        x=df.index[sell_signals],
-        y=df.loc[sell_signals, "signal_plot"],
+        x=signal_plot.index[sell_signals],
+        y=signal_plot[sell_signals],
         marker_color="red",
         name="Sell Signal",
     )