feat: CSI500 universe + reversal_vol signal + argparse CLI

data/universe.py

@@ -1,4 +1,4 @@
-"""CSI 300 (HS300) universe helpers."""
+"""CSI 300 (HS300) and CSI 500 (ZZ500) universe helpers."""
 import logging
 
 import baostock as bs
@@ -19,6 +19,20 @@ SYMBOLS = [
 ]
 
 
+# First 30 CSI 500 (ZZ500) constituents (mid/small caps) in 'shXXXXXX' /
+# 'szXXXXXX' format. Hardcoded for fast, deterministic smoke tests. Use
+# get_zz500_stocks() for the live, full list. Mean reversion tends to be
+# stronger in these smaller caps than in the HS300 large caps.
+CSI500_SYMBOLS = [
+    "sh600006", "sh600008", "sh600017", "sh600020", "sh600021",
+    "sh600026", "sh600037", "sh600039", "sh600053", "sh600056",
+    "sh600060", "sh600061", "sh600062", "sh600073", "sh600089",
+    "sh600095", "sh600118", "sh600125", "sh600126", "sh600143",
+    "sh600153", "sh600160", "sh600169", "sh600176", "sh600183",
+    "sz000009", "sz000012", "sz000021", "sz000025", "sz000027",
+]
+
+
 def get_hs300_stocks() -> pd.DataFrame:
     """Fetch the current CSI 300 constituents from baostock.
 
@@ -36,3 +50,22 @@ def get_hs300_stocks() -> pd.DataFrame:
     df = pd.DataFrame(stocks, columns=["code", "name", "date"])
     df["code"] = df["code"].str.replace(".", "", regex=False)
     return df
+
+
+def get_zz500_stocks() -> pd.DataFrame:
+    """Fetch the current CSI 500 (ZZ500) constituents from baostock.
+
+    Returns:
+        DataFrame with columns ``code`` (e.g. ``sh600006``), ``name``, ``date``.
+    """
+    bs.login()
+    try:
+        rs = bs.query_zz500_stocks()
+        stocks = []
+        while rs.next():
+            stocks.append(rs.get_row_data())
+    finally:
+        bs.logout()
+    df = pd.DataFrame(stocks, columns=["code", "name", "date"])
+    df["code"] = df["code"].str.replace(".", "", regex=False)
+    return df

run_example.py

@@ -1,6 +1,12 @@
 #!/usr/bin/env python3
-"""End-to-end pipeline: HS300 universe -> reversal signal -> cross-sectional IC
+"""End-to-end pipeline: universe -> signal -> cross-sectional IC
--> multi-stock backtest (AlphaStrategy + RankEqualWeightBuilder) -> reports."""
+-> multi-stock backtest (AlphaStrategy + RankEqualWeightBuilder) -> reports.
+
+Usage:
+    python3 run_example.py --universe hs300 --signal reversal
+    python3 run_example.py --universe csi500 --signal reversal_vol
+"""
+import argparse
 import logging
 
 import pandas as pd
@@ -9,10 +15,11 @@ from analysis.report import generate_report
 from backtest.config import BacktestConfig
 from backtest.runner import BacktestRunner
 from data.downloader import download_batch
-from data.universe import SYMBOLS
+from data.universe import SYMBOLS, CSI500_SYMBOLS
 from eval.metrics import evaluate_cross_sectional
 from portfolio.builder import RankEqualWeightBuilder
 from signals.reversal import ReversalSignal
+from signals.reversal_vol import ReversalVolSignal
 from strategies.alpha_strategy import AlphaStrategy
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
@@ -29,18 +36,27 @@ def _forward_returns(data: dict[str, pd.DataFrame], horizon: int) -> pd.DataFram
     return pd.DataFrame(forward_returns)
 
 
-def main(forward_horizon: int = 5):
+def main(forward_horizon: int = 5, universe: str = "csi500", signal_name: str = "reversal_vol"):
-    symbols = SYMBOLS[:30]
+    universes = {"hs300": SYMBOLS, "csi500": CSI500_SYMBOLS}
+    symbols = universes.get(universe, CSI500_SYMBOLS)[:30]
+
+    signals = {
+        "reversal": ReversalSignal(lookback=5),
+        "reversal_vol": ReversalVolSignal(lookback=5, vol_window=20),
+    }
+    signal = signals.get(signal_name, ReversalVolSignal(lookback=5, vol_window=20))
+
     start, end = "2023-01-01", "2024-12-31"
     initial_cash = 1_000_000
 
+    logger.info(f"Universe: {universe} ({len(symbols)} stocks), Signal: {signal.name}")
+
     # 1-2. Download daily data for the universe.
     data = download_batch(symbols, start, end)
     data = {s: df for s, df in data.items() if df is not None and not df.empty}
     logger.info(f"Downloaded {len(data)}/{len(symbols)} symbols")
 
-    # 3. Compute the reversal signal per stock.
+    # 3. Compute the signal per stock.
-    signal = ReversalSignal(lookback=5)
     signal_series: dict[str, pd.Series] = {}
     for sym, df in data.items():
         sig = signal.compute(df)
@@ -52,7 +68,7 @@ def main(forward_horizon: int = 5):
     returns_df = _forward_returns(data, forward_horizon)
     signal_eval = evaluate_cross_sectional(signals_df, returns_df)
 
-    # 4b. Multi-horizon IC to show which horizon the signal works at.
+    # 4b. Multi-horizon IC.
     horizon_evals = {
         h: evaluate_cross_sectional(signals_df, _forward_returns(data, h))
         for h in (1, 5, 20)
@@ -84,6 +100,7 @@ def main(forward_horizon: int = 5):
     # 8. Print summary.
     print("\nSIGNAL IC")
     print("=" * 50)
+    print(f"Universe: {universe}  |  Signal: {signal.name}")
     print(f"IC mean / std / IR:       {signal_eval['ic_mean']:.4f} / "
           f"{signal_eval['ic_std']:.4f} / {signal_eval['ir']:.4f}")
     print(f"Rank IC mean / std / IR:  {signal_eval['rank_ic_mean']:.4f} / "
@@ -102,4 +119,8 @@ def main(forward_horizon: int = 5):
 
 
 if __name__ == "__main__":
-    main()
+    parser = argparse.ArgumentParser(description="Chinese equity quant backtest")
+    parser.add_argument("--universe", default="csi500", choices=["hs300", "csi500"])
+    parser.add_argument("--signal", default="reversal_vol", choices=["reversal", "reversal_vol"])
+    args = parser.parse_args()
+    main(universe=args.universe, signal_name=args.signal)

signals/__init__.py

@@ -1,5 +1,6 @@
 """Alpha signal abstractions."""
 from signals.base import AlphaSignal
 from signals.reversal import ReversalSignal
+from signals.reversal_vol import ReversalVolSignal
 
-__all__ = ["AlphaSignal", "ReversalSignal"]
+__all__ = ["AlphaSignal", "ReversalSignal", "ReversalVolSignal"]

signals/reversal_vol.py

@@ -0,0 +1,27 @@
+"""Volatility-scaled short-horizon reversal signal."""
+import pandas as pd
+
+from signals.base import AlphaSignal
+
+
+class ReversalVolSignal(AlphaSignal):
+    """Reversal score normalized by trailing volatility.
+
+    The raw reversal ``-close.pct_change(lookback)`` is divided by the rolling
+    standard deviation of daily returns over ``vol_window``. Scaling by
+    volatility damps the score of noisy, high-vol names so the signal favors
+    oversold stocks whose move is large *relative* to their own volatility.
+    """
+
+    def __init__(self, lookback: int = 5, vol_window: int = 20):
+        self.lookback = lookback
+        self.vol_window = vol_window
+
+    def compute(self, df: pd.DataFrame) -> pd.Series:
+        reversal = -df["close"].pct_change(self.lookback)
+        vol = df["close"].pct_change().rolling(self.vol_window).std()
+        return reversal / vol
+
+    @property
+    def name(self) -> str:
+        return f"reversal_vol_{self.lookback}d_{self.vol_window}d"