716 lines
23 KiB
Python
716 lines
23 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
数据准备 V2 - 基于时间片对齐的特征计算（修复版）
|
||
|
||
核心改进：
|
||
1. 时间片对齐：9:35 和历史的 9:35 比，而不是和前30分钟比
|
||
2. Z-Score 特征：相对于同时间片历史分布的偏离程度
|
||
3. 滚动窗口基线：每个日期使用它之前 N 天的数据作为基线（不是固定的最后 N 天！）
|
||
4. 基于 Z-Score 的动量：消除一天内波动率异构性
|
||
|
||
修复：
|
||
- 滚动窗口基线：避免未来数据泄露
|
||
- Z-Score 动量：消除早盘/尾盘波动率差异
|
||
- 进程级数据库单例：避免连接池爆炸
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import numpy as np
|
||
import pandas as pd
|
||
from datetime import datetime, timedelta
|
||
from sqlalchemy import create_engine, text
|
||
from elasticsearch import Elasticsearch
|
||
from clickhouse_driver import Client
|
||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||
from typing import Dict, List, Tuple, Optional
|
||
from tqdm import tqdm
|
||
from collections import defaultdict
|
||
import warnings
|
||
import pickle
|
||
|
||
warnings.filterwarnings('ignore')
|
||
|
||
# ==================== 配置 ====================
|
||
|
||
MYSQL_URL = "mysql+pymysql://root:Zzl5588161!@192.168.1.5:3306/stock"
|
||
ES_HOST = 'http://127.0.0.1:9200'
|
||
ES_INDEX = 'concept_library_v3'
|
||
|
||
CLICKHOUSE_CONFIG = {
|
||
'host': '127.0.0.1',
|
||
'port': 9000,
|
||
'user': 'default',
|
||
'password': 'Zzl33818!',
|
||
'database': 'stock'
|
||
}
|
||
|
||
REFERENCE_INDEX = '000001.SH'
|
||
|
||
# 输出目录
|
||
OUTPUT_DIR = os.path.join(os.path.dirname(__file__), 'data_v2')
|
||
BASELINE_DIR = os.path.join(OUTPUT_DIR, 'baselines')
|
||
RAW_CACHE_DIR = os.path.join(OUTPUT_DIR, 'raw_cache')
|
||
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
||
os.makedirs(BASELINE_DIR, exist_ok=True)
|
||
os.makedirs(RAW_CACHE_DIR, exist_ok=True)
|
||
|
||
# 特征配置
|
||
CONFIG = {
|
||
'baseline_days': 20, # 滚动窗口大小
|
||
'min_baseline_samples': 10, # 最少需要10个样本才算有效基线
|
||
'limit_up_threshold': 9.8,
|
||
'limit_down_threshold': -9.8,
|
||
'zscore_clip': 5.0,
|
||
}
|
||
|
||
# 特征列表
|
||
FEATURES_V2 = [
|
||
'alpha', 'alpha_zscore', 'amt_zscore', 'rank_zscore',
|
||
'momentum_3m', 'momentum_5m', 'limit_up_ratio',
|
||
]
|
||
|
||
# ==================== 进程级单例（避免连接池爆炸）====================
|
||
|
||
# 进程级全局变量
|
||
_process_mysql_engine = None
|
||
_process_es_client = None
|
||
_process_ch_client = None
|
||
|
||
|
||
def init_process_connections():
|
||
"""进程初始化时调用，创建连接（单例）"""
|
||
global _process_mysql_engine, _process_es_client, _process_ch_client
|
||
_process_mysql_engine = create_engine(MYSQL_URL, echo=False, pool_pre_ping=True, pool_size=5)
|
||
_process_es_client = Elasticsearch([ES_HOST])
|
||
_process_ch_client = Client(**CLICKHOUSE_CONFIG)
|
||
|
||
|
||
def get_mysql_engine():
|
||
"""获取进程级 MySQL Engine（单例）"""
|
||
global _process_mysql_engine
|
||
if _process_mysql_engine is None:
|
||
_process_mysql_engine = create_engine(MYSQL_URL, echo=False, pool_pre_ping=True, pool_size=5)
|
||
return _process_mysql_engine
|
||
|
||
|
||
def get_es_client():
|
||
"""获取进程级 ES 客户端（单例）"""
|
||
global _process_es_client
|
||
if _process_es_client is None:
|
||
_process_es_client = Elasticsearch([ES_HOST])
|
||
return _process_es_client
|
||
|
||
|
||
def get_ch_client():
|
||
"""获取进程级 ClickHouse 客户端（单例）"""
|
||
global _process_ch_client
|
||
if _process_ch_client is None:
|
||
_process_ch_client = Client(**CLICKHOUSE_CONFIG)
|
||
return _process_ch_client
|
||
|
||
|
||
# ==================== 工具函数 ====================
|
||
|
||
def code_to_ch_format(code: str) -> str:
|
||
if not code or len(code) != 6 or not code.isdigit():
|
||
return None
|
||
if code.startswith('6'):
|
||
return f"{code}.SH"
|
||
elif code.startswith('0') or code.startswith('3'):
|
||
return f"{code}.SZ"
|
||
else:
|
||
return f"{code}.BJ"
|
||
|
||
|
||
def time_to_slot(ts) -> str:
|
||
"""将时间戳转换为时间片（HH:MM格式）"""
|
||
if isinstance(ts, str):
|
||
return ts
|
||
return ts.strftime('%H:%M')
|
||
|
||
|
||
# ==================== 获取概念列表 ====================
|
||
|
||
def get_all_concepts() -> List[dict]:
|
||
"""从ES获取所有叶子概念"""
|
||
es_client = get_es_client()
|
||
concepts = []
|
||
|
||
query = {
|
||
"query": {"match_all": {}},
|
||
"size": 100,
|
||
"_source": ["concept_id", "concept", "stocks"]
|
||
}
|
||
|
||
resp = es_client.search(index=ES_INDEX, body=query, scroll='2m')
|
||
scroll_id = resp['_scroll_id']
|
||
hits = resp['hits']['hits']
|
||
|
||
while len(hits) > 0:
|
||
for hit in hits:
|
||
source = hit['_source']
|
||
stocks = []
|
||
if 'stocks' in source and isinstance(source['stocks'], list):
|
||
for stock in source['stocks']:
|
||
if isinstance(stock, dict) and 'code' in stock and stock['code']:
|
||
stocks.append(stock['code'])
|
||
|
||
if stocks:
|
||
concepts.append({
|
||
'concept_id': source.get('concept_id'),
|
||
'concept_name': source.get('concept'),
|
||
'stocks': stocks
|
||
})
|
||
|
||
resp = es_client.scroll(scroll_id=scroll_id, scroll='2m')
|
||
scroll_id = resp['_scroll_id']
|
||
hits = resp['hits']['hits']
|
||
|
||
es_client.clear_scroll(scroll_id=scroll_id)
|
||
print(f"获取到 {len(concepts)} 个概念")
|
||
return concepts
|
||
|
||
|
||
# ==================== 获取交易日列表 ====================
|
||
|
||
def get_trading_days(start_date: str, end_date: str) -> List[str]:
|
||
"""获取交易日列表"""
|
||
client = get_ch_client()
|
||
|
||
query = f"""
|
||
SELECT DISTINCT toDate(timestamp) as trade_date
|
||
FROM stock_minute
|
||
WHERE toDate(timestamp) >= '{start_date}'
|
||
AND toDate(timestamp) <= '{end_date}'
|
||
ORDER BY trade_date
|
||
"""
|
||
|
||
result = client.execute(query)
|
||
days = [row[0].strftime('%Y-%m-%d') for row in result]
|
||
if days:
|
||
print(f"找到 {len(days)} 个交易日: {days[0]} ~ {days[-1]}")
|
||
return days
|
||
|
||
|
||
# ==================== 获取昨收价 ====================
|
||
|
||
def get_prev_close(stock_codes: List[str], trade_date: str) -> Dict[str, float]:
|
||
"""获取昨收价（上一交易日的收盘价 F007N）"""
|
||
valid_codes = [c for c in stock_codes if c and len(c) == 6 and c.isdigit()]
|
||
if not valid_codes:
|
||
return {}
|
||
|
||
codes_str = "','".join(valid_codes)
|
||
query = f"""
|
||
SELECT SECCODE, F007N
|
||
FROM ea_trade
|
||
WHERE SECCODE IN ('{codes_str}')
|
||
AND TRADEDATE = (
|
||
SELECT MAX(TRADEDATE) FROM ea_trade WHERE TRADEDATE < '{trade_date}'
|
||
)
|
||
AND F007N IS NOT NULL AND F007N > 0
|
||
"""
|
||
|
||
try:
|
||
engine = get_mysql_engine()
|
||
with engine.connect() as conn:
|
||
result = conn.execute(text(query))
|
||
return {row[0]: float(row[1]) for row in result if row[1]}
|
||
except Exception as e:
|
||
print(f"获取昨收价失败: {e}")
|
||
return {}
|
||
|
||
|
||
def get_index_prev_close(trade_date: str, index_code: str = REFERENCE_INDEX) -> float:
|
||
"""获取指数昨收价"""
|
||
code_no_suffix = index_code.split('.')[0]
|
||
|
||
try:
|
||
engine = get_mysql_engine()
|
||
with engine.connect() as conn:
|
||
result = conn.execute(text("""
|
||
SELECT F006N FROM ea_exchangetrade
|
||
WHERE INDEXCODE = :code AND TRADEDATE < :today
|
||
ORDER BY TRADEDATE DESC LIMIT 1
|
||
"""), {'code': code_no_suffix, 'today': trade_date}).fetchone()
|
||
|
||
if result and result[0]:
|
||
return float(result[0])
|
||
except Exception as e:
|
||
print(f"获取指数昨收失败: {e}")
|
||
|
||
return None
|
||
|
||
|
||
# ==================== 获取分钟数据 ====================
|
||
|
||
def get_daily_stock_data(trade_date: str, stock_codes: List[str]) -> pd.DataFrame:
|
||
"""获取单日所有股票的分钟数据"""
|
||
client = get_ch_client()
|
||
|
||
ch_codes = []
|
||
code_map = {}
|
||
for code in stock_codes:
|
||
ch_code = code_to_ch_format(code)
|
||
if ch_code:
|
||
ch_codes.append(ch_code)
|
||
code_map[ch_code] = code
|
||
|
||
if not ch_codes:
|
||
return pd.DataFrame()
|
||
|
||
ch_codes_str = "','".join(ch_codes)
|
||
|
||
query = f"""
|
||
SELECT code, timestamp, close, volume, amt
|
||
FROM stock_minute
|
||
WHERE toDate(timestamp) = '{trade_date}'
|
||
AND code IN ('{ch_codes_str}')
|
||
ORDER BY code, timestamp
|
||
"""
|
||
|
||
result = client.execute(query)
|
||
if not result:
|
||
return pd.DataFrame()
|
||
|
||
df = pd.DataFrame(result, columns=['ch_code', 'timestamp', 'close', 'volume', 'amt'])
|
||
df['code'] = df['ch_code'].map(code_map)
|
||
df = df.dropna(subset=['code'])
|
||
|
||
return df[['code', 'timestamp', 'close', 'volume', 'amt']]
|
||
|
||
|
||
def get_daily_index_data(trade_date: str, index_code: str = REFERENCE_INDEX) -> pd.DataFrame:
|
||
"""获取单日指数分钟数据"""
|
||
client = get_ch_client()
|
||
|
||
query = f"""
|
||
SELECT timestamp, close, volume, amt
|
||
FROM index_minute
|
||
WHERE toDate(timestamp) = '{trade_date}'
|
||
AND code = '{index_code}'
|
||
ORDER BY timestamp
|
||
"""
|
||
|
||
result = client.execute(query)
|
||
if not result:
|
||
return pd.DataFrame()
|
||
|
||
df = pd.DataFrame(result, columns=['timestamp', 'close', 'volume', 'amt'])
|
||
return df
|
||
|
||
|
||
# ==================== 计算原始概念特征（单日）====================
|
||
|
||
def compute_raw_concept_features(
|
||
trade_date: str,
|
||
concepts: List[dict],
|
||
all_stocks: List[str]
|
||
) -> pd.DataFrame:
|
||
"""计算单日概念的原始特征（alpha, amt, rank_pct, limit_up_ratio）"""
|
||
# 检查缓存
|
||
cache_file = os.path.join(RAW_CACHE_DIR, f'raw_{trade_date}.parquet')
|
||
if os.path.exists(cache_file):
|
||
return pd.read_parquet(cache_file)
|
||
|
||
# 获取数据
|
||
stock_df = get_daily_stock_data(trade_date, all_stocks)
|
||
if stock_df.empty:
|
||
return pd.DataFrame()
|
||
|
||
index_df = get_daily_index_data(trade_date)
|
||
if index_df.empty:
|
||
return pd.DataFrame()
|
||
|
||
# 获取昨收价
|
||
prev_close = get_prev_close(all_stocks, trade_date)
|
||
index_prev_close = get_index_prev_close(trade_date)
|
||
|
||
if not prev_close or not index_prev_close:
|
||
return pd.DataFrame()
|
||
|
||
# 计算涨跌幅
|
||
stock_df['prev_close'] = stock_df['code'].map(prev_close)
|
||
stock_df = stock_df.dropna(subset=['prev_close'])
|
||
stock_df['change_pct'] = (stock_df['close'] - stock_df['prev_close']) / stock_df['prev_close'] * 100
|
||
|
||
index_df['change_pct'] = (index_df['close'] - index_prev_close) / index_prev_close * 100
|
||
index_change_map = dict(zip(index_df['timestamp'], index_df['change_pct']))
|
||
|
||
# 获取所有时间点
|
||
timestamps = sorted(stock_df['timestamp'].unique())
|
||
|
||
# 概念到股票的映射
|
||
concept_stocks = {c['concept_id']: set(c['stocks']) for c in concepts}
|
||
|
||
results = []
|
||
|
||
for ts in timestamps:
|
||
ts_stock_data = stock_df[stock_df['timestamp'] == ts]
|
||
index_change = index_change_map.get(ts, 0)
|
||
|
||
stock_change = dict(zip(ts_stock_data['code'], ts_stock_data['change_pct']))
|
||
stock_amt = dict(zip(ts_stock_data['code'], ts_stock_data['amt']))
|
||
|
||
concept_features = []
|
||
|
||
for concept_id, stocks in concept_stocks.items():
|
||
concept_changes = [stock_change[s] for s in stocks if s in stock_change]
|
||
concept_amts = [stock_amt.get(s, 0) for s in stocks if s in stock_change]
|
||
|
||
if not concept_changes:
|
||
continue
|
||
|
||
avg_change = np.mean(concept_changes)
|
||
total_amt = sum(concept_amts)
|
||
alpha = avg_change - index_change
|
||
|
||
limit_up_count = sum(1 for c in concept_changes if c >= CONFIG['limit_up_threshold'])
|
||
limit_up_ratio = limit_up_count / len(concept_changes)
|
||
|
||
concept_features.append({
|
||
'concept_id': concept_id,
|
||
'alpha': alpha,
|
||
'total_amt': total_amt,
|
||
'limit_up_ratio': limit_up_ratio,
|
||
'stock_count': len(concept_changes),
|
||
})
|
||
|
||
if not concept_features:
|
||
continue
|
||
|
||
concept_df = pd.DataFrame(concept_features)
|
||
concept_df['rank_pct'] = concept_df['alpha'].rank(pct=True)
|
||
concept_df['timestamp'] = ts
|
||
concept_df['time_slot'] = time_to_slot(ts)
|
||
concept_df['trade_date'] = trade_date
|
||
|
||
results.append(concept_df)
|
||
|
||
if not results:
|
||
return pd.DataFrame()
|
||
|
||
result_df = pd.concat(results, ignore_index=True)
|
||
|
||
# 保存缓存
|
||
result_df.to_parquet(cache_file, index=False)
|
||
|
||
return result_df
|
||
|
||
|
||
# ==================== 滚动窗口基线计算 ====================
|
||
|
||
def compute_rolling_baseline(
|
||
historical_data: pd.DataFrame,
|
||
concept_id: str
|
||
) -> Dict[str, Dict]:
|
||
"""
|
||
计算单个概念的滚动基线
|
||
|
||
返回: {time_slot: {alpha_mean, alpha_std, amt_mean, amt_std, rank_mean, rank_std, sample_count}}
|
||
"""
|
||
if historical_data.empty:
|
||
return {}
|
||
|
||
concept_data = historical_data[historical_data['concept_id'] == concept_id]
|
||
if concept_data.empty:
|
||
return {}
|
||
|
||
baseline_dict = {}
|
||
|
||
for time_slot, group in concept_data.groupby('time_slot'):
|
||
if len(group) < CONFIG['min_baseline_samples']:
|
||
continue
|
||
|
||
alpha_std = group['alpha'].std()
|
||
amt_std = group['total_amt'].std()
|
||
rank_std = group['rank_pct'].std()
|
||
|
||
baseline_dict[time_slot] = {
|
||
'alpha_mean': group['alpha'].mean(),
|
||
'alpha_std': max(alpha_std if pd.notna(alpha_std) else 1.0, 0.1),
|
||
'amt_mean': group['total_amt'].mean(),
|
||
'amt_std': max(amt_std if pd.notna(amt_std) else group['total_amt'].mean() * 0.5, 1.0),
|
||
'rank_mean': group['rank_pct'].mean(),
|
||
'rank_std': max(rank_std if pd.notna(rank_std) else 0.2, 0.05),
|
||
'sample_count': len(group),
|
||
}
|
||
|
||
return baseline_dict
|
||
|
||
|
||
# ==================== 计算单日 Z-Score 特征（带滚动基线）====================
|
||
|
||
def compute_zscore_features_rolling(
|
||
trade_date: str,
|
||
concepts: List[dict],
|
||
all_stocks: List[str],
|
||
historical_raw_data: pd.DataFrame # 该日期之前 N 天的原始数据
|
||
) -> pd.DataFrame:
|
||
"""
|
||
计算单日的 Z-Score 特征（使用滚动窗口基线）
|
||
|
||
关键改进：
|
||
1. 基线只使用 trade_date 之前的数据（无未来泄露）
|
||
2. 动量基于 Z-Score 计算（消除波动率异构性）
|
||
"""
|
||
# 计算当日原始特征
|
||
raw_df = compute_raw_concept_features(trade_date, concepts, all_stocks)
|
||
|
||
if raw_df.empty:
|
||
return pd.DataFrame()
|
||
|
||
zscore_records = []
|
||
|
||
for concept_id, group in raw_df.groupby('concept_id'):
|
||
# 计算该概念的滚动基线（只用历史数据）
|
||
baseline_dict = compute_rolling_baseline(historical_raw_data, concept_id)
|
||
|
||
if not baseline_dict:
|
||
continue
|
||
|
||
# 按时间排序
|
||
group = group.sort_values('timestamp').reset_index(drop=True)
|
||
|
||
# Z-Score 历史（用于计算基于 Z-Score 的动量）
|
||
zscore_history = []
|
||
|
||
for idx, row in group.iterrows():
|
||
time_slot = row['time_slot']
|
||
|
||
if time_slot not in baseline_dict:
|
||
continue
|
||
|
||
bl = baseline_dict[time_slot]
|
||
|
||
# 计算 Z-Score
|
||
alpha_zscore = (row['alpha'] - bl['alpha_mean']) / bl['alpha_std']
|
||
amt_zscore = (row['total_amt'] - bl['amt_mean']) / bl['amt_std']
|
||
rank_zscore = (row['rank_pct'] - bl['rank_mean']) / bl['rank_std']
|
||
|
||
# 截断极端值
|
||
clip = CONFIG['zscore_clip']
|
||
alpha_zscore = np.clip(alpha_zscore, -clip, clip)
|
||
amt_zscore = np.clip(amt_zscore, -clip, clip)
|
||
rank_zscore = np.clip(rank_zscore, -clip, clip)
|
||
|
||
# 记录 Z-Score 历史
|
||
zscore_history.append(alpha_zscore)
|
||
|
||
# 基于 Z-Score 计算动量（消除波动率异构性）
|
||
momentum_3m = 0.0
|
||
momentum_5m = 0.0
|
||
|
||
if len(zscore_history) >= 3:
|
||
recent_3 = zscore_history[-3:]
|
||
older_3 = zscore_history[-6:-3] if len(zscore_history) >= 6 else [zscore_history[0]]
|
||
momentum_3m = np.mean(recent_3) - np.mean(older_3)
|
||
|
||
if len(zscore_history) >= 5:
|
||
recent_5 = zscore_history[-5:]
|
||
older_5 = zscore_history[-10:-5] if len(zscore_history) >= 10 else [zscore_history[0]]
|
||
momentum_5m = np.mean(recent_5) - np.mean(older_5)
|
||
|
||
zscore_records.append({
|
||
'concept_id': concept_id,
|
||
'timestamp': row['timestamp'],
|
||
'time_slot': time_slot,
|
||
'trade_date': trade_date,
|
||
# 原始特征
|
||
'alpha': row['alpha'],
|
||
'total_amt': row['total_amt'],
|
||
'limit_up_ratio': row['limit_up_ratio'],
|
||
'stock_count': row['stock_count'],
|
||
'rank_pct': row['rank_pct'],
|
||
# Z-Score 特征
|
||
'alpha_zscore': alpha_zscore,
|
||
'amt_zscore': amt_zscore,
|
||
'rank_zscore': rank_zscore,
|
||
# 基于 Z-Score 的动量
|
||
'momentum_3m': momentum_3m,
|
||
'momentum_5m': momentum_5m,
|
||
})
|
||
|
||
if not zscore_records:
|
||
return pd.DataFrame()
|
||
|
||
return pd.DataFrame(zscore_records)
|
||
|
||
|
||
# ==================== 多进程处理 ====================
|
||
|
||
def process_single_day_v2(args) -> Tuple[str, bool]:
|
||
"""处理单个交易日（多进程版本）"""
|
||
trade_date, day_index, concepts, all_stocks, all_trading_days = args
|
||
output_file = os.path.join(OUTPUT_DIR, f'features_v2_{trade_date}.parquet')
|
||
|
||
if os.path.exists(output_file):
|
||
return (trade_date, True)
|
||
|
||
try:
|
||
# 计算滚动窗口范围（该日期之前的 N 天）
|
||
baseline_days = CONFIG['baseline_days']
|
||
|
||
# 找出 trade_date 之前的交易日
|
||
start_idx = max(0, day_index - baseline_days)
|
||
end_idx = day_index # 不包含当天
|
||
|
||
if end_idx <= start_idx:
|
||
# 没有足够的历史数据
|
||
return (trade_date, False)
|
||
|
||
historical_days = all_trading_days[start_idx:end_idx]
|
||
|
||
# 加载历史原始数据
|
||
historical_dfs = []
|
||
for hist_date in historical_days:
|
||
cache_file = os.path.join(RAW_CACHE_DIR, f'raw_{hist_date}.parquet')
|
||
if os.path.exists(cache_file):
|
||
historical_dfs.append(pd.read_parquet(cache_file))
|
||
else:
|
||
# 需要计算
|
||
hist_df = compute_raw_concept_features(hist_date, concepts, all_stocks)
|
||
if not hist_df.empty:
|
||
historical_dfs.append(hist_df)
|
||
|
||
if not historical_dfs:
|
||
return (trade_date, False)
|
||
|
||
historical_raw_data = pd.concat(historical_dfs, ignore_index=True)
|
||
|
||
# 计算当日 Z-Score 特征（使用滚动基线）
|
||
df = compute_zscore_features_rolling(trade_date, concepts, all_stocks, historical_raw_data)
|
||
|
||
if df.empty:
|
||
return (trade_date, False)
|
||
|
||
df.to_parquet(output_file, index=False)
|
||
return (trade_date, True)
|
||
|
||
except Exception as e:
|
||
print(f"[{trade_date}] 处理失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return (trade_date, False)
|
||
|
||
|
||
# ==================== 主流程 ====================
|
||
|
||
def main():
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description='准备训练数据 V2（滚动窗口基线 + Z-Score 动量）')
|
||
parser.add_argument('--start', type=str, default='2022-01-01', help='开始日期')
|
||
parser.add_argument('--end', type=str, default=None, help='结束日期（默认今天）')
|
||
parser.add_argument('--workers', type=int, default=18, help='并行进程数')
|
||
parser.add_argument('--baseline-days', type=int, default=20, help='滚动基线窗口大小')
|
||
parser.add_argument('--force', action='store_true', help='强制重新计算（忽略缓存）')
|
||
|
||
args = parser.parse_args()
|
||
|
||
end_date = args.end or datetime.now().strftime('%Y-%m-%d')
|
||
CONFIG['baseline_days'] = args.baseline_days
|
||
|
||
print("=" * 60)
|
||
print("数据准备 V2 - 滚动窗口基线 + Z-Score 动量")
|
||
print("=" * 60)
|
||
print(f"日期范围: {args.start} ~ {end_date}")
|
||
print(f"并行进程数: {args.workers}")
|
||
print(f"滚动基线窗口: {args.baseline_days} 天")
|
||
|
||
# 初始化主进程连接
|
||
init_process_connections()
|
||
|
||
# 1. 获取概念列表
|
||
concepts = get_all_concepts()
|
||
all_stocks = list(set(s for c in concepts for s in c['stocks']))
|
||
print(f"股票总数: {len(all_stocks)}")
|
||
|
||
# 2. 获取交易日列表
|
||
trading_days = get_trading_days(args.start, end_date)
|
||
|
||
if not trading_days:
|
||
print("无交易日数据")
|
||
return
|
||
|
||
# 3. 第一阶段：预计算所有原始特征（用于缓存）
|
||
print(f"\n{'='*60}")
|
||
print("第一阶段：预计算原始特征（用于滚动基线）")
|
||
print(f"{'='*60}")
|
||
|
||
# 如果强制重新计算，删除缓存
|
||
if args.force:
|
||
import shutil
|
||
if os.path.exists(RAW_CACHE_DIR):
|
||
shutil.rmtree(RAW_CACHE_DIR)
|
||
os.makedirs(RAW_CACHE_DIR, exist_ok=True)
|
||
if os.path.exists(OUTPUT_DIR):
|
||
for f in os.listdir(OUTPUT_DIR):
|
||
if f.startswith('features_v2_'):
|
||
os.remove(os.path.join(OUTPUT_DIR, f))
|
||
|
||
# 单线程预计算原始特征（因为需要顺序缓存）
|
||
print(f"预计算 {len(trading_days)} 天的原始特征...")
|
||
for trade_date in tqdm(trading_days, desc="预计算原始特征"):
|
||
cache_file = os.path.join(RAW_CACHE_DIR, f'raw_{trade_date}.parquet')
|
||
if not os.path.exists(cache_file):
|
||
compute_raw_concept_features(trade_date, concepts, all_stocks)
|
||
|
||
# 4. 第二阶段：计算 Z-Score 特征（多进程）
|
||
print(f"\n{'='*60}")
|
||
print("第二阶段：计算 Z-Score 特征（滚动基线）")
|
||
print(f"{'='*60}")
|
||
|
||
# 从第 baseline_days 天开始（前面的没有足够历史）
|
||
start_idx = args.baseline_days
|
||
processable_days = trading_days[start_idx:]
|
||
|
||
if not processable_days:
|
||
print(f"错误：需要至少 {args.baseline_days + 1} 天的数据")
|
||
return
|
||
|
||
print(f"可处理日期: {processable_days[0]} ~ {processable_days[-1]} ({len(processable_days)} 天)")
|
||
print(f"跳过前 {start_idx} 天（基线预热期）")
|
||
|
||
# 构建任务
|
||
tasks = []
|
||
for i, trade_date in enumerate(trading_days):
|
||
if i >= start_idx:
|
||
tasks.append((trade_date, i, concepts, all_stocks, trading_days))
|
||
|
||
print(f"开始处理 {len(tasks)} 个交易日（{args.workers} 进程并行）...")
|
||
|
||
success_count = 0
|
||
failed_dates = []
|
||
|
||
# 使用进程池初始化器
|
||
with ProcessPoolExecutor(max_workers=args.workers, initializer=init_process_connections) as executor:
|
||
futures = {executor.submit(process_single_day_v2, task): task[0] for task in tasks}
|
||
|
||
with tqdm(total=len(futures), desc="处理进度", unit="天") as pbar:
|
||
for future in as_completed(futures):
|
||
trade_date = futures[future]
|
||
try:
|
||
result_date, success = future.result()
|
||
if success:
|
||
success_count += 1
|
||
else:
|
||
failed_dates.append(result_date)
|
||
except Exception as e:
|
||
print(f"\n[{trade_date}] 进程异常: {e}")
|
||
failed_dates.append(trade_date)
|
||
pbar.update(1)
|
||
|
||
print("\n" + "=" * 60)
|
||
print(f"处理完成: {success_count}/{len(tasks)} 个交易日")
|
||
if failed_dates:
|
||
print(f"失败日期: {failed_dates[:10]}{'...' if len(failed_dates) > 10 else ''}")
|
||
print(f"数据保存在: {OUTPUT_DIR}")
|
||
print("=" * 60)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|