scripts/plot_curriculum_analysis.py

#!/usr/bin/env python3
"""Curriculum-strategy comparison figures (validation set).

Produces four panels comparing the three L10 curriculum-management strategies
(backtracking vs. two no-backtracking baselines):

  (A) Stage-wise learning curve: validation accuracy vs training epoch, with the
      curriculum stage drawn as a step line (shows *how* each method climbs the
      10 stages, and where the no-repair baseline deadlocks).
  (B) Per-hop latent depth-identification accuracy on the validation set
      (frontier_acc: does latent slot m decode to a correct depth-m node?).
  (C) Per-hop decoy-arm leakage on the validation set (fraction of latents that
      decode into the UNREACHABLE arm -- the interpretability failure signature).
  (D) Per-hop exact-depth accuracy (confusion-matrix diagonal / N): the latent
      decodes to a node whose TRUE reachable depth == its slot index.

Interpretability metrics are computed via logit-lens on the final checkpoint of
each arm (node id == token id, so argmax of the LM head at a latent position is a
predicted graph node).

Usage:
  CUDA_VISIBLE_DEVICES=7 PYTHONPATH=. python scripts/plot_curriculum_analysis.py
"""
import argparse
import json
import os
import re

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import torch
from transformers import AutoModelForCausalLM, AutoConfig

from stokenizer import STokenizer
from coconut import Coconut
from scripts.probe_latents import build_prefix_tokens, node_depth_maps


ARMS = [
    # (label, ckpt-dir slug, log slug, color)
    ("Backtracking (CE-gated)", "backtrack-ce", "backtrack_ce", "#e08214"),
    ("Backtracking (superposition-gated)", "backtrack-superpos", "backtrack_superpos", "#762a83"),
    ("Backtracking (frontier-gated)", "backtrack", "backtrack", "#1b7837"),
    ("Current-stage-only", "curstage", "curstage", "#2166ac"),
    ("Retention-gated (no repair)", "accstage-nobt", "accstage_nobt", "#b2182b"),
]


def latest_ckpt(slug):
    d = f"ckpts/star-coconut-L10-bfs-{slug}"
    cks = [f for f in os.listdir(d) if f.startswith("checkpoint_")]
    cks.sort(key=lambda x: int(x.split("_")[1]))
    return os.path.join(d, cks[-1])


def parse_log(log_slug):
    """Return (epochs, val_acc, stage_at_eval, promote_epochs).

    We align each 'Accuracy on validation set' line to the most recent
    'train epoch N/.. stage=S' line preceding it.
    """
    path = f"logs/star_coconut_L10_bfs_{log_slug}.log"
    ep_re = re.compile(r"train epoch (\d+)/\d+ stage=(\d+)")
    acc_re = re.compile(r"Accuracy on validation set: \d+ / \d+ = ([0-9.]+)")
    prom_re = re.compile(r"PROMOTE stage (\d+) -> (\d+)")
    cur_ep, cur_stage = 0, 0
    epochs, accs, stages, promotes = [], [], [], []
    with open(path, errors="ignore") as fh:
        for line in fh:
            m = ep_re.search(line)
            if m:
                cur_ep, cur_stage = int(m.group(1)), int(m.group(2))
                continue
            m = acc_re.search(line)
            if m:
                epochs.append(cur_ep)
                accs.append(float(m.group(1)))
                stages.append(cur_stage)
                continue
            m = prom_re.search(line)
            if m:
                promotes.append((cur_ep, int(m.group(2))))
    return epochs, accs, stages, promotes


@torch.no_grad()
def probe(ckpt, val_path, model_id, L, device, batch_size=64):
    tok = STokenizer()
    latent_id = tok.convert_tokens_to_ids("<|latent|>")
    base = AutoModelForCausalLM.from_config(AutoConfig.from_pretrained(model_id))
    model = Coconut(base, latent_id,
                    tok.convert_tokens_to_ids("<|start-latent|>"),
                    tok.convert_tokens_to_ids("<|end-latent|>"),
                    tok.eos_token_id)
    sd = torch.load(ckpt, map_location="cpu")
    model.load_state_dict(sd, strict=False)
    model.to(device).eval()

    data = json.load(open(val_path))
    frontier = [0] * (L + 1)   # decoded node on correct depth-m frontier (either arm)
    exact = [0] * (L + 1)      # decoded node's TRUE reachable depth == m
    neg = [0] * (L + 1)        # decoded node in the UNREACHABLE arm
    superpos = [0] * (L + 1)   # top-|F| logits == the full depth-m frontier SET
    total = 0
    # confusion: row = latent slot m (1..L), col = TRUE role of decoded node
    #   cols 0..L = reachable depth, L+1 = decoy(neg) arm, L+2 = off-graph
    NEG, OFF = L + 1, L + 2
    conf = [[0] * (L + 3) for _ in range(L + 1)]
    for i in range(0, len(data), batch_size):
        batch = data[i:i + batch_size]
        seqs = [build_prefix_tokens(s, tok) + [latent_id] * L for s in batch]
        maxlen = max(len(x) for x in seqs)
        seqs = [x for x in seqs if len(x) == maxlen]  # fixed L => equal already
        input_ids = torch.tensor([build_prefix_tokens(s, tok) + [latent_id] * L
                                  for s in batch], device=device)
        attn = torch.ones_like(input_ids)
        pos = torch.arange(input_ids.shape[1], device=device).unsqueeze(0).expand(len(batch), -1)
        logits = model.forward(input_ids, attn, input_ids.clone(), pos).logits
        for bi, s in enumerate(batch):
            total += 1
            role = node_depth_maps(s, L)
            root_pos = len(build_prefix_tokens(s, tok)) - 1
            for m in range(1, L + 1):
                slot_logits = logits[bi, root_pos + (m - 1)]
                pred = int(torch.argmax(slot_logits).item())
                F = {int(n) for n in s["neighbor_k"].get(str(m), [])}
                if F:
                    top = torch.topk(slot_logits, k=len(F)).indices.tolist()
                    if {int(t) for t in top} == F:
                        superpos[m] += 1
                if pred in s["neighbor_k"].get(str(m), []):
                    frontier[m] += 1
                r = role.get(pred)
                if r is None:
                    conf[m][OFF] += 1
                elif r[0] == "neg":
                    neg[m] += 1
                    conf[m][NEG] += 1
                else:
                    conf[m][r[1]] += 1
                    if r[1] == m:
                        exact[m] += 1
    hops = list(range(1, L + 1))
    return {
        "hops": hops,
        "total": total,
        "frontier": [frontier[m] / total for m in hops],
        "exact": [exact[m] / total for m in hops],
        "neg": [neg[m] / total for m in hops],
        "superposition": [superpos[m] / total for m in hops],
        "conf": conf,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--val", default="data/star_2arm_L10_valid_fo_bfs.json")
    ap.add_argument("--model_id", default="configs/symbol-2layer-8head-768dim-L20.json")
    ap.add_argument("--L", type=int, default=10)
    ap.add_argument("--device", default="cuda:0")
    ap.add_argument("--out", default="figs/curriculum_analysis.png")
    ap.add_argument("--cache", default="figs/curriculum_metrics.json")
    args = ap.parse_args()
    os.makedirs(os.path.dirname(args.out), exist_ok=True)

    results = {}
    for label, ck_slug, log_slug, color in ARMS:
        ckpt = latest_ckpt(ck_slug)
        print(f"[{label}] log={log_slug}  ckpt={ckpt}")
        pr = probe(ckpt, args.val, args.model_id, args.L, args.device)
        ep, acc, stg, prom = parse_log(log_slug)
        results[label] = {"color": color, "ckpt": ckpt, "probe": pr,
                          "epochs": ep, "val_acc": acc, "stage": stg,
                          "promotes": prom}
    json.dump(results, open(args.cache, "w"), indent=2)

    fig, axes = plt.subplots(2, 2, figsize=(15, 11))
    axA, axB, axC, axD = axes[0, 0], axes[0, 1], axes[1, 0], axes[1, 1]

    # (A) stage-wise learning curve
    for label, *_ , color in ARMS:
        r = results[label]
        axA.plot(r["epochs"], r["val_acc"], color=color, lw=1.6, label=label)
    axA.set_xlabel("training epoch"); axA.set_ylabel("validation accuracy")
    axA.set_title("(A) Stage-wise learning: val accuracy vs epoch")
    axA.set_ylim(0, 1); axA.grid(alpha=0.3); axA.legend(fontsize=8, loc="upper left")
    axA2 = axA.twinx()
    for label, *_, color in ARMS:
        r = results[label]
        axA2.step(r["epochs"], r["stage"], color=color, lw=1.0, ls=":", alpha=0.6, where="post")
    axA2.set_ylabel("curriculum stage (dotted)")
    axA2.set_ylim(0, args.L + 0.5)

    hops = list(range(1, args.L + 1))
    # (B) per-hop depth-identification (frontier) accuracy
    for label, *_, color in ARMS:
        axB.plot(hops, results[label]["probe"]["frontier"], "-o", color=color, label=label, ms=4)
    axB.set_xlabel("latent slot / hop m"); axB.set_ylabel("frontier accuracy")
    axB.set_title("(B) Latent depth-ID accuracy per hop (val)")
    axB.set_ylim(0, 1.02); axB.grid(alpha=0.3); axB.legend(fontsize=8, loc="lower left")

    # (C) per-hop decoy-arm leakage
    for label, *_, color in ARMS:
        axC.plot(hops, results[label]["probe"]["neg"], "-o", color=color, label=label, ms=4)
    axC.set_xlabel("latent slot / hop m"); axC.set_ylabel("fraction decoding to decoy arm")
    axC.set_title("(C) Decoy-arm leakage per hop (val)")
    axC.grid(alpha=0.3); axC.legend(fontsize=8, loc="upper left")

    # (D) per-hop SUPERPOSITION accuracy: top-|F| logits == full frontier set
    for label, *_, color in ARMS:
        axD.plot(hops, results[label]["probe"]["superposition"], "-o", color=color, label=label, ms=4)
    axD.set_xlabel("latent slot / hop m"); axD.set_ylabel("superposition accuracy")
    axD.set_title("(D) Superposition per hop: top-2 == BOTH frontier nodes (val)")
    axD.set_ylim(0, 1.02); axD.grid(alpha=0.3); axD.legend(fontsize=8, loc="lower left")

    fig.suptitle("L10 2-arm star: curriculum strategy comparison (validation set)", fontsize=13)
    fig.tight_layout(rect=[0, 0, 1, 0.98])
    fig.savefig(args.out, dpi=150)
    print(f"\nsaved combined figure -> {args.out}")

    outdir = os.path.dirname(args.out)

    # ---- individual panels ----
    def save_line(fname, ykey, ylabel, title, ylim=None, loc="lower left"):
        f, ax = plt.subplots(figsize=(7, 5))
        for label, *_, color in ARMS:
            ax.plot(hops, results[label]["probe"][ykey], "-o", color=color, label=label, ms=5, lw=1.8)
        ax.set_xlabel("latent slot / hop m"); ax.set_ylabel(ylabel); ax.set_title(title)
        if ylim: ax.set_ylim(*ylim)
        ax.grid(alpha=0.3); ax.set_xticks(hops); ax.legend(fontsize=9, loc=loc)
        f.tight_layout(); p = os.path.join(outdir, fname); f.savefig(p, dpi=150); plt.close(f)
        print(f"saved -> {p}")

    # (A) stage-wise learning as its own figure
    fA, ax = plt.subplots(figsize=(9, 5.5))
    for label, *_, color in ARMS:
        r = results[label]
        ax.plot(r["epochs"], r["val_acc"], color=color, lw=1.7, label=label)
    ax.set_xlabel("training epoch"); ax.set_ylabel("validation accuracy")
    ax.set_title("Stage-wise learning: val accuracy vs epoch (dotted = curriculum stage)")
    ax.set_ylim(0, 1); ax.grid(alpha=0.3); ax.legend(fontsize=9, loc="upper left")
    ax2 = ax.twinx()
    for label, *_, color in ARMS:
        r = results[label]
        ax2.step(r["epochs"], r["stage"], color=color, lw=1.1, ls=":", alpha=0.65, where="post")
    ax2.set_ylabel("curriculum stage (dotted)"); ax2.set_ylim(0, args.L + 0.5)
    fA.tight_layout(); pA = os.path.join(outdir, "panelA_stagewise_learning.png")
    fA.savefig(pA, dpi=150); plt.close(fA); print(f"saved -> {pA}")

    save_line("panelB_depth_id_accuracy.png", "frontier", "frontier accuracy",
              "Latent depth-ID accuracy per hop (validation)", ylim=(0, 1.02))
    save_line("panelC_decoy_leakage.png", "neg", "fraction decoding to decoy arm",
              "Decoy-arm leakage per hop (validation)", loc="upper left")
    save_line("panelD_exact_depth.png", "exact", "exact-depth accuracy",
              "Exact BFS-depth identification per hop (validation)", ylim=(0, 1.02))
    save_line("panelE_superposition.png", "superposition", "superposition accuracy",
              "Superposition per hop: top-2 == BOTH frontier nodes (validation)",
              ylim=(0, 1.02))

    # ---- confusion-matrix heatmaps (one per arm) ----
    fh, axes_h = plt.subplots(1, len(ARMS), figsize=(6.3 * len(ARMS), 6))
    for ax, (label, *_, color) in zip(axes_h, ARMS):
        conf = results[label]["probe"]["conf"]
        L = args.L
        mat = [[conf[m][c] for c in range(L + 3)] for m in range(1, L + 1)]
        im = ax.imshow(mat, aspect="auto", cmap="magma")
        ax.set_xticks(range(L + 3))
        ax.set_xticklabels([str(d) for d in range(L + 1)] + ["neg", "off"], fontsize=8)
        ax.set_yticks(range(L)); ax.set_yticklabels(range(1, L + 1))
        ax.set_xlabel("TRUE role of decoded node (reachable depth / decoy / off-graph)")
        ax.set_ylabel("latent slot m")
        ax.set_title(label)
        for mi in range(L):
            for c in range(L + 3):
                v = mat[mi][c]
                if v:
                    ax.text(c, mi, str(v), ha="center", va="center",
                            color="white" if v < results[label]["probe"]["total"] * 0.5 else "black",
                            fontsize=6)
        fh.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    fh.suptitle("Latent confusion matrices: slot m vs. TRUE decoded-node depth (validation, N=256)", fontsize=13)
    fh.tight_layout(rect=[0, 0, 1, 0.96])
    ph = os.path.join(outdir, "confusion_heatmaps.png")
    fh.savefig(ph, dpi=150); plt.close(fh); print(f"saved -> {ph}")

    print(f"\nsaved metrics -> {args.cache}")


if __name__ == "__main__":
    main()