tests/test_stats.py

"""Tests for owmi.stats (Appendix B estimators): dose-response, ICC, power.

This module was written but never exercised by any test before this file.
The tests below are genuine planted-signal recovery checks (not just smoke
tests): synthetic data with a KNOWN ground truth is generated, the estimator
is run on it, and the recovered estimate is checked against the true value
within a tolerance appropriate to the sampling noise of the design. Every
random draw uses a fixed ``numpy.random.Generator`` seed and every fitting
routine here is itself deterministic given its own ``seed`` argument, so
these tests are exactly reproducible.
"""
import math
import unittest
from pathlib import Path

import numpy as np

from owmi.stats import (
    HAVE_STATSMODELS,
    extract_dose_rows,
    extract_score_rows,
    fit_dose_response,
    load_pilot_artifacts,
    pilot_variance_summary,
    power_recommendation,
    variance_components,
)
from owmi.stats import _simulated_ci_width  # private, needed for the monotonicity check

PILOT_DIR = Path(
    "/Users/emiliofe/Dropbox/-CODEX-/claude-developer/owmi-paper/pilot_results/run_2026-08-08"
)


def _expit(x: float) -> float:
    return 1.0 / (1.0 + math.exp(-x))


def _planted_dose_rows(rng, sham_prob, beta0, slope, dose_grid, n_items):
    """Simulate paired (intervention, sham) rows from a known dose-response DGP.

    ``logit P(y=1 | dbar=d) = beta0 + slope * d`` for intervention rows (so
    the true zero-divergence margin is ``expit(beta0) - sham_prob`` exactly,
    since the ramp basis used by the fitter is 0 at d=0 by construction);
    sham rows are independent Bernoulli(sham_prob) draws with dbar pinned to
    0. Returns the rows plus the true per-row intervention probabilities
    (needed by the caller to compute the exact aggregate margin).
    """
    rows = []
    true_int_p = []
    for i in range(n_items):
        d = dose_grid[i % len(dose_grid)]
        p = _expit(beta0 + slope * d)
        true_int_p.append(p)
        rows.append({
            "condition": "intervention", "y": int(rng.random() < p),
            "dbar": d, "item": f"item{i}", "site": "s0",
        })
        rows.append({
            "condition": "sham", "y": int(rng.random() < sham_prob),
            "dbar": 0.0, "item": f"item{i}", "site": "s0",
        })
    return rows, true_int_p


class DoseResponseRecoveryTests(unittest.TestCase):
    def test_recovers_planted_positive_m0_and_lambda(self):
        rng = np.random.default_rng(42)
        sham_prob = 0.10
        beta0 = math.log(0.25 / 0.75)  # expit(beta0) = 0.25 -> true m0 = 0.15
        slope = 2.5
        dose_grid = [0.25, 0.5, 0.75, 1.0]
        rows, true_int_p = _planted_dose_rows(
            rng, sham_prob, beta0, slope, dose_grid, n_items=400
        )

        m0_true = _expit(beta0) - sham_prob
        agg_true = float(np.mean(true_int_p)) - sham_prob
        lambda_true = 1.0 - m0_true / agg_true
        self.assertGreater(m0_true, 0.0)
        self.assertGreater(lambda_true, 0.0)

        result = fit_dose_response(rows, n_boot=100, seed=3)
        self.assertEqual(result["n_intervention"], 400)
        self.assertEqual(result["n_sham"], 400)
        self.assertAlmostEqual(result["m0"]["estimate"], m0_true, delta=0.08)
        self.assertAlmostEqual(result["lambda"]["truncated"], lambda_true, delta=0.15)
        # The point estimate should fall inside its own bootstrap interval.
        self.assertLessEqual(result["m0"]["ci_low"], result["m0"]["estimate"])
        self.assertGreaterEqual(result["m0"]["ci_high"], result["m0"]["estimate"])
        # The fitted curve g must be non-decreasing and g(0) == 0, as designed.
        curve = [v for _, v in result["g"]["curve"]]
        self.assertAlmostEqual(curve[0], 0.0, places=6)
        self.assertTrue(all(b >= a - 1e-9 for a, b in zip(curve, curve[1:])))

    def test_null_data_recovers_zero_m0_and_no_spurious_slope(self):
        rng = np.random.default_rng(11)
        sham_prob = 0.20
        beta0 = math.log(sham_prob / (1.0 - sham_prob))  # expit(beta0) == sham_prob
        dose_grid = [0.25, 0.5, 0.75, 1.0]
        rows, _ = _planted_dose_rows(
            rng, sham_prob, beta0, slope=0.0, dose_grid=dose_grid, n_items=400
        )

        result = fit_dose_response(rows, n_boot=100, seed=3)
        self.assertLess(abs(result["m0"]["estimate"]), 0.08)
        # A true-null CI should straddle zero.
        self.assertLessEqual(result["m0"]["ci_low"], 0.0)
        self.assertGreaterEqual(result["m0"]["ci_high"], 0.0)
        # No spurious dose-response: the fitted curve stays close to flat.
        curve_values = [v for _, v in result["g"]["curve"]]
        self.assertLess(max(curve_values), 0.15)

    def test_extract_dose_rows_matches_flat_row_contract(self):
        """extract_dose_rows must produce the same shape fit_dose_response consumes
        directly, whether fed raw result.json-style artifacts or flat rows."""
        artifact = {
            "condition": "intervention",
            "probe_intervention_score": 0.8,
            "output_divergence": {"mean_js": 0.4},
            "benchmark_example": {"item_id": "it1"},
            "object": {"object_id": "site1"},
        }
        rows = extract_dose_rows([artifact])
        self.assertEqual(len(rows), 1)
        row = rows[0]
        self.assertEqual(row["condition"], "intervention")
        self.assertEqual(row["y"], 1)
        self.assertAlmostEqual(row["dbar"], 0.4)
        self.assertEqual(row["item"], "it1")
        self.assertEqual(row["site"], "site1")


class VarianceComponentsRecoveryTests(unittest.TestCase):
    def test_recovers_known_item_variance_and_nulls_flat_factors(self):
        rng = np.random.default_rng(123)
        n_domains, n_objects, n_items, n_seeds = 2, 2, 30, 3
        sigma_item = 0.25
        sigma_resid = 0.10
        grand_mean = 0.5
        item_effects = rng.normal(0.0, sigma_item, size=n_items)

        rows = []
        for di in range(n_domains):
            for oi in range(n_objects):
                for ii in range(n_items):
                    for si in range(n_seeds):
                        score = grand_mean + item_effects[ii] + rng.normal(0.0, sigma_resid)
                        rows.append({
                            "score": score, "domain": f"d{di}", "object": f"o{oi}",
                            "item": f"i{ii}", "seed": f"s{si}",
                        })

        result = variance_components(rows, n_boot=200, seed=5)
        true_item_var = sigma_item ** 2
        true_resid_var = sigma_resid ** 2

        item_est = result["components"]["item"]["estimate"]
        resid_est = result["components"]["residual"]["estimate"]
        # Item variance recovered to the right order of magnitude (crossed
        # MoM/MixedLM estimators of a variance component from ~30 levels
        # carry real sampling noise; a tight point match is not realistic).
        self.assertGreater(item_est, 0.5 * true_item_var)
        self.assertLess(item_est, 1.5 * true_item_var)
        self.assertAlmostEqual(resid_est, true_resid_var, delta=0.4 * true_resid_var)
        # domain/object/seed carry no planted effect and must stay small
        # relative to the item component that does.
        for null_factor in ("domain", "object", "seed"):
            self.assertLess(result["components"][null_factor]["estimate"], 0.15 * true_item_var)
            self.assertLess(result["icc"][null_factor]["estimate"], 0.1)
        self.assertGreater(result["icc"]["item"]["estimate"], 0.5)

    @unittest.skipUnless(HAVE_STATSMODELS, "regression test targets the statsmodels backend")
    def test_singleton_level_factor_reports_zero_not_degenerate(self):
        """Regression test for the 2026-08-08 bug in _statsmodels_components.

        A factor observed at only one level (e.g. a pilot that only probed
        one intervention site, as the real pilot_results/run_2026-08-08 data
        does for "object") has no between-level variance to estimate. The
        vc_formula previously included it anyway, colliding with the fixed
        intercept; on real pilot data this made the singleton factor's
        "variance component" come out numerically identical to the residual
        scale. Both must be exactly 0.0 now, matching the numpy MoM backend's
        convention (see owmi/stats.py::_statsmodels_components).
        """
        rng = np.random.default_rng(7)
        rows = []
        for ii in range(20):
            for si in range(3):
                score = 0.5 + rng.normal(0.0, 0.1)
                rows.append({
                    "score": score, "domain": "only_domain", "object": "only_object",
                    "item": f"i{ii}", "seed": f"s{si}",
                })
        result = variance_components(rows, n_boot=20, seed=1, backend="statsmodels")
        self.assertEqual(result["backend"], "statsmodels_mixedlm_crossed")
        self.assertEqual(result["components"]["domain"]["estimate"], 0.0)
        self.assertEqual(result["components"]["object"]["estimate"], 0.0)
        self.assertGreater(result["components"]["residual"]["estimate"], 0.001)

    def test_extract_score_rows_prefers_object_id_over_raw_object_mapping(self):
        artifact = {
            "condition": "intervention",
            "probe_intervention_score": 0.5,
            "benchmark": {"name": "toy"},
            "benchmark_example": {"item_id": "it1"},
            "object": {"object_id": "resid_l16", "kind": "residual_stream"},
            "seed": 7,
        }
        rows = extract_score_rows([artifact])
        self.assertEqual(len(rows), 1)
        self.assertEqual(rows[0]["object"], "resid_l16")


class PowerSimulationMonotonicityTests(unittest.TestCase):
    def test_ci_width_is_monotonically_decreasing_in_n(self):
        rng = np.random.default_rng(9)
        hit_rate, hit_tau2 = 0.6, 0.02
        fa_rate, fa_tau2 = 0.2, 0.02
        grid = (8, 32, 128, 512)
        widths = [
            _simulated_ci_width(n, 2, hit_rate, hit_tau2, fa_rate, fa_tau2, 300, 0.05, rng)
            for n in grid
        ]
        for w in widths:
            self.assertGreater(w, 0.0)
        self.assertTrue(
            all(a >= b for a, b in zip(widths, widths[1:])),
            f"expected non-increasing CI width as n grows, got {widths}",
        )

    def test_power_recommendation_curve_is_monotone_and_recommends_larger_n_for_tighter_target(self):
        loose = power_recommendation(
            PILOT_DIR, target_width=2.0, n_grid=(8, 16, 32, 64), n_sim=80, seed=1,
        )
        tight = power_recommendation(
            PILOT_DIR, target_width=0.3, n_grid=(8, 16, 32, 64, 128, 256), n_sim=80, seed=1,
        )
        for result in (loose, tight):
            widths = [entry["expected_ci_width_monotone"] for entry in result["curve"]]
            self.assertTrue(all(a >= b for a, b in zip(widths, widths[1:])))
        if loose["recommended_n_per_cell"] is not None and tight["recommended_n_per_cell"] is not None:
            self.assertGreaterEqual(tight["recommended_n_per_cell"], loose["recommended_n_per_cell"])


class PilotDirectoryEndToEndTests(unittest.TestCase):
    """Runnable entry-point tests against the real pilot data on disk.

    Reads /Users/emiliofe/Dropbox/-CODEX-/claude-developer/owmi-paper/pilot_results/run_2026-08-08,
    which as of 2026-08-08 holds 192 result.json artifacts (single-model
    pilot: Qwen2.5-7B-Instruct and Mistral-7B-Instruct-v0.3, one probed site
    "resid_l16", 3 benchmarks x up to 2 items each, 2 seeds, 4 conditions).
    """

    @classmethod
    def setUpClass(cls):
        if not PILOT_DIR.exists():
            raise unittest.SkipTest(f"pilot directory not found: {PILOT_DIR}")

    def test_load_pilot_artifacts_reads_every_result_json(self):
        artifacts = load_pilot_artifacts(PILOT_DIR)
        expected = len(list(PILOT_DIR.rglob("result.json")))
        self.assertEqual(len(artifacts), expected)
        self.assertGreater(len(artifacts), 0)

    def test_pilot_variance_summary_runs_end_to_end(self):
        artifacts = load_pilot_artifacts(PILOT_DIR)
        summary = pilot_variance_summary(artifacts)
        self.assertGreater(summary["n_runs"], 0)
        self.assertGreater(summary["n_item_cells"], 0)
        self.assertTrue(0.0 <= summary["hit_rate"] <= 1.0)
        self.assertTrue(0.0 <= summary["false_alarm_rate"] <= 1.0)
        self.assertGreaterEqual(summary["hit_item_variance"], 0.0)
        self.assertGreaterEqual(summary["false_alarm_item_variance"], 0.0)

    def test_fit_dose_response_runs_end_to_end_on_real_pilot_data(self):
        artifacts = load_pilot_artifacts(PILOT_DIR)
        result = fit_dose_response(artifacts, n_boot=50)
        self.assertGreater(result["n_intervention"], 0)
        self.assertGreater(result["n_sham"], 0)
        self.assertTrue(math.isfinite(result["m0"]["estimate"]))
        import json
        json.dumps(result, allow_nan=True)  # must not raise structurally

    def test_variance_components_runs_end_to_end_on_real_pilot_data(self):
        artifacts = load_pilot_artifacts(PILOT_DIR)
        rows = extract_score_rows(artifacts)
        self.assertGreater(len(rows), 0)
        result = variance_components(rows, n_boot=30)
        for factor in ("domain", "object", "item", "seed", "residual"):
            self.assertIn(factor, result["components"])
            self.assertGreaterEqual(result["components"][factor]["estimate"], 0.0)
        # The pilot only ever probed one site/object, so it must not be
        # mistaken for a source of variance (see the singleton-factor
        # regression test above).
        self.assertEqual(result["components"]["object"]["estimate"], 0.0)

    def test_power_recommendation_runs_end_to_end_on_real_pilot_data(self):
        result = power_recommendation(PILOT_DIR, n_sim=60, n_grid=(8, 16, 32, 64))
        self.assertIn("pilot", result)
        self.assertIn("curve", result)
        self.assertGreater(len(result["curve"]), 0)
        for entry in result["curve"]:
            self.assertIn("n_items", entry)
            self.assertIn("expected_ci_width", entry)
            self.assertGreater(entry["expected_ci_width"], 0.0)


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