import sys import time import json from datetime import datetime from pathlib import Path import numpy as np import pandas as pd sys.path.insert(0, str(Path(__file__).parent)) from data.sina_fetcher import sina_get_stock_hist, sina_get_market_snapshot from data.preprocessor import prepare_hist, clean_snapshot from features.technical import compute_technical_features from models.ensemble import EnsembleModel from models.lgbm_model import WalkForwardTrainer from config import BASE_DIR, TRAINING_DATA_DIR METRICS_DIR = BASE_DIR / "training_data" METRICS_DIR.mkdir(exist_ok=True) STATIC_POOL = [ "000001", "000002", "000063", "000100", "000157", "000333", "000338", "000425", "000538", "000568", "000596", "000625", "000651", "000661", "000725", "000768", "000776", "000858", "000895", "000938", "000963", "001979", "002027", "002049", "002120", "002142", "002230", "002236", "002271", "002304", "002352", "002371", "002415", "002460", "002475", "002594", "002601", "002607", "002714", "002736", "002841", "002916", "002938", "300015", "300033", "300059", "300122", "300124", "300142", "300347", "600000", "600009", "600010", "600015", "600016", "600019", "600028", "600029", "600030", "600031", "600036", "600048", "600050", "600061", "600085", "600104", "600109", "600111", "600115", "600132", "600150", "600160", "600170", "600176", "600183", "600196", "600276", "600309", "600332", "600346", "600352", "600362", "600406", "600436", "600438", "600489", "600519", "600547", "600570", "600585", "600588", "600600", "600690", "600703", "600741", "600745", "600809", "600837", "600887", "600893", "600900", "601006", "601012", "601088", "601111", "601138", "601155", "601166", "601169", "601186", "601211", "601225", "601228", "601229", "601236", "601288", "601318", "601328", "601336", "601360", "601377", "601390", "601398", "601601", "601628", "601633", "601668", "601669", "601688", "601698", "601766", "601788", "601800", "601818", "601838", "601857", "601877", "601878", "601881", "601888", "601899", "601919", "601939", "601985", "601988", "601989", "601998", "603019", "603160", "603259", "603288", "603369", "603501", "603517", "603568", "603799", "603833", "603899", "603986", "603993", ] def fetch_dynamic_pool(max_stocks: int = 300) -> list: print(" 获取动态选股池（成交额Top300）...") snapshot = sina_get_market_snapshot(max_pages=10) if snapshot.empty: print(" 动态选股池获取失败，使用静态池") return STATIC_POOL[:max_stocks] snapshot = clean_snapshot(snapshot) snapshot = snapshot.sort_values("成交额(元)", ascending=False) codes = snapshot.head(max_stocks)["代码"].tolist() print(f" 动态选股池: {len(codes)} 只") return codes def fetch_training_data(max_stocks: int = 300, days: int = 250, use_dynamic: bool = True) -> pd.DataFrame: if use_dynamic: pool = fetch_dynamic_pool(max_stocks) else: pool = STATIC_POOL[:max_stocks] print(f"获取 {len(pool)} 只股票 {days} 天历史数据...") all_data = [] success = 0 cache_path = TRAINING_DATA_DIR / "training_data.pkl" if cache_path.exists(): cached = pd.read_pickle(str(cache_path)) if not cached.empty: cached_date = cached["date"].max() if "date" in cached.columns else None if cached_date and (pd.Timestamp.now() - cached_date).days < 3: print(f" 使用缓存数据 (截至 {cached_date.date()}, {len(cached)} 条)") return cached for i, code in enumerate(pool): df = sina_get_stock_hist(code, days=days) if df is None or df.empty: continue df = prepare_hist(df) if df.empty or len(df) < 30: continue tech = compute_technical_features(df) if tech.empty: continue tech["code"] = code all_data.append(tech) success += 1 if (i + 1) % 50 == 0: print(f" 已处理 {i + 1}/{len(pool)}, 成功 {success}") time.sleep(0.3) print(f" 数据获取完成: {success} 只股票") if all_data: result = pd.concat(all_data, ignore_index=True) result.to_pickle(str(cache_path)) return result return pd.DataFrame() def generate_labels(df: pd.DataFrame, forward_days: int = 1, threshold: float = 0.02) -> pd.Series: future_ret = df.groupby("code")["close"].pct_change(forward_days).shift(-forward_days) labels = (future_ret >= threshold).astype(int) return labels def save_model_with_version(model: EnsembleModel, metrics: dict, wf_auc: float): today = datetime.now().strftime("%Y%m%d") version_dir = MODEL_DIR / "versions" version_dir.mkdir(exist_ok=True) model.lgbm.save(f"ensemble_lgbm") model.xgb.save(f"ensemble_xgb") version_dir = BASE_DIR / "saved_models" / "versions" version_dir.mkdir(exist_ok=True) import shutil for ext in [".txt", "_meta.json"]: src = MODEL_DIR / f"ensemble_lgbm{ext}" if src.exists(): shutil.copy2(str(src), str(version_dir / f"ensemble_lgbm_{today}{ext}")) for ext in [".json", "_meta.json"]: src = MODEL_DIR / f"ensemble_xgb{ext}" if src.exists(): shutil.copy2(str(src), str(version_dir / f"ensemble_xgb_{today}{ext}")) versions = sorted(version_dir.glob("ensemble_lgbm_*_meta.json")) if len(versions) > 10: for old in versions[:len(versions) - 10]: date_part = old.name.replace("ensemble_lgbm_", "").replace("_meta.json", "") for ext in [".txt", "_meta.json"]: f = version_dir / f"ensemble_lgbm_{date_part}{ext}" if f.exists(): f.unlink() for ext in [".json", "_meta.json"]: f = version_dir / f"ensemble_xgb_{date_part}{ext}" if f.exists(): f.unlink() return today def save_metrics(metrics: dict, wf_auc: float, train_samples: int, positive_ratio: float, today: str): metrics_file = METRICS_DIR / "model_metrics.json" history = [] if metrics_file.exists(): try: history = json.loads(metrics_file.read_text()) except Exception: history = [] entry = { "date": today, "train_samples": train_samples, "positive_ratio": round(positive_ratio, 4), "lgbm_auc": round(metrics["lgbm"]["auc"], 4), "lgbm_accuracy": round(metrics["lgbm"]["accuracy"], 4), "xgb_auc": round(metrics["xgb"]["auc"], 4), "xgb_accuracy": round(metrics["xgb"]["accuracy"], 4), "wf_auc_mean": round(wf_auc, 4), } history.append(entry) history = history[-50:] metrics_file.write_text(json.dumps(history, ensure_ascii=False, indent=2)) MODEL_DIR = BASE_DIR / "saved_models" MODEL_DIR.mkdir(exist_ok=True) def train_model(df: pd.DataFrame): print("\n准备训练数据...") skip_cols = {"date", "code", "open", "close", "high", "low", "volume", "amount"} feature_cols = [c for c in df.columns if c not in skip_cols and df[c].dtype in ["float64", "int64", "float32"]] df = df.dropna(subset=feature_cols, how="all") df[feature_cols] = df[feature_cols].fillna(0) labels = generate_labels(df) valid_mask = labels.notna() df = df[valid_mask] labels = labels[valid_mask] train_samples = len(df) positive_ratio = labels.mean() print(f" 训练样本: {train_samples}, 正样本比例: {positive_ratio:.3f}") print(f" 特征数量: {len(feature_cols)}") X = df[feature_cols] y = labels.astype(int) print("\n训练集成模型...") model = EnsembleModel() metrics = model.train(X, y, eval_pct=0.2) print(f" LightGBM AUC: {metrics['lgbm']['auc']:.4f}") print(f" XGBoost AUC: {metrics['xgb']['auc']:.4f}") print("\nWalk-forward 验证...") trainer = WalkForwardTrainer(n_splits=3, gap=5) wf_results, wf_models = trainer.train(X, y) wf_auc = wf_results["auc"].mean() print(f" WF 平均 AUC: {wf_auc:.4f} ± {wf_results['auc'].std():.4f}") print("\n特征重要性 Top 10:") imp = model.feature_importance() for _, row in imp.head(10).iterrows(): print(f" {row['feature']:<25} {row['importance']:.4f}") today = save_model_with_version(model, metrics, wf_auc) save_metrics(metrics, wf_auc, train_samples, positive_ratio, today) print(f"\n模型已保存 (版本: {today})") print(f"性能指标已保存: training_data/model_metrics.json") return model, metrics def main(): import argparse parser = argparse.ArgumentParser(description="A股超短期选股模型训练") parser.add_argument("--stocks", type=int, default=300, help="训练股票数量 (默认300)") parser.add_argument("--days", type=int, default=250, help="历史天数 (默认250)") parser.add_argument("--static", action="store_true", help="使用静态股票池而非动态") args = parser.parse_args() print("=" * 60) print(" A股超短期选股模型训练") print(f" 股票数: {args.stocks}, 天数: {args.days}") print("=" * 60) df = fetch_training_data(max_stocks=args.stocks, days=args.days, use_dynamic=not args.static) if df.empty: print("ERROR: 无法获取训练数据") return model, metrics = train_model(df) print("\n" + "=" * 60) print(" 训练完成!") print("=" * 60) if __name__ == "__main__": main()