feat(factor): add factor preprocessing pipeline (winsorize/zscore/rank/fill/orthogonalize)

src/easy_tdx/factor/transform.py

@@ -0,0 +1,196 @@
+# src/easy_tdx/factor/transform.py
+"""因子预处理 — 纯函数管道。"""
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+
+
+def winsorize(
+    factor_data: pd.DataFrame,
+    columns: str | list[str],
+    method: str = "mad",
+    threshold: float = 3.0,
+) -> pd.DataFrame:
+    """截面去极值。"""
+    if isinstance(columns, str):
+        columns = [columns]
+    result = factor_data.copy()
+
+    for col in columns:
+        if col not in result.columns:
+            continue
+
+        def _clip_group(group: pd.Series) -> pd.Series:
+            valid = group.dropna()
+            if len(valid) < 3:
+                return group
+            if method == "mad":
+                median = valid.median()
+                mad = (valid - median).abs().median() * 1.4826
+                lower = median - threshold * mad
+                upper = median + threshold * mad
+            elif method == "sigma":
+                mean = valid.mean()
+                std = valid.std()
+                lower = mean - threshold * std
+                upper = mean + threshold * std
+            elif method == "percentile":
+                lower = valid.quantile(0.025)
+                upper = valid.quantile(0.975)
+            else:
+                raise ValueError(f"未知去极值方法: {method!r}")
+            return group.clip(lower, upper)
+
+        if "date" in result.columns:
+            result[col] = result.groupby("date")[col].transform(_clip_group)
+        else:
+            result[col] = _clip_group(result[col])
+
+    return result
+
+
+def zscore(
+    factor_data: pd.DataFrame,
+    columns: str | list[str],
+    cross_section: bool = True,
+) -> pd.DataFrame:
+    """标准化。"""
+    if isinstance(columns, str):
+        columns = [columns]
+    result = factor_data.copy()
+
+    for col in columns:
+        if col not in result.columns:
+            continue
+
+        def _zscore_group(group: pd.Series) -> pd.Series:
+            std = group.std()
+            if std == 0 or pd.isna(std):
+                return group * 0
+            return (group - group.mean()) / std
+
+        if cross_section and "date" in result.columns:
+            result[col] = result.groupby("date")[col].transform(_zscore_group)
+        else:
+            result[col] = _zscore_group(result[col])
+
+    return result
+
+
+def rank_normalize(
+    factor_data: pd.DataFrame,
+    columns: str | list[str],
+) -> pd.DataFrame:
+    """排名归一化 [0, 1]。"""
+    if isinstance(columns, str):
+        columns = [columns]
+    result = factor_data.copy()
+
+    for col in columns:
+        if col not in result.columns:
+            continue
+
+        def _rank_group(group: pd.Series) -> pd.Series:
+            return group.rank(pct=True)
+
+        if "date" in result.columns:
+            result[col] = result.groupby("date")[col].transform(_rank_group)
+        else:
+            result[col] = _rank_group(result[col])
+
+    return result
+
+
+def fill_missing(
+    factor_data: pd.DataFrame,
+    columns: str | list[str],
+    method: str = "cross_mean",
+) -> pd.DataFrame:
+    """缺失值填充。"""
+    if isinstance(columns, str):
+        columns = [columns]
+    result = factor_data.copy()
+
+    for col in columns:
+        if col not in result.columns:
+            continue
+        if method == "cross_mean":
+            if "date" in result.columns:
+
+                def _fill_mean(group: pd.Series) -> pd.Series:
+                    return group.fillna(group.mean())
+
+                result[col] = result.groupby("date")[col].transform(_fill_mean)
+            else:
+                result[col] = result[col].fillna(result[col].mean())
+        elif method == "forward_fill":
+            if "code" in result.columns:
+                result[col] = result.groupby("code")[col].ffill()
+            else:
+                result[col] = result[col].ffill()
+        else:
+            raise ValueError(f"未知填充方法: {method!r}")
+
+    return result
+
+
+def orthogonalize(
+    factor_data: pd.DataFrame,
+    target: str,
+    by: str | list[str],
+) -> pd.DataFrame:
+    """因子正交化。"""
+    if isinstance(by, str):
+        by = [by]
+    result = factor_data.copy()
+
+    if target not in result.columns:
+        return result
+    for b in by:
+        if b not in result.columns:
+            return result
+
+    y = result[target].to_numpy(dtype=np.float64)
+    X_cols = [result[b].to_numpy(dtype=np.float64) for b in by]
+    X = np.column_stack([np.ones(len(y))] + X_cols)
+
+    mask = ~np.isnan(y)
+    for xc in X_cols:
+        mask &= ~np.isnan(xc)
+
+    if mask.sum() < len(by) + 2:
+        return result
+
+    coef, _, _, _ = np.linalg.lstsq(X[mask], y[mask], rcond=None)
+    predicted = X @ coef
+    residual = y - predicted
+    residual[~mask] = np.nan
+    result[target] = residual
+
+    return result
+
+
+def preprocess(
+    factor_data: pd.DataFrame,
+    columns: list[str],
+    steps: list[str] | None = None,
+) -> pd.DataFrame:
+    """一键预处理管道。"""
+    if steps is None:
+        steps = ["winsorize", "zscore", "fill_missing"]
+
+    result = factor_data.copy()
+    for step in steps:
+        if step == "winsorize":
+            result = winsorize(result, columns)
+        elif step == "zscore":
+            result = zscore(result, columns)
+        elif step == "rank_normalize":
+            result = rank_normalize(result, columns)
+        elif step == "fill_missing":
+            result = fill_missing(result, columns)
+        else:
+            raise ValueError(f"未知预处理步骤: {step!r}")
+
+    return result

tests/unit/test_factor_transform.py

@@ -0,0 +1,117 @@
+# tests/unit/test_factor_transform.py
+"""Test factor preprocessing functions."""
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from easy_tdx.factor.transform import (
+    fill_missing,
+    orthogonalize,
+    preprocess,
+    rank_normalize,
+    winsorize,
+    zscore,
+)
+
+
+def _make_cross_section(n_dates: int = 20, n_stocks: int = 30, seed: int = 42) -> pd.DataFrame:
+    rng = np.random.default_rng(seed)
+    rows = []
+    for d in range(n_dates):
+        for s in range(n_stocks):
+            rows.append({
+                "date": 20240101 + d,
+                "code": f"{s:06d}",
+                "momentum_20d": rng.normal(0.02, 0.05),
+                "volatility_20d": abs(rng.normal(0.02, 0.01)),
+            })
+    df = pd.DataFrame(rows)
+    df.loc[0, "momentum_20d"] = 10.0
+    df.loc[1, "momentum_20d"] = -10.0
+    df.loc[2, "momentum_20d"] = np.nan
+    return df
+
+
+class TestWinsorize:
+    def test_mad_clips_extremes(self):
+        df = _make_cross_section()
+        result = winsorize(df, ["momentum_20d"], method="mad", threshold=3.0)
+        assert result["momentum_20d"].max() < 10.0
+        assert result["momentum_20d"].min() > -10.0
+
+    def test_preserves_shape(self):
+        df = _make_cross_section()
+        result = winsorize(df, ["momentum_20d"])
+        assert len(result) == len(df)
+        assert list(result.columns) == list(df.columns)
+
+
+class TestZscore:
+    def test_cross_section_standardization(self):
+        df = _make_cross_section()
+        result = zscore(df, ["momentum_20d"], cross_section=True)
+        for date in result["date"].unique():
+            sub = result[result["date"] == date]["momentum_20d"].dropna()
+            if len(sub) > 2:
+                assert abs(sub.mean()) < 0.5
+
+    def test_preserves_nan(self):
+        df = _make_cross_section()
+        result = zscore(df, ["momentum_20d"])
+        assert result["momentum_20d"].isna().sum() >= 1
+
+
+class TestRankNormalize:
+    def test_output_range(self):
+        df = _make_cross_section()
+        result = rank_normalize(df, ["momentum_20d"])
+        valid = result["momentum_20d"].dropna()
+        assert valid.min() >= 0
+        assert valid.max() <= 1
+
+
+class TestFillMissing:
+    def test_cross_mean_fills(self):
+        df = _make_cross_section()
+        na_before = df["momentum_20d"].isna().sum()
+        result = fill_missing(df, ["momentum_20d"], method="cross_mean")
+        na_after = result["momentum_20d"].isna().sum()
+        assert na_after < na_before
+
+    def test_forward_fill(self):
+        df = _make_cross_section()
+        result = fill_missing(df, ["momentum_20d"], method="forward_fill")
+        assert len(result) == len(df)
+
+
+class TestOrthogonalize:
+    def test_residual_differs(self):
+        df = _make_cross_section()
+        df = zscore(df, ["momentum_20d", "volatility_20d"])
+        df = fill_missing(df, ["momentum_20d", "volatility_20d"], method="cross_mean")
+        result = orthogonalize(df, target="momentum_20d", by="volatility_20d")
+        assert "momentum_20d" in result.columns
+        assert not result["momentum_20d"].equals(df["momentum_20d"])
+
+
+class TestPreprocess:
+    def test_default_pipeline(self):
+        df = _make_cross_section()
+        result = preprocess(df, ["momentum_20d"])
+        assert len(result) == len(df)
+        assert "momentum_20d" in result.columns
+        assert result["momentum_20d"].isna().sum() <= df["momentum_20d"].isna().sum()
+
+    def test_custom_steps(self):
+        df = _make_cross_section()
+        result = preprocess(df, ["momentum_20d"], steps=["winsorize", "zscore"])
+        assert len(result) == len(df)
+
+    def test_preserves_other_columns(self):
+        df = _make_cross_section()
+        result = preprocess(df, ["momentum_20d"])
+        assert "date" in result.columns
+        assert "code" in result.columns
+        assert "volatility_20d" in result.columns