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

#!/usr/bin/env python3
# pyright: reportAttributeAccessIssue=false
"""Corrected 6-metric robustness recompute (top of tab:selectivity-robust), reading
gamma DIRECTLY from gamma_olmes_tidy.csv. Mirrors pilot_bootstrap_olmes.py ingestion
(expA = bin-targeted, expC = empirical naive top-K) and selectivity_M1 verbatim.

Six selectivity metrics, headline bin-rule = highest-attribution topic (argmax raw z-bin):
  M1 mean-collateral      sel = |g_t| / mean_{b!=t}|g_b|
  M2 worst-case collateral sel = |g_t| / max_{b!=t}|g_b|
  M3 excluding-shared      sel = |g_t| / mean_{b!=t, b!=shared(t)}|g_b|
  M4 signed damage ratio   sel = |g_t| / mean_{b!=t} g_b      (signed denominator)
  M5 cosine-to-ideal       cos([g_b]_b, onehot_t)             (win if bin > naive)
  M6 neg-entropy concentr. -sum p log p, p=|g|/sum|g|         (win if bin > naive)
For M1-M4: ratio = sel_bin / sel_naive, win if >1. For M5-M6: win if bin metric > naive.

--baseline old        : recompute g = OLD_BASE[b] - unlearned_acc (validation vs published table)
--baseline corrected  : use the tidy g (config-matched :mc::olmes baselines)   [default]
"""

import sys
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",
}
# Pre-correction (OLMo-3 Tech-Report) baselines — for reproducing the published table.
OLD_BASE = {
    "socialiqa": 0.8029,
    "mmlu_social_science": 0.7508,
    "mmlu_stem": 0.5979,
    "arc_challenge": 0.892,
}
# "Shared" benchmark excluded by M3 (same MMLU family pairing; socialiqa/arc pair to their
# nearest neighbor). Pinned by validation against the published table.
SHARED = {
    "mmlu_stem": "mmlu_social_science",
    "mmlu_social_science": "mmlu_stem",
    "socialiqa": "mmlu_social_science",
    "arc_challenge": "mmlu_stem",
}


def _gamma_col(df, mode):
    if mode == "old":
        return df.apply(lambda r: OLD_BASE[r.eval_benchmark] - r.unlearned_acc, axis=1)
    return df.gamma


def load(mode):
    df = pd.read_csv(TIDY)
    df = df[df.eval_family == "primary"].copy()
    df["g"] = _gamma_col(df, mode)

    def cells(cond):
        out = {}
        for key, g in df[df.condition == cond].groupby("target_or_topic"):
            out[str(key)] = dict(zip(g.eval_benchmark, g.g))
        return out

    expA_raw, expC = cells("expA"), cells("expC")
    gamma = {}
    for key, cell in expA_raw.items():
        if "__" not in key:
            continue
        topic, target = key.rsplit("__", 1)
        gamma.setdefault(target, {})[topic] = cell
    return gamma, expC


def top1_topic(target):
    z = pd.read_csv(ZSCORED_DIR / f"zscored_{target}.csv")
    return z.loc[z.zscore.idxmax(), "topic_label"]


# ---- the six metrics ----
def _offs(target):
    return [b for b in BENCHMARKS if b != target]


def m_mean(d, t):
    o = [abs(d[b]) for b in _offs(t) if b in d]
    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 m_worst(d, t):
    o = [abs(d[b]) for b in _offs(t) if b in d]
    m = max(o) if o else 0.0
    return abs(d.get(t, 0.0)) / m if m > 0 else float("nan")


def m_exclshared(d, t):
    o = [abs(d[b]) for b in _offs(t) if b in d and b != SHARED.get(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 m_signed(d, t):
    o = [d[b] for b in _offs(t) if b in d]
    m = float(np.mean(o)) if o else 0.0
    return abs(d.get(t, 0.0)) / abs(m) if m != 0 else float("nan")


def _vec(d, t):
    return np.array([d.get(b, 0.0) for b in BENCHMARKS], float), BENCHMARKS.index(t)


def m_cosine(d, t):
    v, ti = _vec(d, t)
    n = np.linalg.norm(v)
    if n == 0:
        return float("nan")
    ideal = np.zeros_like(v)
    ideal[ti] = 1.0
    return float(np.dot(v, ideal) / n)  # = |v_t|/||v|| (sign via v_t)


def m_negentropy(d, t):
    a = np.abs([d.get(b, 0.0) for b in BENCHMARKS])
    s = a.sum()
    if s == 0:
        return float("nan")
    p = a / s
    nz = p[p > 0]
    return float(
        np.sum(nz * np.log(nz))
    )  # negative entropy (higher = more concentrated)


RATIO_METRICS = [
    ("mean", m_mean),
    ("worst", m_worst),
    ("excl_shared", m_exclshared),
    ("signed", m_signed),
]
WIN_METRICS = [("cosine", m_cosine), ("neg_entropy", m_negentropy)]


def main():
    mode = "corrected"
    if "--baseline" in sys.argv:
        mode = sys.argv[sys.argv.index("--baseline") + 1]
    gamma, expC = load(mode)
    print(f"=== 6-metric robustness (highest-attribution topic), baseline={mode} ===\n")
    rows = []
    for name, fn in RATIO_METRICS:
        cells, wins = [], 0
        for t in BENCHMARKS:
            k = top1_topic(t)
            dbin = gamma.get(t, {}).get(k, {})
            dnaive = expC.get(t, {})
            if not dbin or not dnaive:
                cells.append("---")
                continue
            sb, sn = fn(dbin, t), fn(dnaive, t)
            r = sb / sn if (sn and not np.isnan(sn) and sn != 0) else float("nan")
            win = (not np.isnan(r)) and r > 1.0
            wins += int(win)
            cells.append(f"{r:.2f}x{'*' if win else ' '}")
        rows.append((name, cells, wins))
        print(
            f"  {name:14s} "
            + "  ".join(f"{DISPLAY[b]}={c}" for b, c in zip(BENCHMARKS, cells))
            + f"   wins={wins}/4"
        )
    for name, fn in WIN_METRICS:
        cells, wins = [], 0
        for t in BENCHMARKS:
            k = top1_topic(t)
            dbin, dnaive = gamma.get(t, {}).get(k, {}), expC.get(t, {})
            if not dbin or not dnaive:
                cells.append("---")
                continue
            vb, vn = fn(dbin, t), fn(dnaive, t)
            win = (not np.isnan(vb)) and (not np.isnan(vn)) and vb > vn
            wins += int(win)
            cells.append("win" if win else "loss")
        rows.append((name, cells, wins))
        print(
            f"  {name:14s} "
            + "  ".join(f"{DISPLAY[b]}={c}" for b, c in zip(BENCHMARKS, cells))
            + f"   wins={wins}/4"
        )

    out = pd.DataFrame(
        [
            {"metric": n, **{DISPLAY[b]: c for b, c in zip(BENCHMARKS, cs)}, "wins": w}
            for n, cs, w in rows
        ]
    )
    dest = OUT_DIR / f"robustness_6metric_{mode}.csv"
    out.to_csv(dest, index=False)
    print(f"\nwrote {dest}")


if __name__ == "__main__":
    main()