Update utils/plotting.py

utils/plotting.py

@@ -1,51 +1,49 @@
-# utils/plotting.py
-
 import matplotlib.pyplot as plt
-import pandas as pd
+import matplotlib.dates as mdates
 
-def plot_stock_data(data, buy_signals=None, sell_signals=None, title="Stock Data with Indicators"):
+def plot_stock_data_with_signals(stock_data):
     """
-    Plots stock data with SMAs, Bollinger Bands, and buy/sell signals.
+    Plots the stock data along with SMAs, Bollinger Bands, and buy/sell signals.
 
     Parameters:
-    - data: DataFrame containing the stock data with 'Close', 'SMA_21', 'SMA_50', 'BB_Upper', and 'BB_Lower' columns.
-    - buy_signals: DataFrame or Series with buy signals. Must contain a 'Date' or similar index for plotting.
-    - sell_signals: DataFrame or Series with sell signals. Must contain a 'Date' or similar index for plotting.
-    - title: Title of the plot.
+    - stock_data (pd.DataFrame): The stock data with 'Close', 'SMA_21', 'SMA_50', 'BB_Upper', 'BB_Lower', 'Buy_Signal', and 'Sell_Signal' columns.
     """
-    # Create a new figure and set the size
-    plt.figure(figsize=(14, 7))
+    # Setting up the plot
+    fig, ax = plt.subplots(figsize=(14, 7))
+
+    # Plotting the closing prices
+    ax.plot(stock_data.index, stock_data['Close'], label='Close Price', color='blue', alpha=0.5)
 
-    # Plot the closing price
-    plt.plot(data.index, data['Close'], label='Close Price', color='skyblue', linewidth=2)
+    # Plotting the SMAs
+    ax.plot(stock_data.index, stock_data['SMA_21'], label='21-Period SMA', color='orange', alpha=0.75)
+    ax.plot(stock_data.index, stock_data['SMA_50'], label='50-Period SMA', color='green', alpha=0.75)
 
-    # Plot SMAs
-    plt.plot(data.index, data['SMA_21'], label='21-period SMA', color='orange', linewidth=1.5)
-    plt.plot(data.index, data['SMA_50'], label='50-period SMA', color='green', linewidth=1.5)
+    # Plotting the Bollinger Bands
+    ax.plot(stock_data.index, stock_data['BB_Upper'], label='Upper Bollinger Band', color='red', linestyle='--', alpha=0.5)
+    ax.plot(stock_data.index, stock_data['BB_Lower'], label='Lower Bollinger Band', color='cyan', linestyle='--', alpha=0.5)
 
-    # Plot Bollinger Bands
-    plt.plot(data.index, data['BB_Upper'], label='Upper Bollinger Band', color='grey', linestyle='--', linewidth=1)
-    plt.plot(data.index, data['BB_Lower'], label='Lower Bollinger Band', color='grey', linestyle='--', linewidth=1)
+    # Highlighting buy signals
+    buy_signals = stock_data[stock_data['Buy_Signal']]
+    ax.scatter(buy_signals.index, buy_signals['Close'], label='Buy Signal', marker='^', color='green', alpha=1, s=100)
 
-    # Highlight buy signals
-    if buy_signals is not None:
-        plt.scatter(buy_signals.index, data.loc[buy_signals.index]['Close'], marker='^', color='green', label='Buy Signal', alpha=1)
+    # Highlighting sell signals
+    sell_signals = stock_data[stock_data['Sell_Signal']]
+    ax.scatter(sell_signals.index, sell_signals['Close'], label='Sell Signal', marker='v', color='red', alpha=1, s=100)
 
-    # Highlight sell signals
-    if sell_signals is not None:
-        plt.scatter(sell_signals.index, data.loc[sell_signals.index]['Close'], marker='v', color='red', label='Sell Signal', alpha=1)
+    # Beautifying the plot
+    ax.set_title("Stock Price with Indicators and Signals")
+    ax.set_xlabel("Date")
+    ax.set_ylabel("Price")
+    ax.legend()
 
-    # Customize the plot
-    plt.title(title)
-    plt.xlabel('Date')
-    plt.ylabel('Price')
-    plt.legend()
-    plt.grid(True)
+    # Formatting the x-axis to show dates nicely
+    ax.xaxis.set_major_locator(mdates.WeekdayLocator())
+    ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
+    fig.autofmt_xdate() # Rotate date labels to prevent overlap
 
-    # Show the plot
     plt.show()
 
-# Example usage:
-# This assumes `data` DataFrame is already loaded with the required columns including 'Close', 'SMA_21', 'SMA_50', 'BB_Upper', 'BB_Lower'.
-# `buy_signals` and `sell_signals` DataFrames/Series should have the dates of signals.
-# Due to the nature of this example, actual data is not provided here. To test this function, ensure you have a DataFrame with the appropriate structure.
+if __name__ == "__main__":
+    # Example usage:
+    # Generate or fetch your stock_data DataFrame with necessary columns before calling this function
+    pass