social-data-attribution-w2/scripts/analysis/multiseed/cross_method_comparison.py

#!/usr/bin/env python3
# pyright: reportAttributeAccessIssue=false
"""Cross-method selectivity comparison for the COLM #3137 rebuttal.

Compares per-target pilot-validated selectivity under two methods:
  (1) bin-level seed-1 (corrected OLMES gamma matrix: bin-level z-scored forget-sets)
  (2) faithful 3-seed (per-doc influence top-200 forget-sets)

Both use the uniform 4-target pool for the unique-z shortlist (transportation /
politics / finance for MMLU-SS, sports / food / health for ARC, etc.). Reports
sel(expA pilot pick), sel(expC), ratio, and gamma-on-target so the reader can
diagnose where divergence reflects real method differences vs near-zero-gamma noise.

Outputs: results/cross_method_comparison.csv (tidy) + a printed markdown summary.
"""

from pathlib import Path
import numpy as np
import pandas as pd

HOME = Path.home()
TIDY_BIN = HOME / "scratch/n16_selectivity/results/gamma_olmes_tidy.csv"
TIDY_FAITHFUL = HOME / "scratch/n16_selectivity/results/faithful_gamma_tidy.csv"
ZS = HOME / "dev/data-attribution/artifacts/zscored_bin_scores/aggregated"
OUT_CSV = HOME / "scratch/n16_selectivity/results/cross_method_comparison.csv"

BENCHES = ["socialiqa", "mmlu_social_science", "mmlu_stem", "arc_challenge"]
DISPLAY = {
    "socialiqa": "SocialIQA",
    "mmlu_social_science": "MMLU-SS",
    "mmlu_stem": "MMLU-STEM",
    "arc_challenge": "ARC-C",
}


def sel(d, t):
    o = [abs(d.get(b, 0.0)) for b in BENCHES if b != t]
    m = float(np.mean(o)) if o else 0.0
    return abs(d.get(t, 0.0)) / m if m > 0 else float("nan")


def unique_z_top3(target):
    def tm(b):
        return pd.read_csv(ZS / f"zscored_{b}.csv").groupby("topic_label").zscore.mean()

    t = tm(target)
    others = pd.concat([tm(b) for b in BENCHES if b != target], axis=1).max(axis=1)
    return (t - others).sort_values(ascending=False).head(3).index.tolist()


def bin_level_pilot(target, shortlist):
    """Bin-level seed-1: find shortlist topic with best sel(expA) and report cells."""
    df = pd.read_csv(TIDY_BIN)
    df = df[df.eval_family == "primary"]
    expA = {}
    for k, g in df[df.condition == "expA"].groupby("target_or_topic"):
        if "__" not in str(k):
            continue
        topic, tgt = str(k).rsplit("__", 1)
        expA.setdefault(tgt, {})[topic] = dict(zip(g.eval_benchmark, g.gamma))
    eC = df[(df.condition == "expC") & (df.target_or_topic == target)]
    expC_d = dict(zip(eC.eval_benchmark, eC.gamma))
    picks = []
    for top in shortlist:
        if top not in expA.get(target, {}):
            continue
        s = sel(expA[target][top], target)
        picks.append((top, s, expA[target][top]))
    if not picks:
        return None
    best_topic, best_sel, best_cell = max(picks, key=lambda x: x[1])
    sC = sel(expC_d, target)
    return dict(
        method="bin_seed1",
        target=DISPLAY[target],
        pick=best_topic,
        sel_expA=round(best_sel, 3),
        sel_expC=round(sC, 3),
        ratio=f"{best_sel / sC:.2f}" if sC > 0 else "n/a",
        gamma_on_target=round(best_cell.get(target, 0), 5),
        n_seeds=1,
    )


def faithful_pilot(target, shortlist):
    """Faithful 3-seed: per-seed best-of-shortlist, then mean+/-std across seeds."""
    df = pd.read_csv(TIDY_FAITHFUL)
    ratios, sels, gtargs = [], [], []
    picks = set()
    for seed in sorted(df.seed.unique()):
        eC = df[(df.condition == "expC") & (df.target == target) & (df.seed == seed)]
        if eC.empty:
            continue
        eC_d = dict(zip(eC.eval_benchmark, eC.gamma))
        sC = sel(eC_d, target)
        if not sC or sC == 0:
            continue
        cands = []
        for top in shortlist:
            cell = df[
                (df.condition == "expA")
                & (df.target == target)
                & (df.topic == top)
                & (df.seed == seed)
            ]
            if cell.empty:
                continue
            d = dict(zip(cell.eval_benchmark, cell.gamma))
            cands.append((top, sel(d, target), d.get(target, 0)))
        if not cands:
            continue
        best = max(cands, key=lambda x: x[1])
        picks.add(best[0])
        sels.append(best[1])
        ratios.append(best[1] / sC)
        gtargs.append(best[2])
    if not ratios:
        return None
    a, g = np.array(ratios), np.array(gtargs)
    s = np.array(sels)
    return dict(
        method="faithful_3seed",
        target=DISPLAY[target],
        pick=",".join(sorted(picks)),
        sel_expA=f"{s.mean():.2f}+/-{s.std():.2f}",
        sel_expC=f"{(s / a).mean():.2f}+/-{(s / a).std():.2f}",
        ratio=f"{a.mean():.2f}+/-{a.std():.2f}",
        gamma_on_target=f"{g.mean():.4f}+/-{g.std():.4f}",
        n_seeds=len(a),
    )


def main():
    rows = []
    print("# Cross-method selectivity comparison (uniform 4-target pool pilot)\n")
    print(
        "| Target | Method | Pick | sel(expA) | sel(expC) | Ratio | gamma-on-target | n_seeds |"
    )
    print("|---|---|---|---|---|---|---|---|")
    for t in BENCHES:
        sh = unique_z_top3(t)
        for fn in (bin_level_pilot, faithful_pilot):
            r = fn(t, sh)
            if r is None:
                continue
            rows.append(r)
            print(
                f"| {r['target']} | {r['method']} | {r['pick']} | "
                f"{r['sel_expA']} | {r['sel_expC']} | {r['ratio']} | "
                f"{r['gamma_on_target']} | {r['n_seeds']} |"
            )
    pd.DataFrame(rows).to_csv(OUT_CSV, index=False)
    print(f"\nwrote {OUT_CSV}")

    # Verdict per target
    print("\n## Verdict per target\n")
    for t in BENCHES:
        b = next(
            (
                r
                for r in rows
                if r["target"] == DISPLAY[t] and r["method"] == "bin_seed1"
            ),
            None,
        )
        f = next(
            (
                r
                for r in rows
                if r["target"] == DISPLAY[t] and r["method"] == "faithful_3seed"
            ),
            None,
        )
        if not (b and f):
            continue
        fg = abs(float(f["gamma_on_target"].split("+/-")[0]))
        f_mean = float(f["ratio"].split("+/-")[0])
        b_ratio = float(b["ratio"]) if b["ratio"] != "n/a" else float("nan")
        # heuristic verdict
        agree = (
            "AGREE"
            if (b_ratio > 1) == (f_mean > 1) and abs(b_ratio - f_mean) < 2
            else "DIVERGE"
        )
        nz_flag = " (faithful gamma ~= 0 -> noise-dominated)" if fg < 0.01 else ""
        print(
            f"- **{DISPLAY[t]}**: bin ratio {b_ratio:.2f} vs faithful {f_mean:.2f} -> {agree}{nz_flag}"
        )


if __name__ == "__main__":
    main()