cyb50-quant/index-rotation/tests/test_phase2_signals.py

from __future__ import annotations

import unittest

import pandas as pd

from src.portfolio.allocator import allocate_weights
from src.signals.selector import build_signal_panel
from src.signals.trend import apply_trend_filter


def make_signal_input(trade_date: str = "2020-01-10") -> pd.DataFrame:
    return pd.DataFrame(
        [
            {
                "instrument": "sse50",
                "trade_date": pd.Timestamp(trade_date),
                "close": 110,
                "daily_return": 0.01,
                "ret_5d": 0.08,
                "ret_10d": 0.10,
                "ret_20d": 0.12,
                "ret_60d": 0.20,
                "ma_20": 100,
                "ma_60": 95,
                "vol_10d": 0.10,
                "vol_20d": 0.11,
            },
            {
                "instrument": "hs300",
                "trade_date": pd.Timestamp(trade_date),
                "close": 108,
                "daily_return": 0.01,
                "ret_5d": 0.05,
                "ret_10d": 0.06,
                "ret_20d": 0.08,
                "ret_60d": 0.12,
                "ma_20": 102,
                "ma_60": 101,
                "vol_10d": 0.12,
                "vol_20d": 0.13,
            },
            {
                "instrument": "chinext50",
                "trade_date": pd.Timestamp(trade_date),
                "close": 96,
                "daily_return": -0.01,
                "ret_5d": -0.02,
                "ret_10d": -0.01,
                "ret_20d": -0.03,
                "ret_60d": 0.02,
                "ma_20": 98,
                "ma_60": 100,
                "vol_10d": 0.18,
                "vol_20d": 0.20,
            },
            {
                "instrument": "star50",
                "trade_date": pd.Timestamp(trade_date),
                "close": 102,
                "daily_return": 0.00,
                "ret_5d": 0.02,
                "ret_10d": 0.03,
                "ret_20d": 0.01,
                "ret_60d": 0.04,
                "ma_20": 101,
                "ma_60": 103,
                "vol_10d": 0.08,
                "vol_20d": 0.09,
            },
        ]
    )


class SignalLayerTests(unittest.TestCase):
    def test_trend_filter_requires_at_least_two_rules(self) -> None:
        frame = apply_trend_filter(make_signal_input())
        outcome = frame.set_index("instrument")["trend_pass"].to_dict()
        self.assertEqual(outcome["sse50"], True)
        self.assertEqual(outcome["hs300"], True)
        self.assertEqual(outcome["star50"], True)
        self.assertEqual(outcome["chinext50"], False)

    def test_trend_filter_threshold_can_be_raised(self) -> None:
        frame = apply_trend_filter(make_signal_input(), min_rules=3)
        outcome = frame.set_index("instrument")["trend_pass"].to_dict()
        self.assertEqual(outcome["sse50"], True)
        self.assertEqual(outcome["hs300"], True)
        self.assertEqual(outcome["star50"], False)
        self.assertEqual(outcome["chinext50"], False)

    def test_ranking_and_scoring_only_selects_trend_pass_members(self) -> None:
        signals = build_signal_panel(make_signal_input(), top_n=2).set_index("instrument")
        self.assertEqual(signals.loc["sse50", "selection_rank"], 1)
        self.assertEqual(signals.loc["hs300", "selection_rank"], 2)
        self.assertTrue(pd.isna(signals.loc["chinext50", "selection_rank"]))
        self.assertGreater(signals.loc["sse50", "final_score"], signals.loc["hs300", "final_score"])

    def test_custom_risk_penalty_can_change_selection_order(self) -> None:
        frame = pd.DataFrame(
            [
                {
                    "instrument": "high_beta",
                    "trade_date": pd.Timestamp("2020-01-10"),
                    "close": 110,
                    "daily_return": 0.01,
                    "ret_5d": 0.12,
                    "ret_10d": 0.12,
                    "ret_20d": 0.12,
                    "ret_60d": 0.12,
                    "ma_20": 100,
                    "ma_60": 95,
                    "vol_10d": 0.30,
                    "vol_20d": 0.30,
                },
                {
                    "instrument": "defensive",
                    "trade_date": pd.Timestamp("2020-01-10"),
                    "close": 108,
                    "daily_return": 0.01,
                    "ret_5d": 0.11,
                    "ret_10d": 0.11,
                    "ret_20d": 0.11,
                    "ret_60d": 0.11,
                    "ma_20": 100,
                    "ma_60": 95,
                    "vol_10d": 0.05,
                    "vol_20d": 0.05,
                },
                {
                    "instrument": "filler_a",
                    "trade_date": pd.Timestamp("2020-01-10"),
                    "close": 90,
                    "daily_return": -0.01,
                    "ret_5d": 0.02,
                    "ret_10d": 0.02,
                    "ret_20d": 0.02,
                    "ret_60d": 0.02,
                    "ma_20": 100,
                    "ma_60": 101,
                    "vol_10d": 0.10,
                    "vol_20d": 0.10,
                },
                {
                    "instrument": "filler_b",
                    "trade_date": pd.Timestamp("2020-01-10"),
                    "close": 89,
                    "daily_return": -0.01,
                    "ret_5d": 0.01,
                    "ret_10d": 0.01,
                    "ret_20d": 0.01,
                    "ret_60d": 0.01,
                    "ma_20": 100,
                    "ma_60": 101,
                    "vol_10d": 0.11,
                    "vol_20d": 0.11,
                },
            ]
        )
        low_penalty = build_signal_panel(frame, top_n=1, risk_penalty_multiplier=0.0).set_index("instrument")
        high_penalty = build_signal_panel(frame, top_n=1, risk_penalty_multiplier=1.0).set_index("instrument")
        self.assertEqual(low_penalty.loc["high_beta", "selection_rank"], 1)
        self.assertEqual(high_penalty.loc["defensive", "selection_rank"], 1)

    def test_custom_momentum_weights_can_change_selection_order(self) -> None:
        frame = pd.DataFrame(
            [
                {
                    "instrument": "short_burst",
                    "trade_date": pd.Timestamp("2020-01-10"),
                    "close": 110,
                    "daily_return": 0.01,
                    "ret_5d": 0.20,
                    "ret_10d": 0.20,
                    "ret_20d": 0.05,
                    "ret_60d": 0.01,
                    "ma_20": 100,
                    "ma_60": 95,
                    "vol_10d": 0.10,
                    "vol_20d": 0.10,
                },
                {
                    "instrument": "long_runner",
                    "trade_date": pd.Timestamp("2020-01-10"),
                    "close": 108,
                    "daily_return": 0.01,
                    "ret_5d": 0.10,
                    "ret_10d": 0.10,
                    "ret_20d": 0.15,
                    "ret_60d": 0.18,
                    "ma_20": 100,
                    "ma_60": 95,
                    "vol_10d": 0.10,
                    "vol_20d": 0.10,
                },
                {
                    "instrument": "filler_a",
                    "trade_date": pd.Timestamp("2020-01-10"),
                    "close": 90,
                    "daily_return": -0.01,
                    "ret_5d": 0.01,
                    "ret_10d": 0.01,
                    "ret_20d": 0.01,
                    "ret_60d": 0.01,
                    "ma_20": 100,
                    "ma_60": 95,
                    "vol_10d": 0.10,
                    "vol_20d": 0.10,
                },
                {
                    "instrument": "filler_b",
                    "trade_date": pd.Timestamp("2020-01-10"),
                    "close": 89,
                    "daily_return": -0.01,
                    "ret_5d": 0.05,
                    "ret_10d": 0.04,
                    "ret_20d": 0.03,
                    "ret_60d": 0.02,
                    "ma_20": 100,
                    "ma_60": 95,
                    "vol_10d": 0.10,
                    "vol_20d": 0.10,
                },
            ]
        )
        default_panel = build_signal_panel(frame, top_n=1, risk_penalty_multiplier=0.0).set_index("instrument")
        short_tilt_panel = build_signal_panel(
            frame,
            top_n=1,
            risk_penalty_multiplier=0.0,
            momentum_weights={"ret_5d": 0.50, "ret_10d": 0.30, "ret_20d": 0.15, "ret_60d": 0.05},
        ).set_index("instrument")
        self.assertEqual(default_panel.loc["long_runner", "selection_rank"], 1)
        self.assertEqual(short_tilt_panel.loc["short_burst", "selection_rank"], 1)

    def test_top1_top2_and_empty_allocation(self) -> None:
        base_signals = build_signal_panel(make_signal_input(), top_n=2)

        top2 = allocate_weights(base_signals, top_n=2)
        top2_weights = top2.set_index("instrument")["target_weight"].to_dict()
        self.assertEqual(top2_weights["sse50"], 0.5)
        self.assertEqual(top2_weights["hs300"], 0.5)
        self.assertEqual(top2["cash_weight"].iloc[0], 0.0)

        top1 = allocate_weights(build_signal_panel(make_signal_input(), top_n=1), top_n=1)
        top1_weights = top1.set_index("instrument")["target_weight"].to_dict()
        self.assertEqual(top1_weights["sse50"], 1.0)
        self.assertEqual(sum(top1_weights.values()), 1.0)

        empty = make_signal_input()
        empty[["close", "ret_20d"]] = [90, -0.10]
        empty["ma_20"] = 100
        empty["ma_60"] = 110
        allocated = allocate_weights(build_signal_panel(empty, top_n=2), top_n=2)
        self.assertTrue((allocated["target_weight"] == 0.0).all())
        self.assertEqual(allocated["cash_weight"].iloc[0], 1.0)


if __name__ == "__main__":
    unittest.main()