1 files changed, 74 insertions, 0 deletions
diff --git a/scripts/plot_teacher_test_gap.py b/scripts/plot_teacher_test_gap.py
new file mode 100644
index 0000000..8c4e0de
--- /dev/null
+++ b/scripts/plot_teacher_test_gap.py
@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+"""Two-panel teacher-task figure: train gap (soft ramp) vs test gap (sign flip).
+
+Reads width_summary.csv from a teacher-task sweep produced by
+downstream_capacity_sweep.py and plots, side by side:
+  left  - FA/BP train gap vs width (log y): the optimization cost, soft ramp;
+  right - FA/BP test gap vs width (linear y, error bars): task-dependent,
+          crosses zero (FA can generalize better at intermediate width).
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Teacher-task train vs test gap figure.")
+    p.add_argument("--run-dir", type=Path,
+                   default=Path("outputs/teacher_test_gap_sgd_T8000"))
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    rows = []
+    with (args.run_dir / "width_summary.csv").open() as fh:
+        for row in csv.DictReader(fh):
+            rows.append({k: float(v) for k, v in row.items()})
+    rows.sort(key=lambda r: r["width"])
+
+    widths = np.array([r["width"] for r in rows])
+    n_fa = np.array([r["runs_fa"] for r in rows])
+    train_gap = np.array([r["train_gap_mean"] for r in rows])
+    train_se = np.array([r["train_gap_std"] for r in rows]) / np.sqrt(n_fa)
+    test_gap = np.array([r["gap_mean"] for r in rows])
+    test_se = np.array([r["gap_std"] for r in rows]) / np.sqrt(n_fa)
+
+    fig, axes = plt.subplots(1, 2, figsize=(11.2, 4.4), dpi=180)
+    axes[0].errorbar(widths, train_gap, yerr=train_se, fmt="o-", color="#c65f16",
+                     capsize=2.5)
+    axes[0].set_xscale("log", base=2)
+    axes[0].set_yscale("log")
+    axes[0].set_xlabel("hidden width")
+    axes[0].set_ylabel("FA train MSE - BP train MSE")
+    axes[0].set_title("Optimization gap: soft, monotone, never zero")
+    axes[0].grid(alpha=0.18, which="both")
+
+    axes[1].errorbar(widths, test_gap, yerr=test_se, fmt="o-", color="#174f91",
+                     capsize=2.5)
+    axes[1].axhline(0.0, color="black", linewidth=1)
+    axes[1].set_xscale("log", base=2)
+    axes[1].set_xlabel("hidden width")
+    axes[1].set_ylabel("FA test MSE - BP test MSE")
+    axes[1].set_title("Generalization gap: task-dependent, changes sign")
+    axes[1].grid(alpha=0.18)
+
+    fig.suptitle("Teacher task (MLP teacher, SGD, T=8000): optimization cost vs test effect",
+                 y=1.02)
+    fig.tight_layout()
+    out = args.run_dir / "teacher_train_vs_test_gap.png"
+    fig.savefig(out, bbox_inches="tight")
+    print(f"plot: {out}")
+
+
+if __name__ == "__main__":
+    main()