social-data-attribution-w2/scripts/analysis/tom_figure_helpers.py

from __future__ import annotations

import json
import math
import statistics as st
from pathlib import Path
from typing import cast

import matplotlib

matplotlib.use("Agg")
matplotlib.rcParams.update({"pdf.fonttype": 42, "ps.fonttype": 42})

import matplotlib.pyplot as plt  # noqa: E402
from matplotlib.patches import Rectangle  # noqa: E402
import numpy as np  # noqa: E402
import pandas as pd  # noqa: E402

PRAGMATIC = [
    f"tombench_{name}"
    for name in "hinting_task_test faux_pas_recognition_test strange_story_task".split()
]
SWEET = [
    f"tombench_{name}"
    for name in "false_belief_task faux_pas_recognition_test strange_story_task hinting_task_test multiple_desires".split()
]
ALL20 = [
    f"tombench_{name}"
    for name in (
        "ambiguous_story_task completion_of_failed_actions discrepant_desires "
        "discrepant_emotions discrepant_intentions emotion_regulation false_belief_task "
        "faux_pas_recognition_test hidden_emotions hinting_task_test knowledge_attention_links "
        "knowledge_pretend_play_links moral_emotions multiple_desires percepts_knowledge_links "
        "persuasion_story_task prediction_of_actions scalar_implicature_test strange_story_task "
        "unexpected_outcome_test"
    ).split()
]
TCRIT = {2: 4.303, 8: 2.306}


def load_gamma_cells(gamma_ci_path: Path) -> dict:
    gamma_ci = json.loads(gamma_ci_path.read_text())
    return gamma_ci["results"]["acc_uncond"]["absolute"]["net"]


def pooled(cells: dict, probes: list[str], topic: str) -> list[float]:
    vals: list[float] = []
    for probe in probes:
        vals += cells[probe][topic].get("vals", [])
    return vals


def short_probe(probe: str) -> str:
    return (
        probe.replace("tombench_", "")
        .replace("_task", "")
        .replace("_recognition_test", "")
        .replace("_test", "")
    )


def save_current_figure(out_dir: Path, stem: str, dpi: int = 300) -> None:
    plt.tight_layout()
    plt.savefig(out_dir / f"{stem}.pdf")
    plt.savefig(out_dir / f"{stem}.png", dpi=dpi)
    plt.close()


def plot_pragmatic_bar(cells: dict, topics: list[str], out_dir: Path) -> None:
    rows = []
    for topic in topics:
        vals = pooled(cells, PRAGMATIC, topic)
        if len(vals) >= 2:
            mean = st.mean(vals)
            ci95 = TCRIT[len(vals) - 1] * st.stdev(vals) / math.sqrt(len(vals))
            rows.append((topic, mean, ci95))
    rows.sort(key=lambda row: row[1])

    labels = [row[0] for row in rows]
    means = [row[1] for row in rows]
    cis = [row[2] for row in rows]
    colors = ["#c0392b" if label == "social_life" else "#95a5a6" for label in labels]

    _, ax = plt.subplots(figsize=(7, 6.5))
    ax.barh(
        range(len(labels)),
        means,
        xerr=cis,
        color=colors,
        height=0.7,
        error_kw={"elinewidth": 0.8, "capsize": 2},
    )
    ax.set_yticks(range(len(labels)))
    ax.set_yticklabels(labels, fontsize=7)
    ax.axvline(0, color="black", lw=0.9)
    ax.set_xlabel("net $\\gamma$ on social-pragmatic ToM (acc_uncond, 3 seeds)")
    ax.set_title(
        "Topic-unlearning damage on social-pragmatic ToM\n"
        "(hinting + faux-pas + strange-stories; 95% CI)",
        fontsize=10,
    )
    save_current_figure(out_dir, "fig_tom_pragmatic_bar")


def plot_social_life_forest(cells: dict, out_dir: Path) -> None:
    items = []
    for probe in SWEET:
        cell = cells[probe]["social_life"]
        if cell["mean"] is not None:
            items.append((short_probe(probe), cell["mean"], cell.get("ci95") or 0.0))

    vals = pooled(cells, PRAGMATIC, "social_life")
    mean = st.mean(vals)
    ci95 = TCRIT[8] * st.stdev(vals) / math.sqrt(9)
    items.append(("PRAGMATIC POOL", mean, ci95))

    _, ax = plt.subplots(figsize=(6.5, 3.8))
    ys = range(len(items))
    ax.errorbar(
        [item[1] for item in items],
        list(ys),
        xerr=[item[2] for item in items],
        fmt="o",
        color="#c0392b",
        ecolor="#c0392b",
        capsize=3,
    )
    ax.set_yticks(list(ys))
    ax.set_yticklabels([item[0] for item in items], fontsize=8)
    ax.axvline(0, color="black", lw=0.9)
    ax.invert_yaxis()
    ax.set_xlabel("social_life net $\\gamma$ (acc_uncond, 95% CI)")
    ax.set_title(
        "Unlearning social_life: per-subtask effect on held-out ToM", fontsize=10
    )
    save_current_figure(out_dir, "fig_tom_social_life_forest")


def plot_gamma_heatmap(cells: dict, topics: list[str], out_dir: Path) -> None:
    matrix = np.full((len(topics), len(ALL20)), np.nan)
    for col_idx, probe in enumerate(ALL20):
        for row_idx, topic in enumerate(topics):
            mean = cells[probe][topic]["mean"]
            if mean is not None:
                matrix[row_idx, col_idx] = mean

    _, ax = plt.subplots(figsize=(10, 8))
    vmax = np.nanmax(np.abs(matrix))
    image = ax.imshow(matrix, aspect="auto", cmap="RdBu", vmin=-vmax, vmax=vmax)
    ax.set_xticks(range(len(ALL20)))
    ax.set_xticklabels([short_probe(probe) for probe in ALL20], rotation=90, fontsize=6)
    ax.set_yticks(range(len(topics)))
    ax.set_yticklabels(topics, fontsize=6)
    social_life_idx = topics.index("social_life")
    ax.add_patch(
        Rectangle(
            (-0.5, social_life_idx - 0.5),
            len(ALL20),
            1,
            fill=False,
            edgecolor="black",
            lw=1.5,
        )
    )
    plt.colorbar(image, ax=ax, label="net $\\gamma$ (acc_uncond)", shrink=0.6)
    ax.set_title("Held-out ToMBench: net $\\gamma$ by topic x subtask", fontsize=10)
    save_current_figure(out_dir, "fig_tom_gamma_heatmap")


def plot_attribution_bar(cells: dict, bin_scores_dir: Path, out_dir: Path) -> None:
    topics = list(cells[PRAGMATIC[0]].keys())
    attr = {topic: 0.0 for topic in topics}
    for probe in PRAGMATIC:
        df = pd.read_csv(bin_scores_dir / f"queries_{probe}_bin_scores.csv")
        df["mass"] = df["mean_score"] * df["doc_count"]
        grouped = cast(pd.Series, df.groupby("topic_label")["mass"].sum())
        for topic in topics:
            if topic in grouped.index:
                attr[topic] += float(grouped.loc[topic]) / len(PRAGMATIC)

    rows = sorted(attr.items(), key=lambda row: row[1])
    labels = [row[0] for row in rows]
    values = [row[1] for row in rows]
    colors = ["#c0392b" if label == "social_life" else "#95a5a6" for label in labels]

    _, ax = plt.subplots(figsize=(7, 6.5))
    ax.barh(range(len(labels)), values, color=colors, height=0.7)
    ax.set_yticks(range(len(labels)))
    ax.set_yticklabels(labels, fontsize=7)
    ax.axvline(0, color="black", lw=0.9)
    ax.set_xlabel("mean signed attribution mass (social-pragmatic probes)")
    ax.set_title(
        "Attribution: topic influence on social-pragmatic ToM\n"
        "(social_life does not match the unlearning ranking)",
        fontsize=9,
    )
    save_current_figure(out_dir, "fig_tom_attribution_bar")