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

"""Corrected pilot-validated (R14) + best-in-hindsight (R6) + bootstrap CIs + LOBO,
reading gamma DIRECTLY from gamma_olmes_tidy.csv (config-matched :mc::olmes baselines).

Port of bootstrap_and_pilot_sensitivity.py: only the gamma/naive INGESTION changes
(no BASELINE dict, no dual-CSV stitch). selectivity_M1 / find_best / unique_z ranking /
bootstrap / leave-one-out / R14_with_K are byte-for-byte the validated logic.
"""

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

HOME = Path.home()
ROOT = HOME / "scratch" / "n16_selectivity"
TIDY = ROOT / "results" / "gamma_olmes_tidy.csv"
ZSCORED_DIR = (
    HOME
    / "dev"
    / "data-attribution"
    / "artifacts"
    / "zscored_bin_scores"
    / "aggregated"
)
OUT_DIR = ROOT / "results" / "robustness"
OUT_DIR.mkdir(parents=True, exist_ok=True)

BENCHMARKS = ["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 _cell(g):
    return dict(zip(g.eval_benchmark, g.gamma))


def build_gamma_matrix():
    """gamma[target][topic] = {bench: gamma}, from expA cells of the tidy matrix."""
    df = pd.read_csv(TIDY)
    df = df[(df.eval_family == "primary") & (df.condition == "expA")]
    gamma = {}
    for key, g in df.groupby("target_or_topic"):
        if "__" not in str(key):
            continue
        topic, target = str(key).rsplit("__", 1)
        gamma.setdefault(target, {})[topic] = _cell(g)
    return gamma


def build_empirical_naive():
    """naive[target] = {bench: gamma}, from expC (corpus-wide influence-targeted)."""
    df = pd.read_csv(TIDY)
    df = df[(df.eval_family == "primary") & (df.condition == "expC")]
    return {str(k): _cell(g) for k, g in df.groupby("target_or_topic")}


def selectivity_M1(damage, target, off_targets=None):
    if off_targets is None:
        off_targets = [b for b in BENCHMARKS if b != target]
    t = abs(damage.get(target, 0.0))
    others = [abs(damage[b]) for b in off_targets if b in damage]
    if not others:
        return float("nan")
    m = float(np.mean(others))
    return t / m if m > 0 else (float("inf") if t > 0 else float("nan"))


def find_best_selectivity_topic_among(gamma_matrix, target, candidate_topics):
    best = None
    best_sel = -float("inf")
    for topic in candidate_topics:
        cell = gamma_matrix.get(target, {}).get(topic, {})
        if not cell:
            continue
        single = {sb: cell.get(sb, 0.0) for sb in BENCHMARKS}
        s = selectivity_M1(single, target)
        if not np.isnan(s) and not np.isinf(s) and s > best_sel:
            best_sel = s
            best = topic
    return best, best_sel


def topic_marginal_z(z_df):
    return z_df.groupby("topic_label")["zscore"].mean()


def get_unique_z_ranking(z_all, target):
    tm = topic_marginal_z(z_all[target])
    other_means = pd.concat(
        [topic_marginal_z(z_all[b]) for b in BENCHMARKS if b != target], axis=1
    )
    other_max = other_means.max(axis=1)
    unique = tm - other_max
    return unique.sort_values(ascending=False)


def bootstrap_R6_ratio(gamma_matrix, target, naive_damage, B=2000, rng_seed=42):
    rng = np.random.default_rng(rng_seed)
    topics = list(gamma_matrix.get(target, {}).keys())
    if not topics:
        return float("nan"), float("nan"), float("nan")
    sel_naive = selectivity_M1(naive_damage.get(target, {}), target)
    if not sel_naive or np.isnan(sel_naive):
        return float("nan"), float("nan"), float("nan")
    ratios = []
    for _ in range(B):
        resampled = rng.choice(topics, size=len(topics), replace=True)
        bt, bs = find_best_selectivity_topic_among(
            gamma_matrix, target, list(resampled)
        )
        if bt is None or np.isnan(bs) or np.isinf(bs):
            continue
        ratios.append(bs / sel_naive)
    r = np.array(ratios)
    return (
        float(np.median(r)),
        float(np.quantile(r, 0.025)),
        float(np.quantile(r, 0.975)),
    )


def bootstrap_R14_ratio(
    gamma_matrix, z_all, target, naive_damage, K=3, B=2000, rng_seed=42
):
    rng = np.random.default_rng(rng_seed)
    sel_naive = selectivity_M1(naive_damage.get(target, {}), target)
    if not sel_naive or np.isnan(sel_naive):
        return float("nan"), float("nan"), float("nan")
    topics = list(gamma_matrix.get(target, {}).keys())
    if not topics:
        return float("nan"), float("nan"), float("nan")
    unique_z_rank = get_unique_z_ranking(z_all, target)
    ratios = []
    for _ in range(B):
        noise = rng.normal(0, 0.5, size=len(unique_z_rank))
        noisy = unique_z_rank.copy() + noise
        noisy_top = noisy.nlargest(K).index.tolist()
        bt, bs = find_best_selectivity_topic_among(gamma_matrix, target, noisy_top)
        if bt is None or np.isnan(bs) or np.isinf(bs):
            continue
        ratios.append(bs / sel_naive)
    r = np.array(ratios)
    return (
        float(np.median(r)),
        float(np.quantile(r, 0.025)),
        float(np.quantile(r, 0.975)),
    )


def leave_one_out_M1(gamma_matrix, target, topic_weights, naive_damage):
    bd = {}
    for sb in BENCHMARKS:
        acc = 0.0
        wsum = 0.0
        for topic, w in topic_weights.items():
            g = gamma_matrix.get(target, {}).get(topic, {}).get(sb)
            if g is None:
                continue
            acc += w * g
            wsum += w
        if wsum > 0:
            bd[sb] = acc / wsum
    nd = naive_damage.get(target, {})
    result = {}
    for excluded in BENCHMARKS:
        if excluded == target:
            continue
        off_targets = [b for b in BENCHMARKS if b != target and b != excluded]
        sel_bin = selectivity_M1(bd, target, off_targets=off_targets)
        sel_naive = selectivity_M1(nd, target, off_targets=off_targets)
        ratio = (
            sel_bin / sel_naive
            if (sel_naive and not np.isnan(sel_naive) and sel_naive != 0)
            else float("nan")
        )
        result[excluded] = (sel_bin, sel_naive, ratio)
    return result


def R14_with_K(gamma_matrix, z_all, target, naive_damage, K=3):
    sel_naive = selectivity_M1(naive_damage.get(target, {}), target)
    if not sel_naive or np.isnan(sel_naive):
        return None, float("nan"), float("nan")
    unique_z_rank = get_unique_z_ranking(z_all, target)
    topK = unique_z_rank.nlargest(K).index.tolist()
    bt, bs = find_best_selectivity_topic_among(gamma_matrix, target, topK)
    if bt is None or np.isnan(bs) or np.isinf(bs):
        return None, float("nan"), float("nan")
    return bt, bs, bs / sel_naive


def main():
    gamma = build_gamma_matrix()
    naive_emp = build_empirical_naive()
    z_all = {b: pd.read_csv(ZSCORED_DIR / f"zscored_{b}.csv") for b in BENCHMARKS}

    print("=" * 78)
    print(
        "HEAD-TO-HEAD (top-1 attribution topic): bin expA vs empirical naive expC -- CORRECTED"
    )
    print("=" * 78)
    h2h = 0
    for target in BENCHMARKS:
        z = z_all[target]
        k1 = z.loc[z["zscore"].idxmax(), "topic_label"]
        sb = selectivity_M1(gamma.get(target, {}).get(k1, {}), target)
        sn = selectivity_M1(naive_emp.get(target, {}), target)
        ratio = sb / sn if (sn and not np.isnan(sn) and sn != 0) else float("nan")
        win = ratio > 1.0
        h2h += int(win)
        print(
            f"  {DISPLAY[target]:12s} k*={k1:<28s} sel_bin={sb:7.3f} sel_naive={sn:7.3f} ratio={ratio:6.2f} {'WIN' if win else 'lose'}"
        )
    print(f"  >>> head-to-head bin wins: {h2h}/4")

    print("\n" + "=" * 78)
    print(
        "PILOT-VALIDATED (R14, top-3 unique-z shortlist -> most selective) -- CORRECTED"
    )
    print("=" * 78)
    pilot = 0
    for target in BENCHMARKS:
        bt, bs, ratio = R14_with_K(gamma, z_all, target, naive_emp, K=3)
        win = (not np.isnan(ratio)) and ratio > 1.0
        pilot += int(win)
        print(
            f"  {DISPLAY[target]:12s} picked={str(bt):<28s} sel_bin={bs:7.3f} ratio_vs_naive={ratio:6.2f} {'WIN' if win else 'lose'}"
        )
    print(f"  >>> pilot-validated bin wins: {pilot}/4")

    print("\n" + "=" * 78)
    print("BEST-IN-HINDSIGHT (R6, best of all 24 topics) -- CORRECTED")
    print("=" * 78)
    hind = 0
    for target in BENCHMARKS:
        bt, bs = find_best_selectivity_topic_among(
            gamma, target, list(gamma.get(target, {}).keys())
        )
        sn = selectivity_M1(naive_emp.get(target, {}), target)
        ratio = bs / sn if (sn and not np.isnan(sn) and sn != 0) else float("nan")
        win = (not np.isnan(ratio)) and ratio > 1.0
        hind += int(win)
        print(
            f"  {DISPLAY[target]:12s} best={str(bt):<28s} sel_bin={bs:7.3f} ratio_vs_naive={ratio:6.2f} {'WIN' if win else 'lose'}"
        )
    print(f"  >>> best-in-hindsight bin wins: {hind}/4")

    print("\n" + "=" * 78)
    print(
        "BOOTSTRAP CIs (B=2000): R6 (hindsight) and R14 (pilot) ratio vs naive -- CORRECTED"
    )
    print("=" * 78)
    rows = []
    for target in BENCHMARKS:
        m6, lo6, hi6 = bootstrap_R6_ratio(gamma, target, naive_emp)
        m14, lo14, hi14 = bootstrap_R14_ratio(gamma, z_all, target, naive_emp, K=3)
        rows.append(
            {
                "target": target,
                "display": DISPLAY[target],
                "R6_median": m6,
                "R6_lo": lo6,
                "R6_hi": hi6,
                "R14_median": m14,
                "R14_lo": lo14,
                "R14_hi": hi14,
            }
        )
        print(
            f"  {DISPLAY[target]:12s} R6 median={m6:5.2f} CI=[{lo6:5.2f},{hi6:5.2f}]   R14 median={m14:5.2f} CI=[{lo14:5.2f},{hi14:5.2f}]"
        )
    pd.DataFrame(rows).to_csv(OUT_DIR / "bootstrap_ci_olmes.csv", index=False)

    print("\n" + "=" * 78)
    print("LEAVE-ONE-BENCHMARK-OUT (pilot R14 picks) -- CORRECTED")
    print("=" * 78)
    for target in BENCHMARKS:
        bt, _, _ = R14_with_K(gamma, z_all, target, naive_emp, K=3)
        if bt is None:
            continue
        loo = leave_one_out_M1(gamma, target, {bt: 1.0}, naive_emp)
        ratios = [r for (_, _, r) in loo.values() if not np.isnan(r)]
        mn = min(ratios) if ratios else float("nan")
        s = " | ".join(
            [f"-{b[:3]}={loo[b][2]:5.2f}" for b in BENCHMARKS if b != target]
        )
        print(
            f"  {DISPLAY[target]:12s} {s}   min={mn:5.2f}  {'robust_win' if (ratios and mn > 1.0) else 'NOT robust'}"
        )

    print(
        f"\nSUMMARY: head-to-head {h2h}/4, pilot {pilot}/4, best-in-hindsight {hind}/4"
    )


if __name__ == "__main__":
    main()