mirror of
https://ghfast.top/https://github.com/aeroxw/easy-tdx.git
synced 2026-09-12 15:44:15 +08:00
118 lines
3.8 KiB
Python
118 lines
3.8 KiB
Python
# 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
|