217 lines
7.9 KiB
Python
217 lines
7.9 KiB
Python
"""
|
||
grid_seeker v6.6 三级模型推理管道
|
||
Rank → Top → Stacking → stacking_probability (最终排序)
|
||
"""
|
||
import pickle
|
||
import numpy as np
|
||
import pandas as pd
|
||
from pathlib import Path
|
||
from datetime import date
|
||
from core.scoring.config import (
|
||
get_model_path, RANK_MODEL, TOP_MODEL, STACKING_MODEL,
|
||
STACKING_THRESHOLD,
|
||
)
|
||
from core.scoring.features.pipeline import FeaturePipeline
|
||
from core.scoring.models import ScoringResult
|
||
from core.database import db
|
||
from core.logger import LogLevel, PrintLog
|
||
|
||
|
||
# ============================================================
|
||
# Rank 模型输入特征 (52维, v3.4, 直接从模型文件的 selected_features 读取)
|
||
# ============================================================
|
||
def _get_rank_features() -> list:
|
||
import pickle
|
||
from core.scoring.config import get_model_path
|
||
path = get_model_path(RANK_MODEL)
|
||
with open(path, 'rb') as f:
|
||
obj = pickle.load(f)
|
||
if isinstance(obj, dict):
|
||
sf = obj.get('selected_features', [])
|
||
if sf:
|
||
return sf
|
||
raise RuntimeError("无法从 rank.pkl 读取 selected_features")
|
||
|
||
RANK_FEATURE_COLS = _get_rank_features()
|
||
|
||
|
||
class GridSeekerPipeline:
|
||
"""
|
||
grid_seeker v6.6 三级模型评分管道。
|
||
|
||
Usage:
|
||
engine = GridSeekerPipeline()
|
||
rankings = engine.run(trade_date=date.today())
|
||
# 返回 DataFrame: stock_code, stacking_probability, rank 等
|
||
"""
|
||
|
||
def __init__(self, model_dir: Path = None):
|
||
self._rank_model = None
|
||
self._top_model = None
|
||
self._stacking_model = None
|
||
|
||
# ---- 模型加载 ----
|
||
|
||
def _load_model(self, name: str):
|
||
"""加载单个 .pkl 模型"""
|
||
path = get_model_path(name)
|
||
if not path.exists():
|
||
raise FileNotFoundError(f'模型文件不存在: {path}')
|
||
with open(path, 'rb') as f:
|
||
obj = pickle.load(f)
|
||
# 支持 dict 格式 {"model": lgbm_model, ...} 或直接返回模型对象
|
||
if isinstance(obj, dict):
|
||
return obj.get('model', obj)
|
||
return obj
|
||
|
||
@property
|
||
def rank_model(self):
|
||
if self._rank_model is None:
|
||
self._rank_model = self._load_model(RANK_MODEL)
|
||
return self._rank_model
|
||
|
||
@property
|
||
def top_model(self):
|
||
if self._top_model is None:
|
||
self._top_model = self._load_model(TOP_MODEL)
|
||
return self._top_model
|
||
|
||
@property
|
||
def stacking_model(self):
|
||
if self._stacking_model is None:
|
||
self._stacking_model = self._load_model(STACKING_MODEL)
|
||
return self._stacking_model
|
||
|
||
# ---- 预测 ----
|
||
|
||
def _predict_with_model(self, model, X: pd.DataFrame, feature_cols: list) -> np.ndarray:
|
||
"""
|
||
使用模型预测。自动选择特征子集，兼容 sklearn API (predict/predict_proba)。
|
||
"""
|
||
available = [c for c in feature_cols if c in X.columns]
|
||
missing = set(feature_cols) - set(available)
|
||
if missing:
|
||
PrintLog(LogLevel.WARNING,
|
||
f'[scorer] 缺少特征列 ({len(missing)}): {list(missing)[:5]}...')
|
||
|
||
X_sub = X[available].fillna(0).values
|
||
|
||
try:
|
||
if hasattr(model, 'predict_proba'):
|
||
proba = model.predict_proba(X_sub)
|
||
if proba.shape[1] >= 2:
|
||
return proba[:, 1]
|
||
return proba[:, 0]
|
||
elif hasattr(model, 'predict'):
|
||
return model.predict(X_sub)
|
||
else:
|
||
return model.predict(X_sub)
|
||
except Exception as e:
|
||
PrintLog(LogLevel.ERROR, f'[scorer] 模型预测失败: {e}')
|
||
raise
|
||
|
||
# ---- 主流程 ----
|
||
|
||
def run(self, trade_date: date) -> pd.DataFrame:
|
||
"""
|
||
执行完整的 3 级评分管道。
|
||
|
||
Returns:
|
||
DataFrame indexed by stock_code, 含 stacking_probability / rank 等列,
|
||
按 stacking_probability 降序排列
|
||
"""
|
||
PrintLog(LogLevel.INFO, f'[scorer] ===== grid_seeker v6.6 评分开始 ({trade_date}) =====')
|
||
|
||
# 1. 特征工程
|
||
pipeline = FeaturePipeline(trade_date)
|
||
feature_df = pipeline.run()
|
||
|
||
if feature_df.empty:
|
||
PrintLog(LogLevel.WARNING, '[scorer] 无股票通过特征工程, 终止')
|
||
return pd.DataFrame()
|
||
|
||
PrintLog(LogLevel.INFO,
|
||
f'[scorer] 特征工程完成: {len(feature_df)} stocks, '
|
||
f'{len(feature_df.columns)} dims')
|
||
|
||
# 2. Stage 1: Rank 模型 → rank_predicted_rounds (52维)
|
||
PrintLog(LogLevel.INFO, '[scorer] Stage 1/3: Rank 模型...')
|
||
feature_df['rank_predicted_rounds'] = self._predict_with_model(
|
||
self.rank_model, feature_df, RANK_FEATURE_COLS
|
||
)
|
||
|
||
# 3. Stage 2: Top 模型 → top_elite_prob (53维 = 52 + rank_predicted_rounds)
|
||
PrintLog(LogLevel.INFO, '[scorer] Stage 2/3: Top 模型...')
|
||
top_cols = RANK_FEATURE_COLS + ['rank_predicted_rounds']
|
||
feature_df['top_elite_prob'] = self._predict_with_model(
|
||
self.top_model, feature_df, top_cols
|
||
)
|
||
|
||
# 4. Stage 3: Stacking 模型 → stacking_probability (54维 = 52 + rank + top)
|
||
PrintLog(LogLevel.INFO, '[scorer] Stage 3/3: Stacking 模型...')
|
||
stk_cols = RANK_FEATURE_COLS + ['rank_predicted_rounds', 'top_elite_prob']
|
||
feature_df['stacking_probability'] = self._predict_with_model(
|
||
self.stacking_model, feature_df, stk_cols
|
||
)
|
||
|
||
# 5. 排序（直接用 stacking_probability）
|
||
feature_df['score_rank'] = feature_df['stacking_probability'].rank(
|
||
ascending=False, method='min'
|
||
).astype(int)
|
||
feature_df['candidate_count'] = len(feature_df)
|
||
feature_df = feature_df.sort_values('score_rank')
|
||
|
||
n_above = (feature_df['stacking_probability'] >= STACKING_THRESHOLD).sum()
|
||
PrintLog(LogLevel.INFO,
|
||
f'[scorer] 评分完成: {len(feature_df)} 只候选, '
|
||
f'{n_above} 只高于阈值 {STACKING_THRESHOLD}')
|
||
PrintLog(LogLevel.INFO,
|
||
f'[scorer] Top-5: '
|
||
f'{feature_df.head(5)[["stacking_probability", "rank_predicted_rounds"]].to_dict("index")}')
|
||
|
||
return feature_df
|
||
|
||
def persist(self, rankings: pd.DataFrame, trade_date: date):
|
||
"""将评分结果持久化到 ScoringResult 表"""
|
||
if rankings.empty:
|
||
return
|
||
|
||
records = []
|
||
for code, row in rankings.iterrows():
|
||
records.append({
|
||
'stock_code': str(code),
|
||
'trade_date': trade_date,
|
||
'predicted_profit': float(row.get('stacking_probability', 0)),
|
||
'rank_predicted_rounds': float(row.get('rank_predicted_rounds', 0))
|
||
if 'rank_predicted_rounds' in row else None,
|
||
'top_elite_prob': float(row.get('top_elite_prob', 0))
|
||
if 'top_elite_prob' in row else None,
|
||
'stacking_probability': float(row.get('stacking_probability', 0))
|
||
if 'stacking_probability' in row else None,
|
||
'score_rank': int(row.get('score_rank', 0)),
|
||
'candidate_count': int(row.get('candidate_count', 0)),
|
||
})
|
||
|
||
with db.atomic():
|
||
for batch in _chunked(records, 500):
|
||
ScoringResult.insert_many(batch).on_conflict_replace().execute()
|
||
|
||
PrintLog(LogLevel.INFO,
|
||
f'[scorer] 评分结果已持久化: {len(records)} 条')
|
||
|
||
def get_top_n(self, trade_date: date, n: int = 50) -> list[dict]:
|
||
"""查询历史评分 Top-N"""
|
||
rows = (ScoringResult
|
||
.select()
|
||
.where(
|
||
(ScoringResult.trade_date == trade_date) &
|
||
(ScoringResult.score_rank <= n)
|
||
)
|
||
.order_by(ScoringResult.score_rank)
|
||
.dicts())
|
||
return list(rows)
|
||
|
||
|
||
def _chunked(lst: list, n: int):
|
||
for i in range(0, len(lst), n):
|
||
yield lst[i:i + n] |