tests/test_ensemble.py

import math
import sys
import unittest
from pathlib import Path

SERVICE_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(SERVICE_ROOT))

import ensemble  # noqa: E402


def _blend(member_tf, member_ch, qtype="event", prices=None, horizon=5):
    """Blend two flat member forecasts. `prices` defaults to a flat history at the TimesFM
    member's level so `last` (the base-rate anchor) matches the forecast level."""
    if prices is None:
        prices = [member_tf] * 15
    return ensemble.blend(
        prices=prices,
        timesfm_fc=[member_tf] * horizon,
        chronos_fc=[member_ch] * horizon,
        question_type=qtype,
        horizon=horizon,
    )


class EnsembleCoherenceTests(unittest.TestCase):
    """The ensemble must be a COMBINATION of its members: its level stays within/near the member
    range. Regression tests for the 2026-07-23 live audit, where the extremise step drove a market
    at 3.4% to 0% and one at 17% to 7% — an "ensemble" that no longer combined its members and that
    the Next.js client now has to withhold (app/lib/ensemble.ts guard). TimesFM and Chronos-2
    forecast the SAME market price series, so extremising their aggregate away from 0.5 is unsound.
    """

    def test_confident_low_event_forecast_is_not_collapsed_to_zero(self):
        fc = _blend(0.0342, 0.034)["ensemble"]["forecast"]
        # members ~3.4%; the ensemble must stay near that, never the old clamped 0.0.
        self.assertGreater(min(fc), 0.0)
        for p in fc:
            self.assertAlmostEqual(p, 0.034, delta=0.01)

    def test_confident_mid_low_event_forecast_is_not_halved(self):
        out = _blend(0.17, 0.17)
        fc = out["ensemble"]["forecast"]
        # members 17%; the old extremise produced ~7% ("falling"). Must stay ~17%, flat.
        for p in fc:
            self.assertAlmostEqual(p, 0.17, delta=0.01)
        self.assertEqual(out["ensemble"]["trend"], "flat")

    def test_confident_high_event_forecast_is_not_clamped_to_one(self):
        fc = _blend(0.9, 0.9)["ensemble"]["forecast"]
        for p in fc:
            self.assertAlmostEqual(p, 0.9, delta=0.01)

    def test_ensemble_level_stays_within_member_band_across_the_range(self):
        # The Next.js coherence guard withholds an ensemble whose mean falls outside the member
        # range by > 0.05; the service must never emit such an ensemble on coherent members.
        for p in [0.01, 0.03, 0.05, 0.1, 0.17, 0.3, 0.45, 0.55, 0.7, 0.9, 0.97, 0.99]:
            fc = _blend(p, p)["ensemble"]["forecast"]
            mean_fc = sum(fc) / len(fc)
            self.assertGreaterEqual(mean_fc, p - 0.05, f"underflow at p={p}: {mean_fc}")
            self.assertLessEqual(mean_fc, p + 0.05, f"overflow at p={p}: {mean_fc}")

    def test_brier_estimate_reflects_the_coherent_probability(self):
        # brier = p*(1-p) on the last point; at 3.4% it is ~0.033, not the old 0.0 (from p=0).
        brier = _blend(0.0342, 0.034)["ensemble"]["brier_estimate"]
        self.assertAlmostEqual(brier, 0.03, delta=0.01)


class EnsembleKeptBehaviourTests(unittest.TestCase):
    """Behaviour the fix must NOT regress."""

    def test_base_rate_prior_still_pulls_coin_flip_toward_half(self):
        # last=0.45 is within the base-rate band; the ensemble regresses mildly toward 0.5.
        fc = _blend(0.45, 0.45)["ensemble"]["forecast"]
        mean_fc = sum(fc) / len(fc)
        self.assertGreater(mean_fc, 0.45)  # moved toward 0.5
        self.assertLess(mean_fc, 0.50)  # but not past it
        self.assertAlmostEqual(mean_fc, 0.4575, delta=0.005)  # gentle, not a jump

    def test_numeric_path_is_a_weighted_blend(self):
        # Chronos-2 favoured 0.6/0.4; distinct members so the weighting is observable.
        fc = _blend(0.10, 0.20, qtype="numeric")["ensemble"]["forecast"]
        for p in fc:
            self.assertAlmostEqual(p, 0.4 * 0.10 + 0.6 * 0.20, delta=1e-6)

    def test_forecast_values_are_bounded_and_finite(self):
        for p in [0.0, 0.02, 0.5, 0.98, 1.0]:
            for x in _blend(p, p)["ensemble"]["forecast"]:
                self.assertTrue(math.isfinite(x))
                self.assertGreaterEqual(x, 0.0)
                self.assertLessEqual(x, 1.0)


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