chinese-equity-quant/tests/test_alpha.py

"""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)
    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_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)
    '''))

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