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

#!/usr/bin/env python3
"""Experiment 2: Multi-Bin Targeted Unlearning — Figures 5-7."""
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
import seaborn as sns

from scripts.visualization._shared import (
    ACCURACY_METRICS, BASELINE, BENCH_COLORS, EXP_COLORS,
    METRIC_LABELS, METRIC_LABELS_WRAP,
    gamma, paper_rc, save_fig,
)

RESULTS = [
    ("GSM8K-A", "gsm8k", 1800, 0.7582, 0.747414, 0.590482, 0.797, 10.775, 12.7943),
    ("GSM8K-B", "gsm8k", 2000, 0.6914, 0.733589, 0.574171, 0.7916, 15.092, 33.0875),
    ("GSM8K-C", "gsm8k", 2000, 0.7354, 0.738941, 0.576476, 0.7995, 13.018, 22.1878),
    ("SocSci-A", "mmlu_socsci", 2000, 0.7437, 0.722862, 0.577522, 0.7746, 20.260, 36.4323),
    ("SocSci-B", "mmlu_socsci", 2200, 0.7210, 0.691554, 0.583471, 0.7596, 26.683, 50.0129),
    ("SocSci-C", "mmlu_socsci", 2000, 0.7324, 0.695046, 0.581136, 0.7505, 19.067, 39.3390),
    ("STEM-A", "mmlu_stem", 1600, 0.7612, 0.747794, 0.598986, 0.7991, 10.538, 12.8607),
    ("STEM-B", "mmlu_stem", 1800, 0.7718, 0.750188, 0.592849, 0.7998, 10.855, 12.8944),
    ("STEM-C", "mmlu_stem", 1600, 0.7680, 0.746359, 0.590859, 0.7973, 11.307, 12.9272),
    ("SIQa-A", "socialiqa", 2600, 0.6823, 0.723957, 0.575314, 0.7284, 12.630, 17.9116),
    ("SIQa-B", "socialiqa", 2400, 0.7453, 0.725527, 0.580134, 0.7411, 11.228, 12.5594),
    ("SIQa-C", "socialiqa", 1800, 0.7650, 0.750120, 0.592634, 0.7912, 10.126, 11.5622),
]

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

HEATMAP_COL_ORDER = ["socialiqa", "mmlu_socsci", "gsm8k", "mmlu_stem"]

VARIANT_ORDER = {"A": 0, "B": 1, "C": 2}


def build_dataframe() -> pd.DataFrame:
    df = pd.DataFrame(RESULTS, columns=COLS)
    for m in ACCURACY_METRICS:
        df[f"{m}_gamma"] = (df[m] - BASELINE[m]) / abs(BASELINE[m])
    base_ppl = BASELINE["wikitext_ppl"]
    df["wikitext_ppl_gamma"] = (df["wikitext_ppl"] - base_ppl) / abs(base_ppl)
    df["mean_acc_gamma"] = df[[f"{m}_gamma" for m in ACCURACY_METRICS]].mean(axis=1)
    df["variant"] = df["run"].str[-1]
    df["variant_order"] = df["variant"].map(VARIANT_ORDER)
    return df


def fig5_heatmap(df: pd.DataFrame, output_dir: Path) -> None:
    gamma_cols = [f"{m}_gamma" for m in HEATMAP_COL_ORDER]
    col_labels = [METRIC_LABELS_WRAP[m] for m in HEATMAP_COL_ORDER]
    target_order = ["gsm8k", "mmlu_socsci", "mmlu_stem", "socialiqa"]

    ordered = pd.concat([
        df[df["target"] == t].sort_values("variant_order") for t in target_order
    ])
    runs = ordered["run"].values

    fig, (ax_a, ax_p) = plt.subplots(
        1, 2, figsize=(8, 5.5),
        gridspec_kw={"width_ratios": [4, 1.2], "wspace": 0.15},
    )

    acc = ordered[gamma_cols].values
    vmax_a = max(abs(acc.min()), abs(acc.max()), 0.01)
    sns.heatmap(
        acc, ax=ax_a,
        xticklabels=col_labels,
        yticklabels=runs,
        cmap="RdBu", center=0, vmin=-vmax_a, vmax=vmax_a,
        annot=True, fmt=".1%", annot_kws={"fontsize": 7},
        linewidths=0.4, linecolor="white",
        cbar_kws={"label": r"Accuracy $\gamma$", "shrink": 0.4, "aspect": 18, "pad": 0.03},
    )
    ax_a.axvline(x=2, color="black", linewidth=1.0)
    for b in [3, 6, 9]:
        ax_a.axhline(y=b, color="black", linewidth=1.0)
    ax_a.set_yticklabels(ax_a.get_yticklabels(), fontsize=9, rotation=0)
    ax_a.tick_params(axis="both", length=0)
    ax_a.set_xlabel("")
    ax_a.set_ylabel("")

    ppl = ordered[["wikitext_ppl_gamma"]].values
    sns.heatmap(
        ppl, ax=ax_p,
        xticklabels=["Wikitext\nPPL"],
        yticklabels=False,
        cmap="Purples", vmin=0, vmax=ppl.max() * 1.05,
        annot=True, fmt=".1%", annot_kws={"fontsize": 7},
        linewidths=0.4, linecolor="white",
        cbar_kws={"label": r"PPL $\gamma$", "shrink": 0.4, "aspect": 18, "pad": 0.15},
    )
    for b in [3, 6, 9]:
        ax_p.axhline(y=b, color="black", linewidth=1.0)
    ax_p.set_yticks([])
    ax_p.tick_params(axis="x", length=0)
    fig.subplots_adjust(left=0.12)
    save_fig(fig, output_dir, "fig5_multi_bin_heatmap")


def fig6_selectivity(df: pd.DataFrame, output_dir: Path) -> None:
    targets = ACCURACY_METRICS
    fig, axes = plt.subplots(2, 2, figsize=(8, 6.5))

    for idx, target in enumerate(targets):
        ax = axes.flat[idx]
        sub = df[df["target"] == target].sort_values("variant_order")
        x = np.arange(len(sub))
        width = 0.18

        for j, m in enumerate(ACCURACY_METRICS):
            col = f"{m}_gamma"
            is_target = m == target
            ax.bar(
                x + j * width, sub[col], width,
                color=BENCH_COLORS[m],
                alpha=1.0 if is_target else 0.25,
                edgecolor="black" if is_target else "none",
                linewidth=0.6 if is_target else 0,
                label=METRIC_LABELS[m],
            )

        ax.set_xticks(x + 1.5 * width)
        ax.set_xticklabels(sub["run"], fontsize=7)
        ax.axhline(0, color="black", linewidth=0.4)
        ax.set_ylabel(r"$\gamma$")
        ax.set_title(f"Target: {METRIC_LABELS[target]}", fontweight="bold")
        ax.grid(True, axis="y", alpha=0.2, linewidth=0.3)
        ax.yaxis.set_major_formatter(mticker.PercentFormatter(xmax=1, decimals=1))
        ax.legend(fontsize=5.5, ncol=2, loc="lower left")

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


def fig7_comparison(df_multi: pd.DataFrame, output_dir: Path) -> None:
    from scripts.visualization.experiment1_single_bin import build_dataframe as build_df1

    df1 = build_df1()

    fig, ax = plt.subplots(figsize=(7, 4))
    n_metrics = len(ACCURACY_METRICS)
    group_positions = np.arange(n_metrics)
    box_width = 0.3
    offset = 0.17

    bp1 = ax.boxplot(
        [df1[f"{m}_gamma"].values for m in ACCURACY_METRICS],
        positions=group_positions - offset,
        widths=box_width, patch_artist=True, showmeans=True,
        meanprops=dict(marker="D", markerfacecolor="white", markersize=3),
    )
    bp2 = ax.boxplot(
        [df_multi[f"{m}_gamma"].values for m in ACCURACY_METRICS],
        positions=group_positions + offset,
        widths=box_width, patch_artist=True, showmeans=True,
        meanprops=dict(marker="D", markerfacecolor="white", markersize=3),
    )
    for box in bp1["boxes"]:
        box.set_facecolor(EXP_COLORS["exp1"])
    for box in bp2["boxes"]:
        box.set_facecolor(EXP_COLORS["exp2"])

    rng = np.random.default_rng(42)
    for i, m in enumerate(ACCURACY_METRICS):
        d1 = df1[f"{m}_gamma"].values
        d2 = df_multi[f"{m}_gamma"].values
        j1 = rng.uniform(-0.06, 0.06, len(d1))
        j2 = rng.uniform(-0.06, 0.06, len(d2))
        ax.scatter(np.full(len(d1), i - offset) + j1, d1, color="black", s=6, alpha=0.3, zorder=3)
        ax.scatter(np.full(len(d2), i + offset) + j2, d2, color="black", s=6, alpha=0.3, zorder=3)

    ax.axhline(0, color="gray", linestyle="--", linewidth=0.5, alpha=0.5)
    ax.set_xticks(group_positions)
    ax.set_xticklabels([METRIC_LABELS[m] for m in ACCURACY_METRICS])
    ax.set_ylabel(r"$\gamma$")
    ax.yaxis.set_major_formatter(mticker.PercentFormatter(xmax=1, decimals=1))
    ax.grid(True, axis="y", alpha=0.2, linewidth=0.3)

    from matplotlib.patches import Patch
    ax.legend(
        handles=[
            Patch(facecolor=EXP_COLORS["exp1"], label="Exp 1: Single-Bin (n=24)"),
            Patch(facecolor=EXP_COLORS["exp2"], label="Exp 2: Multi-Bin (n=12)"),
        ],
        fontsize=7, loc="lower left",
    )
    save_fig(fig, output_dir, "fig7_exp1_vs_exp2_comparison")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=Path, default=Path("artifacts/experiment2_figures"))
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    paper_rc()
    df = build_dataframe()
    print("Generating Experiment 2 figures...")
    fig5_heatmap(df, args.output_dir)
    fig6_selectivity(df, args.output_dir)
    try:
        fig7_comparison(df, args.output_dir)
    except Exception as e:
        print(f"  fig7 skipped: {e}")
    print(f"All figures saved to {args.output_dir}/")


if __name__ == "__main__":
    main()