225 lines
7.3 KiB
Python
225 lines
7.3 KiB
Python
"""Chart plotting utilities for creating standard, beautiful OHLC charts."""
|
|
|
|
import pandas as pd
|
|
import matplotlib.pyplot as plt
|
|
from typing import Optional, Tuple
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def plot_ohlc(
|
|
df: pd.DataFrame,
|
|
title: Optional[str] = None,
|
|
volume: bool = True,
|
|
figsize: Tuple[int, int] = (14, 8),
|
|
style: str = 'seaborn-v0_8-darkgrid',
|
|
**kwargs
|
|
) -> plt.Figure:
|
|
"""Create a beautiful standard OHLC candlestick chart.
|
|
|
|
This is a convenience function that generates a professional-looking candlestick
|
|
chart with consistent styling across all generated charts. It uses mplfinance
|
|
with seaborn aesthetics for a polished appearance.
|
|
|
|
Args:
|
|
df: pandas DataFrame with DatetimeIndex and columns: open, high, low, close, volume
|
|
title: Optional chart title. If None, uses symbol from chart context
|
|
volume: Whether to include volume subplot (default: True)
|
|
figsize: Figure size as (width, height) in inches (default: (14, 8))
|
|
style: Base matplotlib style to use (default: 'seaborn-v0_8-darkgrid')
|
|
**kwargs: Additional arguments to pass to mplfinance.plot()
|
|
|
|
Returns:
|
|
matplotlib.figure.Figure: The created figure object
|
|
|
|
Example:
|
|
```python
|
|
# Basic usage in analyze_chart_data script
|
|
fig = plot_ohlc(df, title='BTC/USDT 15min')
|
|
|
|
# Customize with additional indicators
|
|
fig = plot_ohlc(df, volume=True, title='Price Action')
|
|
|
|
# Add custom overlays after calling plot_ohlc
|
|
df['SMA20'] = df['close'].rolling(20).mean()
|
|
fig = plot_ohlc(df, title='With SMA')
|
|
# Note: For mplfinance overlays, use the mav or addplot parameters
|
|
```
|
|
|
|
Note:
|
|
The DataFrame must have a DatetimeIndex and the standard OHLCV columns.
|
|
Column names should be lowercase: open, high, low, close, volume
|
|
"""
|
|
try:
|
|
import mplfinance as mpf
|
|
except ImportError:
|
|
raise ImportError(
|
|
"mplfinance is required for plot_ohlc(). "
|
|
"Install it with: pip install mplfinance"
|
|
)
|
|
|
|
# Validate DataFrame structure
|
|
required_cols = ['open', 'high', 'low', 'close']
|
|
missing_cols = [col for col in required_cols if col not in df.columns]
|
|
if missing_cols:
|
|
raise ValueError(
|
|
f"DataFrame missing required columns: {missing_cols}. "
|
|
f"Required: {required_cols}"
|
|
)
|
|
|
|
if not isinstance(df.index, pd.DatetimeIndex):
|
|
raise ValueError(
|
|
"DataFrame must have a DatetimeIndex. "
|
|
"Convert with: df.index = pd.to_datetime(df.index)"
|
|
)
|
|
|
|
# Ensure volume column exists for volume plot
|
|
if volume and 'volume' not in df.columns:
|
|
logger.warning("volume=True but 'volume' column not found in DataFrame. Disabling volume.")
|
|
volume = False
|
|
|
|
# Create custom style with seaborn aesthetics
|
|
# Using a professional color scheme: green for up candles, red for down candles
|
|
mc = mpf.make_marketcolors(
|
|
up='#26a69a', # Teal green (calmer than bright green)
|
|
down='#ef5350', # Coral red (softer than pure red)
|
|
edge='inherit', # Match candle color for edges
|
|
wick='inherit', # Match candle color for wicks
|
|
volume='in', # Volume bars colored by price direction
|
|
alpha=0.9 # Slight transparency for elegance
|
|
)
|
|
|
|
s = mpf.make_mpf_style(
|
|
base_mpf_style='charles', # Clean base style
|
|
marketcolors=mc,
|
|
rc={
|
|
'font.size': 10,
|
|
'axes.labelsize': 11,
|
|
'axes.titlesize': 12,
|
|
'xtick.labelsize': 9,
|
|
'ytick.labelsize': 9,
|
|
'legend.fontsize': 10,
|
|
'figure.facecolor': '#f0f0f0',
|
|
'axes.facecolor': '#ffffff',
|
|
'axes.grid': True,
|
|
'grid.alpha': 0.3,
|
|
'grid.linestyle': '--',
|
|
}
|
|
)
|
|
|
|
# Prepare plot parameters
|
|
plot_params = {
|
|
'type': 'candle',
|
|
'style': s,
|
|
'volume': volume,
|
|
'figsize': figsize,
|
|
'tight_layout': True,
|
|
'returnfig': True,
|
|
'warn_too_much_data': 1000, # Warn if > 1000 candles for performance
|
|
}
|
|
|
|
# Add title if provided
|
|
if title:
|
|
plot_params['title'] = title
|
|
|
|
# Merge any additional kwargs
|
|
plot_params.update(kwargs)
|
|
|
|
# Create the plot
|
|
logger.info(
|
|
f"Creating OHLC chart with {len(df)} candles, "
|
|
f"date range: {df.index.min()} to {df.index.max()}, "
|
|
f"volume: {volume}"
|
|
)
|
|
|
|
fig, axes = mpf.plot(df, **plot_params)
|
|
|
|
return fig
|
|
|
|
|
|
def add_indicators_to_plot(
|
|
df: pd.DataFrame,
|
|
indicators: dict,
|
|
**plot_kwargs
|
|
) -> plt.Figure:
|
|
"""Create an OHLC chart with technical indicators overlaid.
|
|
|
|
This extends plot_ohlc() to include common technical indicators using
|
|
mplfinance's addplot functionality for proper overlay on candlestick charts.
|
|
|
|
Args:
|
|
df: pandas DataFrame with OHLCV data and indicator columns
|
|
indicators: Dictionary mapping indicator names to parameters
|
|
Example: {
|
|
'SMA_20': {'color': 'blue', 'width': 1.5},
|
|
'EMA_50': {'color': 'orange', 'width': 1.5}
|
|
}
|
|
**plot_kwargs: Additional arguments for plot_ohlc()
|
|
|
|
Returns:
|
|
matplotlib.figure.Figure: The created figure object
|
|
|
|
Example:
|
|
```python
|
|
# Calculate indicators
|
|
df['SMA_20'] = df['close'].rolling(20).mean()
|
|
df['SMA_50'] = df['close'].rolling(50).mean()
|
|
|
|
# Plot with indicators
|
|
fig = add_indicators_to_plot(
|
|
df,
|
|
indicators={
|
|
'SMA_20': {'color': 'blue', 'width': 1.5, 'label': '20 SMA'},
|
|
'SMA_50': {'color': 'red', 'width': 1.5, 'label': '50 SMA'}
|
|
},
|
|
title='BTC/USDT with Moving Averages'
|
|
)
|
|
```
|
|
"""
|
|
try:
|
|
import mplfinance as mpf
|
|
except ImportError:
|
|
raise ImportError(
|
|
"mplfinance is required. Install it with: pip install mplfinance"
|
|
)
|
|
|
|
# Build addplot list for indicators
|
|
addplots = []
|
|
for indicator_col, params in indicators.items():
|
|
if indicator_col not in df.columns:
|
|
logger.warning(f"Indicator column '{indicator_col}' not found in DataFrame. Skipping.")
|
|
continue
|
|
|
|
color = params.get('color', 'blue')
|
|
width = params.get('width', 1.0)
|
|
panel = params.get('panel', 0) # 0 = main panel with candles
|
|
ylabel = params.get('ylabel', '')
|
|
|
|
addplots.append(
|
|
mpf.make_addplot(
|
|
df[indicator_col],
|
|
color=color,
|
|
width=width,
|
|
panel=panel,
|
|
ylabel=ylabel
|
|
)
|
|
)
|
|
|
|
# Pass addplot to plot_ohlc via kwargs
|
|
if addplots:
|
|
plot_kwargs['addplot'] = addplots
|
|
|
|
return plot_ohlc(df, **plot_kwargs)
|
|
|
|
|
|
# Convenience presets for common chart types
|
|
def plot_price_volume(df: pd.DataFrame, title: Optional[str] = None) -> plt.Figure:
|
|
"""Create a standard price + volume chart."""
|
|
return plot_ohlc(df, title=title, volume=True, figsize=(14, 8))
|
|
|
|
|
|
def plot_price_only(df: pd.DataFrame, title: Optional[str] = None) -> plt.Figure:
|
|
"""Create a price-only candlestick chart without volume."""
|
|
return plot_ohlc(df, title=title, volume=False, figsize=(14, 6))
|