Yuxuan Yan
0a6f367fbf
Weights formed from close[t] now earn the t→t+1 return: forward returns are computed on the full market calendar before selecting signal dates, so a sparse signal grid earns the next available return rather than the next signal date, and the final signal date (no forward return) is dropped. Signal pct_change uses fill_method=None so suspended names do not inherit stale prices. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
258 lines
8.8 KiB
Python
258 lines
8.8 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
|
|
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_period_returns():
|
|
dates = pd.date_range("2024-01-01", periods=4)
|
|
data = pd.concat([
|
|
pd.DataFrame({
|
|
"symbol_id": "sh600000",
|
|
"symbol_name": "sh600000",
|
|
"date": dates,
|
|
"open": [100.0, 200.0, 200.0, 200.0],
|
|
"high": [100.0, 200.0, 200.0, 200.0],
|
|
"low": [100.0, 200.0, 200.0, 200.0],
|
|
"close": [100.0, 200.0, 200.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, 200.0, 200.0],
|
|
"high": [100.0, 100.0, 200.0, 200.0],
|
|
"low": [100.0, 100.0, 200.0, 200.0],
|
|
"close": [100.0, 100.0, 200.0, 200.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"], 0.5)
|
|
|
|
|
|
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[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"] == 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()
|
|
|
|
|
|
# --- registry / factory -----------------------------------------------------
|
|
|
|
def test_builtins_are_registered():
|
|
assert {"reversal", "reversal_vol", "momentum"} <= 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()
|
|
|