"""训练管道 - 包含数据处理、模型训练和预测功能 本模块提供： 1. 数据准备：使用 FactorEngine 从因子计算结果中准备训练/测试数据 2. 数据处理：Fillna(0) -> Dropna 3. 模型训练：使用LightGBM训练分类模型 4. 预测和选股：输出每日top5股票池 注意：本模块使用 src.factors 框架进行因子计算，确保防泄露机制生效。 因子配置示例： from src.factors import MovingAverageFactor, ReturnRankFactor from src.training.pipeline import prepare_data # 直接传入因子实例列表 - 简单直观 factors = [ MovingAverageFactor(period=5), MovingAverageFactor(period=10), ReturnRankFactor(period=5), ] train_data, val_data, test_data, factor_config = prepare_data( factors=factors, ... ) # 或者使用 FactorConfig 包装（支持链式添加） from src.training.pipeline import FactorConfig config = FactorConfig() .add(MovingAverageFactor(period=5)) .add(MovingAverageFactor(period=10)) .add(ReturnRankFactor(period=5)) """ from datetime import datetime from pathlib import Path from typing import List, Optional import numpy as np import polars as pl from src.factors import DataLoader, FactorEngine, BaseFactor from src.factors.data_spec import DataSpec from src.factors.momentum import MovingAverageFactor, ReturnRankFactor from src.pipeline import ( DropNAProcessor, FillNAProcessor, LightGBMModel, PipelineStage, ProcessingPipeline, TaskType, ) # ========== 因子配置类 ========== class FactorConfig: """因子配置类 - 管理因子列表 用于包装因子实例列表，支持链式添加。 示例： # 方式1：初始化时传入列表 config = FactorConfig([ MovingAverageFactor(period=5), ReturnRankFactor(period=5), ]) # 方式2：链式添加 config = FactorConfig() .add(MovingAverageFactor(period=5)) .add(ReturnRankFactor(period=5)) # 获取因子实例列表 factors = config.get_factors() # 获取特征列名 feature_cols = config.get_feature_names() """ def __init__(self, factors: Optional[List[BaseFactor]] = None): """初始化因子配置 Args: factors: 因子实例列表 """ self._factors: List[BaseFactor] = factors or [] def add(self, factor: BaseFactor) -> "FactorConfig": """添加因子到配置 支持链式调用： config = FactorConfig() .add(MovingAverageFactor(period=5)) .add(ReturnRankFactor(period=5)) Args: factor: 因子实例 Returns: self，支持链式调用 """ if not isinstance(factor, BaseFactor): raise ValueError(f"必须是 BaseFactor 实例， got {type(factor)}") self._factors.append(factor) return self def get_factors(self) -> List[BaseFactor]: """获取因子实例列表 Returns: 因子实例列表 """ return self._factors def get_feature_names(self) -> List[str]: """获取所有因子的特征列名 Returns: 特征列名列表 """ return [f.name for f in self._factors] def get_max_lookback(self) -> int: """获取所有因子中最大的 lookback 天数 Returns: 最大 lookback 天数 """ max_lookback = 0 for factor in self._factors: for spec in factor.data_specs: max_lookback = max(max_lookback, spec.lookback_days) return max_lookback def __len__(self) -> int: return len(self._factors) def __repr__(self) -> str: names = [f.name for f in self._factors] return f"FactorConfig({names})" def prepare_data( factors: Optional[List[BaseFactor]] = None, data_dir: str = "data", train_start: str = "20180101", train_end: str = "20230101", val_start: str = "20230101", val_end: str = "20230601", test_start: str = "20230601", test_end: str = "20240101", ) -> tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame, FactorConfig]: """准备训练、验证和测试数据 使用 FactorEngine 计算因子，确保防泄露机制生效。 Args: factors: 因子实例列表，默认为 None（使用 MA5, MA10, ReturnRank5） data_dir: 数据目录 train_start: 训练集开始日期 train_end: 训练集结束日期 val_start: 验证集开始日期 val_end: 验证集结束日期 test_start: 测试集开始日期 test_end: 测试集结束日期 Returns: (train_data, val_data, test_data, factor_config): 训练集、验证集、测试集的DataFrame，以及使用的因子配置 """ from src.data.storage import Storage storage = Storage() # 1. 处理因子配置 if factors is None: # 默认因子配置 factor_config = FactorConfig( [ MovingAverageFactor(period=5), MovingAverageFactor(period=10), ReturnRankFactor(period=5), ] ) elif isinstance(factors, FactorConfig): factor_config = factors elif isinstance(factors, list): # 转换为 FactorConfig factor_config = FactorConfig(factors) else: raise ValueError( f"factors 必须是 List[BaseFactor] 或 FactorConfig， got {type(factors)}" ) factors_list = factor_config.get_factors() feature_cols = factor_config.get_feature_names() if not factors_list: raise ValueError("至少需要提供一个因子") print(f"[PrepareData] 使用因子: {feature_cols}") # 2. 初始化 FactorEngine loader = DataLoader(data_dir=data_dir) engine = FactorEngine(loader) # 3. 计算需要的回溯天数 max_lookback = factor_config.get_max_lookback() start_with_lookback = str(int(train_start) - 10000) # 往前取一年 # 获取所有股票列表 start_dt = f"{start_with_lookback[:4]}-{start_with_lookback[4:6]}-{start_with_lookback[6:8]}" end_dt = f"{test_end[:4]}-{test_end[4:6]}-{test_end[6:8]}" all_stocks_query = f""" SELECT DISTINCT ts_code FROM daily WHERE trade_date >= '{start_dt}' AND trade_date <= '{end_dt}' """ all_stocks_df = storage._connection.sql(all_stocks_query).pl() all_stocks = all_stocks_df["ts_code"].to_list() print(f"[PrepareData] 股票数量: {len(all_stocks)}") # 4. 计算所有因子并合并 all_features = None for factor in factors_list: print(f"[PrepareData] 计算因子: {factor.name} ({factor.factor_type})") if factor.factor_type == "time_series": # 时序因子需要传入股票列表 result = engine.compute( factor, stock_codes=all_stocks, start_date=start_with_lookback, end_date=test_end, ) else: # 截面因子不需要股票列表 result = engine.compute( factor, start_date=start_with_lookback, end_date=test_end, ) # 合并结果 if all_features is None: all_features = result else: # 确保没有重复的 _right 列 result = result.select([ c for c in result.columns if not c.endswith("_right") ]) all_features = all_features.select([ c for c in all_features.columns if not c.endswith("_right") ]) all_features = all_features.join( result, on=["trade_date", "ts_code"], how="outer" ) if all_features is None: raise ValueError("没有计算任何因子") # 5. 计算标签（未来5日收益率） all_data = _compute_label(all_features, start_date=train_start, end_date=test_end) # 6. 过滤不符合条件的股票 all_data = _filter_invalid_stocks(all_data) print(f"[PrepareData] After filtering: total={len(all_data)}") # 转换日期格式用于比较 train_start_fmt = f"{train_start[:4]}-{train_start[4:6]}-{train_start[6:8]}" train_end_fmt = f"{train_end[:4]}-{train_end[4:6]}-{train_end[6:8]}" val_start_fmt = f"{val_start[:4]}-{val_start[4:6]}-{val_start[6:8]}" val_end_fmt = f"{val_end[:4]}-{val_end[4:6]}-{val_end[6:8]}" test_start_fmt = f"{test_start[:4]}-{test_start[4:6]}-{test_start[6:8]}" test_end_fmt = f"{test_end[:4]}-{test_end[4:6]}-{test_end[6:8]}" # 拆分数据 train_data = all_data.filter( (pl.col("trade_date") >= train_start_fmt) & (pl.col("trade_date") <= train_end_fmt) ) val_data = all_data.filter( (pl.col("trade_date") >= val_start_fmt) & (pl.col("trade_date") <= val_end_fmt) ) test_data = all_data.filter( (pl.col("trade_date") >= test_start_fmt) & (pl.col("trade_date") <= test_end_fmt) ) print( f"[PrepareData] Split result: train={len(train_data)}, val={len(val_data)}, test={len(test_data)}" ) return train_data, val_data, test_data, factor_config def _filter_invalid_stocks(df: pl.DataFrame) -> pl.DataFrame: """过滤不符合条件的股票 过滤规则： 1. 过滤北交所股票（ts_code 以 BJ 结尾） 2. 过滤创业板股票（ts_code 以 30 开头） 3. 过滤科创板股票（ts_code 以 68 开头） 4. 过滤退市/风险股票（ts_code 以 8 开头） Args: df: 原始数据 Returns: 过滤后的数据 """ ts_code_col = pl.col("ts_code") return df.filter( ~ts_code_col.str.ends_with("BJ") & ~ts_code_col.str.starts_with("30") & ~ts_code_col.str.starts_with("68") & ~ts_code_col.str.starts_with("8") ) def _compute_label( features_df: pl.DataFrame, start_date: str, end_date: str, ) -> pl.DataFrame: """计算标签（未来5日收益率） 标签定义：未来5日收益率大于0为1，否则为0 Args: features_df: 包含因子的DataFrame start_date: 开始日期 end_date: 结束日期 Returns: 包含因子和标签的DataFrame """ from src.data.storage import Storage storage = Storage() # 从数据库获取收盘价数据用于计算标签 start_dt = f"{start_date[:4]}-{start_date[4:6]}-{start_date[6:8]}" end_dt = f"{end_date[:4]}-{end_date[4:6]}-{end_date[6:8]}" # 需要多取5天数据来计算未来收益率 end_dt_extended = f"{end_date[:4]}-{end_date[4:6]}-{int(end_date[6:8]) + 5}" price_query = f""" SELECT ts_code, trade_date, close FROM daily WHERE trade_date >= '{start_dt}' AND trade_date <= '{end_dt_extended}' ORDER BY ts_code, trade_date """ price_data = storage._connection.sql(price_query).pl() price_data = price_data.with_columns( pl.col("trade_date").dt.strftime("%Y-%m-%d").alias("trade_date") ) # 按股票计算未来5日收益率 result_list = [] for ts_code in price_data["ts_code"].unique(): stock_data = price_data.filter(pl.col("ts_code") == ts_code).sort("trade_date") if len(stock_data) < 6: continue # 计算未来5日收益率 future_return = stock_data["close"].shift(-5) - stock_data["close"] future_return_pct = future_return / stock_data["close"] stock_data = stock_data.with_columns( [ future_return_pct.alias("future_return_5"), ] ) # 生成标签：收益率>0为1，否则为0 stock_data = stock_data.with_columns( [ (pl.col("future_return_5") > 0).cast(pl.Int8).alias("label"), ] ) result_list.append(stock_data.select(["trade_date", "ts_code", "label"])) if not result_list: return pl.DataFrame() label_df = pl.concat(result_list) # 将标签合并到因子数据 result = features_df.join(label_df, on=["trade_date", "ts_code"], how="inner") # 过滤有效日期范围 result = result.filter( (pl.col("trade_date") >= start_dt) & (pl.col("trade_date") <= end_dt) ) return result def create_pipeline() -> ProcessingPipeline: """创建数据处理流水线 处理流程： 1. FillNA(0): 将缺失值填充为0 注意：不使用 Dropna，因为会导致训练和预测时的行数不匹配 Returns: 配置好的ProcessingPipeline """ processors = [ FillNAProcessor(method="zero"), # 缺失值填充为0 ] return ProcessingPipeline(processors) def train_model( train_data: pl.DataFrame, val_data: Optional[pl.DataFrame], feature_cols: List[str], label_col: str = "label", model_params: Optional[dict] = None, ) -> tuple[LightGBMModel, ProcessingPipeline]: """训练LightGBM分类模型 Args: train_data: 训练数据 val_data: 验证数据（用于早停） feature_cols: 特征列名列表 label_col: 标签列名 model_params: 模型参数字典 Returns: (训练好的模型, 处理流水线) """ # 创建处理流水线 pipeline = create_pipeline() print("[TrainModel] Pipeline created: FillNA(0)") # 准备训练特征和标签 X_train = train_data.select(feature_cols) y_train = train_data[label_col] print(f"[TrainModel] Train samples: {len(X_train)}, features: {feature_cols}") # 处理训练数据 X_train_processed = pipeline.fit_transform(X_train, stage=PipelineStage.TRAIN) print(f"[TrainModel] After processing: {len(X_train_processed)} samples") # 过滤训练集有效标签（排除-1等无效值） valid_mask = y_train.is_in([0, 1]) X_train_processed = X_train_processed.filter(valid_mask) y_train = y_train.filter(valid_mask) print( f"[TrainModel] After filtering valid labels: {len(X_train_processed)} samples" ) print( f"[TrainModel] Train label distribution: {dict(y_train.value_counts().sort('label').iter_rows())}" ) # 准备验证集 X_val_processed = None y_val = None if val_data is not None and len(val_data) > 0: X_val = val_data.select(feature_cols) y_val = val_data[label_col] print(f"[TrainModel] Val samples: {len(X_val)}") # 处理验证集数据（使用训练集的参数） X_val_processed = pipeline.transform(X_val, stage=PipelineStage.TEST) # 过滤验证集有效标签 val_valid_mask = y_val.is_in([0, 1]) X_val_processed = X_val_processed.filter(val_valid_mask) y_val = y_val.filter(val_valid_mask) print(f"[TrainModel] Val after filtering: {len(X_val_processed)} samples") print( f"[TrainModel] Val label distribution: {dict(y_val.value_counts().sort('label').iter_rows())}" ) # 创建模型 params = model_params or { "n_estimators": 100, "learning_rate": 0.05, "max_depth": 5, "num_leaves": 31, } print(f"[TrainModel] Model params: {params}") model = LightGBMModel( task_type="classification", params=params, ) # 训练模型（使用验证集早停） print("[TrainModel] Training LightGBM...") if X_val_processed is not None and y_val is not None: print("[TrainModel] Using validation set for early stopping") model.fit(X_train_processed, y_train, X_val_processed, y_val) else: model.fit(X_train_processed, y_train) print("[TrainModel] Training completed!") return model, pipeline def predict_top_stocks( model: LightGBMModel, pipeline: ProcessingPipeline, test_data: pl.DataFrame, feature_cols: List[str], top_n: int = 5, ) -> pl.DataFrame: """预测并选出每日top N股票 Args: model: 训练好的模型 pipeline: 数据处理流水线 test_data: 测试数据 feature_cols: 特征列名 top_n: 每日选出的股票数量 Returns: 包含日期和股票代码的DataFrame """ # 准备特征和必要列 X_test = test_data.select(feature_cols) key_cols = ["trade_date", "ts_code"] key_data = test_data.select(key_cols) print(f"[Predict] Test samples: {len(X_test)}, top_n: {top_n}") # 处理数据（使用训练阶段的参数） X_test_processed = pipeline.transform(X_test, stage=PipelineStage.TEST) print(f"[Predict] Data processed, shape: {X_test_processed.shape}") # 预测概率 probs = model.predict_proba(X_test_processed) print(f"[Predict] Predictions generated, probability shape: {probs.shape}") # 使用 key_data 添加预测结果，保持行数一致 result = key_data.with_columns( pl.Series( name="pred_prob", values=probs[:, 1] if len(probs.shape) > 1 and probs.shape[1] > 1 else probs.flatten(), ), ) # 每日选出top N top_stocks = [] for date in result["trade_date"].unique().sort(): day_data = result.filter(pl.col("trade_date") == date) # 按概率降序排序，选出top N day_top = day_data.sort("pred_prob", descending=True).head(top_n) top_stocks.append( day_top.select(["trade_date", "pred_prob", "ts_code"]).rename( {"pred_prob": "score"} ) ) return pl.concat(top_stocks) def save_top_stocks(top_stocks: pl.DataFrame, output_path: str) -> None: """保存选股结果到TSV文件 Args: top_stocks: 选股结果 output_path: 输出文件路径 """ # 转换为pandas并保存为TSV df = top_stocks.to_pandas() df.to_csv(output_path, sep="\t", index=False) print(f"[Training] Top stocks saved to: {output_path}") def run_training( factors: Optional[List[BaseFactor]] = None, data_dir: str = "data", output_path: str = "output/top_stocks.tsv", train_start: str = "20180101", train_end: str = "20230101", val_start: str = "20230101", val_end: str = "20230601", test_start: str = "20230601", test_end: str = "20240101", top_n: int = 5, ) -> pl.DataFrame: """运行完整训练流程 Args: factors: 因子实例列表，默认为 None（使用 MA5, MA10, ReturnRank5） data_dir: 数据目录 output_path: 输出文件路径 train_start: 训练集开始日期 train_end: 训练集结束日期 val_start: 验证集开始日期 val_end: 验证集结束日期 test_start: 测试集开始日期 test_end: 测试集结束日期 top_n: 每日选股数量 Returns: 选股结果DataFrame """ print(f"[Training] Starting training pipeline...") print(f"[Training] Train period: {train_start} -> {train_end}") print(f"[Training] Val period: {val_start} -> {val_end}") print(f"[Training] Test period: {test_start} -> {test_end}") # 1. 准备数据 print("[Training] Preparing data...") train_data, val_data, test_data, factor_config = prepare_data( factors=factors, data_dir=data_dir, train_start=train_start, train_end=train_end, val_start=val_start, val_end=val_end, test_start=test_start, test_end=test_end, ) print(f"[Training] Train samples: {len(train_data)}") print(f"[Training] Val samples: {len(val_data)}") print(f"[Training] Test samples: {len(test_data)}") # 2. 获取特征列名 feature_cols = factor_config.get_feature_names() label_col = "label" print(f"[Training] Feature columns: {feature_cols}") # 3. 训练模型 print("[Training] Training model...") model, pipeline = train_model( train_data=train_data, val_data=val_data, feature_cols=feature_cols, label_col=label_col, ) # 4. 测试集预测 print("[Training] Predicting on test set...") top_stocks = predict_top_stocks( model=model, pipeline=pipeline, test_data=test_data, feature_cols=feature_cols, top_n=top_n, ) # 5. 保存结果 print(f"[Training] Saving results to {output_path}...") Path(output_path).parent.mkdir(parents=True, exist_ok=True) save_top_stocks(top_stocks, output_path) print("[Training] Training completed!") return top_stocks