diff --git a/src/easy_tdx/factor/__init__.py b/src/easy_tdx/factor/__init__.py index e6a4db6..768287f 100644 --- a/src/easy_tdx/factor/__init__.py +++ b/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"] diff --git a/src/easy_tdx/factor/engine.py b/src/easy_tdx/factor/engine.py new file mode 100644 index 0000000..3d6caee --- /dev/null +++ b/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 diff --git a/tests/unit/test_factor_engine.py b/tests/unit/test_factor_engine.py new file mode 100644 index 0000000..7f3ff5f --- /dev/null +++ b/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