qlib/backend/app/quant/selection.py

"""Selection Engine（ARCHITECTURE_v2 §14）—— 纯 pandas 执行层。

当前实现 method=score：因子加权复合分 → TopN/Top% 截断，输出 SelectionResult。
M6.2 在同一模块加入 method=condition（结构化条件选股）。

未来函数纪律：面板只在 <= observation_date 的数据上计算；observation_date 是
<= as_of 的最近可用交易日（as_of 显式传入即历史选股，None 则到数据最新）。
data 长表由 Service 装配（已按 universe 过滤 symbol、含预热窗口）。
"""

from __future__ import annotations

from datetime import date

import pandas as pd

from app.domain.entities.selection import (
    SelectionCandidate,
    SelectionQuery,
    SelectionResult,
    SelectionStatistics,
)
from app.quant.factors import FactorError, compute_factor, get_factor
from app.quant.local_engine import build_factor_panels, composite_score

_UNIMPLEMENTED_DEFAULT = [
    "exclude_suspended 依赖停牌数据，当前未建模（结果可能包含停牌股）",
]


def resolve_observation_date(daily: pd.DataFrame, as_of: date | None) -> pd.Timestamp | None:
    """<= as_of 的最近可用交易日；as_of=None 取数据最新一日。"""
    if daily.empty:
        return None
    dates = pd.to_datetime(daily["trade_date"])
    if as_of is None:
        return dates.max()
    avail = dates[dates <= pd.Timestamp(as_of)]
    return avail.max() if len(avail) else None


def factor_columns(query: SelectionQuery) -> set[str]:
    """score 模式所需行情数值列（数据装配裁剪用）。"""
    needed = {"close"}
    for fs in query.factors:
        try:
            defn, _fn = get_factor(fs.name)
        except FactorError:
            continue  # 未知因子由执行期统一报错（score_selection 中 build_factor_panels）
        needed.update(defn.requires)
    return needed


def run_score_selection(
    daily: pd.DataFrame,
    query: SelectionQuery,
    as_of: date | None,
) -> SelectionResult:
    """因子评分选股（v2 §14.1B）：复合分 → 排序 → TopN/Top%。"""
    if query.method != "score":
        raise ValueError(f"run_score_selection 需要 method=score，当前 {query.method}")
    obs = resolve_observation_date(daily, as_of)
    if obs is None:
        resolved = as_of or date.today()
        return SelectionResult(
            as_of_date=resolved,
            method=query.method,
            statistics=SelectionStatistics(),
            candidates=[],
            unimplemented=list(_UNIMPLEMENTED_DEFAULT),
            config_snapshot=query.model_dump(mode="json"),
        )
    resolved = obs.date()
    # 只允许使用 <= obs 的数据（面板计算在截断后数据上进行）
    view = daily[pd.to_datetime(daily["trade_date"]) <= obs]
    if view.empty:
        return SelectionResult(
            as_of_date=resolved,
            method=query.method,
            statistics=SelectionStatistics(),
            candidates=[],
            unimplemented=list(_UNIMPLEMENTED_DEFAULT),
            config_snapshot=query.model_dump(mode="json"),
        )

    panels = build_factor_panels(view, query.factors)  # 未知因子在此抛 FactorError
    score = composite_score(panels).loc[obs].dropna().sort_values(ascending=False)

    # 每因子在 obs 行的原始值（factor_values 供展示与解释；与 build_factor_panels 同数据）
    raw: dict[str, pd.Series] = {}
    for fs in query.factors:
        _defn, panel = compute_factor(fs.name, view)
        if obs in panel.index:
            raw[fs.name] = panel.loc[obs]

    candidates_df = _truncate(score, query)
    evaluated = int(len(score))  # score 已 dropna，长度即有分股票数
    candidates: list[SelectionCandidate] = []
    for rank, (sym, sc) in enumerate(candidates_df.items(), start=1):
        factor_values = {
            name: _to_float(series.get(sym))
            for name, series in raw.items()
            if isinstance(series, pd.Series)
        }
        factor_values = {k: v for k, v in factor_values.items() if v is not None}
        candidates.append(
            SelectionCandidate(
                symbol=sym,
                rank=rank,
                score=round(float(sc), 6),
                factor_values=factor_values,
                selection_reason=_score_reason(query, sym, raw),
            )
        )

    return SelectionResult(
        as_of_date=resolved,
        method=query.method,
        statistics=SelectionStatistics(
            universe_size=_symbol_count(view),
            evaluated=evaluated,
            selected=len(candidates),
        ),
        candidates=candidates,
        unimplemented=list(_UNIMPLEMENTED_DEFAULT),
        config_snapshot=query.model_dump(mode="json"),
    )


def _truncate(score: pd.Series, query: SelectionQuery) -> pd.Series:
    """按 top_n / top_pct / min_score 截断（入参已按分数降序）。"""
    s = score
    if query.min_score is not None:
        s = s[s >= query.min_score]
    if query.top_pct is not None:
        n = max(int(round(len(s) * query.top_pct)), 1)
        s = s.head(n)
    elif query.top_n is not None:
        s = s.head(query.top_n)
    return s


def _score_reason(query: SelectionQuery, symbol: str, raw: dict[str, pd.Series]) -> list[str]:
    """生成可读的入选理由：列每个因子的观测值与权重。"""
    reasons: list[str] = []
    for fs in query.factors:
        try:
            defn, _fn = get_factor(fs.name)
        except FactorError:
            continue
        series = raw.get(fs.name)
        val = _to_float(series.get(symbol)) if isinstance(series, pd.Series) else None
        if val is None:
            continue
        good = defn.direction == "higher_is_better"
        reasons.append(
            f"{fs.name}={val:.4f}（权重 {fs.weight}，{'越高越好' if good else '越低越好'}）"
        )
    return reasons


def _symbol_count(daily: pd.DataFrame) -> int:
    return int(daily["symbol"].nunique()) if not daily.empty and "symbol" in daily else 0


def _to_float(v) -> float | None:
    if v is None:
        return None
    try:
        f = float(v)
    except (TypeError, ValueError):
        return None
    if f != f:  # NaN
        return None
    return f