Update utils/plotting.py

utils/plotting.py

@@ -1,60 +1,44 @@
 import matplotlib.pyplot as plt
 import matplotlib.dates as mdates
+import pandas as pd
 from pandas.plotting import register_matplotlib_converters
+
 register_matplotlib_converters()
 
 def plot_stock_data_with_signals(data):
- """
- Enhanced plotting function to display stock data, indicators, and buy/sell signals with improvements.
- 
- Parameters:
- - data (DataFrame): DataFrame containing stock 'Close' prices, optional indicator values, and signals.
- """
+ # First, ensure the DataFrame's index is in datetime format.
+ data.index = pd.to_datetime(data.index, errors='coerce')
+
  # Create a new figure and set the size.
- plt.figure(figsize=(14, 10))
- 
- # 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, 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 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:
+ plt.figure(figsize=(14, 7))
+
+ # Plotting stock 'Close' prices
+ plt.plot(data.index, data['Close'], label='Close Price', color='black', lw=2)
+
+ # Check and plot EMAs if they exist
+ if 'EMA_20' in data.columns and 'EMA_50' in data.columns:
+ plt.plot(data.index, data['EMA_20'], label='EMA 20', color='blue', lw=1.5)
+ plt.plot(data.index, data['EMA_50'], label='EMA 50', color='red', lw=1.5)
+
+ # Check and plot Bollinger Bands if they exist
+ if 'BB_Upper' in data.columns and 'BB_Lower' in data.columns:
+ plt.fill_between(data.index, data['BB_Lower'], data['BB_Upper'], color='grey', alpha=0.1, label='Bollinger Bands')
+
+ # Highlight buy/sell signals if they exist
+ if 'Combined_Signal' in data.columns:
  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')
- 
- ax1.legend()
-
- # 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()
- 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()
+ plt.scatter(buy_signals.index, buy_signals['Close'], label='Buy Signal', marker='^', color='green', alpha=1)
+ plt.scatter(sell_signals.index, sell_signals['Close'], label='Sell Signal', marker='v', color='red', alpha=1)
+
+ plt.title('Stock Price with Buy/Sell Signals')
+ plt.xlabel('Date')
+ plt.ylabel('Price')
+ plt.legend()
+
+ # Improve the x-axis date format
+ plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
+ plt.gca().xaxis.set_major_locator(mdates.MonthLocator())
+ plt.gcf().autofmt_xdate() # Rotation
+
+ plt.show()