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

#!/usr/bin/env python3
"""Figures for the 3-variation ARC rescore unlearning.

A. 3-way on-target gamma bars (easy / challenge / combined) with sd error bars.
B. 3-panel 24x24 influence z-score heatmaps (topic x format) per variation.

Run: uv run --with matplotlib --with pandas python scripts/visualization/arc_rescore_3var_figures.py
"""
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import pandas as pd

ART = Path("artifacts")
FIG = Path("figures")
FIG.mkdir(exist_ok=True)
VARIATIONS = ["easy", "challenge", "combined"]


def fig_gamma_bars():
    df = pd.read_csv(ART / "unlearning_arc_rescore_3var_summary.csv")
    df = df.set_index("variation").reindex(VARIATIONS)
    fig, ax = plt.subplots(figsize=(5.2, 3.6))
    x = range(len(VARIATIONS))
    ax.bar(x, df["on_target_gamma_mean"], yerr=df["on_target_gamma_std"],
           capsize=5, color=["#4C72B0", "#DD8452", "#55A868"], width=0.6)
    ax.axhline(0, color="black", lw=0.8)
    ax.set_xticks(list(x))
    ax.set_xticklabels([v.capitalize() for v in VARIATIONS])
    ax.set_ylabel("On-target gamma (base - unlearned acc)")
    ax.set_title("ARC unlearning on-target effect (5.68M rescore)")
    ax.set_ylim(-0.005, 0.03)
    for i, v in enumerate(VARIATIONS):
        ax.text(i, df.loc[v, "on_target_gamma_mean"] + df.loc[v, "on_target_gamma_std"] + 0.001,
                f"{df.loc[v, 'on_target_gamma_mean']:+.4f}", ha="center", fontsize=8)
    fig.tight_layout()
    out = FIG / "arc_rescore_3var_gamma_bars.png"
    fig.savefig(out, dpi=200)
    plt.close(fig)
    print(f"wrote {out}")


def fig_heatmaps():
    mats = {}
    vmax = 0.0
    for v in VARIATIONS:
        d = pd.read_csv(ART / f"zscored_bin_scores/aggregated/zscored_arc_{v}.csv")
        m = d.pivot_table(index="topic_label", columns="format_label", values="zscore")
        mats[v] = m
        vmax = max(vmax, float(m.abs().max().max()))
    fig, axes = plt.subplots(1, 3, figsize=(16, 6), constrained_layout=True)
    im = None
    for ax, v in zip(axes, VARIATIONS):
        m = mats[v]
        im = ax.imshow(m.values, aspect="auto", cmap="RdBu_r", vmin=-vmax, vmax=vmax)
        ax.set_title(f"ARC-{v.capitalize()}")
        ax.set_xticks(range(len(m.columns)))
        ax.set_xticklabels(m.columns, rotation=90, fontsize=5)
        if v == "easy":
            ax.set_yticks(range(len(m.index)))
            ax.set_yticklabels(m.index, fontsize=5)
        else:
            ax.set_yticks([])
    fig.colorbar(im, ax=axes, shrink=0.7, label="influence z-score")
    fig.suptitle("ARC influence z-scores by topic x format (5.68M rescore)", fontsize=13)
    out = FIG / "arc_rescore_3var_influence_heatmaps.png"
    fig.savefig(out, dpi=200, bbox_inches="tight")
    plt.close(fig)
    print(f"wrote {out}")


if __name__ == "__main__":
    fig_gamma_bars()
    fig_heatmaps()