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

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

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