"""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()