social-data-attribution-w2/scripts/analysis/tombench_sensitivity.py

"""Sensitivity checks for the held-out ToMBench social_life result."""

from __future__ import annotations

import argparse
import json
import math
import statistics as st
from pathlib import Path

PRAGMATIC = [
    "tombench_hinting_task_test",
    "tombench_faux_pas_recognition_test",
    "tombench_strange_story_task",
]
TARGET = "social_life"
TCRIT = {2: 4.303, 5: 2.571, 8: 2.306, 14: 2.145}


def vals(cells: dict, probes: list[str], topic: str) -> list[float]:
    out: list[float] = []
    for probe in probes:
        out += cells[probe].get(topic, {}).get("vals", [])
    return out


def stats(values: list[float]) -> tuple[int, float, float, bool, int, float]:
    n = len(values)
    mean = st.mean(values)
    sd = st.stdev(values) if n > 1 else 0.0
    tcrit = TCRIT.get(n - 1)
    ci95 = tcrit * sd / math.sqrt(n) if tcrit else float("nan")
    neg = sum(1 for value in values if value < 0)
    k = min(neg, n - neg)
    p_value = min(2 * sum(math.comb(n, idx) for idx in range(k + 1)) * 0.5**n, 1.0)
    excludes_zero = (mean + ci95) < 0 or (mean - ci95) > 0
    return n, mean, ci95, excludes_zero, neg, p_value


def print_metric_robustness(gamma_ci: dict) -> None:
    print("(1) METRIC ROBUSTNESS: social_life x social-pragmatic, net, absolute")
    for metric in ["acc_per_char", "acc_raw", "acc_uncond", "primary_score"]:
        cells = gamma_ci["results"][metric]["absolute"]["net"]
        n, mean, ci95, excludes_zero, neg, p_value = stats(
            vals(cells, PRAGMATIC, TARGET)
        )
        status = "EXCL0" if excludes_zero else "incl0"
        print(
            f"   {metric:14s} mean={mean:+.3f} "
            f"95%CI=[{mean - ci95:+.3f},{mean + ci95:+.3f}] "
            f"{status} sign {neg}/{n} p={p_value:.3f}"
        )


def print_null_correction(gamma_ci: dict) -> None:
    print("\n(2) NULL-CORRECTION: raw vs net, acc_uncond, pragmatic")
    for kind in ["raw", "net"]:
        cells = gamma_ci["results"]["acc_uncond"]["absolute"][kind]
        n, mean, ci95, excludes_zero, _, _ = stats(vals(cells, PRAGMATIC, TARGET))
        status = "EXCL0" if excludes_zero else "incl0"
        print(
            f"   {kind:4s} mean={mean:+.3f} 95%CI=[{mean - ci95:+.3f},{mean + ci95:+.3f}] {status}"
        )


def print_seed_ranks(gamma_ci: dict) -> None:
    print("\n(3) PER-SEED: social_life pragmatic rank among 24 topics, acc_uncond net")
    cells = gamma_ci["results"]["acc_uncond"]["absolute"]["net"]
    topics = list(cells[PRAGMATIC[0]].keys())
    for seed_idx, seed in enumerate(gamma_ci["seeds"]):
        per_seed_mean = {}
        for topic in topics:
            topic_vals = [
                cells[probe][topic]["vals"][seed_idx]
                for probe in PRAGMATIC
                if len(cells[probe][topic].get("vals", [])) > seed_idx
            ]
            if len(topic_vals) == len(PRAGMATIC):
                per_seed_mean[topic] = st.mean(topic_vals)
        ranked = sorted(per_seed_mean, key=lambda topic: per_seed_mean[topic])
        rank = ranked.index(TARGET) + 1
        print(
            f"   seed {seed}: social_life rank {rank}/{len(ranked)}  (mean {per_seed_mean[TARGET]:+.3f})"
        )


def print_jackknife(gamma_ci: dict) -> None:
    print("\n(4) JACKKNIFE: drop one pragmatic subtask, acc_uncond net")
    cells = gamma_ci["results"]["acc_uncond"]["absolute"]["net"]
    for drop_probe in PRAGMATIC:
        keep = [probe for probe in PRAGMATIC if probe != drop_probe]
        n, mean, ci95, excludes_zero, _, _ = stats(vals(cells, keep, TARGET))
        status = "EXCL0" if excludes_zero else "incl0"
        label = drop_probe.replace("tombench_", "")
        print(
            f"   drop {label:26s} mean={mean:+.3f} 95%CI=[{mean - ci95:+.3f},{mean + ci95:+.3f}] {status}"
        )


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gamma-ci", type=Path, required=True)
    args = parser.parse_args(argv)

    gamma_ci = json.loads(args.gamma_ci.read_text())
    print_metric_robustness(gamma_ci)
    print_null_correction(gamma_ci)
    print_seed_ranks(gamma_ci)
    print_jackknife(gamma_ci)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())