"""
生成回测可视化图表
显示K线、布林带、开仓/平仓位置
"""
import pandas as pd
import json
from pathlib import Path
from peewee import *
import time
from loguru import logger
# 数据库配置
DB_PATH = Path(__file__).parent / 'models' / 'database.db'
db = SqliteDatabase(str(DB_PATH))
class BitMartETH5M(Model):
"""5分钟K线模型"""
id = BigIntegerField(primary_key=True)
open = FloatField(null=True)
high = FloatField(null=True)
low = FloatField(null=True)
close = FloatField(null=True)
class Meta:
database = db
table_name = 'bitmart_eth_5m'
def calculate_bollinger_bands(df, period=10, std_dev=2.5):
"""计算布林带（右移1根，与回测口径一致）"""
df['sma'] = df['close'].rolling(window=period).mean()
df['std'] = df['close'].rolling(window=period).std()
df['bb_upper'] = (df['sma'] + std_dev * df['std']).shift(1)
df['bb_mid'] = df['sma'].shift(1)
df['bb_lower'] = (df['sma'] - std_dev * df['std']).shift(1)
return df
def generate_chart_data(start_date, end_date, trades_file):
"""生成图表数据"""
# 1. 加载K线数据
start_ts = int(pd.Timestamp(start_date).timestamp() * 1000)
end_ts = int(pd.Timestamp(end_date).timestamp() * 1000) + 86400000 # 加一天确保包含end_date当天的数据
query = BitMartETH5M.select().where(
(BitMartETH5M.id >= start_ts) & (BitMartETH5M.id <= end_ts)
).order_by(BitMartETH5M.id)
data = []
for row in query:
data.append({
'timestamp': row.id,
'open': row.open,
'high': row.high,
'low': row.low,
'close': row.close
})
df = pd.DataFrame(data)
df['datetime'] = pd.to_datetime(df['timestamp'], unit='ms')
# 2. 计算布林带
df = calculate_bollinger_bands(df)
# 3. 加载交易记录
trades_df = pd.read_csv(trades_file)
trades_df['datetime'] = pd.to_datetime(trades_df['timestamp'])
# 4. 准备图表数据
chart_data = []
for idx, row in df.iterrows():
chart_data.append({
'timestamp': int(row['timestamp']),
'datetime': row['datetime'].strftime('%Y-%m-%d %H:%M'),
'open': float(row['open']) if pd.notna(row['open']) else None,
'high': float(row['high']) if pd.notna(row['high']) else None,
'low': float(row['low']) if pd.notna(row['low']) else None,
'close': float(row['close']) if pd.notna(row['close']) else None,
'bb_upper': float(row['bb_upper']) if pd.notna(row['bb_upper']) else None,
'bb_mid': float(row['bb_mid']) if pd.notna(row['bb_mid']) else None,
'bb_lower': float(row['bb_lower']) if pd.notna(row['bb_lower']) else None,
})
# 5. 准备交易标记数据
trades_markers = []
for idx, trade in trades_df.iterrows():
action = trade['action']
price = trade['price']
timestamp = trade['datetime']
reason = trade['reason']
# 找到对应的K线索引（使用最近的K线）
kline_idx = df[df['datetime'] == timestamp].index
if len(kline_idx) == 0:
# 如果找不到完全匹配的，找最近的K线
time_diff = abs(df['datetime'] - timestamp)
min_diff = time_diff.min()
# 如果时间差超过10分钟，跳过这个标记
if min_diff > pd.Timedelta(minutes=10):
print(f"跳过标记 {timestamp}，找不到匹配的K线（最小时间差: {min_diff}）")
continue
kline_idx = time_diff.idxmin()
else:
kline_idx = kline_idx[0]
# 图上显示真实成交价，避免“总在最高/最低点成交”的错觉
kline = df.loc[kline_idx]
if pd.notna(price):
display_price = float(price)
else:
display_price = float(kline['close'])
marker = {
'timestamp': timestamp.strftime('%Y-%m-%d %H:%M'),
'price': display_price, # 使用K线实际价格
'action': action,
'reason': reason,
'index': int(kline_idx)
}
# 关键操作直接显示在图上，便于快速理解“为什么做这笔交易”
# 新增：开仓/加仓也展示标签
marker['show_reason_label'] = (
('延迟反转' in reason)
or ('止损' in reason)
or ('开long' in action)
or ('开short' in action)
or ('加long' in action)
or ('加short' in action)
)
if '止损' in reason:
marker['short_reason'] = '止损'
elif '延迟反转' in reason:
marker['short_reason'] = '延迟反转'
elif '触中轨平50%' in reason:
marker['short_reason'] = '中轨平半'
elif '回开仓价全平' in reason:
marker['short_reason'] = '回本全平'
elif '触上轨开空' in reason:
marker['short_reason'] = '上轨开空'
elif '触下轨开多' in reason:
marker['short_reason'] = '下轨开多'
elif '触上轨加空' in reason:
marker['short_reason'] = '上轨加空'
elif '触下轨加多' in reason:
marker['short_reason'] = '下轨加多'
else:
marker['short_reason'] = reason
# 分类标记
if '开long' in action or '加long' in action:
marker['type'] = 'open_long'
marker['color'] = '#00ff00'
marker['symbol'] = 'triangle'
elif '开short' in action or '加short' in action:
marker['type'] = 'open_short'
marker['color'] = '#ff0000'
marker['symbol'] = 'triangle'
elif '平仓' in action:
if '50%' in action:
marker['type'] = 'close_half'
marker['color'] = '#ffff00'
marker['symbol'] = 'diamond'
else:
marker['type'] = 'close_all'
marker['color'] = '#ff00ff'
marker['symbol'] = 'circle'
else:
marker['type'] = 'other'
marker['color'] = '#ffffff'
marker['symbol'] = 'circle'
trades_markers.append(marker)
return chart_data, trades_markers
def generate_html(chart_data, trades_markers, output_file):
"""生成HTML文件"""
html_content = f"""
布林带策略回测可视化 - 2026年3月
"""
with open(output_file, 'w', encoding='utf-8') as f:
f.write(html_content)
print(f"✅ 图表已生成: {output_file}")
if __name__ == '__main__':
db.connect(reuse_if_open=True)
try:
print("正在生成回测可视化图表...")
output_dir = Path(__file__).parent / 'backtest_outputs' / 'charts'
output_dir.mkdir(parents=True, exist_ok=True)
# 生成图表数据
chart_data, trades_markers = generate_chart_data(
start_date='2026-03-01',
end_date='2026-03-03',
trades_file=str(Path(__file__).parent / 'backtest_outputs' / 'trades' / 'bb_backtest_march_2026_trades.csv')
)
print(f"📊 K线数据: {len(chart_data)} 根")
print(f"📍 交易标记: {len(trades_markers)} 个")
# 生成HTML
output_file = output_dir / 'bb_backtest_visualization.html'
generate_html(chart_data, trades_markers, str(output_file))
print(f"\n🎉 完成！请在浏览器中打开 {output_file} 查看图表")
finally:
db.close()