backend redesign
This commit is contained in:
224
backend.old/src/agent/tools/chart_utils.py
Normal file
224
backend.old/src/agent/tools/chart_utils.py
Normal file
@@ -0,0 +1,224 @@
|
||||
"""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))
|
||||
Reference in New Issue
Block a user