myquant/finance/factors/engine.py

"""
FactorEngine — 因子计算引擎。

批量计算因子，处理技术/基本面/情绪因子的不同数据需求。
"""

import copy

import pandas as pd

from factors.base import BaseFactor
from factors.fundamental.roe import ROEFactor
from factors.fundamental.pe_pb import PEFactor, PBFactor, EPFactor

FUNDAMENTAL_FACTOR_TYPES = (ROEFactor, PEFactor, PBFactor, EPFactor)


def _is_sentiment(factor: BaseFactor) -> bool:
    return getattr(factor, "category", "") == "sentiment"


class FactorEngine:
    """因子计算引擎。"""

    def __init__(self, data_manager, sentiment_engine=None):
        """
        参数:
            data_manager: DataManager 实例。
            sentiment_engine: SentimentEngine 实例（可选，启用情绪因子时需提供）。
        """
        self._dm = data_manager
        self._sentiment_engine = sentiment_engine
        self._financial_cache: dict[str, pd.DataFrame] = {}

    def _get_financial(self, ts_code: str) -> pd.DataFrame:
        """获取财务数据（带缓存）。"""
        if ts_code not in self._financial_cache:
            df = self._dm.get_financial(ts_code)
            self._financial_cache[ts_code] = df
        return self._financial_cache[ts_code]

    def _resolve_factors(
        self, factors: list[BaseFactor], fina: pd.DataFrame
    ) -> list[BaseFactor]:
        """为每个股票 clone 基本面因子并注入财务数据。"""
        resolved = []
        for f in factors:
            if isinstance(f, FUNDAMENTAL_FACTOR_TYPES):
                f = copy.copy(f)
                f._financial_df = fina
            resolved.append(f)
        return resolved

    def compute(
        self,
        ts_code: str,
        factors: list[BaseFactor],
    ) -> pd.DataFrame:
        """
        对单只股票计算多个因子。

        参数:
            ts_code: 如 '000001.SZ'
            factors: 因子实例列表

        返回:
            DataFrame，index=trade_date，columns=因子名
        """
        if not factors:
            return pd.DataFrame()

        # 分离情绪因子（通过 SentimentEngine 处理）
        sent_factors = [f for f in factors if _is_sentiment(f)]
        other_factors = [f for f in factors if not _is_sentiment(f)]

        # 收集所有需要的列
        required_cols = set()
        has_fundamental = False
        for f in other_factors:
            required_cols.update(f.get_required_columns())
            if isinstance(f, FUNDAMENTAL_FACTOR_TYPES):
                has_fundamental = True

        # 获取日线数据
        daily = self._dm.get_daily(ts_code)
        if daily.empty:
            return pd.DataFrame()

        daily = daily.set_index("trade_date").sort_index()

        # 获取财务数据（如有基本面因子）
        fina = self._dm.get_financial(ts_code) if has_fundamental else pd.DataFrame()

        # 为当前股票解析因子（clone 基本面因子注入财务数据）
        resolved_factors = self._resolve_factors(other_factors, fina)

        # 逐因子计算
        results = {}
        for factor in resolved_factors:
            try:
                series = factor.calculate(daily)
                results[factor.name] = series.astype("float64")
            except Exception as e:
                print(f"[WARN] 因子 {factor.name} 计算失败 ({ts_code}): {e}")
                results[factor.name] = pd.Series(float("nan"), index=daily.index)

        # 情绪因子：通过 SentimentEngine 计算后合并
        if sent_factors and self._sentiment_engine:
            try:
                sent_df = self._sentiment_engine.compute(ts_code)
                for f in sent_factors:
                    if f.name in sent_df.columns:
                        results[f.name] = sent_df[f.name]
                    else:
                        results[f.name] = pd.Series(float("nan"), index=daily.index)
            except Exception as e:
                print(f"[WARN] 情绪因子计算失败 ({ts_code}): {e}")
                for f in sent_factors:
                    results[f.name] = pd.Series(float("nan"), index=daily.index)

        factor_df = pd.DataFrame(results)
        factor_df.index.name = "trade_date"
        return factor_df

    def compute_batch(
        self,
        ts_codes: list[str],
        factors: list[BaseFactor],
    ) -> dict[str, pd.DataFrame]:
        """
        批量计算多只股票的因子。

        返回:
            {ts_code: factor_df}
        """
        results = {}
        total = len(ts_codes)
        for i, ts_code in enumerate(ts_codes):
            try:
                results[ts_code] = self.compute(ts_code, factors)
            except Exception as e:
                print(f"[WARN] {ts_code} 因子计算失败: {e}")
                results[ts_code] = pd.DataFrame()
            if (i + 1) % 50 == 0:
                print(f"[FactorEngine] 进度: {i + 1}/{total}")
        return results

    def compute_universe(
        self,
        factors: list[BaseFactor],
        date: str,
        ts_codes: list[str] | None = None,
    ) -> pd.DataFrame:
        """
        计算全市场某一天的因子截面。

        参数:
            factors: 因子列表
            date: 目标日期 'YYYYMMDD'
            ts_codes: 股票列表，None 表示全部

        返回:
            DataFrame，index=ts_code，columns=因子名
        """
        if ts_codes is None:
            stocks = self._dm.get_stock_list()
            ts_codes = list(stocks.index)

        rows = []
        for ts_code in ts_codes:
            daily = self._dm.get_daily(ts_code)
            if daily.empty:
                continue
            daily = daily.set_index("trade_date")
            if date not in daily.index:
                continue

            row = {"ts_code": ts_code}
            fina = self._get_financial(ts_code)
            resolved = self._resolve_factors(factors, fina)
            for factor in resolved:
                try:
                    series = factor.calculate(daily)
                    row[factor.name] = series.get(date, float("nan"))
                except Exception:
                    row[factor.name] = float("nan")
            rows.append(row)

        if not rows:
            return pd.DataFrame()
        result = pd.DataFrame(rows).set_index("ts_code")
        return result