Update utils/plotting.py

utils/plotting.py

@@ -1,53 +1,60 @@
-# utils/plotting.py
-
 import matplotlib.pyplot as plt
 import matplotlib.dates as mdates
+from pandas.plotting import register_matplotlib_converters
+register_matplotlib_converters()
 
 def plot_stock_data_with_signals(data):
     """
-    Plots stock data, indicators, and buy/sell signals.
-
+    Enhanced plotting function to display stock data, indicators, and buy/sell signals with improvements.
+    
     Parameters:
-    - data (DataFrame): DataFrame containing stock 'Close' prices, indicator values, and signals.
+    - data (DataFrame): DataFrame containing stock 'Close' prices, optional indicator values, and signals.
     """
     # Create a new figure and set the size.
     plt.figure(figsize=(14, 10))
-
-    # Plot closing prices and EMAs
+    
+    # Plotting stock prices and EMAs if they exist
     ax1 = plt.subplot(311)  # 3 rows, 1 column, 1st subplot
-    data['Close'].plot(ax=ax1, color='black', lw=2., legend=True)
-    if 'EMA_Short' in data.columns and 'EMA_Long' in data.columns:
+    data['Close'].plot(ax=ax1, color='black', lw=2., legend=True, label='Close')
+    if 'EMA_Short' in data and 'EMA_Long' in data:
         data[['EMA_Short', 'EMA_Long']].plot(ax=ax1, lw=1.5, legend=True)
-    ax1.set_title('Stock Price, EMAs, and Bollinger Bands')
-    ax1.fill_between(data.index, data['BB_Lower'], data['BB_Upper'], color='grey', alpha=0.3)
+    ax1.set_title('Stock Price and Indicators')
+    
+    # Plotting Bollinger Bands if they exist
+    if 'BB_Upper' in data and 'BB_Lower' in data:
+        ax1.fill_between(data.index, data['BB_Lower'], data['BB_Upper'], color='grey', alpha=0.3, label='Bollinger Bands')
+    
+    # Highlight buy/sell signals if Combined_Signal column exists
+    if 'Combined_Signal' in data:
+        buy_signals = data[data['Combined_Signal'] == 'buy']
+        sell_signals = data[data['Combined_Signal'] == 'sell']
+        ax1.plot(buy_signals.index, buy_signals['Close'], '^', markersize=10, color='g', lw=0, label='Buy Signal')
+        ax1.plot(sell_signals.index, sell_signals['Close'], 'v', markersize=10, color='r', lw=0, label='Sell Signal')
     
-    # Highlight buy/sell signals
-    buy_signals = data[data['Combined_Signal'] == 'buy']
-    sell_signals = data[data['Combined_Signal'] == 'sell']
-    ax1.plot(buy_signals.index, data.loc[buy_signals.index]['Close'], '^', markersize=10, color='g', lw=0, label='Buy Signal')
-    ax1.plot(sell_signals.index, data.loc[sell_signals.index]['Close'], 'v', markersize=10, color='r', lw=0, label='Sell Signal')
     ax1.legend()
 
-    # Plot MACD and Signal Line
-    ax2 = plt.subplot(312, sharex=ax1)  # Share x-axis with ax1
-    data['MACD'].plot(ax=ax2, color='blue', label='MACD', legend=True)
-    data['MACD_Signal_Line'].plot(ax=ax2, color='red', label='Signal Line', legend=True)
-    ax2.fill_between(data.index, data['MACD'] - data['MACD_Signal_Line'], color='grey', alpha=0.3)
-    ax2.set_title('MACD')
-
-    # Plot RSI
-    ax3 = plt.subplot(313, sharex=ax1)  # Share x-axis with ax1
-    data['RSI'].plot(ax=ax3, color='purple', legend=True)
-    ax3.axhline(70, linestyle='--', alpha=0.5, color='red')
-    ax3.axhline(30, linestyle='--', alpha=0.5, color='green')
-    ax3.set_title('RSI')
-
-    # Improve layout and x-axis date format
+    # Plotting MACD and Signal Line if they exist
+    if 'MACD' in data and 'MACD_Signal_Line' in data:
+        ax2 = plt.subplot(312, sharex=ax1)  # Share x-axis with ax1
+        data['MACD'].plot(ax=ax2, color='blue', label='MACD', legend=True)
+        data['MACD_Signal_Line'].plot(ax=ax2, color='red', label='Signal Line', legend=True)
+        ax2.set_title('MACD')
+        ax2.legend()
+
+    # Plotting RSI if it exists
+    if 'RSI' in data:
+        ax3 = plt.subplot(313, sharex=ax1)  # Share x-axis with ax1
+        data['RSI'].plot(ax=ax3, color='purple', legend=True, label='RSI')
+        ax3.axhline(70, linestyle='--', alpha=0.5, color='red', label='Overbought')
+        ax3.axhline(30, linestyle='--', alpha=0.5, color='green', label='Oversold')
+        ax3.set_title('RSI')
+        ax3.legend()
+
+    # Improving layout, setting x-axis format for better date handling
     plt.tight_layout()
-    plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
-    plt.gca().xaxis.set_major_locator(mdates.DayLocator())
-    plt.xticks(rotation=45)
-
+    for ax in [ax1, ax2, ax3]:
+        ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
+        ax.xaxis.set_major_locator(mdates.AutoDateLocator())
+        plt.setp(ax.xaxis.get_majorticklabels(), rotation=45)
+    
     plt.show()
-
-# Note: This is a basic example for visualization. You may need to adjust it based on your actual 'data' DataFrame structure and the specific indicators you are plotting.