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

#!/usr/bin/env python3
"""Main paper figures — single-bin unlearning results (all 4 benchmarks).

Color scheme: purple (#984EA3) = social, green (#1b7837) = math/STEM.
Solid fill = reasoning, hatched = knowledge.
Layout: horizontal orientation for paper readability.
"""
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, BENCH_HATCHES,
    EXP_COLORS, HEATMAP_COL_ORDER,
    METRIC_LABELS, METRIC_LABELS_WRAP,
    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

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


def fig1_single_bin_heatmap(df: pd.DataFrame, out: Path) -> None:
    null_gammas = {m: gamma(NULL_BIN[3 + i], m) for i, m in enumerate(ACCURACY_METRICS)}

    gamma_cols = [f"{m}_gamma" for m in HEATMAP_COL_ORDER]
    row_labels = [METRIC_LABELS[m] for m in HEATMAP_COL_ORDER]
    df_s = df.sort_values("mean_acc_gamma", ascending=False)

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

    data = df_plot[gamma_cols].values.T

    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": 7},
        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)

    ax.set_xticklabels(ax.get_xticklabels(), fontsize=7.5, 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=9, 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")


def fig2_per_metric_bars(df: pd.DataFrame, out: Path) -> None:
    panel_order = [BARS_ORDER[0][0], BARS_ORDER[0][1],
                   BARS_ORDER[1][0], BARS_ORDER[1][1]]

    fig, axes = plt.subplots(2, 2, figsize=(16, 8))
    for idx, m in enumerate(panel_order):
        ax = axes.flat[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)
        ax.set_xticklabels(topics, fontsize=6, rotation=45, ha="right")
        ax.axhline(0, color="black", linewidth=0.5)
        ax.set_ylabel(r"$\gamma$", fontsize=8)
        ax.set_title(METRIC_LABELS[m], fontweight="bold", fontsize=10)
        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")


def fig3_null_bin_control(out: Path) -> None:
    df1 = build_exp1_df()
    null_gammas = {m: gamma(NULL_BIN[3 + i], m) for i, m in enumerate(ACCURACY_METRICS)}

    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 = df1[col].values
        null_v = null_gammas[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["exp3"], marker="*", s=150,
            edgecolors="black", linewidth=0.5, zorder=4,
        )

        ax.set_xticks([0, 1])
        ax.set_xticklabels(["Single-Bin\n(n=24)", "Global\nRandom"], 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, out, "fig3_null_bin_control")


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

    df = build_exp1_df()

    print("Generating main paper figures...")
    fig1_single_bin_heatmap(df, args.output_dir)
    fig2_per_metric_bars(df, args.output_dir)
    fig3_null_bin_control(args.output_dir)
    print(f"All main figures saved to {args.output_dir}/")


if __name__ == "__main__":
    main()