social-data-attribution-w2/scripts/visualization/paper_main_figures_arc.py

#!/usr/bin/env python3
"""Main paper figures — single-bin unlearning results with ARC-Challenge replacing GSM8K.

Adapted from scripts/visualization/paper_main_figures.py (commit e76a748).
Data source: unlearning_analysis.csv (exp1 + targeted baselines) +
             results/all_eval_results.csv (arc_challenge values from soc171 TrackStar).

Figures produced (all saved with _arc suffix):
  1. fig1_single_bin_heatmap_arc  — 24 topics × 4 benchmarks heatmap + Global Random col
  2. fig2_per_metric_gamma_bars_arc — 2×2 per-benchmark gamma bars
  3. fig3_null_bin_control_arc    — boxplot exp1 vs expC_arc null-bin star

Usage:
  python paper_main_figures_arc.py --output-dir artifacts/paper_main_arc
  python paper_main_figures_arc.py --data-dir /path/to/csvs --output-dir /path/to/out
"""
from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import numpy as np
import pandas as pd
import seaborn as sns

from dolma.distribution_report.style import mpl_font, style_size

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

BASELINE = {
    "arc_challenge": 0.8344709897610921,
    "mmlu_socsci":   0.750846,
    "mmlu_stem":     0.597871,
    "socialiqa":     0.802900,
    "wikitext_ppl":  9.1559,
}

# arc_challenge replaces gsm8k; same reasoning category → same green
ACCURACY_METRICS = ["arc_challenge", "mmlu_socsci", "mmlu_stem", "socialiqa"]
HEATMAP_COL_ORDER = ["socialiqa", "mmlu_socsci", "arc_challenge", "mmlu_stem"]

METRIC_LABELS = {
    "arc_challenge": "ARC-Challenge",
    "mmlu_socsci":   "MMLU Social Science",
    "mmlu_stem":     "MMLU STEM",
    "socialiqa":     "SocialIQA",
    "wikitext_ppl":  "Wikitext PPL",
}
METRIC_LABELS_WRAP = {
    "arc_challenge": "ARC\nChallenge",
    "mmlu_socsci":   "MMLU\nSocial Sci.",
    "mmlu_stem":     "MMLU\nSTEM",
    "socialiqa":     "SocialIQA",
    "wikitext_ppl":  "Wikitext\nPPL",
}

BENCH_COLORS = {
    "socialiqa":     "#984EA3",
    "mmlu_socsci":   "#984EA3",
    "arc_challenge": "#1b7837",   # same green as gsm8k (reasoning)
    "mmlu_stem":     "#1b7837",
    "wikitext_ppl":  "#ff7f00",
}
BENCH_HATCHES = {
    "socialiqa":     "",
    "mmlu_socsci":   "//",
    "arc_challenge": "",          # solid fill (reasoning)
    "mmlu_stem":     "//",
}

EXP_COLORS = {
    "exp1": "#a6cee3",
    "expC_arc": "#fdbf6f",
}

# Snake-case bin name → display label (matching experiment1_single_bin.py)
TOPIC_DISPLAY = {
    "adult_content":             "Adult",
    "art_and_design":            "Art & Design",
    "crime_and_law":             "Crime & Law",
    "education_and_jobs":        "Education & Jobs",
    "electronics_and_hardware":  "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",
    "home_and_hobbies":          "Home & Hobbies",
    "industrial":                "Industrial",
    "literature":                "Literature",
    "politics":                  "Politics",
    "religion":                  "Religion",
    "science_math_and_technology": "Science & Tech.",
    "social_life":               "Social Life",
    "software":                  "Software",
    "software_development":      "Software Dev.",
    "sports_and_fitness":        "Sports & Fitness",
    "transportation":            "Transportation",
    "travel_and_tourism":        "Travel",
}


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

def gamma(score: float, metric: str) -> float:
    return (score - BASELINE[metric]) / abs(BASELINE[metric])


def paper_rc():
    """Apply unified role-based fonts (mirrors Fig 4 / plot_paired.py style)."""
    plt.rcParams.update({
        "font.family": "serif",
        "font.serif": ["Liberation Serif", "DejaVu Serif", "Times New Roman"],
        "mathtext.fontset": "dejavuserif",
        "font.size":        style_size("TICK"),
        "axes.titlesize":   style_size("SUBPLOT_TITLE"),
        "axes.labelsize":   style_size("AXIS_TITLE"),
        "xtick.labelsize":  style_size("TICK"),
        "ytick.labelsize":  style_size("TICK"),
        "legend.fontsize":  style_size("LEGEND"),
        "figure.dpi": 300,
        "savefig.dpi": 300,
        "savefig.bbox": "tight",
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.linewidth": 0.5,
        "xtick.major.width": 0.5,
        "ytick.major.width": 0.5,
        "grid.linewidth": 0.3,
        "grid.alpha": 0.2,
    })


def save_fig(fig, output_dir: Path, name: str) -> None:
    for ext in ("pdf", "svg", "png"):
        fig.savefig(output_dir / f"{name}.{ext}", dpi=300, bbox_inches="tight")
    plt.close(fig)
    print(f"  {name} (pdf/svg/png)")


# ---------------------------------------------------------------------------
# Data loading
# ---------------------------------------------------------------------------

def build_dataframe(data_dir: Path) -> tuple[pd.DataFrame, dict]:
    """
    Returns:
        df   — exp1 (single-bin unlearning) dataframe with gamma cols
        null — dict of metric → gamma for the expC_arc null-bin run
    """
    orig = pd.read_csv(data_dir / "unlearning_analysis.csv")
    arc_df = pd.read_csv(data_dir / "results" / "all_eval_results.csv")

    # --- exp1 rows from unlearning_analysis.csv ---
    exp1 = orig[orig["experiment"] == "singe-topic_unlearning"].copy()

    # Join arc_challenge from all_eval_results (exp=='exp1')
    arc_exp1 = arc_df[arc_df["exp"] == "exp1"][["topic", "arc_challenge"]].rename(
        columns={"topic": "bin"}
    )
    exp1 = exp1.merge(arc_exp1, on="bin", how="left")

    # Display name
    exp1["topic"] = exp1["bin"].map(TOPIC_DISPLAY).fillna(exp1["bin"])

    # Gamma columns
    for m in ACCURACY_METRICS:
        exp1[f"{m}_gamma"] = exp1[m].apply(lambda v: gamma(v, m))
    exp1["wikitext_ppl_gamma"] = exp1["wikitext_ppl"].apply(
        lambda v: gamma(v, "wikitext_ppl")
    )
    exp1["mean_acc_gamma"] = exp1[[f"{m}_gamma" for m in ACCURACY_METRICS]].mean(axis=1)

    # --- Null-bin column: exp3 (global random, matches original paper_main_figures.py) ---
    # socialiqa / mmlu_socsci / mmlu_stem come from unlearning_analysis.csv null_bin row
    # arc_challenge comes from exp3 row in all_eval_results.csv (soc171 eval of that same run)
    null_row = orig[orig["experiment"] == "null_bin"].iloc[0]
    exp3_row = arc_df[arc_df["exp"] == "exp3"].iloc[0]
    null = {
        "arc_challenge": gamma(float(exp3_row["arc_challenge"]), "arc_challenge"),
        "mmlu_socsci":   gamma(float(null_row["mmlu_socsci"]),   "mmlu_socsci"),
        "mmlu_stem":     gamma(float(null_row["mmlu_stem"]),      "mmlu_stem"),
        "socialiqa":     gamma(float(null_row["socialiqa"]),      "socialiqa"),
    }

    return exp1.reset_index(drop=True), null


# ---------------------------------------------------------------------------
# Figure 1: Cross-benchmark heatmap
# ---------------------------------------------------------------------------

def fig1_single_bin_heatmap_arc(df: pd.DataFrame, null: dict, out: Path) -> None:
    gamma_cols = [f"{m}_gamma" for m in HEATMAP_COL_ORDER]
    row_labels  = [METRIC_LABELS_WRAP[m] for m in HEATMAP_COL_ORDER]

    df_s = df.sort_values("mean_acc_gamma", ascending=False)
    n_topics = len(df_s)

    null_col = {f"{m}_gamma": null[m] for m in HEATMAP_COL_ORDER}
    null_col["topic"] = "Global\nRandom"
    df_plot = pd.concat([df_s, pd.DataFrame([null_col])], ignore_index=True)
    topics = df_plot["topic"].values

    data = df_plot[gamma_cols].values.T  # shape: (4 metrics, n_topics+1)

    fig, ax = plt.subplots(figsize=(16, 3.8))

    vmax_a = max(abs(data[:, :n_topics].min()), abs(data[:, :n_topics].max()), 0.01)
    sns.heatmap(
        data, ax=ax,
        xticklabels=topics,
        yticklabels=row_labels,
        cmap="RdBu", center=0, vmin=-vmax_a, vmax=vmax_a,
        annot=True, fmt=".1%",
        annot_kws={"fontsize": style_size("ANNOTATION")},
        linewidths=0.4, linecolor="white",
        cbar_kws={
            "label": r"$\gamma$",
            "shrink": 0.8, "aspect": 15, "pad": 0.01,
        },
    )
    ax.axhline(y=2, color="black", linewidth=1.0)
    ax.axvline(x=n_topics, color="black", linewidth=1.5)

    # Dense 24-topic x-axis at 45° → small tick role; sparse 4-metric y-axis → standard tick role.
    ax.set_xticklabels(ax.get_xticklabels(),
                       fontsize=style_size("COLORBAR_TICK"),
                       rotation=45, ha="right")
    xtick_labels = ax.get_xticklabels()
    xtick_labels[-1].set_fontweight("bold")
    xtick_labels[-1].set_fontstyle("italic")
    ax.set_yticklabels(ax.get_yticklabels(),
                       fontsize=style_size("TICK"), rotation=0)
    ax.tick_params(axis="both", length=0)
    ax.set_xlabel("")
    ax.set_ylabel("")
    fig.subplots_adjust(bottom=0.30)
    save_fig(fig, out, "fig1_single_bin_heatmap_arc")


# ---------------------------------------------------------------------------
# Figure 2: Per-metric gamma bars
# ---------------------------------------------------------------------------

def fig2_per_metric_bars_arc(df: pd.DataFrame, out: Path) -> None:
    panel_order = ["socialiqa", "arc_challenge", "mmlu_socsci", "mmlu_stem"]

    # Round 8: 4x1 vertical stack instead of 2x2 grid (user feedback:
    # "we should be converting them into stacked. So like all four
    # charts on top of each other so that we can see it easier").
    fig, axes = plt.subplots(4, 1, figsize=(10, 14))
    for idx, m in enumerate(panel_order):
        ax = axes[idx]
        col = f"{m}_gamma"
        ds = df.sort_values(col, ascending=True)
        vals = ds[col].values
        topics = ds["topic"].values
        x = np.arange(len(ds))

        ax.bar(
            x, vals,
            color=BENCH_COLORS[m],
            hatch=BENCH_HATCHES[m],
            edgecolor="black" if BENCH_HATCHES[m] else "none",
            linewidth=0.3 if BENCH_HATCHES[m] else 0,
            width=0.7,
        )
        ax.set_xticks(x)
        # 24 topics rotated 45° in an 8" panel → use dense tick role.
        ax.set_xticklabels(topics,
                           fontsize=style_size("COLORBAR_TICK"),
                           rotation=45, ha="right")
        ax.axhline(0, color="black", linewidth=0.5)
        ax.set_ylabel(r"$\gamma$", fontdict=mpl_font("AXIS_TITLE"))
        ax.set_title(METRIC_LABELS[m], fontdict=mpl_font("SUBPLOT_TITLE"))
        ax.yaxis.set_major_formatter(mticker.PercentFormatter(xmax=1, decimals=1))
        ax.grid(True, axis="y", alpha=0.2, linewidth=0.3)
        lim = max(abs(vals.min()), abs(vals.max()), 0.005) * 1.3
        ax.set_ylim(-lim, lim)

    fig.tight_layout()
    save_fig(fig, out, "fig2_per_metric_gamma_bars_arc")


# ---------------------------------------------------------------------------
# Figure 3: Null-bin control (exp1 box + expC_arc star)
# ---------------------------------------------------------------------------

def fig3_null_bin_control_arc(df: pd.DataFrame, null: dict, out: Path) -> None:
    fig, axes = plt.subplots(1, 4, figsize=(10, 3))
    for idx, m in enumerate(ACCURACY_METRICS):
        ax = axes[idx]
        col = f"{m}_gamma"
        exp1_v = df[col].values
        null_v = null[m]

        bp = ax.boxplot(
            [exp1_v], positions=[0], widths=0.4, patch_artist=True,
            showmeans=True,
            meanprops=dict(marker="D", markerfacecolor="white", markersize=3),
        )
        bp["boxes"][0].set_facecolor(EXP_COLORS["exp1"])

        rng = np.random.default_rng(42)
        jitter = rng.uniform(-0.08, 0.08, len(exp1_v))
        ax.scatter(
            np.full(len(exp1_v), 0) + jitter, exp1_v,
            color="black", s=6, alpha=0.35, zorder=3,
        )

        ax.scatter(
            1, null_v, color=EXP_COLORS["expC_arc"], marker="*", s=150,
            edgecolors="black", linewidth=0.5, zorder=4,
        )

        ax.set_xticks([0, 1])
        # 2 labels in a 2.5" panel → sparse tick role.
        ax.set_xticklabels(["Single-Bin\n(n=24)", "Global\nRandom"],
                           fontsize=style_size("TICK"))
        ax.axhline(0, color="gray", linestyle="--", linewidth=0.4, alpha=0.5)
        ax.set_title(METRIC_LABELS[m], fontdict=mpl_font("SUBPLOT_TITLE"))
        ax.grid(True, axis="y", alpha=0.2, linewidth=0.3)
        ax.yaxis.set_major_formatter(mticker.PercentFormatter(xmax=1, decimals=0))
        if idx == 0:
            ax.set_ylabel(r"$\gamma$", fontdict=mpl_font("AXIS_TITLE"))

    fig.tight_layout()
    save_fig(fig, out, "fig3_null_bin_control_arc")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--data-dir", type=Path,
        default=Path(__file__).parent,
        help="Directory containing unlearning_analysis.csv and results/all_eval_results.csv",
    )
    parser.add_argument(
        "--output-dir", type=Path,
        default=Path("artifacts/paper_main_arc"),
    )
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)

    paper_rc()
    df, null = build_dataframe(args.data_dir)

    print("Generating ARC-variant paper figures...")
    fig1_single_bin_heatmap_arc(df, null, args.output_dir)
    fig2_per_metric_bars_arc(df, args.output_dir)
    fig3_null_bin_control_arc(df, null, args.output_dir)
    print(f"All figures saved to {args.output_dir}/")


if __name__ == "__main__":
    main()