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

"""Aggregate per-seed γ grids into per-cell mean +/- std / 95% CI.

Reuses build_figure4_table.py (same dir) to compute per-(probe,topic) γ for each seed's
eval grid against the shared (seed-independent) baselines, then stacks seeds {42,43,44}.

Parity with the paper's random-baseline Single-Bin experiment:
  - reports the net causal impact net_γ = γ_topic - γ_null, where γ_null is that
    same seed's null-bin control checkpoint (random cross-topic forget set),
    read from <eval_grid>/__null__/. The raw (un-corrected) γ is reported alongside.
  - reports γ in both modes: 'relative' = (post-pre)/pre (the paper's definition),
    and 'absolute' = post-pre (robust where the relative ratio is undefined on
    floor/saturated ToM baselines).

Usage:
  python aggregate_seeds_ci.py \
    --baselines /path/to/parallel_5shot/allenai-Olmo-3-1025-7B \
    --seed-grid 42:/path/to/n10b/eval_grid \
    --seed-grid 43:/path/to/n10b/eval_grid_seed43 \
    --seed-grid 44:/path/to/n10b/eval_grid_seed44 \
    --out-dir /path/to/n10b/figure4_ci
"""

from __future__ import annotations

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

import build_figure4_table as bf

METRICS = ["acc_uncond", "primary_score"]
MODES = ["relative", "absolute"]
TCRIT = {2: 12.706, 3: 4.303, 4: 3.182, 5: 2.776}  # two-sided t, df=n-1


def _stack(vals_with_none: list[float | None]) -> dict:
    """Per-cell across-seed summary: mean, sample std, 95% CI half-width, stability flag."""
    vals = [v for v in vals_with_none if v is not None]
    n = len(vals)
    mean = st.mean(vals) if n >= 1 else None
    sd = st.stdev(vals) if n >= 2 else None
    tcrit = TCRIT.get(n)
    ci95 = (tcrit * sd / math.sqrt(n)) if (sd is not None and tcrit) else None
    stable = (
        n >= 2
        and sd is not None
        and mean is not None
        and abs(mean) > sd
        and all((v > 0) == (mean > 0) for v in vals)
    )
    return {
        "n": n,
        "mean": mean,
        "std": sd,
        "ci95": ci95,
        "vals": vals,
        "stable": bool(stable),
    }


def main(argv: list[str] | None = None) -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--baselines", type=Path, required=True)
    p.add_argument(
        "--seed-grid",
        action="append",
        required=True,
        help="SEED:/path/to/eval_grid (repeatable). Each grid may hold a __null__ row.",
    )
    p.add_argument("--out-dir", type=Path, required=True)
    p.add_argument("--metrics", nargs="*", default=METRICS)
    p.add_argument("--modes", nargs="*", default=MODES)
    args = p.parse_args(argv)
    args.out_dir.mkdir(parents=True, exist_ok=True)

    seeds = [(s, Path(d)) for s, d in (sg.split(":", 1) for sg in args.seed_grid)]
    baselines = bf.collect_baselines(args.baselines)

    full = {
        "seeds": [s for s, _ in seeds],
        "definition": "net_gamma = gamma_topic - gamma_null (paper's net causal impact); "
        "modes: relative=(post-pre)/pre, absolute=post-pre; raw = uncorrected gamma_topic",
        "results": {},
    }
    # null coverage diagnostic: how many seeds actually have a null row populated
    null_seeds = {
        s: any(bf.collect_null(d).get(pr, {}) for pr, _ in bf.PROBES) for s, d in seeds
    }

    for metric in args.metrics:
        full["results"][metric] = {}
        for mode in args.modes:
            per_seed = {}
            for s, d in seeds:
                post = bf.collect_post(d)
                null_post = bf.collect_null(d)
                per_seed[s] = bf.compute_gamma(
                    baselines, post, metric, mode=mode, null_post=null_post
                )
            kinds = {}
            for kind, key in (("raw", "gamma"), ("net", "net_gamma")):
                cells = {
                    pr: {
                        t: _stack([per_seed[s][pr][t].get(key) for s, _ in seeds])
                        for t in bf.ALL_TOPICS
                    }
                    for pr, _ in bf.PROBES
                }
                kinds[kind] = cells
            full["results"][metric][mode] = kinds

    (args.out_dir / "gamma_ci.json").write_text(json.dumps(full, indent=2, default=str))

    # ---- human-readable summary (net γ is the headline; relative is the paper-parity primary) ----
    n_grid = len(bf.PROBES) * len(bf.ALL_TOPICS)
    lines = [
        f"# Net-corrected γ CIs over seeds {full['seeds']}",
        "",
        "null-bin row present per seed: "
        + ", ".join(f"{s}={'yes' if null_seeds[s] else 'NO'}" for s, _ in seeds),
        "",
    ]
    for metric in args.metrics:
        for mode in args.modes:
            cells = full["results"][metric][mode]["net"]
            n_cells = sum(
                1
                for pr, _ in bf.PROBES
                for t in bf.ALL_TOPICS
                if cells[pr][t]["mean"] is not None
            )
            n_fullseed = sum(
                1
                for pr, _ in bf.PROBES
                for t in bf.ALL_TOPICS
                if cells[pr][t]["n"] == len(seeds)
            )
            n_stable = sum(
                1
                for pr, _ in bf.PROBES
                for t in bf.ALL_TOPICS
                if cells[pr][t]["stable"]
            )
            lines += [
                f"## net γ — metric={metric}, mode={mode}",
                f"- cells with >=1 seed: {n_cells}/{n_grid}",
                f"- cells with all {len(seeds)} seeds: {n_fullseed}",
                f"- cells 'stable' (|mean|>1std, consistent sign): {n_stable}",
                "  most-degrading topics (mean net γ over probes +/- std-of-cell-means):",
            ]
            rows = []
            for t in bf.ALL_TOPICS:
                ms = [
                    cells[pr][t]["mean"]
                    for pr, _ in bf.PROBES
                    if cells[pr][t]["mean"] is not None
                ]
                if ms:
                    rows.append((t, st.mean(ms), st.stdev(ms) if len(ms) > 1 else 0.0))
            for t, m, sd in sorted(rows, key=lambda x: x[1])[:6]:
                lines.append(f"    {t:32s} {m:+.4f} +/- {sd:.4f}")
            lines.append("")
    (args.out_dir / "gamma_ci_summary.md").write_text("\n".join(lines))

    print(
        f"seeds={full['seeds']}  null_rows="
        + ",".join(f"{s}:{'y' if null_seeds[s] else 'N'}" for s, _ in seeds)
    )
    for metric in args.metrics:
        for mode in args.modes:
            cells = full["results"][metric][mode]["net"]
            n_fullseed = sum(
                1
                for pr, _ in bf.PROBES
                for t in bf.ALL_TOPICS
                if cells[pr][t]["n"] == len(seeds)
            )
            n_stable = sum(
                1
                for pr, _ in bf.PROBES
                for t in bf.ALL_TOPICS
                if cells[pr][t]["stable"]
            )
            print(
                f"  net γ {metric}/{mode}: all-seed cells={n_fullseed}  stable={n_stable}"
            )
    print(f"wrote {args.out_dir / 'gamma_ci.json'} and gamma_ci_summary.md")


if __name__ == "__main__":
    main()