feat: math core v3 engine upgrade

2026-06-27 12:21:41 +00:00
parent 365bc5d4b7
commit b4bf0f2361
34 changed files with 11693 additions and 3 deletions
@@ -0,0 +1,290 @@
+"""Unit tests for v3 correlation-aware clustering.
+
+Tests for compute_n_eff, compute_cluster_llr, and cluster_evidence functions.
+
+Requirements validated: 4.1–4.5
+"""
+from __future__ import annotations
+
+from datetime import datetime, timezone
+
+import pytest
+
+from services.aggregation.scoring import EvidenceUnit
+from services.aggregation.worker import (
+    cluster_evidence,
+    compute_cluster_llr,
+    compute_n_eff,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_unit(cluster_id: str = "cluster_a", symbol: str = "AAPL") -> EvidenceUnit:
+    """Create a minimal EvidenceUnit for testing."""
+    return EvidenceUnit(
+        symbol=symbol,
+        layer="company",
+        event_type="earnings",
+        source_id="doc_1",
+        source_group="reuters",
+        timestamp=datetime(2024, 1, 15, 12, 0, tzinfo=timezone.utc),
+        horizon="7d",
+        direction=1,
+        sentiment_strength=0.8,
+        impact=0.7,
+        extraction_conf=0.9,
+        source_cred=0.85,
+        novelty=0.6,
+        event_base_rate=0.25,
+        cluster_id=cluster_id,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Test: 3 identical articles from same source → n_eff < 3
+# Requirement: 4.2, 4.3
+# ---------------------------------------------------------------------------
+
+
+class TestNEffIdenticalArticles:
+    """3 identical articles from same wire (default rho=0.80) → n_eff < 3."""
+
+    def test_n_eff_less_than_cluster_size(self):
+        llrs = [1.0, 1.0, 1.0]
+        # Default correlations: rho=0.80 for all pairs (same wire/source)
+        n_eff = compute_n_eff(llrs)
+
+        # Formula: (3)² / (3 + 2×3×0.80×1×1) = 9 / (3 + 4.8) = 9/7.8 ≈ 1.154
+        expected = 9.0 / 7.8
+        assert n_eff < 3.0
+        assert n_eff == pytest.approx(expected, rel=1e-6)
+
+    def test_n_eff_discounts_correlated_signals(self):
+        """Higher correlation → lower n_eff."""
+        llrs = [1.0, 1.0, 1.0]
+        n_eff_correlated = compute_n_eff(llrs)  # default rho=0.80
+
+        # Independent: rho=0.0
+        identity = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
+        n_eff_independent = compute_n_eff(llrs, correlations=identity)
+
+        assert n_eff_correlated < n_eff_independent
+
+
+# ---------------------------------------------------------------------------
+# Test: 3 independent articles → n_eff ≈ 3
+# Requirement: 4.2, 4.3
+# ---------------------------------------------------------------------------
+
+
+class TestNEffIndependentArticles:
+    """3 independent articles (rho=0.0) → n_eff = 3.0."""
+
+    def test_n_eff_equals_cluster_size(self):
+        llrs = [1.0, 1.0, 1.0]
+        # Zero off-diagonal correlations
+        correlations = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
+        n_eff = compute_n_eff(llrs, correlations=correlations)
+
+        # n_eff = (3)² / (3 + 0) = 3.0
+        assert n_eff == pytest.approx(3.0, rel=1e-6)
+
+    def test_n_eff_with_varying_magnitudes(self):
+        """Independent signals with different magnitudes still give n <= cluster size."""
+        llrs = [0.5, 1.0, 2.0]
+        correlations = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
+        n_eff = compute_n_eff(llrs, correlations=correlations)
+
+        # With zero correlations, n_eff = (sum |w|)^2 / sum(w^2)
+        # = (0.5+1.0+2.0)^2 / (0.25+1.0+4.0) = 12.25 / 5.25 ≈ 2.333
+        expected = (3.5**2) / (0.25 + 1.0 + 4.0)
+        assert n_eff == pytest.approx(expected, rel=1e-6)
+        assert n_eff <= 3.0
+
+
+# ---------------------------------------------------------------------------
+# Test: Single signal cluster → n_eff = 1.0
+# Requirement: 4.2
+# ---------------------------------------------------------------------------
+
+
+class TestNEffSingleSignal:
+    """Single signal in a cluster → n_eff = 1.0."""
+
+    def test_single_signal(self):
+        llrs = [0.5]
+        n_eff = compute_n_eff(llrs)
+        assert n_eff == 1.0
+
+    def test_empty_cluster(self):
+        llrs: list[float] = []
+        n_eff = compute_n_eff(llrs)
+        assert n_eff == 1.0
+
+
+# ---------------------------------------------------------------------------
+# Test: Cluster LLR clamp at ±2.5
+# Requirement: 4.4, 4.5
+# ---------------------------------------------------------------------------
+
+
+class TestClusterLLRClamp:
+    """Cluster LLR is clamped to [-2.5, 2.5]."""
+
+    def test_positive_clamp(self):
+        """Very large positive LLRs with high n_eff → clamped to 2.5."""
+        llrs = [2.0, 2.0, 2.0, 2.0, 2.0]
+        # Use independent correlations for max n_eff
+        correlations = [
+            [1, 0, 0, 0, 0],
+            [0, 1, 0, 0, 0],
+            [0, 0, 1, 0, 0],
+            [0, 0, 0, 1, 0],
+            [0, 0, 0, 0, 1],
+        ]
+        n_eff = compute_n_eff(llrs, correlations=correlations)
+        cluster_llr = compute_cluster_llr(llrs, n_eff)
+
+        # weighted_mean = 2.0, sqrt(5) ≈ 2.236, raw = 4.47 → clamp to 2.5
+        assert cluster_llr == pytest.approx(2.5, rel=1e-6)
+
+    def test_negative_clamp(self):
+        """Very negative LLRs with high n_eff → clamped to -2.5."""
+        llrs = [-2.0, -2.0, -2.0, -2.0, -2.0]
+        correlations = [
+            [1, 0, 0, 0, 0],
+            [0, 1, 0, 0, 0],
+            [0, 0, 1, 0, 0],
+            [0, 0, 0, 1, 0],
+            [0, 0, 0, 0, 1],
+        ]
+        n_eff = compute_n_eff(llrs, correlations=correlations)
+        cluster_llr = compute_cluster_llr(llrs, n_eff)
+
+        assert cluster_llr == pytest.approx(-2.5, rel=1e-6)
+
+    def test_within_bounds_no_clamp(self):
+        """Small LLRs with low n_eff → no clamping needed."""
+        llrs = [0.3, 0.4]
+        n_eff = compute_n_eff(llrs)
+        cluster_llr = compute_cluster_llr(llrs, n_eff)
+
+        assert -2.5 <= cluster_llr <= 2.5
+        # Should NOT be at the clamp boundary
+        assert abs(cluster_llr) < 2.5
+
+    def test_single_signal_clamp(self):
+        """Single signal beyond clamp → clamped."""
+        llrs = [3.0]
+        cluster_llr = compute_cluster_llr(llrs, n_eff=1.0)
+        assert cluster_llr == pytest.approx(2.5, rel=1e-6)
+
+    def test_single_signal_negative_clamp(self):
+        """Single negative signal beyond clamp → clamped to -2.5."""
+        llrs = [-3.0]
+        cluster_llr = compute_cluster_llr(llrs, n_eff=1.0)
+        assert cluster_llr == pytest.approx(-2.5, rel=1e-6)
+
+
+# ---------------------------------------------------------------------------
+# Test: All-zero LLRs → cluster_llr = 0.0
+# Requirement: 4.4
+# ---------------------------------------------------------------------------
+
+
+class TestClusterLLRZero:
+    """All-zero LLRs produce zero cluster LLR."""
+
+    def test_all_zeros(self):
+        llrs = [0.0, 0.0, 0.0]
+        n_eff = compute_n_eff(llrs)
+        cluster_llr = compute_cluster_llr(llrs, n_eff)
+        assert cluster_llr == 0.0
+
+    def test_empty_llrs(self):
+        """Empty LLR list → 0.0."""
+        cluster_llr = compute_cluster_llr([], n_eff=1.0)
+        assert cluster_llr == 0.0
+
+
+# ---------------------------------------------------------------------------
+# Test: Grouping by correct key dimensions (cluster_id)
+# Requirement: 4.1
+# ---------------------------------------------------------------------------
+
+
+class TestClusterEvidence:
+    """cluster_evidence groups EvidenceUnits by cluster_id."""
+
+    def test_grouping_by_cluster_id(self):
+        """Units with same cluster_id are grouped together."""
+        unit_a1 = _make_unit(cluster_id="cluster_a")
+        unit_a2 = _make_unit(cluster_id="cluster_a")
+        unit_b1 = _make_unit(cluster_id="cluster_b")
+
+        units = [unit_a1, unit_a2, unit_b1]
+        llrs = [1.0, 0.5, -0.3]
+
+        clusters = cluster_evidence(units, llrs)
+
+        assert len(clusters) == 2
+
+        # Find clusters by id
+        cluster_map = {c.cluster_id: c for c in clusters}
+
+        assert "cluster_a" in cluster_map
+        assert "cluster_b" in cluster_map
+
+        # Cluster A has 2 units
+        assert len(cluster_map["cluster_a"].units) == 2
+        assert cluster_map["cluster_a"].llrs == [1.0, 0.5]
+
+        # Cluster B has 1 unit
+        assert len(cluster_map["cluster_b"].units) == 1
+        assert cluster_map["cluster_b"].llrs == [-0.3]
+
+    def test_single_cluster(self):
+        """All units with same cluster_id → one cluster."""
+        units = [_make_unit(cluster_id="only") for _ in range(4)]
+        llrs = [0.1, 0.2, 0.3, 0.4]
+
+        clusters = cluster_evidence(units, llrs)
+
+        assert len(clusters) == 1
+        assert clusters[0].cluster_id == "only"
+        assert len(clusters[0].units) == 4
+        assert clusters[0].llrs == [0.1, 0.2, 0.3, 0.4]
+
+    def test_each_unit_different_cluster(self):
+        """Each unit in its own cluster → N clusters."""
+        units = [_make_unit(cluster_id=f"c_{i}") for i in range(5)]
+        llrs = [0.1 * i for i in range(5)]
+
+        clusters = cluster_evidence(units, llrs)
+
+        assert len(clusters) == 5
+        for c in clusters:
+            assert len(c.units) == 1
+
+    def test_empty_input(self):
+        """No units → no clusters."""
+        clusters = cluster_evidence([], [])
+        assert clusters == []
+
+    def test_llrs_parallel_to_units(self):
+        """LLRs are correctly associated with their units."""
+        unit_x = _make_unit(cluster_id="x")
+        unit_y = _make_unit(cluster_id="y")
+        unit_x2 = _make_unit(cluster_id="x")
+
+        units = [unit_x, unit_y, unit_x2]
+        llrs = [1.5, -0.7, 2.3]
+
+        clusters = cluster_evidence(units, llrs)
+        cluster_map = {c.cluster_id: c for c in clusters}
+
+        assert cluster_map["x"].llrs == [1.5, 2.3]
+        assert cluster_map["y"].llrs == [-0.7]