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

#!/usr/bin/env python3
"""Experiment 3: Null-Bin Control — Figures 8-10."""
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,
)

NULL_BIN = ("Null-Bin Control", "null", 2800, 0.41395, 0.744879, 0.591468, 0.7977, 14.77, 20.2)

COLS = [
    "run", "target", "checkpoint", "gsm8k", "mmlu_socsci",
    "mmlu_stem", "socialiqa", "wikitext_ppl", "forget_ppl",
]


def _load_exp_gammas():
    from scripts.visualization.experiment1_single_bin import build_dataframe as build_df1
    from scripts.visualization.experiment2_multi_bin import build_dataframe as build_df2
    return build_df1(), build_df2()


def fig8_control(output_dir: Path) -> None:
    df1, df2 = _load_exp_gammas()
    null_gammas = {m: gamma(NULL_BIN[3 + i], m) for i, m in enumerate(ACCURACY_METRICS)}

    fig, axes = plt.subplots(1, 4, figsize=(7, 3.5))

    for idx, m in enumerate(ACCURACY_METRICS):
        ax = axes[idx]
        col = f"{m}_gamma"
        exp1_v = df1[col].values
        exp2_v = df2[col].values
        null_v = null_gammas[m]

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

        rng = np.random.default_rng(42)
        for i, d in enumerate([exp1_v, exp2_v]):
            jitter = rng.uniform(-0.08, 0.08, len(d))
            ax.scatter(
                np.full(len(d), i) + jitter, d,
                color="black", s=6, alpha=0.35, zorder=3,
            )

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

        ax.set_xticks([0, 1, 2])
        ax.set_xticklabels(["Exp 1\n(24)", "Exp 2\n(12)", "Null\nCtrl"], fontsize=7)
        ax.axhline(0, color="gray", linestyle="--", linewidth=0.4, alpha=0.5)
        ax.set_title(METRIC_LABELS[m], fontweight="bold", fontsize=9)
        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$")

    fig.tight_layout()
    save_fig(fig, output_dir, "fig8_control_vs_targeted")


def fig9_anomaly(output_dir: 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([METRIC_LABELS[m] for m in ACCURACY_METRICS], 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, output_dir, "fig9_gsm8k_anomaly")


def fig10_efficiency(output_dir: Path) -> None:
    from scripts.visualization.experiment1_single_bin import (
        RESULTS as R1, COLS as C1,
    )
    from scripts.visualization.experiment2_multi_bin import (
        RESULTS as R2, COLS as C2,
    )

    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)

    df2 = pd.DataFrame(R2, columns=C2)
    for m in ACCURACY_METRICS:
        df2[f"{m}_gamma"] = (df2[m] - BASELINE[m]) / abs(BASELINE[m])
    df2["mean_acc_gamma"] = df2[[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="Exp 1: Single-Bin", zorder=3, alpha=0.7,
    )
    ax.scatter(
        df2["checkpoint"], df2["mean_acc_gamma"],
        color=EXP_COLORS["exp2"], edgecolors="black", linewidth=0.3,
        s=35, marker="s", label="Exp 2: Multi-Bin", zorder=3, alpha=0.8,
    )
    ax.scatter(
        NULL_BIN[2], null_gamma,
        color=EXP_COLORS["exp3"], marker="*", s=180,
        edgecolors="black", linewidth=0.5, zorder=4, label="Exp 3: Null-Bin",
    )

    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, output_dir, "fig10_checkpoint_efficiency")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=Path, default=Path("artifacts/experiment3_figures"))
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    paper_rc()
    print("Generating Experiment 3 figures...")
    fig8_control(args.output_dir)
    fig9_anomaly(args.output_dir)
    fig10_efficiency(args.output_dir)
    print(f"All figures saved to {args.output_dir}/")


if __name__ == "__main__":
    main()