860 lines
29 KiB
Python
860 lines
29 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
快速融合异动回测脚本
|
||
|
||
优化策略：
|
||
1. 预先构建所有序列（向量化），避免循环内重复切片
|
||
2. 批量 ML 推理（一次推理所有候选）
|
||
3. 使用 NumPy 向量化操作替代 Python 循环
|
||
|
||
性能对比：
|
||
- 原版：5分钟/天
|
||
- 优化版：预计 10-30秒/天
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import argparse
|
||
import json
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Dict, List, Optional, Tuple
|
||
from collections import defaultdict
|
||
|
||
import numpy as np
|
||
import pandas as pd
|
||
import torch
|
||
from tqdm import tqdm
|
||
from sqlalchemy import create_engine, text
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
||
|
||
# ==================== 配置 ====================
|
||
|
||
MYSQL_ENGINE = create_engine(
|
||
"mysql+pymysql://root:Zzl5588161!@192.168.1.5:3306/stock",
|
||
echo=False
|
||
)
|
||
|
||
FEATURES = ['alpha', 'alpha_delta', 'amt_ratio', 'amt_delta', 'rank_pct', 'limit_up_ratio']
|
||
|
||
CONFIG = {
|
||
'seq_len': 15, # 序列长度（支持跨日后可从 9:30 检测）
|
||
'min_alpha_abs': 0.3, # 最小 alpha 过滤
|
||
'cooldown_minutes': 8,
|
||
'max_alerts_per_minute': 20,
|
||
'clip_value': 10.0,
|
||
# === 融合权重：均衡 ===
|
||
'rule_weight': 0.5,
|
||
'ml_weight': 0.5,
|
||
# === 触发阈值 ===
|
||
'rule_trigger': 65, # 60 -> 65，略提高规则门槛
|
||
'ml_trigger': 70, # 75 -> 70，略降低 ML 门槛
|
||
'fusion_trigger': 45,
|
||
}
|
||
|
||
|
||
# ==================== 规则评分（向量化版）====================
|
||
|
||
def get_size_adjusted_thresholds(stock_count: np.ndarray) -> dict:
|
||
"""
|
||
根据概念股票数量计算动态阈值
|
||
|
||
设计思路：
|
||
- 小概念（<10 只）：波动大是正常的，需要更高阈值
|
||
- 中概念（10-50 只）：标准阈值
|
||
- 大概念（>50 只）：能有明显波动说明是真异动，降低阈值
|
||
|
||
返回各指标的调整系数（乘以基准阈值）
|
||
"""
|
||
n = len(stock_count)
|
||
|
||
# 基于股票数量的调整系数
|
||
# 小概念：系数 > 1（提高阈值，更难触发）
|
||
# 大概念：系数 < 1（降低阈值，更容易触发）
|
||
size_factor = np.ones(n)
|
||
|
||
# 微型概念（<5 只）：阈值 × 1.8
|
||
tiny = stock_count < 5
|
||
size_factor[tiny] = 1.8
|
||
|
||
# 小概念（5-10 只）：阈值 × 1.4
|
||
small = (stock_count >= 5) & (stock_count < 10)
|
||
size_factor[small] = 1.4
|
||
|
||
# 中小概念（10-20 只）：阈值 × 1.2
|
||
medium_small = (stock_count >= 10) & (stock_count < 20)
|
||
size_factor[medium_small] = 1.2
|
||
|
||
# 中概念（20-50 只）：标准阈值 × 1.0
|
||
medium = (stock_count >= 20) & (stock_count < 50)
|
||
size_factor[medium] = 1.0
|
||
|
||
# 大概念（50-100 只）：阈值 × 0.85
|
||
large = (stock_count >= 50) & (stock_count < 100)
|
||
size_factor[large] = 0.85
|
||
|
||
# 超大概念（>100 只）：阈值 × 0.7
|
||
xlarge = stock_count >= 100
|
||
size_factor[xlarge] = 0.7
|
||
|
||
return size_factor
|
||
|
||
|
||
def score_rules_batch(df: pd.DataFrame) -> Tuple[np.ndarray, List[List[str]]]:
|
||
"""
|
||
批量计算规则得分（向量化）- 考虑概念规模版
|
||
|
||
设计原则：
|
||
- 规则作为辅助信号，不应单独主导决策
|
||
- 根据概念股票数量动态调整阈值
|
||
- 大概念异动更有价值，小概念需要更大波动才算异动
|
||
|
||
Args:
|
||
df: DataFrame，包含所有特征列（必须包含 stock_count）
|
||
Returns:
|
||
scores: (n,) 规则得分数组
|
||
triggered_rules: 每行触发的规则列表
|
||
"""
|
||
n = len(df)
|
||
scores = np.zeros(n)
|
||
triggered = [[] for _ in range(n)]
|
||
|
||
alpha = df['alpha'].values
|
||
alpha_delta = df['alpha_delta'].values
|
||
amt_ratio = df['amt_ratio'].values
|
||
amt_delta = df['amt_delta'].values
|
||
rank_pct = df['rank_pct'].values
|
||
limit_up_ratio = df['limit_up_ratio'].values
|
||
stock_count = df['stock_count'].values if 'stock_count' in df.columns else np.full(n, 20)
|
||
|
||
alpha_abs = np.abs(alpha)
|
||
alpha_delta_abs = np.abs(alpha_delta)
|
||
|
||
# 获取基于规模的调整系数
|
||
size_factor = get_size_adjusted_thresholds(stock_count)
|
||
|
||
# ========== Alpha 规则（动态阈值）==========
|
||
# 基准阈值：极强 5%，强 4%，中等 3%
|
||
# 实际阈值 = 基准 × size_factor
|
||
|
||
# 极强信号
|
||
alpha_extreme_thresh = 5.0 * size_factor
|
||
mask = alpha_abs >= alpha_extreme_thresh
|
||
scores[mask] += 20
|
||
for i in np.where(mask)[0]: triggered[i].append('alpha_extreme')
|
||
|
||
# 强信号
|
||
alpha_strong_thresh = 4.0 * size_factor
|
||
mask = (alpha_abs >= alpha_strong_thresh) & (alpha_abs < alpha_extreme_thresh)
|
||
scores[mask] += 15
|
||
for i in np.where(mask)[0]: triggered[i].append('alpha_strong')
|
||
|
||
# 中等信号
|
||
alpha_medium_thresh = 3.0 * size_factor
|
||
mask = (alpha_abs >= alpha_medium_thresh) & (alpha_abs < alpha_strong_thresh)
|
||
scores[mask] += 10
|
||
for i in np.where(mask)[0]: triggered[i].append('alpha_medium')
|
||
|
||
# ========== Alpha 加速度规则（动态阈值）==========
|
||
delta_strong_thresh = 2.0 * size_factor
|
||
mask = alpha_delta_abs >= delta_strong_thresh
|
||
scores[mask] += 15
|
||
for i in np.where(mask)[0]: triggered[i].append('alpha_delta_strong')
|
||
|
||
delta_medium_thresh = 1.5 * size_factor
|
||
mask = (alpha_delta_abs >= delta_medium_thresh) & (alpha_delta_abs < delta_strong_thresh)
|
||
scores[mask] += 10
|
||
for i in np.where(mask)[0]: triggered[i].append('alpha_delta_medium')
|
||
|
||
# ========== 成交额规则（不受规模影响，放量就是放量）==========
|
||
mask = amt_ratio >= 10.0
|
||
scores[mask] += 20
|
||
for i in np.where(mask)[0]: triggered[i].append('volume_extreme')
|
||
|
||
mask = (amt_ratio >= 6.0) & (amt_ratio < 10.0)
|
||
scores[mask] += 12
|
||
for i in np.where(mask)[0]: triggered[i].append('volume_strong')
|
||
|
||
# ========== 排名规则 ==========
|
||
mask = rank_pct >= 0.98
|
||
scores[mask] += 15
|
||
for i in np.where(mask)[0]: triggered[i].append('rank_top')
|
||
|
||
mask = rank_pct <= 0.02
|
||
scores[mask] += 15
|
||
for i in np.where(mask)[0]: triggered[i].append('rank_bottom')
|
||
|
||
# ========== 涨停规则（动态阈值）==========
|
||
# 大概念有涨停更有意义
|
||
limit_high_thresh = 0.30 * size_factor
|
||
mask = limit_up_ratio >= limit_high_thresh
|
||
scores[mask] += 20
|
||
for i in np.where(mask)[0]: triggered[i].append('limit_up_high')
|
||
|
||
limit_medium_thresh = 0.20 * size_factor
|
||
mask = (limit_up_ratio >= limit_medium_thresh) & (limit_up_ratio < limit_high_thresh)
|
||
scores[mask] += 12
|
||
for i in np.where(mask)[0]: triggered[i].append('limit_up_medium')
|
||
|
||
# ========== 概念规模加分（大概念异动更有价值）==========
|
||
# 大概念（50+）额外加分
|
||
large_concept = stock_count >= 50
|
||
has_signal = scores > 0 # 至少触发了某个规则
|
||
mask = large_concept & has_signal
|
||
scores[mask] += 10
|
||
for i in np.where(mask)[0]: triggered[i].append('large_concept_bonus')
|
||
|
||
# 超大概念（100+）再加分
|
||
xlarge_concept = stock_count >= 100
|
||
mask = xlarge_concept & has_signal
|
||
scores[mask] += 10
|
||
for i in np.where(mask)[0]: triggered[i].append('xlarge_concept_bonus')
|
||
|
||
# ========== 组合规则（动态阈值）==========
|
||
combo_alpha_thresh = 3.0 * size_factor
|
||
|
||
# Alpha + 放量 + 排名（三重验证）
|
||
mask = (alpha_abs >= combo_alpha_thresh) & (amt_ratio >= 5.0) & ((rank_pct >= 0.95) | (rank_pct <= 0.05))
|
||
scores[mask] += 20
|
||
for i in np.where(mask)[0]: triggered[i].append('triple_signal')
|
||
|
||
# Alpha + 涨停（强组合）
|
||
mask = (alpha_abs >= combo_alpha_thresh) & (limit_up_ratio >= 0.15 * size_factor)
|
||
scores[mask] += 15
|
||
for i in np.where(mask)[0]: triggered[i].append('alpha_with_limit')
|
||
|
||
# ========== 小概念惩罚（过滤噪音）==========
|
||
# 微型概念（<5 只）如果只有单一信号，减分
|
||
tiny_concept = stock_count < 5
|
||
single_rule = np.array([len(t) <= 1 for t in triggered])
|
||
mask = tiny_concept & single_rule & (scores > 0)
|
||
scores[mask] *= 0.5 # 减半
|
||
for i in np.where(mask)[0]: triggered[i].append('tiny_concept_penalty')
|
||
|
||
scores = np.clip(scores, 0, 100)
|
||
return scores, triggered
|
||
|
||
|
||
# ==================== ML 评分器 ====================
|
||
|
||
class FastMLScorer:
|
||
"""快速 ML 评分器"""
|
||
|
||
def __init__(self, checkpoint_dir: str = 'ml/checkpoints', device: str = 'auto'):
|
||
self.checkpoint_dir = Path(checkpoint_dir)
|
||
|
||
if device == 'auto':
|
||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
elif device == 'cuda' and not torch.cuda.is_available():
|
||
print("警告: CUDA 不可用，使用 CPU")
|
||
self.device = torch.device('cpu')
|
||
else:
|
||
self.device = torch.device(device)
|
||
|
||
self.model = None
|
||
self.thresholds = None
|
||
self._load_model()
|
||
|
||
def _load_model(self):
|
||
model_path = self.checkpoint_dir / 'best_model.pt'
|
||
thresholds_path = self.checkpoint_dir / 'thresholds.json'
|
||
config_path = self.checkpoint_dir / 'config.json'
|
||
|
||
if not model_path.exists():
|
||
print(f"警告: 模型不存在 {model_path}")
|
||
return
|
||
|
||
try:
|
||
from model import LSTMAutoencoder
|
||
|
||
config = {}
|
||
if config_path.exists():
|
||
with open(config_path) as f:
|
||
config = json.load(f).get('model', {})
|
||
|
||
# 处理旧配置键名
|
||
if 'd_model' in config:
|
||
config['hidden_dim'] = config.pop('d_model') // 2
|
||
for key in ['num_encoder_layers', 'num_decoder_layers', 'nhead', 'dim_feedforward', 'max_seq_len', 'use_instance_norm']:
|
||
config.pop(key, None)
|
||
if 'num_layers' not in config:
|
||
config['num_layers'] = 1
|
||
|
||
checkpoint = torch.load(model_path, map_location='cpu')
|
||
self.model = LSTMAutoencoder(**config)
|
||
self.model.load_state_dict(checkpoint['model_state_dict'])
|
||
self.model.to(self.device)
|
||
self.model.eval()
|
||
|
||
if thresholds_path.exists():
|
||
with open(thresholds_path) as f:
|
||
self.thresholds = json.load(f)
|
||
|
||
print(f"ML模型加载成功 (设备: {self.device})")
|
||
except Exception as e:
|
||
print(f"ML模型加载失败: {e}")
|
||
self.model = None
|
||
|
||
def is_ready(self):
|
||
return self.model is not None
|
||
|
||
@torch.no_grad()
|
||
def score_batch(self, sequences: np.ndarray) -> np.ndarray:
|
||
"""
|
||
批量计算 ML 得分
|
||
|
||
Args:
|
||
sequences: (batch, seq_len, n_features)
|
||
Returns:
|
||
scores: (batch,) 0-100 分数
|
||
"""
|
||
if not self.is_ready() or len(sequences) == 0:
|
||
return np.zeros(len(sequences))
|
||
|
||
x = torch.FloatTensor(sequences).to(self.device)
|
||
output, _ = self.model(x)
|
||
mse = ((output - x) ** 2).mean(dim=-1)
|
||
errors = mse[:, -1].cpu().numpy()
|
||
|
||
p95 = self.thresholds.get('p95', 0.1) if self.thresholds else 0.1
|
||
scores = np.clip(errors / p95 * 50, 0, 100)
|
||
return scores
|
||
|
||
|
||
# ==================== 快速回测 ====================
|
||
|
||
def build_sequences_fast(
|
||
df: pd.DataFrame,
|
||
seq_len: int = 30,
|
||
prev_df: pd.DataFrame = None
|
||
) -> Tuple[np.ndarray, pd.DataFrame]:
|
||
"""
|
||
快速构建所有有效序列
|
||
|
||
支持跨日序列：用前一天收盘数据 + 当天开盘数据拼接，实现 9:30 就能检测
|
||
|
||
Args:
|
||
df: 当天数据
|
||
seq_len: 序列长度
|
||
prev_df: 前一天数据（可选，用于构建开盘时的序列）
|
||
|
||
返回：
|
||
sequences: (n_valid, seq_len, n_features) 所有有效序列
|
||
info_df: 对应的元信息 DataFrame
|
||
"""
|
||
# 确保按概念和时间排序
|
||
df = df.sort_values(['concept_id', 'timestamp']).reset_index(drop=True)
|
||
|
||
# 如果有前一天数据，按概念构建尾部缓存（取每个概念最后 seq_len-1 条）
|
||
prev_cache = {}
|
||
if prev_df is not None and len(prev_df) > 0:
|
||
prev_df = prev_df.sort_values(['concept_id', 'timestamp'])
|
||
for concept_id, gdf in prev_df.groupby('concept_id'):
|
||
tail_data = gdf.tail(seq_len - 1)
|
||
if len(tail_data) > 0:
|
||
feat_matrix = tail_data[FEATURES].values
|
||
feat_matrix = np.nan_to_num(feat_matrix, nan=0.0, posinf=0.0, neginf=0.0)
|
||
feat_matrix = np.clip(feat_matrix, -CONFIG['clip_value'], CONFIG['clip_value'])
|
||
prev_cache[concept_id] = feat_matrix
|
||
|
||
# 按概念分组
|
||
groups = df.groupby('concept_id')
|
||
|
||
sequences = []
|
||
infos = []
|
||
|
||
for concept_id, gdf in groups:
|
||
gdf = gdf.reset_index(drop=True)
|
||
|
||
# 获取特征矩阵
|
||
feat_matrix = gdf[FEATURES].values
|
||
feat_matrix = np.nan_to_num(feat_matrix, nan=0.0, posinf=0.0, neginf=0.0)
|
||
feat_matrix = np.clip(feat_matrix, -CONFIG['clip_value'], CONFIG['clip_value'])
|
||
|
||
# 如果有前一天缓存，拼接到当天数据前面
|
||
if concept_id in prev_cache:
|
||
prev_data = prev_cache[concept_id]
|
||
combined_matrix = np.vstack([prev_data, feat_matrix])
|
||
# 计算偏移量：前一天数据的长度
|
||
offset = len(prev_data)
|
||
else:
|
||
combined_matrix = feat_matrix
|
||
offset = 0
|
||
|
||
# 滑动窗口构建序列
|
||
n_total = len(combined_matrix)
|
||
if n_total < seq_len:
|
||
continue
|
||
|
||
for i in range(n_total - seq_len + 1):
|
||
seq = combined_matrix[i:i + seq_len]
|
||
|
||
# 计算对应当天数据的索引
|
||
# 序列最后一个点的位置 = i + seq_len - 1
|
||
# 对应当天数据的索引 = (i + seq_len - 1) - offset
|
||
today_idx = i + seq_len - 1 - offset
|
||
|
||
# 只要序列的最后一个点是当天的数据，就记录
|
||
if today_idx < 0 or today_idx >= len(gdf):
|
||
continue
|
||
|
||
sequences.append(seq)
|
||
|
||
# 记录最后一个时间步的信息（当天的）
|
||
row = gdf.iloc[today_idx]
|
||
infos.append({
|
||
'concept_id': concept_id,
|
||
'timestamp': row['timestamp'],
|
||
'alpha': row['alpha'],
|
||
'alpha_delta': row.get('alpha_delta', 0),
|
||
'amt_ratio': row.get('amt_ratio', 1),
|
||
'amt_delta': row.get('amt_delta', 0),
|
||
'rank_pct': row.get('rank_pct', 0.5),
|
||
'limit_up_ratio': row.get('limit_up_ratio', 0),
|
||
'stock_count': row.get('stock_count', 0),
|
||
'total_amt': row.get('total_amt', 0),
|
||
})
|
||
|
||
if not sequences:
|
||
return np.array([]), pd.DataFrame()
|
||
|
||
return np.array(sequences), pd.DataFrame(infos)
|
||
|
||
|
||
def backtest_single_day_fast(
|
||
ml_scorer: FastMLScorer,
|
||
df: pd.DataFrame,
|
||
date: str,
|
||
config: Dict,
|
||
prev_df: pd.DataFrame = None
|
||
) -> List[Dict]:
|
||
"""
|
||
快速回测单天（向量化版本）
|
||
|
||
Args:
|
||
ml_scorer: ML 评分器
|
||
df: 当天数据
|
||
date: 日期
|
||
config: 配置
|
||
prev_df: 前一天数据（用于 9:30 开始检测）
|
||
"""
|
||
seq_len = config.get('seq_len', 30)
|
||
|
||
# 1. 构建所有序列（支持跨日）
|
||
sequences, info_df = build_sequences_fast(df, seq_len, prev_df)
|
||
|
||
if len(sequences) == 0:
|
||
return []
|
||
|
||
# 2. 过滤小波动
|
||
alpha_abs = np.abs(info_df['alpha'].values)
|
||
valid_mask = alpha_abs >= config['min_alpha_abs']
|
||
|
||
sequences = sequences[valid_mask]
|
||
info_df = info_df[valid_mask].reset_index(drop=True)
|
||
|
||
if len(sequences) == 0:
|
||
return []
|
||
|
||
# 3. 批量规则评分
|
||
rule_scores, triggered_rules = score_rules_batch(info_df)
|
||
|
||
# 4. 批量 ML 评分（分批处理避免显存溢出）
|
||
batch_size = 2048
|
||
ml_scores = []
|
||
for i in range(0, len(sequences), batch_size):
|
||
batch_seq = sequences[i:i+batch_size]
|
||
batch_scores = ml_scorer.score_batch(batch_seq)
|
||
ml_scores.append(batch_scores)
|
||
ml_scores = np.concatenate(ml_scores) if ml_scores else np.zeros(len(sequences))
|
||
|
||
# 5. 融合得分
|
||
w1, w2 = config['rule_weight'], config['ml_weight']
|
||
final_scores = w1 * rule_scores + w2 * ml_scores
|
||
|
||
# 6. 判断异动
|
||
is_anomaly = (
|
||
(rule_scores >= config['rule_trigger']) |
|
||
(ml_scores >= config['ml_trigger']) |
|
||
(final_scores >= config['fusion_trigger'])
|
||
)
|
||
|
||
# 7. 应用冷却期（按概念+时间排序后处理）
|
||
info_df['rule_score'] = rule_scores
|
||
info_df['ml_score'] = ml_scores
|
||
info_df['final_score'] = final_scores
|
||
info_df['is_anomaly'] = is_anomaly
|
||
info_df['triggered_rules'] = triggered_rules
|
||
|
||
# 只保留异动
|
||
anomaly_df = info_df[info_df['is_anomaly']].copy()
|
||
|
||
if len(anomaly_df) == 0:
|
||
return []
|
||
|
||
# 应用冷却期
|
||
anomaly_df = anomaly_df.sort_values(['concept_id', 'timestamp'])
|
||
cooldown = {}
|
||
keep_mask = []
|
||
|
||
for _, row in anomaly_df.iterrows():
|
||
cid = row['concept_id']
|
||
ts = row['timestamp']
|
||
|
||
if cid in cooldown:
|
||
try:
|
||
diff = (ts - cooldown[cid]).total_seconds() / 60
|
||
except:
|
||
diff = config['cooldown_minutes'] + 1
|
||
|
||
if diff < config['cooldown_minutes']:
|
||
keep_mask.append(False)
|
||
continue
|
||
|
||
cooldown[cid] = ts
|
||
keep_mask.append(True)
|
||
|
||
anomaly_df = anomaly_df[keep_mask]
|
||
|
||
# 8. 按时间分组，每分钟最多 max_alerts_per_minute 个
|
||
alerts = []
|
||
for ts, group in anomaly_df.groupby('timestamp'):
|
||
group = group.nlargest(config['max_alerts_per_minute'], 'final_score')
|
||
|
||
for _, row in group.iterrows():
|
||
alpha = row['alpha']
|
||
if alpha >= 1.5:
|
||
atype = 'surge_up'
|
||
elif alpha <= -1.5:
|
||
atype = 'surge_down'
|
||
elif row['amt_ratio'] >= 3.0:
|
||
atype = 'volume_spike'
|
||
else:
|
||
atype = 'unknown'
|
||
|
||
rule_score = row['rule_score']
|
||
ml_score = row['ml_score']
|
||
final_score = row['final_score']
|
||
|
||
if rule_score >= config['rule_trigger']:
|
||
trigger = f'规则强信号({rule_score:.0f}分)'
|
||
elif ml_score >= config['ml_trigger']:
|
||
trigger = f'ML强信号({ml_score:.0f}分)'
|
||
else:
|
||
trigger = f'融合触发({final_score:.0f}分)'
|
||
|
||
alerts.append({
|
||
'concept_id': row['concept_id'],
|
||
'alert_time': row['timestamp'],
|
||
'trade_date': date,
|
||
'alert_type': atype,
|
||
'final_score': final_score,
|
||
'rule_score': rule_score,
|
||
'ml_score': ml_score,
|
||
'trigger_reason': trigger,
|
||
'triggered_rules': row['triggered_rules'],
|
||
'alpha': alpha,
|
||
'alpha_delta': row['alpha_delta'],
|
||
'amt_ratio': row['amt_ratio'],
|
||
'amt_delta': row['amt_delta'],
|
||
'rank_pct': row['rank_pct'],
|
||
'limit_up_ratio': row['limit_up_ratio'],
|
||
'stock_count': row['stock_count'],
|
||
'total_amt': row['total_amt'],
|
||
})
|
||
|
||
return alerts
|
||
|
||
|
||
# ==================== 数据加载 ====================
|
||
|
||
def load_daily_features(data_dir: str, date: str) -> Optional[pd.DataFrame]:
|
||
file_path = Path(data_dir) / f"features_{date}.parquet"
|
||
if not file_path.exists():
|
||
return None
|
||
return pd.read_parquet(file_path)
|
||
|
||
|
||
def get_available_dates(data_dir: str, start: str, end: str) -> List[str]:
|
||
data_path = Path(data_dir)
|
||
dates = []
|
||
for f in sorted(data_path.glob("features_*.parquet")):
|
||
d = f.stem.replace('features_', '')
|
||
if start <= d <= end:
|
||
dates.append(d)
|
||
return dates
|
||
|
||
|
||
def get_prev_trading_day(data_dir: str, date: str) -> Optional[str]:
|
||
"""获取给定日期之前最近的有数据的交易日"""
|
||
data_path = Path(data_dir)
|
||
all_dates = sorted([f.stem.replace('features_', '') for f in data_path.glob("features_*.parquet")])
|
||
|
||
for i, d in enumerate(all_dates):
|
||
if d == date and i > 0:
|
||
return all_dates[i - 1]
|
||
return None
|
||
|
||
|
||
def export_to_csv(alerts: List[Dict], path: str):
|
||
if alerts:
|
||
pd.DataFrame(alerts).to_csv(path, index=False, encoding='utf-8-sig')
|
||
print(f"已导出: {path}")
|
||
|
||
|
||
# ==================== 数据库写入 ====================
|
||
|
||
def init_db_table():
|
||
"""
|
||
初始化数据库表（如果不存在则创建）
|
||
|
||
表结构说明：
|
||
- concept_id: 概念ID
|
||
- alert_time: 异动时间（精确到分钟）
|
||
- trade_date: 交易日期
|
||
- alert_type: 异动类型（surge_up/surge_down/volume_spike/unknown）
|
||
- final_score: 最终得分（0-100）
|
||
- rule_score: 规则得分（0-100）
|
||
- ml_score: ML得分（0-100）
|
||
- trigger_reason: 触发原因
|
||
- alpha: 超额收益率
|
||
- alpha_delta: alpha变化速度
|
||
- amt_ratio: 成交额放大倍数
|
||
- rank_pct: 排名百分位
|
||
- stock_count: 概念股票数量
|
||
- triggered_rules: 触发的规则列表（JSON）
|
||
"""
|
||
create_sql = text("""
|
||
CREATE TABLE IF NOT EXISTS concept_anomaly_hybrid (
|
||
id INT AUTO_INCREMENT PRIMARY KEY,
|
||
concept_id VARCHAR(64) NOT NULL,
|
||
alert_time DATETIME NOT NULL,
|
||
trade_date DATE NOT NULL,
|
||
alert_type VARCHAR(32) NOT NULL,
|
||
final_score FLOAT NOT NULL,
|
||
rule_score FLOAT NOT NULL,
|
||
ml_score FLOAT NOT NULL,
|
||
trigger_reason VARCHAR(64),
|
||
alpha FLOAT,
|
||
alpha_delta FLOAT,
|
||
amt_ratio FLOAT,
|
||
amt_delta FLOAT,
|
||
rank_pct FLOAT,
|
||
limit_up_ratio FLOAT,
|
||
stock_count INT,
|
||
total_amt FLOAT,
|
||
triggered_rules JSON,
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||
UNIQUE KEY uk_concept_time (concept_id, alert_time, trade_date),
|
||
INDEX idx_trade_date (trade_date),
|
||
INDEX idx_concept_id (concept_id),
|
||
INDEX idx_final_score (final_score),
|
||
INDEX idx_alert_type (alert_type)
|
||
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='概念异动检测结果（融合版）'
|
||
""")
|
||
|
||
with MYSQL_ENGINE.begin() as conn:
|
||
conn.execute(create_sql)
|
||
print("数据库表已就绪: concept_anomaly_hybrid")
|
||
|
||
|
||
def save_alerts_to_mysql(alerts: List[Dict], dry_run: bool = False) -> int:
|
||
"""
|
||
保存异动到 MySQL
|
||
|
||
Args:
|
||
alerts: 异动列表
|
||
dry_run: 是否只模拟，不实际写入
|
||
|
||
Returns:
|
||
实际保存的记录数
|
||
"""
|
||
if not alerts:
|
||
return 0
|
||
|
||
if dry_run:
|
||
print(f" [Dry Run] 将写入 {len(alerts)} 条异动")
|
||
return len(alerts)
|
||
|
||
saved = 0
|
||
skipped = 0
|
||
|
||
with MYSQL_ENGINE.begin() as conn:
|
||
for alert in alerts:
|
||
try:
|
||
# 检查是否已存在（使用 INSERT IGNORE 更高效）
|
||
insert_sql = text("""
|
||
INSERT IGNORE INTO concept_anomaly_hybrid
|
||
(concept_id, alert_time, trade_date, alert_type,
|
||
final_score, rule_score, ml_score, trigger_reason,
|
||
alpha, alpha_delta, amt_ratio, amt_delta,
|
||
rank_pct, limit_up_ratio, stock_count, total_amt,
|
||
triggered_rules)
|
||
VALUES
|
||
(:concept_id, :alert_time, :trade_date, :alert_type,
|
||
:final_score, :rule_score, :ml_score, :trigger_reason,
|
||
:alpha, :alpha_delta, :amt_ratio, :amt_delta,
|
||
:rank_pct, :limit_up_ratio, :stock_count, :total_amt,
|
||
:triggered_rules)
|
||
""")
|
||
|
||
result = conn.execute(insert_sql, {
|
||
'concept_id': alert['concept_id'],
|
||
'alert_time': alert['alert_time'],
|
||
'trade_date': alert['trade_date'],
|
||
'alert_type': alert['alert_type'],
|
||
'final_score': alert['final_score'],
|
||
'rule_score': alert['rule_score'],
|
||
'ml_score': alert['ml_score'],
|
||
'trigger_reason': alert['trigger_reason'],
|
||
'alpha': alert.get('alpha', 0),
|
||
'alpha_delta': alert.get('alpha_delta', 0),
|
||
'amt_ratio': alert.get('amt_ratio', 1),
|
||
'amt_delta': alert.get('amt_delta', 0),
|
||
'rank_pct': alert.get('rank_pct', 0.5),
|
||
'limit_up_ratio': alert.get('limit_up_ratio', 0),
|
||
'stock_count': alert.get('stock_count', 0),
|
||
'total_amt': alert.get('total_amt', 0),
|
||
'triggered_rules': json.dumps(alert.get('triggered_rules', []), ensure_ascii=False),
|
||
})
|
||
|
||
if result.rowcount > 0:
|
||
saved += 1
|
||
else:
|
||
skipped += 1
|
||
|
||
except Exception as e:
|
||
print(f" 保存失败: {alert['concept_id']} @ {alert['alert_time']} - {e}")
|
||
|
||
if skipped > 0:
|
||
print(f" 跳过 {skipped} 条重复记录")
|
||
|
||
return saved
|
||
|
||
|
||
def clear_alerts_by_date(trade_date: str) -> int:
|
||
"""清除指定日期的异动记录（用于重新回测）"""
|
||
with MYSQL_ENGINE.begin() as conn:
|
||
result = conn.execute(
|
||
text("DELETE FROM concept_anomaly_hybrid WHERE trade_date = :trade_date"),
|
||
{'trade_date': trade_date}
|
||
)
|
||
return result.rowcount
|
||
|
||
|
||
def analyze_alerts(alerts: List[Dict]):
|
||
if not alerts:
|
||
print("无异动")
|
||
return
|
||
|
||
df = pd.DataFrame(alerts)
|
||
print(f"\n总异动: {len(alerts)}")
|
||
print(f"\n类型分布:\n{df['alert_type'].value_counts()}")
|
||
print(f"\n得分统计:")
|
||
print(f" 最终: {df['final_score'].mean():.1f} (max: {df['final_score'].max():.1f})")
|
||
print(f" 规则: {df['rule_score'].mean():.1f} (max: {df['rule_score'].max():.1f})")
|
||
print(f" ML: {df['ml_score'].mean():.1f} (max: {df['ml_score'].max():.1f})")
|
||
|
||
trigger_type = df['trigger_reason'].apply(
|
||
lambda x: '规则' if '规则' in x else ('ML' if 'ML' in x else '融合')
|
||
)
|
||
print(f"\n触发来源:\n{trigger_type.value_counts()}")
|
||
|
||
|
||
# ==================== 主函数 ====================
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description='快速融合异动回测')
|
||
parser.add_argument('--data_dir', default='ml/data')
|
||
parser.add_argument('--checkpoint_dir', default='ml/checkpoints')
|
||
parser.add_argument('--start', required=True)
|
||
parser.add_argument('--end', default=None)
|
||
parser.add_argument('--dry-run', action='store_true', help='模拟运行，不写入数据库')
|
||
parser.add_argument('--export-csv', default=None, help='导出 CSV 文件路径')
|
||
parser.add_argument('--save-db', action='store_true', help='保存结果到数据库')
|
||
parser.add_argument('--clear-first', action='store_true', help='写入前先清除该日期的旧数据')
|
||
parser.add_argument('--device', default='auto')
|
||
|
||
args = parser.parse_args()
|
||
if args.end is None:
|
||
args.end = args.start
|
||
|
||
print("=" * 60)
|
||
print("快速融合异动回测")
|
||
print("=" * 60)
|
||
print(f"日期: {args.start} ~ {args.end}")
|
||
print(f"设备: {args.device}")
|
||
print(f"保存数据库: {args.save_db}")
|
||
print("=" * 60)
|
||
|
||
# 初始化数据库表（如果需要保存）
|
||
if args.save_db and not args.dry_run:
|
||
init_db_table()
|
||
|
||
# 初始化 ML 评分器
|
||
ml_scorer = FastMLScorer(args.checkpoint_dir, args.device)
|
||
|
||
# 获取日期
|
||
dates = get_available_dates(args.data_dir, args.start, args.end)
|
||
if not dates:
|
||
print("无数据")
|
||
return
|
||
|
||
print(f"找到 {len(dates)} 天数据\n")
|
||
|
||
# 回测（支持跨日序列）
|
||
all_alerts = []
|
||
total_saved = 0
|
||
prev_df = None # 缓存前一天数据
|
||
|
||
for i, date in enumerate(tqdm(dates, desc="回测")):
|
||
df = load_daily_features(args.data_dir, date)
|
||
if df is None or df.empty:
|
||
prev_df = None # 当天无数据，清空缓存
|
||
continue
|
||
|
||
# 第一天需要加载前一天数据（如果存在）
|
||
if i == 0 and prev_df is None:
|
||
prev_date = get_prev_trading_day(args.data_dir, date)
|
||
if prev_date:
|
||
prev_df = load_daily_features(args.data_dir, prev_date)
|
||
if prev_df is not None:
|
||
tqdm.write(f" 加载前一天数据: {prev_date}")
|
||
|
||
alerts = backtest_single_day_fast(ml_scorer, df, date, CONFIG, prev_df)
|
||
all_alerts.extend(alerts)
|
||
|
||
# 保存到数据库
|
||
if args.save_db and alerts:
|
||
if args.clear_first and not args.dry_run:
|
||
cleared = clear_alerts_by_date(date)
|
||
if cleared > 0:
|
||
tqdm.write(f" 清除 {date} 旧数据: {cleared} 条")
|
||
|
||
saved = save_alerts_to_mysql(alerts, dry_run=args.dry_run)
|
||
total_saved += saved
|
||
tqdm.write(f" {date}: {len(alerts)} 个异动, 保存 {saved} 条")
|
||
elif alerts:
|
||
tqdm.write(f" {date}: {len(alerts)} 个异动")
|
||
|
||
# 当天数据成为下一天的 prev_df
|
||
prev_df = df
|
||
|
||
# 导出 CSV
|
||
if args.export_csv:
|
||
export_to_csv(all_alerts, args.export_csv)
|
||
|
||
# 分析
|
||
analyze_alerts(all_alerts)
|
||
|
||
print(f"\n总计: {len(all_alerts)} 个异动")
|
||
if args.save_db:
|
||
print(f"已保存到数据库: {total_saved} 条")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|