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

#!/usr/bin/env python3
# pyright: reportAttributeAccessIssue=false
"""Multi-seed paired-bar figure for Finding 9 / R4 C2.

Reads:
  - ~/scratch/n16_selectivity/results/faithful_gamma_tidy.csv (3-seed gamma)
  - ~/scratch/n16_selectivity/results/paired_significance.csv  (BH-adj Wilcoxon p)

Produces 4-panel dumbbell plot (one panel per primary benchmark):
  - Each row = one WebOrganizer topic
  - White circle = random-in-topic (exp1) seed-mean gamma on that benchmark
  - Colored circle = influence-targeted (expA, target == panel benchmark) seed-mean gamma
  - Horizontal whiskers on each marker = +/- 1 SD across seeds {42, 43, 44}
  - Grey line connects the within-bin pair
  - Panel title carries the BH-adjusted Wilcoxon significance star

Output: ~/scratch/n16_selectivity/results/figures/unlearning_paired_arc.{pdf,png}
"""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

HOME = Path.home()
TIDY = HOME / "scratch/n16_selectivity/results/faithful_gamma_tidy.csv"
SIG = HOME / "scratch/n16_selectivity/results/paired_significance.csv"

# Panel order (row, col) -> benchmark key as it appears in the tidy CSV.
_PANEL_ORDER = [
    ("socialiqa", 0, 0),
    ("mmlu_social_science", 0, 1),
    ("mmlu_stem", 1, 0),
    ("arc_challenge", 1, 1),
]

_BENCH_DISPLAY = {
    "socialiqa": "SocialIQA",
    "mmlu_social_science": "MMLU Social Sciences",
    "mmlu_stem": "MMLU STEM",
    "arc_challenge": "ARC-Challenge",
}

# Match the manuscript's existing aesthetic (from plot_paired.py).
_COLOR_RANDOM = "#bcbddc"  # light purple = neutral baseline
_COLOR_TARGETED = {
    "socialiqa": "#54278f",
    "mmlu_social_science": "#3182bd",
    "mmlu_stem": "#31a354",
    "arc_challenge": "#e6550d",
}

# Compact display labels for the 24 WebOrganizer topics.
_TOPIC_LABEL = {
    "adult_content": "Adult Content",
    "art_and_design": "Art & Design",
    "crime_and_law": "Crime & Law",
    "education_and_jobs": "Education & Jobs",
    "electronics_and_hardware": "Electronics & Hardware",
    "entertainment": "Entertainment",
    "fashion_and_beauty": "Fashion & Beauty",
    "finance_and_business": "Finance & Business",
    "food_and_dining": "Food & Dining",
    "games": "Games",
    "health": "Health",
    "history_and_geography": "History & Geography",
    "home_and_hobbies": "Home & Hobbies",
    "industrial": "Industrial",
    "literature": "Literature",
    "politics": "Politics",
    "religion": "Religion",
    "science_math_and_technology": "Sci, Math & Tech",
    "social_life": "Social Life",
    "software": "Software",
    "software_development": "Software Dev.",
    "sports_and_fitness": "Sports & Fitness",
    "transportation": "Transportation",
    "travel_and_tourism": "Travel & Tourism",
}


def _stars(p_bh: float) -> str:
    if pd.isna(p_bh):
        return ""
    if p_bh < 0.001:
        return "***"
    if p_bh < 0.01:
        return "**"
    if p_bh < 0.05:
        return "*"
    return "ns"


def _build_panel(df: pd.DataFrame, benchmark: str) -> pd.DataFrame:
    """For one benchmark, return per-topic seed-mean and seed-std gamma for
    both expA (target == benchmark) and exp1 (random in same topic).

    Columns: topic, gamma_expA_mean, gamma_expA_std, gamma_exp1_mean,
             gamma_exp1_std, n_seeds.
    """
    a = df[
        (df.condition == "expA")
        & (df.target == benchmark)
        & (df.eval_benchmark == benchmark)
    ][["topic", "seed", "gamma"]]
    e = df[(df.condition == "exp1") & (df.eval_benchmark == benchmark)][
        ["topic", "seed", "gamma"]
    ]

    a_agg = a.groupby("topic").gamma.agg(["mean", "std", "count"]).reset_index()
    a_agg.columns = ["topic", "gamma_expA_mean", "gamma_expA_std", "n_expA"]
    e_agg = e.groupby("topic").gamma.agg(["mean", "std", "count"]).reset_index()
    e_agg.columns = ["topic", "gamma_exp1_mean", "gamma_exp1_std", "n_exp1"]

    merged = pd.merge(a_agg, e_agg, on="topic", how="inner")
    return merged


def plot(
    tidy_path: Path,
    sig_path: Path,
    out_dir: Path,
    out_stem: str = "unlearning_paired_arc",
) -> None:
    df = pd.read_csv(tidy_path)
    sig = pd.read_csv(sig_path)
    sig_wilcoxon = sig[sig.test == "wilcoxon_pooled"].set_index("target")

    panels = {b: _build_panel(df, b) for b, _, _ in _PANEL_ORDER}
    # Sort each panel by influence-targeted seed-mean gamma descending,
    # then REVERSE so the largest-gamma topic renders at the top of the axis
    # (matplotlib y=0 is at the bottom; this matches the original plot_paired.py
    # visual convention: most-damaged topic at the top).
    for b in panels:
        panels[b] = (
            panels[b]
            .sort_values("gamma_expA_mean", ascending=False)
            .iloc[::-1]
            .reset_index(drop=True)
        )

    # Shared x-range from min/max across all panels' gamma values
    # (including +/- std for whisker reach).
    all_lows, all_highs = [], []
    for d in panels.values():
        for col_mean, col_std in [
            ("gamma_expA_mean", "gamma_expA_std"),
            ("gamma_exp1_mean", "gamma_exp1_std"),
        ]:
            mean = d[col_mean].to_numpy()
            std = d[col_std].fillna(0).to_numpy()
            all_lows.append(np.nanmin(mean - std))
            all_highs.append(np.nanmax(mean + std))
    x_lo_raw = float(np.nanmin(all_lows))
    x_hi_raw = float(np.nanmax(all_highs))
    pad = (x_hi_raw - x_lo_raw) * 0.05
    x_lo = min(0.0, x_lo_raw - pad)
    x_hi = x_hi_raw + pad

    # 0.18 in per row matches the original; max(3.5, ...) protects sparse panels.
    n_rows = max(len(d) for d in panels.values())
    panel_height_in = max(3.5, 0.18 * n_rows)

    fig, axes = plt.subplots(
        2, 2, figsize=(13, panel_height_in * 2 + 1.0), sharey=False, sharex=False
    )

    for bench, r, c in _PANEL_ORDER:
        ax = axes[r, c]
        d = panels[bench]
        if d.empty:
            ax.set_title(_BENCH_DISPLAY[bench])
            ax.text(
                0.5, 0.5, "No data", ha="center", va="center", transform=ax.transAxes
            )
            continue
        y_pos = np.arange(len(d))
        c_targeted = _COLOR_TARGETED[bench]

        # Connecting line for each pair.
        for i, (mr, mt) in enumerate(zip(d["gamma_exp1_mean"], d["gamma_expA_mean"])):
            ax.plot([mr, mt], [i, i], color="#999999", linewidth=1.0, zorder=1)

        # Whiskers (1 SD) — drawn as thin caps.
        ax.errorbar(
            d["gamma_exp1_mean"],
            y_pos,
            xerr=d["gamma_exp1_std"].fillna(0),
            fmt="none",
            ecolor="#888888",
            elinewidth=0.8,
            capsize=2,
            zorder=2,
        )
        ax.errorbar(
            d["gamma_expA_mean"],
            y_pos,
            xerr=d["gamma_expA_std"].fillna(0),
            fmt="none",
            ecolor=c_targeted,
            elinewidth=0.8,
            capsize=2,
            zorder=2,
        )

        # Markers.
        ax.scatter(
            d["gamma_exp1_mean"],
            y_pos,
            facecolors="white",
            edgecolors="#666666",
            s=55,
            linewidths=1.1,
            label="In-topic random",
            zorder=3,
        )
        ax.scatter(
            d["gamma_expA_mean"],
            y_pos,
            facecolors=c_targeted,
            edgecolors=c_targeted,
            s=65,
            label="In-topic influence-targeted",
            zorder=3,
        )

        ax.set_yticks(y_pos)
        ax.set_yticklabels(
            [_TOPIC_LABEL.get(t, t) for t in d["topic"]],
            fontsize=8,
            family="serif",
        )
        ax.set_xlim(x_lo, x_hi)
        if r == 1:
            ax.set_xlabel(
                r"$\gamma$ = baseline accuracy $-$ unlearned accuracy",
                fontsize=10,
                family="serif",
            )
        for label in ax.get_xticklabels():
            label.set_fontsize(8)
            label.set_family("serif")

        # Panel title: benchmark + BH-adjusted Wilcoxon star + n.
        if bench in sig_wilcoxon.index:
            row = sig_wilcoxon.loc[bench]
            star = _stars(row["p_bh"])
            n = int(row["n_cells"])
            title = f"{_BENCH_DISPLAY[bench]}  {star}  (n={n} cells, p$_{{\\rm BH}}$={row['p_bh']:.3g})"
        else:
            title = _BENCH_DISPLAY[bench]
        ax.set_title(title, fontsize=11, family="serif")

        ax.axvline(0.0, color="#bbbbbb", linewidth=0.8, zorder=0)
        ax.grid(True, axis="x", linestyle=":", color="#dddddd", zorder=0)
        ax.set_axisbelow(True)
        if r == 0 and c == 0:
            ax.legend(loc="lower right", frameon=True, fontsize=8)

    fig.subplots_adjust(
        left=0.18, right=0.98, top=0.965, bottom=0.05, wspace=0.35, hspace=0.18
    )

    out_dir.mkdir(parents=True, exist_ok=True)
    pdf_path = out_dir / f"{out_stem}.pdf"
    png_path = out_dir / f"{out_stem}.png"
    fig.savefig(pdf_path)
    fig.savefig(png_path, dpi=200)
    plt.close(fig)
    print(f"Wrote {pdf_path}")
    print(f"Wrote {png_path}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--tidy", type=Path, default=TIDY)
    parser.add_argument("--paired", type=Path, default=SIG)
    parser.add_argument(
        "--out-dir",
        type=Path,
        default=HOME / "scratch/n16_selectivity/results/figures",
    )
    parser.add_argument("--out-stem", default="unlearning_paired_arc")
    args = parser.parse_args()
    plot(args.tidy, args.paired, args.out_dir, args.out_stem)


if __name__ == "__main__":
    main()