n10b/analysis/make_tom_figures.py

"""Generate the held-out ToM result figures from figure4_ci/gamma_ci.json.

Outputs (PDF + PNG) to <n10b>/figures/:
  fig_tom_pragmatic_bar      per-topic social-pragmatic net gamma, 95% CI, social_life highlighted
  fig_tom_social_life_forest social_life net gamma across the 5 usable ToMBench subtasks + pool
  fig_tom_gamma_heatmap      20 ToMBench probes x 24 topics net gamma (acc_uncond)
"""
from __future__ import annotations
import json, math, statistics as st
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

N = "/storage/ice-shared/cs7634/staff/TDA/outputs/n10b"
OUT = Path(f"{N}/figures"); OUT.mkdir(parents=True, exist_ok=True)
J = json.loads(Path(f"{N}/figure4_ci/gamma_ci.json").read_text())
C = J["results"]["acc_uncond"]["absolute"]["net"]
PRAG = ["tombench_hinting_task_test", "tombench_faux_pas_recognition_test", "tombench_strange_story_task"]
SWEET = ["tombench_false_belief_task", "tombench_faux_pas_recognition_test",
         "tombench_strange_story_task", "tombench_hinting_task_test", "tombench_multiple_desires"]
ALL20 = ["tombench_" + s for s 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"]]
TCRIT = {2: 4.303, 8: 2.306}
TOPICS = list(C[PRAG[0]].keys())


def pooled(probes, topic):
    v = []
    for p in probes:
        v += C[p][topic].get("vals", [])
    return v


def short(p):
    return p.replace("tombench_", "").replace("_task", "").replace("_recognition_test", "").replace("_test", "")


# ---- Fig 1: per-topic pragmatic-pooled net gamma, sorted, 95% CI ----
rows = []
for t in TOPICS:
    v = pooled(PRAG, t)
    if len(v) >= 2:
        m, sd = st.mean(v), st.stdev(v)
        rows.append((t, m, TCRIT[len(v) - 1] * sd / math.sqrt(len(v))))
rows.sort(key=lambda x: x[1])
labels = [r[0] for r in rows]; means = [r[1] for r in rows]; cis = [r[2] for r in rows]
colors = ["#c0392b" if l == "social_life" else "#95a5a6" for l in labels]
fig, 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)
plt.tight_layout(); plt.savefig(OUT / "fig_tom_pragmatic_bar.pdf"); plt.savefig(OUT / "fig_tom_pragmatic_bar.png", dpi=150); plt.close()

# ---- Fig 2: social_life forest across the 5 usable subtasks + pragmatic pool ----
items = []
for p in SWEET:
    c = C[p]["social_life"]
    if c["mean"] is not None:
        items.append((short(p), c["mean"], c.get("ci95") or 0.0))
v = pooled(PRAG, "social_life"); m, sd = st.mean(v), st.stdev(v)
items.append(("PRAGMATIC POOL", m, TCRIT[8] * sd / math.sqrt(9)))
fig, ax = plt.subplots(figsize=(6.5, 3.8))
ys = range(len(items))
ax.errorbar([it[1] for it in items], list(ys), xerr=[it[2] for it in items],
            fmt="o", color="#c0392b", ecolor="#c0392b", capsize=3)
ax.set_yticks(list(ys)); ax.set_yticklabels([it[0] for it 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)
plt.tight_layout(); plt.savefig(OUT / "fig_tom_social_life_forest.pdf"); plt.savefig(OUT / "fig_tom_social_life_forest.png", dpi=150); plt.close()

# ---- Fig 3: heatmap 20 ToMBench probes x 24 topics ----
M = np.full((len(TOPICS), len(ALL20)), np.nan)
for j, p in enumerate(ALL20):
    for i, t in enumerate(TOPICS):
        m = C[p][t]["mean"]
        if m is not None:
            M[i, j] = m
fig, ax = plt.subplots(figsize=(10, 8))
vmax = np.nanmax(np.abs(M))
im = ax.imshow(M, aspect="auto", cmap="RdBu", vmin=-vmax, vmax=vmax)
ax.set_xticks(range(len(ALL20))); ax.set_xticklabels([short(p) for p in ALL20], rotation=90, fontsize=6)
ax.set_yticks(range(len(TOPICS))); ax.set_yticklabels(TOPICS, fontsize=6)
sl = TOPICS.index("social_life")
ax.add_patch(plt.Rectangle((-0.5, sl - 0.5), len(ALL20), 1, fill=False, edgecolor="black", lw=1.5))
fig.colorbar(im, ax=ax, label="net $\\gamma$ (acc_uncond)", shrink=0.6)
ax.set_title("Held-out ToMBench: net $\\gamma$ by topic (rows) x subtask (cols)", fontsize=10)
plt.tight_layout(); plt.savefig(OUT / "fig_tom_gamma_heatmap.pdf"); plt.savefig(OUT / "fig_tom_gamma_heatmap.png", dpi=150); plt.close()

# ---- Fig 4: attribution — topic signed influence on the social-pragmatic probes ----
attr = {t: 0.0 for t in TOPICS}
for p in PRAG:
    df = pd.read_csv(f"{N}/bin_scores/queries_{p}_bin_scores.csv")
    df["mass"] = df["mean_score"] * df["doc_count"]
    g = df.groupby("topic_label")["mass"].sum()
    for t in TOPICS:
        if t in g.index:
            attr[t] += g[t] / len(PRAG)
arows = sorted(attr.items(), key=lambda x: x[1])
albls = [r[0] for r in arows]; avals = [r[1] for r in arows]
acol = ["#c0392b" if l == "social_life" else "#95a5a6" for l in albls]
fig, ax = plt.subplots(figsize=(7, 6.5))
ax.barh(range(len(albls)), avals, color=acol, height=0.7)
ax.set_yticks(range(len(albls))); ax.set_yticklabels(albls, 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 is mid-pack -- does NOT match the unlearning ranking)", fontsize=9)
plt.tight_layout(); plt.savefig(OUT / "fig_tom_attribution_bar.pdf"); plt.savefig(OUT / "fig_tom_attribution_bar.png", dpi=150); plt.close()

print("wrote:", [p.name for p in sorted(OUT.glob("fig_tom_*.png"))])