diff --git a/src/easy_tdx/factor/transform.py b/src/easy_tdx/factor/transform.py new file mode 100644 index 0000000..94271ec --- /dev/null +++ b/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 diff --git a/tests/unit/test_factor_transform.py b/tests/unit/test_factor_transform.py new file mode 100644 index 0000000..0a02a22 --- /dev/null +++ b/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