analyze_tom_unlearning.py

"""Paired-selectivity analysis for the Comma-2T ToMBench unlearning experiment.

Reads the per-cell ToMBench eval JSONs written by
``scripts/unlearning/eval_tombench_unlearn.py`` (one per adapter), pairs the
influence-targeted arm (expA) with the size-matched random control (exp1) by
(topic, seed) on the ToMBench target, and tests the one-sided hypothesis
d = gamma_influence - gamma_random > 0 (influence targeting damages ToMBench
more than the random control). Mirrors the core-4 paired-selectivity analysis
(findings/comma_2t/analyze_paired.py) so the ToMBench row is directly comparable
to the SocialIQA d_z +0.49 headline.

Usage:
    python findings/comma_2t/analyze_tom_unlearning.py --eval-dir <dir>
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd
from scipy import stats


def load_cells(eval_dir: Path) -> pd.DataFrame:
    rows = []
    for p in sorted(eval_dir.glob("*_tombench_eval.json")):
        d = json.loads(p.read_text(encoding="utf-8"))
        adapter = str(d.get("adapter_dir", ""))
        if "expA" in adapter or "expA" in p.name:
            cond = "influence"
        elif "exp1" in adapter or "exp1" in p.name:
            cond = "random"
        else:
            continue
        rows.append(
            {
                "condition": cond,
                "topic": d.get("topic_bin"),
                "seed": str(d.get("seed")),
                "acc": d.get("tombench_acc"),
                "gamma": d.get("gamma"),
                "file": p.name,
            }
        )
    return pd.DataFrame(rows)


def paired(df: pd.DataFrame) -> pd.DataFrame:
    inf = df[df.condition == "influence"].set_index(["topic", "seed"])["gamma"]
    rnd = df[df.condition == "random"].set_index(["topic", "seed"])["gamma"]
    pairs = pd.DataFrame({"gamma_influence": inf, "gamma_random": rnd}).dropna()
    pairs["d"] = pairs["gamma_influence"] - pairs["gamma_random"]
    return pairs.reset_index()


def summarize(pairs: pd.DataFrame) -> dict:
    d = pairs["d"].to_numpy()
    n = len(d)
    out: dict = {"n_pairs": n, "median_d": float(np.median(d)) if n else None}
    if n:
        out["mean_d"] = float(np.mean(d))
        out["d_z"] = (
            float(np.mean(d) / np.std(d, ddof=1))
            if n > 1 and np.std(d, ddof=1) > 0
            else None
        )
        n_pos = int((d > 0).sum())
        out["sign_pos"] = f"{n_pos}/{n}"
        # exact binomial sign test, one-sided d>0
        out["sign_p_onesided"] = float(
            stats.binomtest(n_pos, n, 0.5, alternative="greater").pvalue
        )
        if n >= 6:
            try:
                w = stats.wilcoxon(d, alternative="greater")
                out["wilcoxon_p_onesided"] = float(w.pvalue)
            except ValueError as exc:
                out["wilcoxon_p_onesided"] = f"n/a ({exc})"
    return out


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--eval-dir", type=Path, required=True)
    ap.add_argument("--out-csv", type=Path, default=None)
    args = ap.parse_args()

    df = load_cells(args.eval_dir)
    if df.empty:
        print(f"no eval JSONs under {args.eval_dir}")
        return 1

    print("=== per-cell ToMBench gamma ===")
    print(df.sort_values(["condition", "topic", "seed"]).to_string(index=False))

    pairs = paired(df)
    print("\n=== paired (influence - random) by (topic, seed) ===")
    print(pairs.to_string(index=False))

    print("\n=== per-topic gamma means (influence arm) ===")
    inf = df[df.condition == "influence"]
    print(inf.groupby("topic")["gamma"].agg(["mean", "count"]).to_string())

    print(
        "\n=== paired-selectivity summary (d = gamma_influence - gamma_random > 0) ==="
    )
    summary = summarize(pairs)
    for k, v in summary.items():
        print(f"  {k}: {v}")

    if args.out_csv:
        pairs.to_csv(args.out_csv, index=False)
        print(f"\nwrote {args.out_csv}")
    return 0


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