Yuxuan Yan
350 lines
13 KiB
Python
350 lines
13 KiB
Python
"""Tests for pipeline alpha computation and combination (no network)."""
|
|
import textwrap
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
from pipeline.alpha.base import BaseAlpha
|
|
from pipeline.alpha.compute import (
|
|
compute_alpha,
|
|
evaluate_alpha,
|
|
investable_universe_mask,
|
|
)
|
|
from pipeline.alpha.registry import (
|
|
available_alphas,
|
|
get_alpha,
|
|
load_alpha_module,
|
|
register_alpha,
|
|
)
|
|
from pipeline.combo.combine import combine_alphas, _equal_weight
|
|
from pipeline.common.schema import ALPHA_COLUMNS, COMBO_COLUMNS
|
|
|
|
|
|
def _make_data(n_days: int = 30, symbols=("sh600000", "sz000001", "sh600519")) -> pd.DataFrame:
|
|
"""Build a synthetic long-format DATA_COLUMNS frame with deterministic prices."""
|
|
dates = pd.date_range("2024-01-01", periods=n_days)
|
|
rng = np.random.default_rng(0)
|
|
frames = []
|
|
for i, sym in enumerate(symbols):
|
|
# Distinct drift per symbol so the cross-section is non-degenerate.
|
|
close = 100.0 + i * 5 + np.cumsum(rng.standard_normal(n_days))
|
|
frames.append(pd.DataFrame({
|
|
"symbol_id": sym,
|
|
"symbol_name": sym,
|
|
"date": dates,
|
|
"open": close,
|
|
"high": close,
|
|
"low": close,
|
|
"close": close,
|
|
"volume": 1_000.0,
|
|
"amount": 1_000.0 * close,
|
|
}))
|
|
return pd.concat(frames, ignore_index=True)
|
|
|
|
|
|
def test_compute_alpha_schema_and_naming():
|
|
alpha = compute_alpha(_make_data(), "rev5", "reversal", lookback=5)
|
|
assert list(alpha.columns) == ALPHA_COLUMNS
|
|
assert (alpha["alpha_name"] == "rev5").all()
|
|
|
|
|
|
def test_reversal_sign_matches_negative_trailing_return():
|
|
# Cross-sectional z-score preserves the sign relative to the cross-section,
|
|
# so the stock with the most negative trailing return ranks highest.
|
|
data = _make_data()
|
|
alpha = compute_alpha(data, "rev5", "reversal", lookback=5)
|
|
close = data.pivot_table(index="date", columns="symbol_id", values="close").sort_index()
|
|
raw = -close.pct_change(5, fill_method=None)
|
|
last = raw.index[-1]
|
|
expected_top = raw.loc[last].idxmax()
|
|
got = alpha[alpha["date"] == last].set_index("symbol_id")["weight"].idxmax()
|
|
assert got == expected_top
|
|
|
|
|
|
def test_weights_are_cross_sectional_zscore():
|
|
# Each date's weights are a z-score, so the per-date mean is ~0.
|
|
alpha = compute_alpha(_make_data(), "rev5", "reversal", lookback=5)
|
|
per_date_mean = alpha.groupby("date")["weight"].mean().abs()
|
|
assert (per_date_mean < 1e-9).all()
|
|
|
|
|
|
def test_evaluate_alpha_keys():
|
|
data = _make_data()
|
|
alpha = compute_alpha(data, "rev5", "reversal", lookback=5)
|
|
metrics = evaluate_alpha(alpha, data)
|
|
for key in ("cumulative_return", "sharpe_annual", "turnover_annual",
|
|
"max_drawdown", "hit_rate", "n_dates"):
|
|
assert key in metrics
|
|
|
|
|
|
def test_evaluate_alpha_uses_next_open_to_next_open_returns():
|
|
dates = pd.date_range("2024-01-01", periods=5)
|
|
data = pd.concat([
|
|
pd.DataFrame({
|
|
"symbol_id": "sh600000",
|
|
"symbol_name": "sh600000",
|
|
"date": dates,
|
|
"open": [100.0, 100.0, 100.0, 100.0, 200.0],
|
|
"high": [100.0, 1000.0, 1000.0, 1000.0, 1000.0],
|
|
"low": [100.0, 1000.0, 1000.0, 1000.0, 1000.0],
|
|
"close": [100.0, 1000.0, 1000.0, 1000.0, 1000.0],
|
|
"volume": 1_000.0,
|
|
"amount": 1_000.0,
|
|
}),
|
|
pd.DataFrame({
|
|
"symbol_id": "sz000001",
|
|
"symbol_name": "sz000001",
|
|
"date": dates,
|
|
"open": [100.0, 100.0, 100.0, 200.0, 200.0],
|
|
"high": [100.0, 10.0, 10.0, 10.0, 10.0],
|
|
"low": [100.0, 10.0, 10.0, 10.0, 10.0],
|
|
"close": [100.0, 10.0, 10.0, 10.0, 10.0],
|
|
"volume": 1_000.0,
|
|
"amount": 1_000.0,
|
|
}),
|
|
], ignore_index=True)
|
|
alpha = pd.DataFrame({
|
|
"symbol_id": ["sh600000", "sz000001", "sh600000", "sz000001"],
|
|
"date": [dates[1], dates[1], dates[2], dates[2]],
|
|
"alpha_name": ["toy"] * 4,
|
|
"weight": [-1.0, 1.0, 1.0, -1.0],
|
|
})
|
|
|
|
metrics = evaluate_alpha(alpha, data)
|
|
|
|
assert metrics["n_dates"] == 2
|
|
assert np.isclose(metrics["cumulative_return"], 1.25)
|
|
|
|
|
|
def test_evaluate_alpha_excludes_signal_without_forward_return():
|
|
dates = pd.date_range("2024-01-01", periods=3)
|
|
data = pd.concat([
|
|
pd.DataFrame({
|
|
"symbol_id": "sh600000",
|
|
"symbol_name": "sh600000",
|
|
"date": dates,
|
|
"open": [100.0, 100.0, 200.0],
|
|
"high": [100.0, 100.0, 200.0],
|
|
"low": [100.0, 100.0, 200.0],
|
|
"close": [100.0, 100.0, 200.0],
|
|
"volume": 1_000.0,
|
|
"amount": 1_000.0,
|
|
}),
|
|
pd.DataFrame({
|
|
"symbol_id": "sz000001",
|
|
"symbol_name": "sz000001",
|
|
"date": dates,
|
|
"open": [100.0, 100.0, 100.0],
|
|
"high": [100.0, 100.0, 100.0],
|
|
"low": [100.0, 100.0, 100.0],
|
|
"close": [100.0, 100.0, 100.0],
|
|
"volume": 1_000.0,
|
|
"amount": 1_000.0,
|
|
}),
|
|
], ignore_index=True)
|
|
alpha = pd.DataFrame({
|
|
"symbol_id": ["sh600000", "sz000001", "sh600000", "sz000001"],
|
|
"date": [dates[0], dates[0], dates[1], dates[1]],
|
|
"alpha_name": ["toy"] * 4,
|
|
"weight": [1.0, -1.0, -1.0, 1.0],
|
|
})
|
|
|
|
metrics = evaluate_alpha(alpha, data)
|
|
|
|
assert metrics["n_dates"] == 1
|
|
|
|
|
|
def test_equal_weight_is_mean_of_alphas():
|
|
data = _make_data()
|
|
a = compute_alpha(data, "rev", "reversal", lookback=5)
|
|
b = compute_alpha(data, "mom", "momentum", lookback=5)
|
|
combo = _equal_weight([a, b])
|
|
# reversal = -momentum before z-scoring, but after independent per-date
|
|
# z-scoring they are exact negatives, so the equal-weight mean is ~0.
|
|
assert combo["weight"].abs().max() < 1e-9
|
|
|
|
|
|
def test_combine_alphas_schema(tmp_path):
|
|
data = _make_data()
|
|
a_path = tmp_path / "a.pq"
|
|
b_path = tmp_path / "b.pq"
|
|
compute_alpha(data, "rev", "reversal", lookback=5).to_parquet(a_path, index=False)
|
|
compute_alpha(data, "revvol", "reversal_vol", lookback=5, vol_window=10).to_parquet(b_path, index=False)
|
|
combo = combine_alphas([str(a_path), str(b_path)], "eq", method="equal_weight")
|
|
assert list(combo.columns) == COMBO_COLUMNS
|
|
assert (combo["combo_name"] == "eq").all()
|
|
|
|
|
|
def test_combine_single_alpha_is_identity(tmp_path):
|
|
data = _make_data()
|
|
a = compute_alpha(data, "rev", "reversal", lookback=5)
|
|
a_path = tmp_path / "a.pq"
|
|
a.to_parquet(a_path, index=False)
|
|
|
|
combo = combine_alphas([str(a_path)], "rev_combo", method="equal_weight")
|
|
|
|
expected = a[["symbol_id", "date", "weight"]].reset_index(drop=True)
|
|
got = combo[["symbol_id", "date", "weight"]].reset_index(drop=True)
|
|
pd.testing.assert_frame_equal(got, expected)
|
|
assert list(combo.columns) == COMBO_COLUMNS
|
|
assert (combo["combo_name"] == "rev_combo").all()
|
|
|
|
|
|
# --- registry / factory -----------------------------------------------------
|
|
|
|
def test_builtins_are_registered():
|
|
assert {"reversal", "reversal_vol", "momentum", "reversal_rank"} <= set(available_alphas())
|
|
|
|
|
|
def test_get_alpha_filters_unaccepted_params():
|
|
# reversal only accepts lookback; passing vol_window too must not error.
|
|
alpha = get_alpha("reversal", lookback=7, vol_window=99)
|
|
assert alpha.name == "reversal"
|
|
assert alpha.lookback == 7
|
|
assert not hasattr(alpha, "vol_window")
|
|
|
|
|
|
def test_get_alpha_unknown_raises():
|
|
with pytest.raises(KeyError):
|
|
get_alpha("does_not_exist")
|
|
|
|
|
|
def test_register_duplicate_name_raises():
|
|
available_alphas() # ensure built-ins loaded
|
|
|
|
with pytest.raises(ValueError):
|
|
@register_alpha
|
|
class Dup(BaseAlpha):
|
|
name = "reversal"
|
|
|
|
def signal(self, close):
|
|
return close
|
|
|
|
|
|
def test_register_rejects_non_basealpha():
|
|
with pytest.raises(TypeError):
|
|
register_alpha(object) # type: ignore[arg-type]
|
|
|
|
|
|
# --- base class --------------------------------------------------------------
|
|
|
|
def test_to_weights_are_per_date_zscore():
|
|
class _Const(BaseAlpha):
|
|
name = "_const_test"
|
|
|
|
def signal(self, close):
|
|
return close # arbitrary finite signal
|
|
|
|
close = _make_data().pivot_table(index="date", columns="symbol_id", values="close")
|
|
weights = _Const().weights(close.sort_index())
|
|
# Each date demeaned to ~0.
|
|
assert (weights.mean(axis=1).abs() < 1e-9).all()
|
|
|
|
|
|
# --- external plugin loading -------------------------------------------------
|
|
|
|
def test_load_external_alpha_module(tmp_path):
|
|
module_path = tmp_path / "my_external_alpha.py"
|
|
module_path.write_text(textwrap.dedent('''
|
|
import pandas as pd
|
|
from pipeline.alpha.base import BaseAlpha
|
|
from pipeline.alpha.registry import register_alpha
|
|
|
|
@register_alpha
|
|
class ExternalDemoAlpha(BaseAlpha):
|
|
name = "external_demo"
|
|
|
|
def __init__(self, span: int = 3):
|
|
self.span = span
|
|
|
|
def signal(self, close: pd.DataFrame) -> pd.DataFrame:
|
|
return -close.pct_change(self.span, fill_method=None)
|
|
'''))
|
|
|
|
load_alpha_module(str(module_path))
|
|
assert "external_demo" in available_alphas()
|
|
|
|
# The factory forwards the external alpha's own param (`span`).
|
|
instance = get_alpha("external_demo", span=4, lookback=99)
|
|
assert instance.span == 4
|
|
|
|
# And it works end-to-end through compute_alpha.
|
|
result = compute_alpha(_make_data(), "ext", "external_demo", span=4)
|
|
assert list(result.columns) == ALPHA_COLUMNS
|
|
assert (result["alpha_name"] == "ext").all()
|
|
|
|
|
|
# --- rank reversal + investable universe filter ------------------------------
|
|
|
|
def _make_rich_data(n_days: int = 70, symbols=("sh600000", "sz000001", "sh600519", "sz300750")):
|
|
"""Long-format data with the columns the universe filter needs."""
|
|
dates = pd.date_range("2024-01-01", periods=n_days)
|
|
rng = np.random.default_rng(1)
|
|
frames = []
|
|
for i, sym in enumerate(symbols):
|
|
close = 100.0 + i * 5 + np.cumsum(rng.standard_normal(n_days))
|
|
frames.append(pd.DataFrame({
|
|
"symbol_id": sym,
|
|
"symbol_name": sym,
|
|
"date": dates,
|
|
"open": close, "high": close, "low": close, "close": close,
|
|
"volume": 1_000.0,
|
|
"amount": (1_000.0 + i * 5_000.0) * close, # higher i = more liquid
|
|
"isST": 0,
|
|
"tradestatus": 1,
|
|
}))
|
|
return pd.concat(frames, ignore_index=True)
|
|
|
|
|
|
def test_reversal_rank_registered_and_bounded():
|
|
data = _make_data(n_days=30)
|
|
alpha = compute_alpha(data, "rr", "reversal_rank", lookback=5)
|
|
assert list(alpha.columns) == ALPHA_COLUMNS
|
|
# Rank-demeaned weights are per-date zero-mean and bounded by the
|
|
# cross-section size, never blowing up the way a z-score outlier can.
|
|
per_date_mean = alpha.groupby("date")["weight"].mean().abs()
|
|
assert (per_date_mean < 1e-9).all()
|
|
assert alpha["weight"].abs().max() <= len(data["symbol_id"].unique())
|
|
|
|
|
|
def test_investable_universe_mask_excludes_st_and_suspended():
|
|
data = _make_rich_data()
|
|
# Flag one name ST throughout, and suspend another on the last date.
|
|
data.loc[data["symbol_id"] == "sh600000", "isST"] = 1
|
|
last = data["date"].max()
|
|
data.loc[(data["symbol_id"] == "sz000001") & (data["date"] == last), "tradestatus"] = 0
|
|
|
|
close = data.pivot_table(index="date", columns="symbol_id", values="close").sort_index()
|
|
mask = investable_universe_mask(data, close, top_n=10, min_history=5)
|
|
|
|
assert not mask["sh600000"].any() # ST excluded on every date
|
|
assert not bool(mask.loc[last, "sz000001"]) # suspended on the last date
|
|
assert bool(mask.loc[last, "sh600519"]) # a normal name stays investable
|
|
|
|
|
|
def test_compute_alpha_universe_filter_zeros_excluded_names():
|
|
data = _make_rich_data()
|
|
data.loc[data["symbol_id"] == "sh600000", "isST"] = 1
|
|
|
|
alpha = compute_alpha(
|
|
data, "rr_liq", "reversal_rank", lookback=5,
|
|
universe={"top_n": 10, "min_history": 5},
|
|
)
|
|
# The ST name is never held; an investable name is.
|
|
st_w = alpha.loc[alpha["symbol_id"] == "sh600000", "weight"]
|
|
assert (st_w.fillna(0.0) == 0.0).all()
|
|
assert alpha.loc[alpha["symbol_id"] == "sz300750", "weight"].abs().sum() > 0.0
|
|
|
|
|
|
def test_universe_filter_does_not_corrupt_signal_history():
|
|
# Masking happens on the signal, not the price history, so weights on
|
|
# investable names match the unfiltered weights restricted to that set.
|
|
data = _make_rich_data()
|
|
universe = {"top_n": 2, "min_history": 5} # keep only the 2 most liquid names
|
|
|
|
filtered = compute_alpha(data, "f", "reversal_rank", lookback=5, universe=universe)
|
|
held = set(filtered.loc[filtered["weight"] != 0.0, "symbol_id"].unique())
|
|
# The two most liquid names (highest amount) are sh600519, sz300750.
|
|
assert held == {"sh600519", "sz300750"}
|