retention.py

"""Just-learned vs. retention decomposition of recall accuracy.

This module derives a new view from existing `events.jsonl` files without
re-training: at each checkpoint N, it splits recall questions into

- **just-learned**: questions whose `source_lessons == {N}` — i.e., the
  lesson whose training just finished.
- **retention**: questions whose `source_lessons` is a non-empty subset
  of `{0, ..., N-1}` — i.e., strictly older facts the model must have
  kept around across at least one additional training step.

The cycle-8 finding (§3.6) is that every architecture in the benchmark
hits 100% just-learned recall at every seed and checkpoint, so the
*only* axis that separates architectures on this benchmark is retention.
This module produces both the numeric table that backs §3.6 and the
`paper/figures/learning_vs_retention.png` figure.

Usage:

    python3 retention.py                          # aggregate + figure
    python3 retention.py --no-figure              # table only
    python3 retention.py --runs 'runs/2026-04-11T03-*'  # custom glob

The default run glob matches the post-cycle-6 multi-seed reruns, which
are the authoritative data for the paper's §3.4 / §3.5 tables.
"""

from __future__ import annotations

import argparse
import glob
import json
import math
import os
import statistics
from collections import defaultdict
from pathlib import Path


def _load_manifest(run_dir: str) -> dict:
    with open(os.path.join(run_dir, "manifest.json")) as f:
        return json.load(f)


def _load_dataset(run_dir: str) -> dict:
    with open(os.path.join(run_dir, "dataset.json")) as f:
        return json.load(f)


def decompose(run_dir: str) -> dict[int, list[tuple[str, str, bool, frozenset[int]]]]:
    """Read events.jsonl and return {checkpoint_lesson: [(qid, kind, correct, source_lessons), ...]}.

    The source_lessons tuple lets callers filter by just-learned / retention
    without re-reading dataset.json.
    """
    ds = _load_dataset(run_dir)
    qid_to_src = {q["qid"]: frozenset(q["source_lessons"]) for q in ds["questions"]}

    cp: dict[int, list] = defaultdict(list)
    with open(os.path.join(run_dir, "events.jsonl")) as f:
        for line in f:
            e = json.loads(line)
            if e.get("event") == "question_answered":
                N = e["checkpoint_lesson"]
                cp[N].append((e["qid"], e["kind"], e["correct"], qid_to_src[e["qid"]]))
    return dict(cp)


def per_checkpoint_split(cp: dict[int, list]) -> list[dict]:
    """For each checkpoint, return {checkpoint, just_new_acc, only_old_acc, ...}.

    Recall kind only. 'just_new' = source_lessons == {N}.
    'only_old' = non-empty source_lessons subset of {0..N-1}.
    only_old is `None` at checkpoint 0 (no older lessons exist yet).
    """
    rows = []
    for N in sorted(cp):
        recall_at_N = [(q, c, s) for (q, k, c, s) in cp[N] if k == "recall"]
        just_new = [c for (q, c, s) in recall_at_N if s == frozenset({N})]
        only_old = (
            [c for (q, c, s) in recall_at_N if s and s.issubset(set(range(N)))]
            if N > 0
            else []
        )
        rows.append(
            {
                "checkpoint": N,
                "just_new_n": len(just_new),
                "just_new_acc": (sum(just_new) / len(just_new)) if just_new else None,
                "only_old_n": len(only_old),
                "only_old_acc": (sum(only_old) / len(only_old)) if only_old else None,
            }
        )
    return rows


def group_key(manifest: dict) -> tuple[str, str]:
    """Stable grouping key: (arch_name, config-suffix-if-relevant)."""
    arch = manifest["arch_name"]
    cfg = manifest.get("arch_config") or {}
    if arch == "replay_mlp":
        return (arch, f"cap={cfg.get('capacity', '?')}")
    if arch == "frozen_embedding_mlp":
        return (arch, f"freeze_after={cfg.get('freeze_after_lesson', '?')}")
    return (arch, "")


def group_label(key: tuple[str, str]) -> str:
    arch, extra = key
    return f"{arch} {extra}".strip()


def _agg_fmt(values: list[float]) -> str:
    if not values:
        return "  -  "
    if len(values) == 1:
        return f"{values[0] * 100:5.1f}%"
    return f"{statistics.mean(values) * 100:5.1f}% ±{statistics.stdev(values) * 100:4.1f}"


def print_report(groups: dict[tuple[str, str], list[str]]) -> dict[tuple[str, str], dict]:
    """Print the per-group decomposition table and return a summary dict.

    Summary shape: {group_key: {'per_n_jn': {N: [vals]}, 'per_n_old': {N: [vals]},
    'cp0_recall': [vals], 'n_runs': int}}.
    """
    # Stable ordering: floor, frozen, novel sprint, replay by capacity, ceiling.
    order = [
        ("catastrophic_mlp", ""),
        ("frozen_embedding_mlp", "freeze_after=0"),
        # Novel-architectures sprint (cycles 9.1 - 9.10).
        ("movefreeze_mlp", ""),
        ("logit_pin_mlp", ""),
        ("dual_timescale_mlp", ""),
        ("pseudo_rehearsal_mlp", ""),
        ("hidden_protect_mlp", ""),
        ("head_expert_mlp", ""),
        ("sdm_head_mlp", ""),
        ("fast_hebbian_mlp", ""),
        ("pinned_expert_mlp", ""),
        ("growing_probe_mlp", ""),
    ]
    for cap in (0, 5, 10, 15, 20, 30, 40, 50, 70, 100):
        order.append(("replay_mlp", f"cap={cap}"))
    order.append(("memory_lookup", ""))

    seen = set()
    summary: dict[tuple[str, str], dict] = {}

    def handle(key: tuple[str, str]) -> None:
        runs = groups[key]
        per_n_jn: dict[int, list[float]] = defaultdict(list)
        per_n_old: dict[int, list[float]] = defaultdict(list)
        cp0_recall: list[float] = []

        for rd in runs:
            cp = decompose(rd)
            # checkpoint 0 lesson-0 recall (did the model actually learn lesson 0?)
            if 0 in cp:
                r0 = [c for (q, k, c, s) in cp[0] if k == "recall" and s == frozenset({0})]
                if r0:
                    cp0_recall.append(sum(r0) / len(r0))

            rows = per_checkpoint_split(cp)
            for r in rows:
                if r["just_new_acc"] is not None:
                    per_n_jn[r["checkpoint"]].append(r["just_new_acc"])
                if r["only_old_acc"] is not None:
                    per_n_old[r["checkpoint"]].append(r["only_old_acc"])

        label = group_label(key)
        print("\n" + "=" * 76)
        print(f"{label}   (n={len(runs)} runs)")
        print("=" * 76)
        if cp0_recall:
            print(
                f"  checkpoint-0 lesson-0 recall: {_agg_fmt(cp0_recall)}"
                "  (did the model actually learn lesson 0?)"
            )
        print("  N | just-learned recall     | retention (old)")
        for N in sorted(per_n_jn):
            jn = per_n_jn[N]
            old = per_n_old.get(N, [])
            print(f"  {N} | {_agg_fmt(jn):22s}  | {_agg_fmt(old)}")

        all_jn = [v for vs in per_n_jn.values() for v in vs]
        all_old = [v for vs in per_n_old.values() for v in vs]
        jn_mean = statistics.mean(all_jn) if all_jn else float("nan")
        old_mean = statistics.mean(all_old) if all_old else float("nan")
        print(
            f"  AGG | just-learned={jn_mean * 100:5.1f}%  "
            f"retention={old_mean * 100:5.1f}%  "
            f"(n={len(all_jn)} jn, {len(all_old)} old)"
        )

        summary[key] = {
            "per_n_jn": dict(per_n_jn),
            "per_n_old": dict(per_n_old),
            "cp0_recall": cp0_recall,
            "n_runs": len(runs),
            "all_jn": all_jn,
            "all_old": all_old,
        }

    for key in order:
        if key in groups:
            handle(key)
            seen.add(key)
    for key in groups:
        if key not in seen:
            handle(key)

    return summary


def make_figure(summary: dict[tuple[str, str], dict], out_path: Path) -> Path:
    """Two-panel figure: retention-over-checkpoints curves (left) and
    aggregate retention bar chart (right).
    """
    import matplotlib.pyplot as plt

    # Curve panel: show a representative set to keep the plot readable.
    curve_keys = [
        ("catastrophic_mlp", "", "#000000", "-"),
        ("frozen_embedding_mlp", "freeze_after=0", "#555555", "--"),
        ("replay_mlp", "cap=10", "#1f77b4", "-"),
        ("replay_mlp", "cap=30", "#2ca02c", "-"),
        ("replay_mlp", "cap=50", "#9467bd", "-"),
        ("replay_mlp", "cap=70", "#ff7f0e", "-"),
        ("replay_mlp", "cap=100", "#d62728", "-"),
        ("memory_lookup", "", "#17becf", "--"),
    ]

    fig, (ax_curve, ax_bar) = plt.subplots(1, 2, figsize=(14, 5))

    # Left: per-checkpoint retention curves for the curated subset.
    for arch, extra, color, ls in curve_keys:
        key = (arch, extra)
        if key not in summary:
            continue
        s = summary[key]
        xs, means, stds = [], [], []
        for N in sorted(s["per_n_old"]):
            vals = s["per_n_old"][N]
            if not vals:
                continue
            xs.append(N)
            means.append(statistics.mean(vals))
            stds.append(statistics.stdev(vals) if len(vals) > 1 else 0.0)
        if not xs:
            continue
        label = group_label(key)
        ax_curve.errorbar(
            xs,
            means,
            yerr=stds,
            marker="o",
            markersize=5,
            linewidth=1.8,
            linestyle=ls,
            color=color,
            capsize=2,
            label=label,
        )

    # 100% just-learned reference line.
    ax_curve.axhline(1.0, linestyle=":", linewidth=1.2, color="#888888", alpha=0.9)
    ax_curve.text(
        1.1,
        1.01,
        "just-learned recall = 100% (all architectures)",
        fontsize=8,
        color="#444444",
        verticalalignment="bottom",
    )

    ax_curve.set_xlabel("checkpoint (lessons taught)")
    ax_curve.set_ylabel("retention: recall on strictly-older facts")
    ax_curve.set_ylim(-0.02, 1.08)
    ax_curve.set_xticks(list(range(1, 10)))
    ax_curve.set_title("retention-over-checkpoints")
    ax_curve.grid(True, alpha=0.3)
    ax_curve.legend(loc="center right", fontsize="small")

    # Right: aggregate retention bar chart, sorted.
    bar_order = [
        ("catastrophic_mlp", ""),
        ("frozen_embedding_mlp", "freeze_after=0"),
        ("replay_mlp", "cap=5"),
        ("replay_mlp", "cap=10"),
        ("replay_mlp", "cap=15"),
        ("replay_mlp", "cap=20"),
        ("replay_mlp", "cap=30"),
        ("replay_mlp", "cap=40"),
        ("replay_mlp", "cap=50"),
        ("replay_mlp", "cap=70"),
        ("replay_mlp", "cap=100"),
        ("memory_lookup", ""),
    ]
    labels: list[str] = []
    means: list[float] = []
    stds: list[float] = []
    colors: list[str] = []
    for key in bar_order:
        if key not in summary:
            continue
        vals = summary[key]["all_old"]
        if not vals:
            continue
        labels.append(group_label(key).replace("_mlp", "").replace("freeze_after=", "frz="))
        means.append(statistics.mean(vals))
        stds.append(statistics.stdev(vals) if len(vals) > 1 else 0.0)
        arch = key[0]
        if arch == "catastrophic_mlp":
            colors.append("#000000")
        elif arch == "frozen_embedding_mlp":
            colors.append("#555555")
        elif arch == "memory_lookup":
            colors.append("#17becf")
        else:
            colors.append("#1f77b4")

    xs = list(range(len(labels)))
    ax_bar.bar(xs, means, yerr=stds, color=colors, capsize=3, alpha=0.85)
    ax_bar.axhline(1.0, linestyle=":", linewidth=1.0, color="#888888", alpha=0.7)
    ax_bar.set_xticks(xs)
    ax_bar.set_xticklabels(labels, rotation=30, ha="right", fontsize=8)
    ax_bar.set_ylabel("mean retention (all checkpoints × seeds)")
    ax_bar.set_ylim(-0.02, 1.08)
    ax_bar.set_title("aggregate retention by architecture")
    ax_bar.grid(True, alpha=0.3, axis="y")

    fig.suptitle(
        "Learning is universal; the benchmark measures retention",
        fontsize=12,
    )
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.tight_layout()
    fig.savefig(out_path, dpi=150)
    plt.close(fig)
    return out_path


def make_novel_figure(summary: dict[tuple[str, str], dict], out_path: Path) -> Path:
    """Novel-architectures-sprint figure. Two panels:
    - Left: aggregate retention (bars), novel + baselines + replay sweep.
    - Right: per-checkpoint retention curves, novel wins vs replay anchor points.
    """
    import matplotlib.pyplot as plt

    novel_keys = [
        ("movefreeze_mlp", ""),
        ("logit_pin_mlp", ""),
        ("dual_timescale_mlp", ""),
        ("pseudo_rehearsal_mlp", ""),
        ("hidden_protect_mlp", ""),
        ("head_expert_mlp", ""),
        ("sdm_head_mlp", ""),
        ("fast_hebbian_mlp", ""),
        ("pinned_expert_mlp", ""),
        ("growing_probe_mlp", ""),
    ]
    baselines = [
        ("catastrophic_mlp", ""),
        ("frozen_embedding_mlp", "freeze_after=0"),
        ("replay_mlp", "cap=5"),
        ("replay_mlp", "cap=10"),
        ("replay_mlp", "cap=20"),
        ("replay_mlp", "cap=30"),
        ("replay_mlp", "cap=50"),
        ("replay_mlp", "cap=70"),
        ("replay_mlp", "cap=100"),
        ("memory_lookup", ""),
    ]

    fig, (ax_bar, ax_curve) = plt.subplots(1, 2, figsize=(16, 6))

    # Left: stacked bar chart grouped by category (baselines + novels interleaved
    # by final-recall order).
    all_rows = []
    for key in baselines + novel_keys:
        if key not in summary:
            continue
        vals = summary[key]["all_old"]
        if not vals:
            continue
        is_novel = key in novel_keys
        all_rows.append((
            statistics.mean(vals),
            group_label(key),
            statistics.stdev(vals) if len(vals) > 1 else 0.0,
            is_novel,
        ))
    all_rows.sort(key=lambda row: row[0])
    xs = list(range(len(all_rows)))
    means = [r[0] for r in all_rows]
    stds = [r[2] for r in all_rows]
    labels = [r[1].replace("_mlp", "").replace("freeze_after=", "frz=") for r in all_rows]
    colors = ["#d62728" if r[3] else "#1f77b4" for r in all_rows]
    ax_bar.bar(xs, means, yerr=stds, color=colors, capsize=3, alpha=0.85)
    ax_bar.axhline(1.0, linestyle=":", linewidth=1.0, color="#888888", alpha=0.7)
    ax_bar.set_xticks(xs)
    ax_bar.set_xticklabels(labels, rotation=45, ha="right", fontsize=7)
    ax_bar.set_ylabel("mean retention (all cps × seeds)")
    ax_bar.set_ylim(-0.02, 1.08)
    ax_bar.set_title("Aggregate retention: novels (red) vs baselines (blue)")
    ax_bar.grid(True, alpha=0.3, axis="y")

    # Right: per-checkpoint retention curves for the top novels + replay anchors.
    curve_keys = [
        ("catastrophic_mlp", "", "#000000", "-"),
        ("logit_pin_mlp", "", "#d62728", "-"),
        ("dual_timescale_mlp", "", "#9467bd", "-"),
        ("head_expert_mlp", "", "#e377c2", "-"),
        ("growing_probe_mlp", "", "#ff7f0e", "-"),
        ("replay_mlp", "cap=30", "#2ca02c", "--"),
        ("replay_mlp", "cap=50", "#17becf", "--"),
        ("replay_mlp", "cap=70", "#1f77b4", "--"),
        ("memory_lookup", "", "#8c564b", ":"),
    ]
    for arch, extra, color, ls in curve_keys:
        key = (arch, extra)
        if key not in summary:
            continue
        s = summary[key]
        xs_c, means_c, stds_c = [], [], []
        for N in sorted(s["per_n_old"]):
            vals = s["per_n_old"][N]
            if not vals:
                continue
            xs_c.append(N)
            means_c.append(statistics.mean(vals))
            stds_c.append(statistics.stdev(vals) if len(vals) > 1 else 0.0)
        if not xs_c:
            continue
        ax_curve.errorbar(
            xs_c, means_c, yerr=stds_c,
            marker="o", markersize=5, linewidth=1.8,
            linestyle=ls, color=color, capsize=2,
            label=group_label(key),
        )
    ax_curve.axhline(1.0, linestyle=":", linewidth=1.0, color="#888888", alpha=0.7)
    ax_curve.set_xlabel("checkpoint (lessons taught)")
    ax_curve.set_ylabel("retention: recall on older facts")
    ax_curve.set_ylim(-0.02, 1.08)
    ax_curve.set_xticks(list(range(1, 10)))
    ax_curve.set_title("Per-checkpoint retention: top novels vs replay anchors")
    ax_curve.grid(True, alpha=0.3)
    ax_curve.legend(loc="center right", fontsize="x-small")

    fig.suptitle("Novel-architectures sprint: 10 experiments vs benchmark baselines", fontsize=13)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.tight_layout()
    fig.savefig(out_path, dpi=150)
    plt.close(fig)
    return out_path


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__ or "")
    ap.add_argument(
        "--runs",
        default="runs/2026-04-11T03-*",
        help="glob pattern for run dirs (default: post-cycle-6 multi-seed reruns)",
    )
    ap.add_argument(
        "--figure",
        default="paper/figures/learning_vs_retention.png",
        help="path for the two-panel figure (set to empty to skip)",
    )
    ap.add_argument(
        "--novel-figure",
        default="",
        help="if set, additionally write a novel-sprint figure to this path",
    )
    ap.add_argument("--no-figure", action="store_true", help="skip figure generation")
    args = ap.parse_args()

    run_dirs = sorted(glob.glob(args.runs))
    groups: dict[tuple[str, str], list[str]] = defaultdict(list)
    for rd in run_dirs:
        mpath = os.path.join(rd, "manifest.json")
        if not os.path.exists(mpath):
            continue
        m = _load_manifest(rd)
        groups[group_key(m)].append(rd)

    if not groups:
        raise SystemExit(f"no runs matched glob {args.runs!r}")

    summary = print_report(groups)

    # Totals
    total_jn = sum(len(s["all_jn"]) for s in summary.values())
    total_old = sum(len(s["all_old"]) for s in summary.values())
    print(
        f"\nTotal checkpoint×seed observations: "
        f"{total_jn} just-learned, {total_old} retention "
        f"across {sum(s['n_runs'] for s in summary.values())} runs"
    )

    if not args.no_figure and args.figure:
        out = make_figure(summary, Path(args.figure))
        print(f"\nwrote figure → {out}")

    if args.novel_figure:
        out2 = make_novel_figure(summary, Path(args.novel_figure))
        print(f"wrote novel-sprint figure → {out2}")


if __name__ == "__main__":
    main()