仿alphaarena项目增加A股ai自动量化
This commit is contained in:
oficcejo
@@ -0,0 +1,484 @@
|
||||
"""
|
||||
智能盯盘 - K线图绘制模块
|
||||
支持AI决策标注、实时更新
|
||||
"""
|
||||
|
||||
import plotly.graph_objects as go
|
||||
from plotly.subplots import make_subplots
|
||||
import pandas as pd
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Dict, List, Optional
|
||||
import logging
|
||||
|
||||
|
||||
class SmartMonitorKline:
|
||||
"""智能盯盘K线图"""
|
||||
|
||||
def __init__(self):
|
||||
"""初始化K线图"""
|
||||
self.logger = logging.getLogger(__name__)
|
||||
|
||||
def create_kline_with_decisions(
|
||||
self,
|
||||
stock_code: str,
|
||||
stock_name: str,
|
||||
kline_data: pd.DataFrame,
|
||||
ai_decisions: List[Dict],
|
||||
show_volume: bool = True,
|
||||
show_ma: bool = True,
|
||||
height: int = 600
|
||||
) -> go.Figure:
|
||||
"""
|
||||
创建带AI决策标注的K线图
|
||||
|
||||
Args:
|
||||
stock_code: 股票代码
|
||||
stock_name: 股票名称
|
||||
kline_data: K线数据（DataFrame）
|
||||
ai_decisions: AI决策列表
|
||||
show_volume: 是否显示成交量
|
||||
show_ma: 是否显示均线
|
||||
height: 图表高度
|
||||
|
||||
Returns:
|
||||
plotly Figure对象
|
||||
"""
|
||||
try:
|
||||
# 确保数据不为空
|
||||
if kline_data is None or kline_data.empty:
|
||||
self.logger.warning(f"K线数据为空 {stock_code}")
|
||||
return self._create_empty_figure(stock_code, stock_name, height)
|
||||
|
||||
# 确保必需的列存在
|
||||
required_cols = ['日期', '开盘', '收盘', '最高', '最低']
|
||||
if not all(col in kline_data.columns for col in required_cols):
|
||||
self.logger.error(f"K线数据缺少必需列 {stock_code}")
|
||||
return self._create_empty_figure(stock_code, stock_name, height)
|
||||
|
||||
# 创建子图
|
||||
if show_volume:
|
||||
fig = make_subplots(
|
||||
rows=2, cols=1,
|
||||
shared_xaxes=True,
|
||||
vertical_spacing=0.03,
|
||||
row_heights=[0.7, 0.3],
|
||||
subplot_titles=(f'{stock_code} {stock_name}', '成交量')
|
||||
)
|
||||
else:
|
||||
fig = make_subplots(
|
||||
rows=1, cols=1,
|
||||
subplot_titles=(f'{stock_code} {stock_name}',)
|
||||
)
|
||||
|
||||
# 1. 添加K线图
|
||||
fig.add_trace(
|
||||
go.Candlestick(
|
||||
x=kline_data['日期'],
|
||||
open=kline_data['开盘'],
|
||||
high=kline_data['最高'],
|
||||
low=kline_data['最低'],
|
||||
close=kline_data['收盘'],
|
||||
name='K线',
|
||||
increasing_line_color='#ef5350', # 红色（涨）
|
||||
decreasing_line_color='#26a69a' # 绿色（跌）
|
||||
),
|
||||
row=1, col=1
|
||||
)
|
||||
|
||||
# 2. 添加均线（如果需要）
|
||||
if show_ma:
|
||||
self._add_moving_averages(fig, kline_data, row=1, col=1)
|
||||
|
||||
# 3. 添加AI决策标注
|
||||
if ai_decisions:
|
||||
self._add_ai_decision_markers(fig, kline_data, ai_decisions, row=1, col=1)
|
||||
|
||||
# 4. 添加成交量（如果需要）
|
||||
if show_volume and '成交量' in kline_data.columns:
|
||||
self._add_volume(fig, kline_data, row=2, col=1)
|
||||
|
||||
# 5. 更新布局
|
||||
fig.update_layout(
|
||||
height=height,
|
||||
xaxis_rangeslider_visible=False,
|
||||
showlegend=True,
|
||||
hovermode='x unified',
|
||||
template='plotly_white',
|
||||
margin=dict(l=50, r=50, t=50, b=50),
|
||||
)
|
||||
|
||||
# 更新x轴
|
||||
fig.update_xaxes(
|
||||
title_text="日期",
|
||||
row=2 if show_volume else 1,
|
||||
col=1
|
||||
)
|
||||
|
||||
# 更新y轴
|
||||
fig.update_yaxes(title_text="价格(元)", row=1, col=1)
|
||||
if show_volume:
|
||||
fig.update_yaxes(title_text="成交量", row=2, col=1)
|
||||
|
||||
return fig
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"创建K线图失败 {stock_code}: {e}")
|
||||
import traceback
|
||||
self.logger.debug(traceback.format_exc())
|
||||
return self._create_empty_figure(stock_code, stock_name, height)
|
||||
|
||||
def _add_moving_averages(self, fig, kline_data: pd.DataFrame, row: int, col: int):
|
||||
"""添加均线"""
|
||||
try:
|
||||
# 计算均线
|
||||
ma_periods = [5, 10, 20, 60]
|
||||
ma_colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#FFA07A']
|
||||
|
||||
for period, color in zip(ma_periods, ma_colors):
|
||||
if len(kline_data) >= period:
|
||||
ma = kline_data['收盘'].rolling(window=period).mean()
|
||||
fig.add_trace(
|
||||
go.Scatter(
|
||||
x=kline_data['日期'],
|
||||
y=ma,
|
||||
name=f'MA{period}',
|
||||
line=dict(color=color, width=1),
|
||||
opacity=0.7
|
||||
),
|
||||
row=row, col=col
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"添加均线失败: {e}")
|
||||
|
||||
def _add_ai_decision_markers(
|
||||
self,
|
||||
fig,
|
||||
kline_data: pd.DataFrame,
|
||||
ai_decisions: List[Dict],
|
||||
row: int,
|
||||
col: int
|
||||
):
|
||||
"""在K线图上添加AI决策标注"""
|
||||
try:
|
||||
# 决策类型映射
|
||||
action_config = {
|
||||
'buy': {
|
||||
'symbol': 'triangle-up',
|
||||
'color': '#ef5350',
|
||||
'text': '买入',
|
||||
'size': 15
|
||||
},
|
||||
'sell': {
|
||||
'symbol': 'triangle-down',
|
||||
'color': '#26a69a',
|
||||
'text': '卖出',
|
||||
'size': 15
|
||||
},
|
||||
'add_position': {
|
||||
'symbol': 'triangle-up',
|
||||
'color': '#ff9800',
|
||||
'text': '加仓',
|
||||
'size': 12
|
||||
},
|
||||
'reduce_position': {
|
||||
'symbol': 'triangle-down',
|
||||
'color': '#9c27b0',
|
||||
'text': '减仓',
|
||||
'size': 12
|
||||
},
|
||||
'hold': {
|
||||
'symbol': 'circle',
|
||||
'color': '#607d8b',
|
||||
'text': '持有',
|
||||
'size': 8
|
||||
}
|
||||
}
|
||||
|
||||
# 将K线数据日期转换为字符串，便于匹配
|
||||
kline_data['日期_str'] = pd.to_datetime(kline_data['日期']).dt.strftime('%Y-%m-%d')
|
||||
|
||||
# 按决策类型分组
|
||||
for action_type, config in action_config.items():
|
||||
decisions_of_type = [d for d in ai_decisions if d.get('action') == action_type]
|
||||
|
||||
if not decisions_of_type:
|
||||
continue
|
||||
|
||||
# 提取决策的日期和价格
|
||||
decision_dates = []
|
||||
decision_prices = []
|
||||
decision_texts = []
|
||||
|
||||
for decision in decisions_of_type:
|
||||
decision_date = decision.get('decision_time', '').split()[0] # 只取日期部分
|
||||
|
||||
# 在K线数据中查找对应日期的收盘价
|
||||
matching_rows = kline_data[kline_data['日期_str'] == decision_date]
|
||||
|
||||
if not matching_rows.empty:
|
||||
price = matching_rows.iloc[0]['收盘']
|
||||
decision_dates.append(decision_date)
|
||||
decision_prices.append(price)
|
||||
|
||||
# 构建hover文本
|
||||
confidence = decision.get('confidence', 0)
|
||||
reasoning = decision.get('reasoning', '无')[:50] # 截断过长的推理
|
||||
hover_text = (
|
||||
f"<b>{config['text']}</b><br>"
|
||||
f"日期: {decision_date}<br>"
|
||||
f"价格: ¥{price:.2f}<br>"
|
||||
f"置信度: {confidence}%<br>"
|
||||
f"推理: {reasoning}..."
|
||||
)
|
||||
decision_texts.append(hover_text)
|
||||
|
||||
# 添加标注
|
||||
if decision_dates:
|
||||
fig.add_trace(
|
||||
go.Scatter(
|
||||
x=decision_dates,
|
||||
y=decision_prices,
|
||||
mode='markers+text',
|
||||
name=config['text'],
|
||||
marker=dict(
|
||||
symbol=config['symbol'],
|
||||
size=config['size'],
|
||||
color=config['color'],
|
||||
line=dict(color='white', width=1)
|
||||
),
|
||||
text=[config['text']] * len(decision_dates),
|
||||
textposition='top center',
|
||||
textfont=dict(size=10, color=config['color']),
|
||||
hovertext=decision_texts,
|
||||
hoverinfo='text',
|
||||
showlegend=True
|
||||
),
|
||||
row=row, col=col
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"添加AI决策标注失败: {e}")
|
||||
import traceback
|
||||
self.logger.debug(traceback.format_exc())
|
||||
|
||||
def _add_volume(self, fig, kline_data: pd.DataFrame, row: int, col: int):
|
||||
"""添加成交量柱状图"""
|
||||
try:
|
||||
# 计算颜色（红涨绿跌）
|
||||
colors = []
|
||||
for i in range(len(kline_data)):
|
||||
if i == 0:
|
||||
colors.append('#ef5350')
|
||||
else:
|
||||
if kline_data.iloc[i]['收盘'] >= kline_data.iloc[i-1]['收盘']:
|
||||
colors.append('#ef5350') # 红色（涨）
|
||||
else:
|
||||
colors.append('#26a69a') # 绿色（跌）
|
||||
|
||||
fig.add_trace(
|
||||
go.Bar(
|
||||
x=kline_data['日期'],
|
||||
y=kline_data['成交量'],
|
||||
name='成交量',
|
||||
marker_color=colors,
|
||||
showlegend=False
|
||||
),
|
||||
row=row, col=col
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"添加成交量失败: {e}")
|
||||
|
||||
def _create_empty_figure(self, stock_code: str, stock_name: str, height: int) -> go.Figure:
|
||||
"""创建空图表"""
|
||||
fig = go.Figure()
|
||||
fig.add_annotation(
|
||||
text=f"暂无 {stock_code} {stock_name} 的K线数据",
|
||||
xref="paper",
|
||||
yref="paper",
|
||||
x=0.5,
|
||||
y=0.5,
|
||||
showarrow=False,
|
||||
font=dict(size=20, color="gray")
|
||||
)
|
||||
fig.update_layout(
|
||||
height=height,
|
||||
xaxis=dict(visible=False),
|
||||
yaxis=dict(visible=False),
|
||||
template='plotly_white'
|
||||
)
|
||||
return fig
|
||||
|
||||
def get_kline_data(self, stock_code: str, days: int = 60, data_fetcher=None) -> Optional[pd.DataFrame]:
|
||||
"""
|
||||
获取K线数据（带Tushare降级机制）
|
||||
|
||||
Args:
|
||||
stock_code: 股票代码
|
||||
days: 获取天数
|
||||
data_fetcher: 数据获取器实例
|
||||
|
||||
Returns:
|
||||
K线数据DataFrame
|
||||
"""
|
||||
try:
|
||||
if data_fetcher is None:
|
||||
from smart_monitor_data import SmartMonitorDataFetcher
|
||||
data_fetcher = SmartMonitorDataFetcher()
|
||||
|
||||
# 计算日期范围
|
||||
end_date = datetime.now().strftime('%Y%m%d')
|
||||
start_date = (datetime.now() - timedelta(days=days + 30)).strftime('%Y%m%d') # 多取30天以确保足够数据
|
||||
|
||||
# 方法1: 尝试使用AKShare获取（只尝试1次，避免IP封禁）
|
||||
try:
|
||||
import akshare as ak
|
||||
df = ak.stock_zh_a_hist(
|
||||
symbol=stock_code,
|
||||
period='daily',
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
adjust='qfq'
|
||||
)
|
||||
|
||||
if df is not None and not df.empty:
|
||||
# 只保留最近days天的数据
|
||||
df = df.tail(days)
|
||||
self.logger.info(f"✅ AKShare获取K线数据成功 {stock_code}，共{len(df)}条")
|
||||
return df
|
||||
else:
|
||||
self.logger.warning(f"AKShare未返回K线数据 {stock_code}，尝试降级到Tushare")
|
||||
except Exception as e:
|
||||
self.logger.warning(f"AKShare获取K线数据失败 {stock_code}: {type(e).__name__}, 尝试降级到Tushare")
|
||||
|
||||
# 方法2: 降级到Tushare
|
||||
if data_fetcher and data_fetcher.ts_pro:
|
||||
self.logger.info(f"降级使用Tushare获取K线数据 {stock_code}")
|
||||
df = self._get_kline_from_tushare(stock_code, days, data_fetcher.ts_pro)
|
||||
if df is not None and not df.empty:
|
||||
self.logger.info(f"✅ Tushare获取K线数据成功 {stock_code}，共{len(df)}条")
|
||||
return df
|
||||
|
||||
self.logger.error(f"所有数据源都无法获取K线数据 {stock_code}")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"获取K线数据失败 {stock_code}: {e}")
|
||||
import traceback
|
||||
self.logger.debug(traceback.format_exc())
|
||||
return None
|
||||
|
||||
def _get_kline_from_tushare(self, stock_code: str, days: int, ts_pro) -> Optional[pd.DataFrame]:
|
||||
"""
|
||||
从Tushare获取K线数据
|
||||
|
||||
Args:
|
||||
stock_code: 股票代码
|
||||
days: 获取天数
|
||||
ts_pro: Tushare API实例
|
||||
|
||||
Returns:
|
||||
K线数据DataFrame
|
||||
"""
|
||||
try:
|
||||
# 转换股票代码格式
|
||||
if stock_code.startswith('6'):
|
||||
ts_code = f"{stock_code}.SH"
|
||||
elif stock_code.startswith(('0', '3')):
|
||||
ts_code = f"{stock_code}.SZ"
|
||||
else:
|
||||
ts_code = stock_code
|
||||
|
||||
# 计算日期范围（多取一些确保足够）
|
||||
end_date = datetime.now().strftime('%Y%m%d')
|
||||
start_date = (datetime.now() - timedelta(days=days + 60)).strftime('%Y%m%d')
|
||||
|
||||
# 获取日K线数据（前复权）
|
||||
df = ts_pro.daily(
|
||||
ts_code=ts_code,
|
||||
start_date=start_date,
|
||||
end_date=end_date,
|
||||
adj='qfq'
|
||||
)
|
||||
|
||||
if df is None or df.empty:
|
||||
self.logger.error(f"Tushare未返回K线数据 {stock_code}")
|
||||
return None
|
||||
|
||||
# Tushare数据是从新到旧，需要反转
|
||||
df = df.sort_values('trade_date', ascending=True).reset_index(drop=True)
|
||||
|
||||
# 统一列名为AKShare格式
|
||||
df = df.rename(columns={
|
||||
'trade_date': '日期',
|
||||
'open': '开盘',
|
||||
'high': '最高',
|
||||
'low': '最低',
|
||||
'close': '收盘',
|
||||
'vol': '成交量',
|
||||
'amount': '成交额'
|
||||
})
|
||||
|
||||
# 转换日期格式（Tushare: 20240115 -> 2024-01-15）
|
||||
df['日期'] = pd.to_datetime(df['日期'])
|
||||
|
||||
# 只保留最近days天的数据
|
||||
df = df.tail(days)
|
||||
|
||||
return df
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Tushare获取K线数据失败 {stock_code}: {type(e).__name__}: {str(e)}")
|
||||
return None
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# 测试代码
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
kline = SmartMonitorKline()
|
||||
|
||||
# 测试获取K线数据
|
||||
df = kline.get_kline_data('600519', days=60)
|
||||
|
||||
if df is not None:
|
||||
print(f"获取到 {len(df)} 条K线数据")
|
||||
print(df.head())
|
||||
|
||||
# 模拟AI决策
|
||||
ai_decisions = [
|
||||
{
|
||||
'decision_time': '2024-01-15 10:00:00',
|
||||
'action': 'buy',
|
||||
'confidence': 85,
|
||||
'reasoning': '技术指标良好，MACD金叉'
|
||||
},
|
||||
{
|
||||
'decision_time': '2024-01-20 14:30:00',
|
||||
'action': 'add_position',
|
||||
'confidence': 75,
|
||||
'reasoning': '突破关键压力位'
|
||||
},
|
||||
{
|
||||
'decision_time': '2024-01-25 11:00:00',
|
||||
'action': 'sell',
|
||||
'confidence': 80,
|
||||
'reasoning': 'RSI超买，建议止盈'
|
||||
}
|
||||
]
|
||||
|
||||
# 创建K线图
|
||||
fig = kline.create_kline_with_decisions(
|
||||
stock_code='600519',
|
||||
stock_name='贵州茅台',
|
||||
kline_data=df,
|
||||
ai_decisions=ai_decisions,
|
||||
show_volume=True,
|
||||
show_ma=True
|
||||
)
|
||||
|
||||
# 保存为HTML
|
||||
fig.write_html('test_kline.html')
|
||||
print("K线图已保存到 test_kline.html")
|
||||
else:
|
||||
print("获取K线数据失败")
|
||||
|
||||
Reference in New Issue
Block a user