feat(factor): add FactorEngine with single/cross-section/forward-return compute

Co-Authored-By: Claude <noreply@anthropic.com>

src/easy_tdx/factor/__init__.py

@@ -2,5 +2,6 @@
 """因子研究模块。"""
 
 from easy_tdx.factor.base import FACTORY_REGISTRY, Factor, register_factor
+from easy_tdx.factor.engine import FactorEngine
 
-__all__ = ["Factor", "register_factor", "FACTORY_REGISTRY"]
+__all__ = ["Factor", "register_factor", "FACTORY_REGISTRY", "FactorEngine"]

src/easy_tdx/factor/engine.py

@@ -0,0 +1,134 @@
+# src/easy_tdx/factor/engine.py
+"""因子计算引擎 — 单股计算与截面批量计算。"""
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+
+from easy_tdx.factor.base import FACTORY_REGISTRY, Factor
+
+
+def _resolve_factor(f: str | Factor) -> Factor:
+    """将因子名或实例解析为 Factor 实例。"""
+    if isinstance(f, Factor):
+        return f
+    name = f.strip()
+    if name not in FACTORY_REGISTRY:
+        raise ValueError(
+            f"未知因子: {name!r}。可用因子: {sorted(FACTORY_REGISTRY.keys())}"
+        )
+    return FACTORY_REGISTRY[name]()
+
+
+def _datetime_to_int(dt_val: object) -> int:
+    """将 datetime 值转为 YYYYMMDD 整数。"""
+    if hasattr(dt_val, "strftime"):
+        return int(dt_val.strftime("%Y%m%d"))  # type: ignore[union-attr]
+    return int(dt_val)
+
+
+_ALL_DATES: object = object()
+
+
+class FactorEngine:
+    """批量因子计算引擎。"""
+
+    def compute_single(
+        self,
+        df: pd.DataFrame,
+        factors: list[str | Factor],
+    ) -> pd.DataFrame:
+        """单股票多因子计算。"""
+        if not factors:
+            return df.copy()
+
+        result = df.copy()
+        for f in factors:
+            factor = _resolve_factor(f)
+            result[factor.name] = factor.compute(df)
+
+        return result
+
+    def compute_cross_section(
+        self,
+        data: dict[str, pd.DataFrame],
+        factors: list[str | Factor],
+        date: int | None = _ALL_DATES,  # type: ignore[assignment]
+    ) -> pd.DataFrame:
+        """多股票截面因子计算。
+
+        Args:
+            date: int 精确匹配日期；None 仅最新一行；默认(不传)全部日期。
+        """
+        if not data:
+            return pd.DataFrame()
+
+        filter_latest = date is None
+        all_frames: list[pd.DataFrame] = []
+
+        for code, df in data.items():
+            if df.empty:
+                continue
+
+            computed = self.compute_single(df, factors)
+            computed["_date_int"] = computed["datetime"].apply(_datetime_to_int)
+
+            if date is not None and date is not _ALL_DATES:
+                computed = computed[computed["_date_int"] == date]
+            elif filter_latest:
+                computed = computed.iloc[[-1]]
+
+            factor_names = [_resolve_factor(f).name for f in factors]
+            keep_cols = ["_date_int"] + factor_names
+            sub = computed[keep_cols].copy()
+            sub["_code"] = code
+            all_frames.append(sub)
+
+        if not all_frames:
+            return pd.DataFrame()
+
+        combined = pd.concat(all_frames, ignore_index=True)
+        combined = combined.rename(columns={"_date_int": "date", "_code": "code"})
+
+        col_order = ["date", "code"] + [_resolve_factor(f).name for f in factors]
+        combined = combined[col_order].sort_values(["date", "code"]).reset_index(drop=True)
+
+        return combined
+
+    def compute_forward_returns(
+        self,
+        data: dict[str, pd.DataFrame],
+        period: int = 5,
+    ) -> pd.DataFrame:
+        """计算远期收益率。"""
+        if not data:
+            return pd.DataFrame()
+
+        col_name = f"forward_{period}d"
+        all_frames: list[pd.DataFrame] = []
+
+        for code, df in data.items():
+            if df.empty or len(df) < period + 1:
+                continue
+
+            close = df["close"].to_numpy()
+            forward = np.full(len(close), np.nan)
+            forward[: len(close) - period] = (
+                close[period:] / close[: len(close) - period] - 1
+            )
+
+            dates = df["datetime"].apply(_datetime_to_int)
+
+            sub = pd.DataFrame({
+                "date": dates,
+                "code": code,
+                col_name: forward,
+            })
+            all_frames.append(sub)
+
+        if not all_frames:
+            return pd.DataFrame(columns=["date", "code", col_name])
+
+        combined = pd.concat(all_frames, ignore_index=True)
+        combined = combined.sort_values(["date", "code"]).reset_index(drop=True)
+        return combined

tests/unit/test_factor_engine.py

@@ -0,0 +1,156 @@
+# tests/unit/test_factor_engine.py
+"""Test FactorEngine."""
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from easy_tdx.factor.base import Factor
+from easy_tdx.factor.engine import FactorEngine
+
+
+def _make_df(n: int = 60, seed: int = 42) -> pd.DataFrame:
+    """生成合成 OHLCV 数据。"""
+    rng = np.random.default_rng(seed)
+    close = 10.0 + np.cumsum(rng.normal(0, 0.5, n))
+    close = np.maximum(close, 1.0)
+    high = close + rng.uniform(0, 0.5, n)
+    low = close - rng.uniform(0, 0.5, n)
+    low = np.maximum(low, 0.1)
+    open_ = low + rng.uniform(0, high - low, n)
+    vol = rng.integers(100_000, 10_000_000, n).astype(float)
+    amount = close * vol
+
+    dates = pd.date_range("2024-01-01", periods=n, freq="D")
+    return pd.DataFrame({
+        "datetime": dates,
+        "open": open_,
+        "high": high,
+        "low": low,
+        "close": close,
+        "vol": vol,
+        "amount": amount,
+    })
+
+
+class _SimpleMomentum(Factor):
+    name = "simple_momentum"
+    category = "momentum"
+    description = "5 日动量"
+    inputs = ("close",)
+
+    def compute(self, df: pd.DataFrame) -> pd.Series:
+        return df["close"].pct_change(5)
+
+
+class _SimpleVolatility(Factor):
+    name = "simple_volatility"
+    category = "volatility"
+    description = "5 日波动率"
+    inputs = ("close",)
+
+    def compute(self, df: pd.DataFrame) -> pd.Series:
+        ret = df["close"].pct_change()
+        return ret.rolling(5).std()
+
+
+class TestComputeSingle:
+    def test_single_factor(self):
+        engine = FactorEngine()
+        df = _make_df()
+        result = engine.compute_single(df, [_SimpleMomentum()])
+        assert "simple_momentum" in result.columns
+        assert len(result) == len(df)
+
+    def test_multiple_factors(self):
+        engine = FactorEngine()
+        df = _make_df()
+        result = engine.compute_single(df, [_SimpleMomentum(), _SimpleVolatility()])
+        assert "simple_momentum" in result.columns
+        assert "simple_volatility" in result.columns
+        assert len(result) == len(df)
+
+    def test_preserves_original_columns(self):
+        engine = FactorEngine()
+        df = _make_df()
+        result = engine.compute_single(df, [_SimpleMomentum()])
+        assert "close" in result.columns
+        assert "datetime" in result.columns
+
+    def test_unknown_factor_name_raises(self):
+        engine = FactorEngine()
+        df = _make_df()
+        with pytest.raises(ValueError, match="未知因子"):
+            engine.compute_single(df, ["nonexistent_factor_xyz"])
+
+    def test_empty_factors_list(self):
+        engine = FactorEngine()
+        df = _make_df()
+        result = engine.compute_single(df, [])
+        assert len(result) == len(df)
+
+
+class TestComputeCrossSection:
+    def test_cross_section_basic(self):
+        engine = FactorEngine()
+        data = {
+            "000001": _make_df(60, seed=1),
+            "000002": _make_df(60, seed=2),
+            "600036": _make_df(60, seed=3),
+        }
+        result = engine.compute_cross_section(data, [_SimpleMomentum()])
+        assert isinstance(result, pd.DataFrame)
+        assert "date" in result.columns
+        assert "code" in result.columns
+        assert "simple_momentum" in result.columns
+        assert len(result) == 180  # 60 days × 3 stocks
+
+    def test_cross_section_latest_date(self):
+        engine = FactorEngine()
+        data = {
+            "000001": _make_df(60, seed=1),
+            "000002": _make_df(60, seed=2),
+        }
+        result = engine.compute_cross_section(data, [_SimpleMomentum()], date=None)
+        assert len(result) == 2
+
+    def test_cross_section_specific_date(self):
+        engine = FactorEngine()
+        df = _make_df(60, seed=1)
+        data = {"000001": df}
+        target_date = int(df["datetime"].iloc[-5].strftime("%Y%m%d"))
+        result = engine.compute_cross_section(data, [_SimpleMomentum()], date=target_date)
+        assert len(result) == 1
+        assert result.iloc[0]["date"] == target_date
+
+    def test_cross_section_empty_data(self):
+        engine = FactorEngine()
+        result = engine.compute_cross_section({}, [_SimpleMomentum()])
+        assert len(result) == 0
+
+
+class TestComputeForwardReturns:
+    def test_forward_returns_basic(self):
+        engine = FactorEngine()
+        data = {
+            "000001": _make_df(60, seed=1),
+            "000002": _make_df(60, seed=2),
+        }
+        result = engine.compute_forward_returns(data, period=5)
+        assert "date" in result.columns
+        assert "code" in result.columns
+        assert "forward_5d" in result.columns
+        code_000001 = result[result["code"] == "000001"]
+        assert np.isnan(code_000001["forward_5d"].iloc[-1])
+
+    def test_forward_returns_period(self):
+        engine = FactorEngine()
+        data = {"000001": _make_df(60, seed=1)}
+        result = engine.compute_forward_returns(data, period=10)
+        assert "forward_10d" in result.columns
+
+    def test_forward_returns_empty(self):
+        engine = FactorEngine()
+        result = engine.compute_forward_returns({}, period=5)
+        assert len(result) == 0