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

#!/usr/bin/env python3
"""Appendix figures — forget PPL, GSM8K anomaly, wikitext PPL, checkpoint efficiency."""
from __future__ import annotations

import argparse
from pathlib import Path

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

from scripts.visualization._shared import (
    ACCURACY_METRICS, BASELINE, BENCH_COLORS,
    EXP_COLORS, METRIC_LABELS,
    gamma, paper_rc, save_fig,
)
from scripts.visualization.experiment1_single_bin import (
    build_dataframe as build_exp1_df,
)
from scripts.visualization.experiment3_null_bin import NULL_BIN


def figA1_forget_ppl(df: pd.DataFrame, out: Path) -> None:
    ds = df.sort_values("forget_ppl", ascending=True)
    topics = ds["topic"].values
    forget_ppl = ds["forget_ppl"].values
    mean_gamma = ds["mean_acc_gamma"].values
    y = np.arange(len(ds))

    fig, (ax_l, ax_r) = plt.subplots(
        1, 2, figsize=(8, 7), sharey=True,
        gridspec_kw={"width_ratios": [1, 1], "wspace": 0.12},
    )

    ax_l.hlines(y, 0, forget_ppl - BASELINE["wikitext_ppl"], color="#636363", linewidth=0.6)
    ax_l.scatter(
        forget_ppl - BASELINE["wikitext_ppl"], y,
        c=forget_ppl, cmap="YlOrRd", s=35,
        edgecolors="black", linewidth=0.3, zorder=3,
        vmin=forget_ppl.min(), vmax=forget_ppl.max(),
    )
    ax_l.axvline(0, color="black", linewidth=0.4)
    ax_l.set_xlabel("Forget PPL increase over baseline")
    ax_l.set_yticks(y)
    ax_l.set_yticklabels(topics, fontsize=6.5)
    ax_l.grid(True, axis="x", alpha=0.2, linewidth=0.3)

    colors = ["#984EA3" if v < -0.005 else "#bdbdbd" for v in mean_gamma]
    ax_r.hlines(y, 0, mean_gamma, color="#636363", linewidth=0.6)
    ax_r.scatter(mean_gamma, y, color=colors, s=35, edgecolors="black", linewidth=0.3, zorder=3)
    ax_r.axvline(0, color="black", linewidth=0.4)
    ax_r.set_xlabel(r"Mean accuracy $\gamma$")
    ax_r.xaxis.set_major_formatter(mticker.PercentFormatter(xmax=1, decimals=1))
    ax_r.grid(True, axis="x", alpha=0.2, linewidth=0.3)

    save_fig(fig, out, "figA1_forget_ppl_vs_retention")


def figA2_gsm8k_anomaly(out: Path) -> None:
    null_scores = {m: NULL_BIN[3 + i] for i, m in enumerate(ACCURACY_METRICS)}

    fig, (ax_l, ax_r) = plt.subplots(
        1, 2, figsize=(7, 3.5), gridspec_kw={"width_ratios": [3, 2]},
    )

    x = np.arange(len(ACCURACY_METRICS))
    width = 0.3
    base_vals = [BASELINE[m] for m in ACCURACY_METRICS]
    null_vals = [null_scores[m] for m in ACCURACY_METRICS]

    ax_l.bar(x - width / 2, base_vals, width, color="#4393c3", label="Baseline")
    ax_l.bar(x + width / 2, null_vals, width, color="#d6604d", label="Null-Bin")
    ax_l.set_xticks(x)
    ax_l.set_xticklabels(
        ["GSM8K", "MMLU\nSocial Sci.", "MMLU\nSTEM", "SocialIQA"], fontsize=7,
    )
    ax_l.set_ylabel("Accuracy")
    ax_l.set_ylim(0, 0.95)
    ax_l.legend(fontsize=7, loc="upper left")
    ax_l.grid(True, axis="y", alpha=0.2, linewidth=0.3)

    for i, (b, n) in enumerate(zip(base_vals, null_vals)):
        g = (n - b) / abs(b)
        color = "#b2182b" if g < -0.05 else "#333"
        ax_l.text(
            i, max(b, n) + 0.02, f"{g:+.1%}",
            ha="center", va="bottom", fontsize=7, fontweight="bold", color=color,
        )

    mc_m = ["mmlu_socsci", "mmlu_stem", "socialiqa"]
    mc_g = np.mean([gamma(null_scores[m], m) for m in mc_m])
    gsm_g = gamma(null_scores["gsm8k"], "gsm8k")

    bars = ax_r.bar(
        ["MC benchmarks\n(mean)", "GSM8K\n(generative)"],
        [mc_g, gsm_g], width=0.45,
        color=["#4393c3", "#d6604d"], edgecolor="black", linewidth=0.4,
    )
    ax_r.axhline(0, color="black", linewidth=0.4)
    ax_r.set_ylabel(r"$\gamma$")
    ax_r.set_title("MC vs Generative", fontweight="bold", fontsize=9)
    ax_r.grid(True, axis="y", alpha=0.2, linewidth=0.3)
    ax_r.yaxis.set_major_formatter(mticker.PercentFormatter(xmax=1, decimals=0))

    for bar, val in zip(bars, [mc_g, gsm_g]):
        xc = bar.get_x() + bar.get_width() / 2
        if abs(val) > 0.05:
            ax_r.text(
                xc, val / 2, f"$\\gamma$={val:+.1%}",
                ha="center", va="center", fontsize=7.5, fontweight="bold", color="white",
            )
        else:
            ax_r.text(
                xc, val - 0.02, f"$\\gamma$={val:+.1%}",
                ha="center", va="top", fontsize=7.5, fontweight="bold",
            )

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


def figA3_wikitext_ppl(df: pd.DataFrame, out: Path) -> None:
    ds = df.sort_values("wikitext_ppl", ascending=True)
    fig, ax = plt.subplots(figsize=(6, 7))
    increase = ds["wikitext_ppl"].values - BASELINE["wikitext_ppl"]
    colors = [
        "#d73027" if v > 1.5 else "#fc8d59" if v > 0.8 else "#fee08b"
        for v in increase
    ]
    ax.barh(
        range(len(ds)), ds["wikitext_ppl"],
        color=colors, height=0.7, edgecolor="none",
    )
    ax.axvline(
        BASELINE["wikitext_ppl"], color="black", linestyle="--", linewidth=0.7,
        label=f"Baseline ({BASELINE['wikitext_ppl']:.2f})",
    )
    ax.set_yticks(range(len(ds)))
    ax.set_yticklabels(ds["topic"], fontsize=7)
    ax.set_xlabel(r"Wikitext-2 perplexity $\longrightarrow$ (higher = worse)")
    ax.legend(loc="lower right", fontsize=7, framealpha=0.9)
    ax.grid(True, axis="x", alpha=0.2, linewidth=0.3)
    for i, (ppl, inc) in enumerate(zip(ds["wikitext_ppl"], increase)):
        ax.text(ppl + 0.05, i, f"+{inc:.2f}", va="center", ha="left", fontsize=6, color="#333")
    ax.set_xlim(BASELINE["wikitext_ppl"] - 0.3, ds["wikitext_ppl"].max() + 0.5)
    save_fig(fig, out, "figA3_wikitext_ppl")


def figA4_checkpoint_efficiency(out: Path) -> None:
    from scripts.visualization.experiment1_single_bin import RESULTS as R1, COLS as C1

    df1 = pd.DataFrame(R1, columns=C1)
    for m in ACCURACY_METRICS:
        df1[f"{m}_gamma"] = (df1[m] - BASELINE[m]) / abs(BASELINE[m])
    df1["mean_acc_gamma"] = df1[[f"{m}_gamma" for m in ACCURACY_METRICS]].mean(axis=1)

    null_gamma = np.mean(
        [gamma(NULL_BIN[3 + i], m) for i, m in enumerate(ACCURACY_METRICS)]
    )

    fig, ax = plt.subplots(figsize=(5, 4))
    ax.scatter(
        df1["checkpoint"], df1["mean_acc_gamma"],
        color=EXP_COLORS["exp1"], edgecolors="black", linewidth=0.3,
        s=25, label="Single-Bin (n=24)", zorder=3, alpha=0.7,
    )
    ax.scatter(
        NULL_BIN[2], null_gamma,
        color=EXP_COLORS["exp3"], marker="*", s=180,
        edgecolors="black", linewidth=0.5, zorder=4, label="Null-Bin Control",
    )

    ax.axhline(0, color="gray", linestyle="--", linewidth=0.4, alpha=0.5)
    ax.set_xlabel("Selected checkpoint (training steps)")
    ax.set_ylabel(r"Mean accuracy $\gamma$")
    ax.yaxis.set_major_formatter(mticker.PercentFormatter(xmax=1, decimals=1))
    ax.grid(True, alpha=0.2, linewidth=0.3)
    ax.legend(fontsize=7, loc="lower right")
    save_fig(fig, out, "figA4_checkpoint_efficiency")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=Path, default=Path("artifacts/paper_appendix"))
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    paper_rc()

    df = build_exp1_df()

    print("Generating appendix figures...")
    figA1_forget_ppl(df, args.output_dir)
    figA2_gsm8k_anomaly(args.output_dir)
    figA3_wikitext_ppl(df, args.output_dir)
    figA4_checkpoint_efficiency(args.output_dir)
    print(f"All appendix figures saved to {args.output_dir}/")


if __name__ == "__main__":
    main()