ProStock/src/experiment/regression.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LightGBM 回归训练示例 - 使用因子字符串表达式\n",
    "\n",
    "使用字符串表达式定义因子，训练 LightGBM 回归模型预测未来5日收益率。\n",
    "\n",
    "**Label**: return_5 = (ts_delay(close, -5) / close) - 1  # 未来5日收益率"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 导入依赖"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:27:36.008131Z",
     "start_time": "2026-03-07T16:27:36.005668Z"
    }
   },
   "source": [
    "import os\n",
    "from datetime import datetime\n",
    "from typing import List\n",
    "\n",
    "import polars as pl\n",
    "\n",
    "from src.factors import FactorEngine\n",
    "from src.training import (\n",
    "    DateSplitter,\n",
    "    LightGBMModel,\n",
    "    STFilter,\n",
    "    StandardScaler,\n",
    "    StockFilterConfig,\n",
    "    StockPoolManager,\n",
    "    Trainer,\n",
    "    Winsorizer,\n",
    "    NullFiller,\n",
    ")\n",
    "from src.training.config import TrainingConfig"
   ],
   "outputs": [],
   "execution_count": 6
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 定义辅助函数"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:27:36.020762Z",
     "start_time": "2026-03-07T16:27:36.011650Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def create_factors_with_strings(engine: FactorEngine, factor_definitions: dict, label_factor: dict) -> List[str]:\n",
    "    print(\"=\" * 80)\n",
    "    print(\"使用字符串表达式定义因子\")\n",
    "    print(\"=\" * 80)\n",
    "\n",
    "    # 注册所有特征因子\n",
    "    print(\"\\n注册特征因子:\")\n",
    "    for name, expr in factor_definitions.items():\n",
    "        engine.add_factor(name, expr)\n",
    "        print(f\"  - {name}: {expr}\")\n",
    "\n",
    "    # 注册 label 因子\n",
    "    print(\"\\n注册 Label 因子:\")\n",
    "    for name, expr in label_factor.items():\n",
    "        engine.add_factor(name, expr)\n",
    "        print(f\"  - {name}: {expr}\")\n",
    "\n",
    "    # 从字典自动获取特征列\n",
    "    feature_cols = list(factor_definitions.keys())\n",
    "\n",
    "    print(f\"\\n特征因子数: {len(feature_cols)}\")\n",
    "    print(f\"Label: {list(label_factor.keys())[0]}\")\n",
    "    print(f\"已注册因子总数: {len(engine.list_registered())}\")\n",
    "\n",
    "    return feature_cols\n",
    "\n",
    "\n",
    "def prepare_data(\n",
    "    engine: FactorEngine,\n",
    "    feature_cols: List[str],\n",
    "    start_date: str,\n",
    "    end_date: str,\n",
    ") -> pl.DataFrame:\n",
    "    print(\"\\n\" + \"=\" * 80)\n",
    "    print(\"准备数据\")\n",
    "    print(\"=\" * 80)\n",
    "\n",
    "    # 计算因子（全市场数据）\n",
    "    print(f\"\\n计算因子: {start_date} - {end_date}\")\n",
    "    factor_names = feature_cols + [\"return_5\"]  # 包含 label\n",
    "\n",
    "    data = engine.compute(\n",
    "        factor_names=factor_names,\n",
    "        start_date=start_date,\n",
    "        end_date=end_date,\n",
    "    )\n",
    "\n",
    "    print(f\"数据形状: {data.shape}\")\n",
    "    print(f\"数据列: {data.columns}\")\n",
    "    print(f\"\\n前5行预览:\")\n",
    "    print(data.head())\n",
    "\n",
    "    return data"
   ],
   "outputs": [],
   "execution_count": 7
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 配置参数\n",
    "\n",
    "### 3.1 因子定义"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:27:36.027366Z",
     "start_time": "2026-03-07T16:27:36.024465Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 特征因子定义字典：新增因子只需在此处添加一行\n",
    "FACTOR_DEFINITIONS = {\n",
    "    # 1. 价格动量因子\n",
    "    \"ma5\": \"ts_mean(close, 5)\",\n",
    "    \"ma10\": \"ts_mean(close, 10)\",\n",
    "    \"ma20\": \"ts_mean(close, 20)\",\n",
    "    \"ma_ratio\": \"ts_mean(close, 5) / ts_mean(close, 20) - 1\",\n",
    "    # 2. 波动率因子\n",
    "    \"volatility_5\": \"ts_std(close, 5)\",\n",
    "    \"volatility_20\": \"ts_std(close, 20)\",\n",
    "    \"vol_ratio\": \"ts_std(close, 5) / (ts_std(close, 20) + 1e-8)\",\n",
    "    # 3. 收益率动量因子\n",
    "    \"return_10\": \"(close / ts_delay(close, 10)) - 1\",\n",
    "    \"return_20\": \"(close / ts_delay(close, 20)) - 1\",\n",
    "    # 4. 收益率变化因子\n",
    "    \"return_diff\": \"(close / ts_delay(close, 5)) - 1 - ((close / ts_delay(close, 10)) - 1)\",\n",
    "    # 5. 成交量因子\n",
    "    \"vol_ma5\": \"ts_mean(vol, 5)\",\n",
    "    \"vol_ma20\": \"ts_mean(vol, 20)\",\n",
    "    \"vol_ratio\": \"ts_mean(vol, 5) / (ts_mean(vol, 20) + 1e-8)\",\n",
    "    # 6. 市值因子（截面排名）\n",
    "    \"market_cap_rank\": \"cs_rank(total_mv)\",\n",
    "    # 7. 价格位置因子\n",
    "    \"high_low_ratio\": \"(close - ts_min(low, 20)) / (ts_max(high, 20) - ts_min(low, 20) + 1e-8)\",\n",
    "    \"n_income\": \"cs_rank(n_income)\",\n",
    "}\n",
    "\n",
    "# Label 因子定义（不参与训练，用于计算目标）\n",
    "LABEL_FACTOR = {\n",
    "    \"return_5\": \"(ts_delay(close, -5) / close) - 1\",  # 未来5日收益率\n",
    "}"
   ],
   "outputs": [],
   "execution_count": 8
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 训练参数配置"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:27:36.035196Z",
     "start_time": "2026-03-07T16:27:36.031430Z"
    }
   },
   "source": [
    "# 日期范围配置\n",
    "TRAIN_START = \"20200101\"\n",
    "TRAIN_END = \"20241231\"\n",
    "TEST_START = \"20250101\"\n",
    "TEST_END = \"20251231\"\n",
    "\n",
    "# 模型参数配置\n",
    "MODEL_PARAMS = {\n",
    "    \"objective\": \"regression\",\n",
    "    \"metric\": \"mae\",  # 改为 MAE，对异常值更稳健\n",
    "    # 树结构控制（防过拟合核心）\n",
    "    \"num_leaves\": 20,  # 从31降为20，降低模型复杂度\n",
    "    \"max_depth\": 4,  # 显式限制深度，防止过度拟合噪声\n",
    "    \"min_child_samples\": 50,  # 叶子最小样本数，防止学习极端样本\n",
    "    \"min_child_weight\": 0.001,\n",
    "    # 学习参数\n",
    "    \"learning_rate\": 0.01,  # 降低学习率，配合更多树\n",
    "    \"n_estimators\": 1000,  # 增加树数量，配合早停\n",
    "    # 采样策略（关键防过拟合）\n",
    "    \"subsample\": 0.8,  # 每棵树随机采样80%数据（行采样）\n",
    "    \"subsample_freq\": 5,  # 每5轮迭代进行一次 subsample\n",
    "    \"colsample_bytree\": 0.8,  # 每棵树随机选择80%特征（列采样）\n",
    "    # 正则化\n",
    "    \"reg_alpha\": 0.1,  # L1正则，增加稀疏性\n",
    "    \"reg_lambda\": 1.0,  # L2正则，平滑权重\n",
    "    # 数值稳定性\n",
    "    \"verbose\": -1,\n",
    "    \"random_state\": 42,\n",
    "}\n",
    "\n",
    "# 数据处理器配置\n",
    "PROCESSOR_CONFIGS = [\n",
    "    {\"name\": \"winsorizer\", \"params\": {\"lower\": 0.01, \"upper\": 0.99}},\n",
    "    {\"name\": \"cs_standard_scaler\", \"params\": {}},\n",
    "]\n",
    "\n",
    "# 股票池筛选配置\n",
    "STOCK_FILTER_CONFIG = {\n",
    "    \"exclude_cyb\": True,  # 排除创业板\n",
    "    \"exclude_kcb\": True,  # 排除科创板\n",
    "    \"exclude_bj\": True,  # 排除北交所\n",
    "    \"exclude_st\": True,  # 排除ST股票\n",
    "}\n",
    "\n",
    "# 输出配置（相对于本文件所在目录）\n",
    "OUTPUT_DIR = \"output\"\n",
    "SAVE_PREDICTIONS = True\n",
    "PERSIST_MODEL = False"
   ],
   "outputs": [],
   "execution_count": 9
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. 训练流程\n",
    "\n",
    "### 4.1 初始化组件"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:27:41.855565Z",
     "start_time": "2026-03-07T16:27:36.040416Z"
    }
   },
   "source": [
    "print(\"\\n\" + \"=\" * 80)\n",
    "print(\"LightGBM 回归模型训练\")\n",
    "print(\"=\" * 80)\n",
    "\n",
    "# 1. 创建 FactorEngine\n",
    "print(\"\\n[1] 创建 FactorEngine\")\n",
    "engine = FactorEngine()\n",
    "\n",
    "# 2. 使用字符串表达式定义因子\n",
    "print(\"\\n[2] 定义因子（字符串表达式）\")\n",
    "feature_cols = create_factors_with_strings(engine, FACTOR_DEFINITIONS, LABEL_FACTOR)\n",
    "target_col = \"return_5\"\n",
    "\n",
    "# 3. 准备数据（使用模块级别的日期配置）\n",
    "print(\"\\n[3] 准备数据\")\n",
    "\n",
    "data = prepare_data(\n",
    "    engine=engine,\n",
    "    feature_cols=feature_cols,\n",
    "    start_date=TRAIN_START,\n",
    "    end_date=TEST_END,\n",
    ")\n",
    "\n",
    "# 4. 打印配置信息\n",
    "print(f\"\\n[配置] 训练期: {TRAIN_START} - {TRAIN_END}\")\n",
    "print(f\"[配置] 测试期: {TEST_START} - {TEST_END}\")\n",
    "print(f\"[配置] 特征数: {len(feature_cols)}\")\n",
    "print(f\"[配置] 目标变量: {target_col}\")\n",
    "\n",
    "# 5. 创建模型\n",
    "model = LightGBMModel(params=MODEL_PARAMS)\n",
    "\n",
    "# 6. 创建数据处理器\n",
    "processors = [\n",
    "    NullFiller(strategy=\"mean\"),\n",
    "    Winsorizer(**PROCESSOR_CONFIGS[0][\"params\"]),\n",
    "    StandardScaler(exclude_cols=[\"ts_code\", \"trade_date\", target_col]),\n",
    "]\n",
    "\n",
    "# 7. 创建数据划分器\n",
    "splitter = DateSplitter(\n",
    "    train_start=TRAIN_START,\n",
    "    train_end=TRAIN_END,\n",
    "    test_start=TEST_START,\n",
    "    test_end=TEST_END,\n",
    ")\n",
    "\n",
    "# 8. 创建股票池管理器\n",
    "pool_manager = StockPoolManager(\n",
    "    filter_config=StockFilterConfig(**STOCK_FILTER_CONFIG),\n",
    "    selector_config=None,  # 暂时不启用市值选择\n",
    "    data_router=engine.router,\n",
    ")\n",
    "\n",
    "# 9. 创建 ST 股票过滤器\n",
    "st_filter = STFilter(\n",
    "    data_router=engine.router,\n",
    ")\n",
    "\n",
    "# 10. 创建训练器\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    pool_manager=pool_manager,\n",
    "    processors=processors,\n",
    "    filters=[st_filter],\n",
    "    splitter=splitter,\n",
    "    target_col=target_col,\n",
    "    feature_cols=feature_cols,\n",
    "    persist_model=PERSIST_MODEL,\n",
    ")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "================================================================================\n",
      "LightGBM 回归模型训练\n",
      "================================================================================\n",
      "\n",
      "[1] 创建 FactorEngine\n",
      "\n",
      "[2] 定义因子（字符串表达式）\n",
      "================================================================================\n",
      "使用字符串表达式定义因子\n",
      "================================================================================\n",
      "\n",
      "注册特征因子:\n",
      "  - ma5: ts_mean(close, 5)\n",
      "  - ma10: ts_mean(close, 10)\n",
      "  - ma20: ts_mean(close, 20)\n",
      "  - ma_ratio: ts_mean(close, 5) / ts_mean(close, 20) - 1\n",
      "  - volatility_5: ts_std(close, 5)\n",
      "  - volatility_20: ts_std(close, 20)\n",
      "  - vol_ratio: ts_mean(vol, 5) / (ts_mean(vol, 20) + 1e-8)\n",
      "  - return_10: (close / ts_delay(close, 10)) - 1\n",
      "  - return_20: (close / ts_delay(close, 20)) - 1\n",
      "  - return_diff: (close / ts_delay(close, 5)) - 1 - ((close / ts_delay(close, 10)) - 1)\n",
      "  - vol_ma5: ts_mean(vol, 5)\n",
      "  - vol_ma20: ts_mean(vol, 20)\n",
      "  - market_cap_rank: cs_rank(total_mv)\n",
      "  - high_low_ratio: (close - ts_min(low, 20)) / (ts_max(high, 20) - ts_min(low, 20) + 1e-8)\n",
      "  - n_income: cs_rank(n_income)\n",
      "\n",
      "注册 Label 因子:\n",
      "  - return_5: (ts_delay(close, -5) / close) - 1\n",
      "\n",
      "特征因子数: 15\n",
      "Label: return_5\n",
      "已注册因子总数: 16\n",
      "\n",
      "[3] 准备数据\n",
      "\n",
      "================================================================================\n",
      "准备数据\n",
      "================================================================================\n",
      "\n",
      "计算因子: 20200101 - 20251231\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\PyProject\\ProStock\\src\\data\\financial_loader.py:123: UserWarning: Sortedness of columns cannot be checked when 'by' groups provided\n",
      "  merged = df_price.join_asof(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "数据形状: (7044952, 24)\n",
      "数据列: ['ts_code', 'trade_date', 'low', 'close', 'vol', 'high', 'total_mv', 'f_ann_date', 'n_income', 'ma5', 'ma10', 'ma20', 'ma_ratio', 'volatility_5', 'volatility_20', 'vol_ratio', 'return_10', 'return_20', 'return_diff', 'vol_ma5', 'vol_ma20', 'market_cap_rank', 'high_low_ratio', 'return_5']\n",
      "\n",
      "前5行预览:\n",
      "shape: (5, 24)\n",
      "┌───────────┬────────────┬─────────┬─────────┬───┬──────────┬─────────────┬────────────┬───────────┐\n",
      "│ ts_code   ┆ trade_date ┆ low     ┆ close   ┆ … ┆ vol_ma20 ┆ market_cap_ ┆ high_low_r ┆ return_5  │\n",
      "│ ---       ┆ ---        ┆ ---     ┆ ---     ┆   ┆ ---      ┆ rank        ┆ atio       ┆ ---       │\n",
      "│ str       ┆ str        ┆ f64     ┆ f64     ┆   ┆ f64      ┆ ---         ┆ ---        ┆ f64       │\n",
      "│           ┆            ┆         ┆         ┆   ┆          ┆ f64         ┆ f64        ┆           │\n",
      "╞═══════════╪════════════╪═════════╪═════════╪═══╪══════════╪═════════════╪════════════╪═══════════╡\n",
      "│ 000001.SZ ┆ 20200102   ┆ 1806.75 ┆ 1841.69 ┆ … ┆ null     ┆ 0.993583    ┆ null       ┆ -0.004746 │\n",
      "│ 000001.SZ ┆ 20200103   ┆ 1847.15 ┆ 1875.53 ┆ … ┆ null     ┆ 0.993585    ┆ null       ┆ -0.02852  │\n",
      "│ 000001.SZ ┆ 20200106   ┆ 1846.05 ┆ 1863.52 ┆ … ┆ null     ┆ 0.993588    ┆ null       ┆ -0.004685 │\n",
      "│ 000001.SZ ┆ 20200107   ┆ 1850.42 ┆ 1872.26 ┆ … ┆ null     ┆ 0.993588    ┆ null       ┆ -0.022743 │\n",
      "│ 000001.SZ ┆ 20200108   ┆ 1815.49 ┆ 1818.76 ┆ … ┆ null     ┆ 0.993586    ┆ null       ┆ -0.008401 │\n",
      "└───────────┴────────────┴─────────┴─────────┴───┴──────────┴─────────────┴────────────┴───────────┘\n",
      "\n",
      "[配置] 训练期: 20200101 - 20241231\n",
      "[配置] 测试期: 20250101 - 20251231\n",
      "[配置] 特征数: 15\n",
      "[配置] 目标变量: return_5\n"
     ]
    }
   ],
   "execution_count": 10
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 执行训练"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:27:48.155192Z",
     "start_time": "2026-03-07T16:27:41.865012Z"
    }
   },
   "source": [
    "print(\"\\n\" + \"=\" * 80)\n",
    "print(\"开始训练\")\n",
    "print(\"=\" * 80)\n",
    "\n",
    "# 步骤 1: 股票池筛选\n",
    "print(\"\\n[步骤 1/6] 股票池筛选\")\n",
    "print(\"-\" * 60)\n",
    "if pool_manager:\n",
    "    print(\"  执行每日独立筛选股票池...\")\n",
    "    filtered_data = pool_manager.filter_and_select_daily(data)\n",
    "    print(f\"  筛选前数据规模: {data.shape}\")\n",
    "    print(f\"  筛选后数据规模: {filtered_data.shape}\")\n",
    "    print(f\"  筛选前股票数: {data['ts_code'].n_unique()}\")\n",
    "    print(f\"  筛选后股票数: {filtered_data['ts_code'].n_unique()}\")\n",
    "    print(f\"  删除记录数: {len(data) - len(filtered_data)}\")\n",
    "else:\n",
    "    filtered_data = data\n",
    "    print(\"  未配置股票池管理器，跳过筛选\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "================================================================================\n",
      "开始训练\n",
      "================================================================================\n",
      "\n",
      "[步骤 1/6] 股票池筛选\n",
      "------------------------------------------------------------\n",
      "  执行每日独立筛选股票池...\n",
      "  筛选前数据规模: (7044952, 24)\n",
      "  筛选后数据规模: (4532198, 24)\n",
      "  筛选前股票数: 5678\n",
      "  筛选后股票数: 3359\n",
      "  删除记录数: 2512754\n"
     ]
    }
   ],
   "execution_count": 11
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:27:48.744530Z",
     "start_time": "2026-03-07T16:27:48.163043Z"
    }
   },
   "source": [
    "# 步骤 2: 划分训练/测试集\n",
    "print(\"\\n[步骤 2/6] 划分训练集和测试集\")\n",
    "print(\"-\" * 60)\n",
    "if splitter:\n",
    "    train_data, test_data = splitter.split(filtered_data)\n",
    "    print(f\"  训练集数据规模: {train_data.shape}\")\n",
    "    print(f\"  测试集数据规模: {test_data.shape}\")\n",
    "    print(f\"  训练集股票数: {train_data['ts_code'].n_unique()}\")\n",
    "    print(f\"  测试集股票数: {test_data['ts_code'].n_unique()}\")\n",
    "    print(\n",
    "        f\"  训练集日期范围: {train_data['trade_date'].min()} - {train_data['trade_date'].max()}\"\n",
    "    )\n",
    "    print(\n",
    "        f\"  测试集日期范围: {test_data['trade_date'].min()} - {test_data['trade_date'].max()}\"\n",
    "    )\n",
    "\n",
    "    print(\"\\n  训练集前5行预览:\")\n",
    "    print(train_data.head())\n",
    "    print(\"\\n  测试集前5行预览:\")\n",
    "    print(test_data.head())\n",
    "else:\n",
    "    train_data = filtered_data\n",
    "    test_data = filtered_data\n",
    "    print(\"  未配置划分器，全部作为训练集\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[步骤 2/6] 划分训练集和测试集\n",
      "------------------------------------------------------------\n",
      "  训练集数据规模: (3760991, 24)\n",
      "  测试集数据规模: (771207, 24)\n",
      "  训练集股票数: 3321\n",
      "  测试集股票数: 3215\n",
      "  训练集日期范围: 20200102 - 20241231\n",
      "  测试集日期范围: 20250102 - 20251231\n",
      "\n",
      "  训练集前5行预览:\n",
      "shape: (5, 24)\n",
      "┌───────────┬────────────┬─────────┬─────────┬───┬──────────┬─────────────┬────────────┬───────────┐\n",
      "│ ts_code   ┆ trade_date ┆ low     ┆ close   ┆ … ┆ vol_ma20 ┆ market_cap_ ┆ high_low_r ┆ return_5  │\n",
      "│ ---       ┆ ---        ┆ ---     ┆ ---     ┆   ┆ ---      ┆ rank        ┆ atio       ┆ ---       │\n",
      "│ str       ┆ str        ┆ f64     ┆ f64     ┆   ┆ f64      ┆ ---         ┆ ---        ┆ f64       │\n",
      "│           ┆            ┆         ┆         ┆   ┆          ┆ f64         ┆ f64        ┆           │\n",
      "╞═══════════╪════════════╪═════════╪═════════╪═══╪══════════╪═════════════╪════════════╪═══════════╡\n",
      "│ 000001.SZ ┆ 20200102   ┆ 1806.75 ┆ 1841.69 ┆ … ┆ null     ┆ 0.993583    ┆ null       ┆ -0.004746 │\n",
      "│ 000002.SZ ┆ 20200102   ┆ 4824.87 ┆ 4832.29 ┆ … ┆ null     ┆ 0.99492     ┆ null       ┆ -0.011057 │\n",
      "│ 000004.SZ ┆ 20200102   ┆ 90.1    ┆ 90.75   ┆ … ┆ null     ┆ 0.057219    ┆ null       ┆ -0.000441 │\n",
      "│ 000005.SZ ┆ 20200102   ┆ 28.82   ┆ 29.1    ┆ … ┆ null     ┆ 0.28984     ┆ null       ┆ 0.022337  │\n",
      "│ 000006.SZ ┆ 20200102   ┆ 190.24  ┆ 191.3   ┆ … ┆ null     ┆ 0.631551    ┆ null       ┆ 0.012964  │\n",
      "└───────────┴────────────┴─────────┴─────────┴───┴──────────┴─────────────┴────────────┴───────────┘\n",
      "\n",
      "  测试集前5行预览:\n",
      "shape: (5, 24)\n",
      "┌───────────┬────────────┬─────────┬─────────┬───┬────────────┬────────────┬───────────┬───────────┐\n",
      "│ ts_code   ┆ trade_date ┆ low     ┆ close   ┆ … ┆ vol_ma20   ┆ market_cap ┆ high_low_ ┆ return_5  │\n",
      "│ ---       ┆ ---        ┆ ---     ┆ ---     ┆   ┆ ---        ┆ _rank      ┆ ratio     ┆ ---       │\n",
      "│ str       ┆ str        ┆ f64     ┆ f64     ┆   ┆ f64        ┆ ---        ┆ ---       ┆ f64       │\n",
      "│           ┆            ┆         ┆         ┆   ┆            ┆ f64        ┆ f64       ┆           │\n",
      "╞═══════════╪════════════╪═════════╪═════════╪═══╪════════════╪════════════╪═══════════╪═══════════╡\n",
      "│ 000001.SZ ┆ 20250102   ┆ 1455.46 ┆ 1460.57 ┆ … ┆ 1.2151e6   ┆ 0.991617   ┆ 0.063478  ┆ -0.002622 │\n",
      "│ 000002.SZ ┆ 20250102   ┆ 1284.65 ┆ 1291.92 ┆ … ┆ 1.2839e6   ┆ 0.970007   ┆ 0.020839  ┆ -0.022509 │\n",
      "│ 000004.SZ ┆ 20250102   ┆ 54.17   ┆ 57.63   ┆ … ┆ 159807.918 ┆ 0.099106   ┆ 0.131858  ┆ -0.064897 │\n",
      "│ 000006.SZ ┆ 20250102   ┆ 285.33  ┆ 288.12  ┆ … ┆ 404264.211 ┆ 0.72392    ┆ 0.028423  ┆ -0.048278 │\n",
      "│ 000007.SZ ┆ 20250102   ┆ 57.49   ┆ 58.15   ┆ … ┆ 88380.2845 ┆ 0.183495   ┆ 0.391829  ┆ 0.015649  │\n",
      "└───────────┴────────────┴─────────┴─────────┴───┴────────────┴────────────┴───────────┴───────────┘\n"
     ]
    }
   ],
   "execution_count": 12
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:27:49.562499Z",
     "start_time": "2026-03-07T16:27:48.753734Z"
    }
   },
   "source": [
    "# 步骤 3: 训练集数据处理\n",
    "print(\"\\n[步骤 3/6] 训练集数据处理\")\n",
    "print(\"-\" * 60)\n",
    "fitted_processors = []\n",
    "if processors:\n",
    "    for i, processor in enumerate(processors, 1):\n",
    "        print(\n",
    "            f\"  [{i}/{len(processors)}] 应用处理器: {processor.__class__.__name__}\"\n",
    "        )\n",
    "        train_data_before = len(train_data)\n",
    "        train_data = processor.fit_transform(train_data)\n",
    "        train_data_after = len(train_data)\n",
    "        fitted_processors.append(processor)\n",
    "        print(f\"      处理前记录数: {train_data_before}\")\n",
    "        print(f\"      处理后记录数: {train_data_after}\")\n",
    "        if train_data_before != train_data_after:\n",
    "            print(f\"      删除记录数: {train_data_before - train_data_after}\")\n",
    "\n",
    "print(\"\\n  训练集处理后前5行预览:\")\n",
    "print(train_data.head())\n",
    "print(f\"\\n  训练集特征统计:\")\n",
    "print(f\"    特征数: {len(feature_cols)}\")\n",
    "print(f\"    样本数: {len(train_data)}\")\n",
    "print(f\"    缺失值统计:\")\n",
    "for col in feature_cols[:5]:  # 只显示前5个特征的缺失值\n",
    "    null_count = train_data[col].null_count()\n",
    "    if null_count > 0:\n",
    "        print(\n",
    "            f\"      {col}: {null_count} ({null_count / len(train_data) * 100:.2f}%)\"\n",
    "        )"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[步骤 3/6] 训练集数据处理\n",
      "------------------------------------------------------------\n",
      "  [1/3] 应用处理器: NullFiller\n",
      "      处理前记录数: 3760991\n",
      "      处理后记录数: 3760991\n",
      "  [2/3] 应用处理器: Winsorizer\n",
      "      处理前记录数: 3760991\n",
      "      处理后记录数: 3760991\n",
      "  [3/3] 应用处理器: StandardScaler\n",
      "      处理前记录数: 3760991\n",
      "      处理后记录数: 3760991\n",
      "\n",
      "  训练集处理后前5行预览:\n",
      "shape: (5, 24)\n",
      "┌───────────┬────────────┬───────────┬──────────┬───┬──────────┬───────────┬───────────┬───────────┐\n",
      "│ ts_code   ┆ trade_date ┆ low       ┆ close    ┆ … ┆ vol_ma20 ┆ market_ca ┆ high_low_ ┆ return_5  │\n",
      "│ ---       ┆ ---        ┆ ---       ┆ ---      ┆   ┆ ---      ┆ p_rank    ┆ ratio     ┆ ---       │\n",
      "│ str       ┆ str        ┆ f64       ┆ f64      ┆   ┆ f64      ┆ ---       ┆ ---       ┆ f64       │\n",
      "│           ┆            ┆           ┆          ┆   ┆          ┆ f64       ┆ f64       ┆           │\n",
      "╞═══════════╪════════════╪═══════════╪══════════╪═══╪══════════╪═══════════╪═══════════╪═══════════╡\n",
      "│ 000001.SZ ┆ 20200102   ┆ 7.215424  ┆ 7.219817 ┆ … ┆ null     ┆ 1.579213  ┆ null      ┆ -0.004746 │\n",
      "│ 000002.SZ ┆ 20200102   ┆ 7.215424  ┆ 7.219817 ┆ … ┆ null     ┆ 1.579213  ┆ null      ┆ -0.011057 │\n",
      "│ 000004.SZ ┆ 20200102   ┆ 0.118662  ┆ 0.111432 ┆ … ┆ null     ┆ -1.671014 ┆ null      ┆ -0.000441 │\n",
      "│ 000005.SZ ┆ 20200102   ┆ -0.292838 ┆ -0.29447 ┆ … ┆ null     ┆ -0.862878 ┆ null      ┆ 0.022337  │\n",
      "│ 000006.SZ ┆ 20200102   ┆ 0.791109  ┆ 0.77345  ┆ … ┆ null     ┆ 0.324248  ┆ null      ┆ 0.012964  │\n",
      "└───────────┴────────────┴───────────┴──────────┴───┴──────────┴───────────┴───────────┴───────────┘\n",
      "\n",
      "  训练集特征统计:\n",
      "    特征数: 15\n",
      "    样本数: 3760991\n",
      "    缺失值统计:\n",
      "      ma5: 11541 (0.31%)\n",
      "      ma10: 25950 (0.69%)\n",
      "      ma20: 54850 (1.46%)\n",
      "      ma_ratio: 54850 (1.46%)\n",
      "      volatility_5: 11541 (0.31%)\n"
     ]
    }
   ],
   "execution_count": 13
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:28:09.621524Z",
     "start_time": "2026-03-07T16:27:49.566330Z"
    }
   },
   "source": [
    "# 步骤 4: 训练模型\n",
    "print(\"\\n[步骤 4/6] 训练模型\")\n",
    "print(\"-\" * 60)\n",
    "print(f\"  模型类型: LightGBM\")\n",
    "print(f\"  训练样本数: {len(train_data)}\")\n",
    "print(f\"  特征数: {len(feature_cols)}\")\n",
    "print(f\"  目标变量: {target_col}\")\n",
    "\n",
    "X_train = train_data.select(feature_cols)\n",
    "y_train = train_data.select(target_col).to_series()\n",
    "\n",
    "print(f\"\\n  目标变量统计:\")\n",
    "print(f\"    均值: {y_train.mean():.6f}\")\n",
    "print(f\"    标准差: {y_train.std():.6f}\")\n",
    "print(f\"    最小值: {y_train.min():.6f}\")\n",
    "print(f\"    最大值: {y_train.max():.6f}\")\n",
    "print(f\"    缺失值: {y_train.null_count()}\")\n",
    "\n",
    "print(\"\\n  开始训练...\")\n",
    "model.fit(X_train, y_train)\n",
    "print(\"  训练完成！\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[步骤 4/6] 训练模型\n",
      "------------------------------------------------------------\n",
      "  模型类型: LightGBM\n",
      "  训练样本数: 3760991\n",
      "  特征数: 15\n",
      "  目标变量: return_5\n",
      "\n",
      "  目标变量统计:\n",
      "    均值: 0.001511\n",
      "    标准差: 0.062515\n",
      "    最小值: -0.165141\n",
      "    最大值: 0.225065\n",
      "    缺失值: 0\n",
      "\n",
      "  开始训练...\n",
      "  训练完成！\n"
     ]
    }
   ],
   "execution_count": 14
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:28:09.687281Z",
     "start_time": "2026-03-07T16:28:09.625395Z"
    }
   },
   "source": [
    "# 步骤 5: 测试集数据处理\n",
    "print(\"\\n[步骤 5/6] 测试集数据处理\")\n",
    "print(\"-\" * 60)\n",
    "if processors and test_data is not train_data:\n",
    "    for i, processor in enumerate(fitted_processors, 1):\n",
    "        print(\n",
    "            f\"  [{i}/{len(fitted_processors)}] 应用处理器: {processor.__class__.__name__}\"\n",
    "        )\n",
    "        test_data_before = len(test_data)\n",
    "        test_data = processor.transform(test_data)\n",
    "        test_data_after = len(test_data)\n",
    "        print(f\"      处理前记录数: {test_data_before}\")\n",
    "        print(f\"      处理后记录数: {test_data_after}\")\n",
    "else:\n",
    "    print(\"  跳过测试集处理\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[步骤 5/6] 测试集数据处理\n",
      "------------------------------------------------------------\n",
      "  [1/3] 应用处理器: NullFiller\n",
      "      处理前记录数: 771207\n",
      "      处理后记录数: 771207\n",
      "  [2/3] 应用处理器: Winsorizer\n",
      "      处理前记录数: 771207\n",
      "      处理后记录数: 771207\n",
      "  [3/3] 应用处理器: StandardScaler\n",
      "      处理前记录数: 771207\n",
      "      处理后记录数: 771207\n"
     ]
    }
   ],
   "execution_count": 15
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:28:10.718020Z",
     "start_time": "2026-03-07T16:28:09.691077Z"
    }
   },
   "source": [
    "# 步骤 6: 生成预测\n",
    "print(\"\\n[步骤 6/6] 生成预测\")\n",
    "print(\"-\" * 60)\n",
    "X_test = test_data.select(feature_cols)\n",
    "print(f\"  测试样本数: {len(X_test)}\")\n",
    "print(\"  预测中...\")\n",
    "predictions = model.predict(X_test)\n",
    "print(f\"  预测完成！\")\n",
    "\n",
    "print(f\"\\n  预测结果统计:\")\n",
    "print(f\"    均值: {predictions.mean():.6f}\")\n",
    "print(f\"    标准差: {predictions.std():.6f}\")\n",
    "print(f\"    最小值: {predictions.min():.6f}\")\n",
    "print(f\"    最大值: {predictions.max():.6f}\")\n",
    "\n",
    "# 保存结果到 trainer\n",
    "trainer.results = test_data.with_columns([pl.Series(\"prediction\", predictions)])"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[步骤 6/6] 生成预测\n",
      "------------------------------------------------------------\n",
      "  测试样本数: 771207\n",
      "  预测中...\n",
      "  预测完成！\n",
      "\n",
      "  预测结果统计:\n",
      "    均值: 0.000635\n",
      "    标准差: 0.006529\n",
      "    最小值: -0.125836\n",
      "    最大值: 0.149914\n"
     ]
    }
   ],
   "execution_count": 16
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.3 查看结果"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:28:10.746425Z",
     "start_time": "2026-03-07T16:28:10.723160Z"
    }
   },
   "source": [
    "print(\"\\n\" + \"=\" * 80)\n",
    "print(\"训练结果\")\n",
    "print(\"=\" * 80)\n",
    "\n",
    "results = trainer.results\n",
    "\n",
    "print(f\"\\n结果数据形状: {results.shape}\")\n",
    "print(f\"结果列: {results.columns}\")\n",
    "print(f\"\\n结果前10行预览:\")\n",
    "print(results.head(10))\n",
    "print(f\"\\n结果后5行预览:\")\n",
    "print(results.tail())\n",
    "\n",
    "print(f\"\\n每日预测样本数统计:\")\n",
    "daily_counts = results.group_by(\"trade_date\").agg(pl.len()).sort(\"trade_date\")\n",
    "print(f\"  最小: {daily_counts['len'].min()}\")\n",
    "print(f\"  最大: {daily_counts['len'].max()}\")\n",
    "print(f\"  平均: {daily_counts['len'].mean():.2f}\")\n",
    "\n",
    "# 展示某一天的前10个预测结果\n",
    "sample_date = results[\"trade_date\"][0]\n",
    "sample_data = results.filter(results[\"trade_date\"] == sample_date).head(10)\n",
    "print(f\"\\n示例日期 {sample_date} 的前10条预测:\")\n",
    "print(sample_data.select([\"ts_code\", \"trade_date\", target_col, \"prediction\"]))"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "================================================================================\n",
      "训练结果\n",
      "================================================================================\n",
      "\n",
      "结果数据形状: (771207, 25)\n",
      "结果列: ['ts_code', 'trade_date', 'low', 'close', 'vol', 'high', 'total_mv', 'f_ann_date', 'n_income', 'ma5', 'ma10', 'ma20', 'ma_ratio', 'volatility_5', 'volatility_20', 'vol_ratio', 'return_10', 'return_20', 'return_diff', 'vol_ma5', 'vol_ma20', 'market_cap_rank', 'high_low_ratio', 'return_5', 'prediction']\n",
      "\n",
      "结果前10行预览:\n",
      "shape: (10, 25)\n",
      "┌───────────┬───────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬──────────┐\n",
      "│ ts_code   ┆ trade_dat ┆ low       ┆ close     ┆ … ┆ market_ca ┆ high_low_ ┆ return_5  ┆ predicti │\n",
      "│ ---       ┆ e         ┆ ---       ┆ ---       ┆   ┆ p_rank    ┆ ratio     ┆ ---       ┆ on       │\n",
      "│ str       ┆ ---       ┆ f64       ┆ f64       ┆   ┆ ---       ┆ ---       ┆ f64       ┆ ---      │\n",
      "│           ┆ str       ┆           ┆           ┆   ┆ f64       ┆ f64       ┆           ┆ f64      │\n",
      "╞═══════════╪═══════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪══════════╡\n",
      "│ 000001.SZ ┆ 20250102  ┆ 7.215424  ┆ 7.219817  ┆ … ┆ 1.575139  ┆ -1.347097 ┆ -0.002622 ┆ -0.00414 │\n",
      "│           ┆           ┆           ┆           ┆   ┆           ┆           ┆           ┆ 1        │\n",
      "│ 000002.SZ ┆ 20250102  ┆ 7.215424  ┆ 7.219817  ┆ … ┆ 1.500066  ┆ -1.494802 ┆ -0.022509 ┆ -0.00644 │\n",
      "│           ┆           ┆           ┆           ┆   ┆           ┆           ┆           ┆ 3        │\n",
      "│ 000004.SZ ┆ 20250102  ┆ -0.122611 ┆ -0.106629 ┆ … ┆ -1.525498 ┆ -1.110228 ┆ -0.064897 ┆ 0.018069 │\n",
      "│ 000006.SZ ┆ 20250102  ┆ 1.429646  ┆ 1.41091   ┆ … ┆ 0.645142  ┆ -1.46853  ┆ -0.048278 ┆ -0.00495 │\n",
      "│           ┆           ┆           ┆           ┆   ┆           ┆           ┆           ┆ 8        │\n",
      "│ 000007.SZ ┆ 20250102  ┆ -0.100317 ┆ -0.103205 ┆ … ┆ -1.232326 ┆ -0.209675 ┆ 0.015649  ┆ -0.00123 │\n",
      "│           ┆           ┆           ┆           ┆   ┆           ┆           ┆           ┆ 1        │\n",
      "│ 000008.SZ ┆ 20250102  ┆ -0.063519 ┆ -0.067059 ┆ … ┆ 0.410216  ┆ -1.48541  ┆ -0.066939 ┆ -0.01023 │\n",
      "│           ┆           ┆           ┆           ┆   ┆           ┆           ┆           ┆ 3        │\n",
      "│ 000009.SZ ┆ 20250102  ┆ 0.045333  ┆ 0.039996  ┆ … ┆ 1.197186  ┆ -1.358101 ┆ -0.036045 ┆ 0.002236 │\n",
      "│ 000010.SZ ┆ 20250102  ┆ -0.293174 ┆ -0.293811 ┆ … ┆ -0.872494 ┆ -1.382331 ┆ 0.092123  ┆ -0.00730 │\n",
      "│           ┆           ┆           ┆           ┆   ┆           ┆           ┆           ┆ 3        │\n",
      "│ 000011.SZ ┆ 20250102  ┆ -0.233141 ┆ -0.235741 ┆ … ┆ -0.098467 ┆ -1.409992 ┆ -0.022094 ┆ -0.00176 │\n",
      "│           ┆           ┆           ┆           ┆   ┆           ┆           ┆           ┆ 3        │\n",
      "│ 000012.SZ ┆ 20250102  ┆ 0.54057   ┆ 0.526749  ┆ … ┆ 1.006268  ┆ -1.411495 ┆ -0.029188 ┆ 0.000613 │\n",
      "└───────────┴───────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴──────────┘\n",
      "\n",
      "结果后5行预览:\n",
      "shape: (5, 25)\n",
      "┌───────────┬───────────┬───────────┬───────────┬───┬───────────┬───────────┬──────────┬───────────┐\n",
      "│ ts_code   ┆ trade_dat ┆ low       ┆ close     ┆ … ┆ market_ca ┆ high_low_ ┆ return_5 ┆ predictio │\n",
      "│ ---       ┆ e         ┆ ---       ┆ ---       ┆   ┆ p_rank    ┆ ratio     ┆ ---      ┆ n         │\n",
      "│ str       ┆ ---       ┆ f64       ┆ f64       ┆   ┆ ---       ┆ ---       ┆ f64      ┆ ---       │\n",
      "│           ┆ str       ┆           ┆           ┆   ┆ f64       ┆ f64       ┆          ┆ f64       │\n",
      "╞═══════════╪═══════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪══════════╪═══════════╡\n",
      "│ 605588.SH ┆ 20251231  ┆ -0.150009 ┆ -0.151005 ┆ … ┆ -1.041064 ┆ 0.083361  ┆ null     ┆ 0.002801  │\n",
      "│ 605589.SH ┆ 20251231  ┆ -0.282833 ┆ -0.286372 ┆ … ┆ 1.082326  ┆ 0.586743  ┆ null     ┆ 0.003757  │\n",
      "│ 605598.SH ┆ 20251231  ┆ 0.163451  ┆ 0.166474  ┆ … ┆ 0.865913  ┆ 1.46704   ┆ null     ┆ -0.003742 │\n",
      "│ 605599.SH ┆ 20251231  ┆ -0.360996 ┆ -0.359651 ┆ … ┆ 0.623404  ┆ 1.560273  ┆ null     ┆ 0.004761  │\n",
      "│ 689009.SH ┆ 20251231  ┆ -0.101257 ┆ -0.108012 ┆ … ┆ 1.304467  ┆ -1.339619 ┆ null     ┆ -0.000987 │\n",
      "└───────────┴───────────┴───────────┴───────────┴───┴───────────┴───────────┴──────────┴───────────┘\n",
      "\n",
      "每日预测样本数统计:\n",
      "  最小: 3147\n",
      "  最大: 3186\n",
      "  平均: 3173.69\n",
      "\n",
      "示例日期 20250102 的前10条预测:\n",
      "shape: (10, 4)\n",
      "┌───────────┬────────────┬───────────┬────────────┐\n",
      "│ ts_code   ┆ trade_date ┆ return_5  ┆ prediction │\n",
      "│ ---       ┆ ---        ┆ ---       ┆ ---        │\n",
      "│ str       ┆ str        ┆ f64       ┆ f64        │\n",
      "╞═══════════╪════════════╪═══════════╪════════════╡\n",
      "│ 000001.SZ ┆ 20250102   ┆ -0.002622 ┆ -0.004141  │\n",
      "│ 000002.SZ ┆ 20250102   ┆ -0.022509 ┆ -0.006443  │\n",
      "│ 000004.SZ ┆ 20250102   ┆ -0.064897 ┆ 0.018069   │\n",
      "│ 000006.SZ ┆ 20250102   ┆ -0.048278 ┆ -0.004958  │\n",
      "│ 000007.SZ ┆ 20250102   ┆ 0.015649  ┆ -0.001231  │\n",
      "│ 000008.SZ ┆ 20250102   ┆ -0.066939 ┆ -0.010233  │\n",
      "│ 000009.SZ ┆ 20250102   ┆ -0.036045 ┆ 0.002236   │\n",
      "│ 000010.SZ ┆ 20250102   ┆ 0.092123  ┆ -0.007303  │\n",
      "│ 000011.SZ ┆ 20250102   ┆ -0.022094 ┆ -0.001763  │\n",
      "│ 000012.SZ ┆ 20250102   ┆ -0.029188 ┆ 0.000613   │\n",
      "└───────────┴────────────┴───────────┴────────────┘\n"
     ]
    }
   ],
   "execution_count": 17
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.4 保存结果"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:28:11.060067Z",
     "start_time": "2026-03-07T16:28:10.751743Z"
    }
   },
   "cell_type": "code",
   "source": [
    "print(\"\\n\" + \"=\" * 80)\n",
    "print(\"保存预测结果\")\n",
    "print(\"=\" * 80)\n",
    "\n",
    "# 确保输出目录存在\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "# 生成时间戳\n",
    "start_dt = datetime.strptime(TEST_START, \"%Y%m%d\")\n",
    "end_dt = datetime.strptime(TEST_END, \"%Y%m%d\")\n",
    "date_str = f\"{start_dt.strftime('%Y%m%d')}_{end_dt.strftime('%Y%m%d')}\"\n",
    "\n",
    "# 保存每日 Top5\n",
    "print(\"\\n[1/1] 保存每日 Top5 股票...\")\n",
    "top5_output_path = os.path.join(OUTPUT_DIR, f\"top5_{date_str}.csv\")\n",
    "\n",
    "# 按日期分组，取每日 top5\n",
    "top5_by_date = []\n",
    "unique_dates = results[\"trade_date\"].unique().sort()\n",
    "for date in unique_dates:\n",
    "    day_data = results.filter(results[\"trade_date\"] == date)\n",
    "    # 按 prediction 降序排序，取前5\n",
    "    top5 = day_data.sort(\"prediction\", descending=True).head(5)\n",
    "    top5_by_date.append(top5)\n",
    "\n",
    "# 合并所有日期的 top5\n",
    "top5_results = pl.concat(top5_by_date)\n",
    "\n",
    "# 格式化日期并调整列顺序：日期、分数、股票\n",
    "top5_to_save = top5_results.select(\n",
    "    [\n",
    "        pl.col(\"trade_date\").str.slice(0, 4)\n",
    "        + \"-\"\n",
    "        + pl.col(\"trade_date\").str.slice(4, 2)\n",
    "        + \"-\"\n",
    "        + pl.col(\"trade_date\").str.slice(6, 2).alias(\"date\"),\n",
    "        pl.col(\"prediction\").alias(\"score\"),\n",
    "        pl.col(\"ts_code\"),\n",
    "    ]\n",
    ")\n",
    "# top5_to_save.write_csv(top5_output_path, include_header=True)\n",
    "print(f\"  保存路径: {top5_output_path}\")\n",
    "print(f\"  保存行数: {len(top5_to_save)}（{len(unique_dates)}个交易日 × 每日top5）\")\n",
    "print(f\"\\n  预览（前15行）:\")\n",
    "print(top5_to_save.head(15))"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "================================================================================\n",
      "保存预测结果\n",
      "================================================================================\n",
      "\n",
      "[1/1] 保存每日 Top5 股票...\n",
      "  保存路径: output\\top5_20250101_20251231.csv\n",
      "  保存行数: 1215（243个交易日 × 每日top5）\n",
      "\n",
      "  预览（前15行）:\n",
      "shape: (15, 3)\n",
      "┌────────────┬──────────┬───────────┐\n",
      "│ trade_date ┆ score    ┆ ts_code   │\n",
      "│ ---        ┆ ---      ┆ ---       │\n",
      "│ str        ┆ f64      ┆ str       │\n",
      "╞════════════╪══════════╪═══════════╡\n",
      "│ 2025-01-02 ┆ 0.129296 ┆ 603007.SH │\n",
      "│ 2025-01-02 ┆ 0.087661 ┆ 603559.SH │\n",
      "│ 2025-01-02 ┆ 0.056859 ┆ 000595.SZ │\n",
      "│ 2025-01-02 ┆ 0.052261 ┆ 600811.SH │\n",
      "│ 2025-01-02 ┆ 0.046187 ┆ 603366.SH │\n",
      "│ …          ┆ …        ┆ …         │\n",
      "│ 2025-01-06 ┆ 0.145456 ┆ 603007.SH │\n",
      "│ 2025-01-06 ┆ 0.066654 ┆ 603959.SH │\n",
      "│ 2025-01-06 ┆ 0.06255  ┆ 002691.SZ │\n",
      "│ 2025-01-06 ┆ 0.054795 ┆ 000638.SZ │\n",
      "│ 2025-01-06 ┆ 0.047319 ┆ 002713.SZ │\n",
      "└────────────┴──────────┴───────────┘\n"
     ]
    }
   ],
   "execution_count": 18
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.5 特征重要性"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:28:11.069209Z",
     "start_time": "2026-03-07T16:28:11.065068Z"
    }
   },
   "source": [
    "importance = model.feature_importance()\n",
    "if importance is not None:\n",
    "    print(\"\\n特征重要性:\")\n",
    "    print(importance.sort_values(ascending=False))\n",
    "\n",
    "print(\"\\n\" + \"=\" * 80)\n",
    "print(\"训练完成！\")\n",
    "print(\"=\" * 80)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "特征重要性:\n",
      "return_10          3245.995642\n",
      "ma_ratio           2818.976538\n",
      "vol_ratio          2701.151667\n",
      "return_20          2681.762756\n",
      "high_low_ratio     2469.737874\n",
      "return_diff        1463.681328\n",
      "volatility_5       1357.026476\n",
      "market_cap_rank     722.088673\n",
      "vol_ma5             653.087969\n",
      "vol_ma20            631.252641\n",
      "ma20                548.830554\n",
      "n_income            495.525785\n",
      "ma5                 469.811824\n",
      "ma10                305.156890\n",
      "volatility_20       203.306830\n",
      "dtype: float64\n",
      "\n",
      "================================================================================\n",
      "训练完成！\n",
      "================================================================================\n"
     ]
    }
   ],
   "execution_count": 19
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. 可视化分析\n",
    "\n",
    "使用训练好的模型直接绘图。\n",
    "- **特征重要性图**：辅助特征选择\n",
    "- **决策树图**：理解决策逻辑"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:28:11.377695Z",
     "start_time": "2026-03-07T16:28:11.076235Z"
    }
   },
   "source": [
    "# 导入可视化库\n",
    "import matplotlib.pyplot as plt\n",
    "import lightgbm as lgb\n",
    "import pandas as pd\n",
    "\n",
    "# 从封装的model中取出底层Booster\n",
    "booster = model.model\n",
    "print(f\"模型类型: {type(booster)}\")\n",
    "print(f\"特征数量: {len(feature_cols)}\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "模型类型: <class 'lightgbm.basic.Booster'>\n",
      "特征数量: 15\n"
     ]
    }
   ],
   "execution_count": 20
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.1 绘制特征重要性（辅助特征选择）\n",
    "\n",
    "**解读**：\n",
    "- 重要性高的特征对模型贡献大\n",
    "- 重要性为0的特征可以考虑删除\n",
    "- 可以帮助理解哪些因子最有效"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-03-07T16:28:11.542731Z",
     "start_time": "2026-03-07T16:28:11.384139Z"
    }
   },
   "cell_type": "code",
   "source": [
    "print(\"绘制特征重要性...\")\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 8))\n",
    "lgb.plot_importance(\n",
    "    booster, \n",
    "    max_num_features=20,\n",
    "    importance_type='gain',\n",
    "    title='Feature Importance (Gain)',\n",
    "    ax=ax\n",
    ")\n",
    "ax.set_xlabel('Importance (Gain)')\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# 打印重要性排名\n",
    "importance_gain = pd.Series(\n",
    "    booster.feature_importance(importance_type='gain'),\n",
    "    index=feature_cols\n",
    ").sort_values(ascending=False)\n",
    "\n",
    "print(\"\\n[特征重要性排名 - Gain]\")\n",
    "print(importance_gain)\n",
    "\n",
    "# 识别低重要性特征\n",
    "zero_importance = importance_gain[importance_gain == 0].index.tolist()\n",
    "if zero_importance:\n",
    "    print(f\"\\n[低重要性特征] 以下{len(zero_importance)}个特征重要性为0，可考虑删除:\")\n",
    "    for feat in zero_importance:\n",
    "        print(f\"  - {feat}\")\n",
    "else:\n",
    "    print(\"\\n所有特征都有一定重要性\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "绘制特征重要性...\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 1000x800 with 1 Axes>"
      ],
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAA9wAAAMWCAYAAAADI47PAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAPYQAAD2EBqD+naQAAw6lJREFUeJzs3QmcjeX///EPQ/Y9WxKFIktEFIVIRNKGiqRVkcpSUkiWiIqS+LYq1be9KEpSlkqWQmQpRUm2kiXCjJn/4319//f5nTlzzphhbjNzz+v5eBxmzrnPfe5zrnPO3J/r87muK1dSUlKSAQAAAACADJU7Y3cHAAAAAAAIuAEAAAAA8AkZbgAAAAAAfEDADQAAAACADwi4AQAAAADwAQE3AAAAAAA+IOAGAAAAAMAHBNwAAAAAAPiAgBsAAAAAAB8QcAMAAGSCTZs2Wf78+e2rr746Lo/XvXt3q1y58lHd9/7777dGjRpl+DEBQNARcAMAAmfKlCmWK1euqBcFDn74+uuvbejQobZr1y7Lqq/H0qVLLbt65pln3PMIkmHDhrkgtkmTJiluW7BggXXq1MkqVKhgJ5xwghUrVsxtq/ts27btuB/rPffcYytWrLDp06cf98cGgOwsT2YfAAAAflFwcuqppya7rlatWr4F3A8//LDLIhYvXtyXx8jJFHCfeOKJ7vUNgh07dtjLL7/sLpGGDBliw4cPt9NOO809X/1/4MAB+/bbb+3xxx939/n555/T/ZjPPfecJSYmHtXxlitXzjp06GCPPfaYXXbZZUe1DwDIiQi4AQCBdckll1iDBg0sO9u3b58VKlTIcqr9+/dbwYIFLWheffVVy5Mnj7Vv3z7Z9W+++aYLtpXdnjp1qstuhxs3bpy7HI28efMe0zHrmDp27Gi//PKL6wQAABwZJeUAgBzr448/tgsuuMAFtEWKFLF27drZDz/8kGyb77//PpRl1HhbZfpuuukm++uvv0LbqJT83nvvdT8ro+6Vr2/cuNFd9HO0cmhdr/uG70fXrV692q677jorUaKEnX/++cmCtPr161uBAgWsZMmSds0117hxwEdDz6lw4cL222+/2aWXXup+VvnyxIkT3e0rV660Fi1auNemUqVK9vrrr0ctU58/f7716NHDSpUqZUWLFrVu3brZ33//HTVDXbNmTcuXL5+ddNJJ1qtXrxTl982bN3cVCMrkNm3a1AXaDzzwgBt3rHaZN29e6LXVtrJz507r37+/1a5d2z0HHYM6WlT+HG7u3Lnufm+99ZaNHDnSTj75ZNeeLVu2tPXr16c43kWLFlnbtm1dG+g1qFOnjj355JPJtlm7dq1dffXVri20L3XupLXk+oMPPnAl4jrmyOy2MvkvvPBCimBbVFoe/p6RadOmufeuXle9vlWqVHFB++HDh1Mdw+29N5W1fvbZZ939dP9zzjnHlixZkuKxL7rootDjAQDShgw3ACCwdu/ebX/++Wey6xTMiLKHN9xwg7Vu3doeffRRl0mdNGmSC3CXLVsWCkxmz57tMno33nijC7YV+Ck40f/ffPONC1iuvPJK+/HHH+2///2vyz56j1G6dGlXOpxeyiJWq1bNHnnkEUtKSnLXKUgcPHiwyzLecsstbr8TJkxwgamO92jK2BWQKTjVPsaMGWOvvfaa3XnnnS7AfPDBB61Lly7uuU2ePNkF0uedd16KEn1tr8dWELhu3Tr3Gv7666+hAFd0m8rtFbDdcccdoe0U1GnCsPDMqzoydEzqTOjatauVLVvWBde9e/d2wamOS3S9qG0UvOo107FpfPN//vMfa9asmeu4UBAabvTo0ZY7d24XpOv9oeet56kA26M2VydE+fLl7e6773btvmbNGvvoo4/c76L219hrdVJoXgC9ZgrmL7/8cnv33XftiiuuiPm6x8fHu+eu1yKc3kO6qH0jA/HUqPND2/ft29f9//nnn7vAfc+ePTZ27Ngj3l+dKXv37nUdJ2ozvSZqd7224W2jYF9BudqsT58+aT4+AMjRkgAACJiXXnpJUWrUi+zduzepePHiSbfeemuy+23dujWpWLFiya7fv39/iv3/97//dfuaP39+6LqxY8e66zZs2JBsW/2u63VMkXT9Qw89FPpdP+u6a6+9Ntl2GzduTIqLi0saOXJksutXrlyZlCdPnhTXx3o9lixZErruhhtucNc98sgjoev+/vvvpAIFCiTlypUr6Y033ghdv3bt2hTH6u2zfv36SYcOHQpdP2bMGHf9tGnT3O/bt29POuGEE5IuvvjipMOHD4e2e/rpp912L774Yui6Zs2auesmT56c4jnUrFnT3R7pwIEDyfbrveb58uVLGjZsWOi6L774wu27Ro0aSQcPHgxd/+STT7rr9VpKQkJC0qmnnppUqVIl93qES0xMDP3csmXLpNq1a7vHD7+9cePGSdWqVUtKzfr1691jTpgwIdn1es10/fjx41M87o4dO5Jd4uPjU32P9ujRI6lgwYLJjk9trucV/jrp8UqVKpW0c+fOFMfx4Ycfptiv2lGvIQAgbSgpBwAElsqjla0Mv4j+Vznztdde6zLg3iUuLs6V+X7xxRehfah826OJq7Tdueee637/7rvvfDnu22+/Pdnv7733npvsStnt8ONV5lWZ8PDjTS9lUz3KVJ9xxhkuW6vH8ug63aaMZ6TbbrstWRZUWVuNTZ45c6b7/bPPPrNDhw65Wa6VWfbceuutrvx7xowZyfankmZVE6SVtvf2q4y9MuTK8uqYo7WP9h1eqq0hBeI9N1ULbNiwwR1vZNWAl7FXGbuyyHqNlBn22kOPrYqJn376yTZv3hzzmL3hCCpXD6eMtERmt5WJV7VE+GX58uVR36Pe8eh5qWpDZe9H0rlz52THEvmahNN2kVUjAIDYKCkHAARWw4YNo06apoBINEY5GgWCHgVXKod+4403bPv27SkCIT9Elm3reJUQV3CdkZNhadyxgrdwKhvW+GYvuAy/PtrY7MhjUrCoUmyNDxaVl4sC4HAKejUu3rvd4y2DlVbqiNDYao0RV6AcPm5Z48ojnXLKKcl+9wJN77l5s3+nNpu9xnyrPVTir0s0eq/ouaTGGy7g0TwC8s8//6R4Tb3Ook8//TRFmbjK2wcNGuQ6AbygPT3v0SO9JpHHHPneAADERsANAMhxvKWRNI5bWeJIytB6lMXUkl+aFK1u3bou+NH927Rpk6YllmIFJ5ETWoULz1h6x6v9aJI3ZeEjpWe8b7ho+0rt+sgA0Q+Rz/1INM5dQa8mstNEYZrATBlvZaijtU9GPDdvvxoHrox2NFWrVo15f68jIDKgrV69uvt/1apVKd6P3oRlv//+e7LbVKmh8erqJNIyeBpjrY4UZfcHDBiQpvdoel4THbM3RwEA4MgIuAEAOY6CEilTpkwokIlGwcWcOXNchluTUEVmyNMSWHvZwsgZuSMzu0c6XgU/ynyffvrplpXotbjwwgtDvys7u2XLFjfDt2iGc9FEaeFLSanMXBnp1F7/tLy+77zzjnt8zeodTq/30QSG3ntDQW+sY/OehyoL0nr8kRlldSzo+YdTFYAqBjQJ3Pjx49O0HJwmp1OJuoYdaPI7T+S+M4r2e9ZZZ/mybwAIIsZwAwByHGUllRFUdlQzRkfyZhb3Mn+RmT4FQ5G84CgysNbjKPDT8lnhVAKdVpoxWseiwD/yWPR7+BJlx5tmbA9/DTX7eEJCgptpXBSQqkT8qaeeSnbsCpBV7qzlrNJCr2/kayt6XSJfk7fffjvVMdSpOfvss13Hhto48vG8x1FHjWZO12zo6lyIdKSZ6RWoa6jD0qVLU9ymGd01Rlpj3KO9NyOfa7T3qDoz0vP+Siu1l0ruGzdunOH7BoCgIsMNAMhxFAQrMLz++utdgKUlqDSWWWtSaxIvLff09NNPu+28JbMU/GhMrsbQRssean1s0bJV2p+Cqvbt27tAUROTaTkq/a9AS8G3ln9KT9Z1xIgRNnDgQDc2WktPabyvjuP99993E5epvDkzKLjTWtYqvVcWW4Gella77LLL3O16XXXc6ixQGb6u97bTes9a+ist9PqqzfQ
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[特征重要性排名 - Gain]\n",
      "return_10          3245.995642\n",
      "ma_ratio           2818.976538\n",
      "vol_ratio          2701.151667\n",
      "return_20          2681.762756\n",
      "high_low_ratio     2469.737874\n",
      "return_diff        1463.681328\n",
      "volatility_5       1357.026476\n",
      "market_cap_rank     722.088673\n",
      "vol_ma5             653.087969\n",
      "vol_ma20            631.252641\n",
      "ma20                548.830554\n",
      "n_income            495.525785\n",
      "ma5                 469.811824\n",
      "ma10                305.156890\n",
      "volatility_20       203.306830\n",
      "dtype: float64\n",
      "\n",
      "所有特征都有一定重要性\n"
     ]
    }
   ],
   "execution_count": 21
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}