Add cluster bootstrap metrics, supplementary finite-time figure, teacher nested SEs

finite_time_supplement.py computes 95% cluster bootstrap intervals for
the predictor metrics (resampling forward initializations within each
cell) and renders the six-panel supplementary figure. The table
generator gains the teacher-task nested-SE fragment, computed from the
raw paired rows pulled from the experiment machine.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

2 files changed, 219 insertions, 0 deletions

diff --git a/scripts/finite_time_supplement.py b/scripts/finite_time_supplement.py
new file mode 100644
index 0000000..cf70491
--- /dev/null
+++ b/scripts/finite_time_supplement.py
@@ -0,0 +1,189 @@
+#!/usr/bin/env python3
+"""Cluster bootstrap metrics and the supplementary finite-time figure.
+
+Both outputs are derived from outputs/aaai_depth_experiments/finite_time_rows.csv.
+
+1. A LaTeX fragment with MAE, relative MAE, and correlation for the three
+   predictors, each with a 95% cluster bootstrap confidence interval that
+   resamples the four forward initializations within every architecture cell
+   (feedback draws stay nested inside their forward initialization).
+2. A six-panel supplementary figure: measured cost against depth, the three
+   predictor calibration scatters, and signed residuals against depth and width.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+from collections import defaultdict
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+
+PREDICTORS = [
+    ("fixed_gap", "Frozen initialization operator", "frozen"),
+    ("velocity_gap", "Linear extrapolation", "linear"),
+    ("retangent_gap", "Relinearization", "relinearized"),
+]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--rows", type=Path, default=Path("outputs/aaai_depth_experiments/finite_time_rows.csv"))
+    parser.add_argument("--table-outdir", type=Path, default=Path("outputs/tables"))
+    parser.add_argument("--figure-outdir", type=Path, default=Path("outputs/aaai_supplementary"))
+    parser.add_argument("--bootstrap-samples", type=int, default=10_000)
+    parser.add_argument("--bootstrap-seed", type=int, default=2027)
+    return parser.parse_args()
+
+
+def load(path: Path) -> list[dict[str, str]]:
+    with path.open() as handle:
+        return list(csv.DictReader(handle))
+
+
+def metrics(measured: np.ndarray, predicted: np.ndarray) -> tuple[float, float, float]:
+    mae = float(np.mean(np.abs(predicted - measured)))
+    rmae = mae / float(np.mean(measured))
+    corr = float(np.corrcoef(predicted, measured)[0, 1])
+    return mae, rmae, corr
+
+
+def bootstrap_table(rows: list[dict[str, str]], args: argparse.Namespace) -> None:
+    # cell -> init -> array of shape (8, 1 + predictors)
+    cells: dict[tuple[int, int], dict[int, list[list[float]]]] = defaultdict(lambda: defaultdict(list))
+    for row in rows:
+        cells[(int(row["depth"]), int(row["width"]))][int(row["init_seed"])].append(
+            [float(row["empirical_gap"])] + [float(row[f]) for f, _, _ in PREDICTORS]
+        )
+    stacked = [
+        np.asarray([by_init[i] for i in sorted(by_init)])  # (4, 8, 4)
+        for _, by_init in sorted(cells.items())
+    ]
+    full = np.concatenate([cell.reshape(-1, cell.shape[-1]) for cell in stacked])
+    point = [metrics(full[:, 0], full[:, 1 + k]) for k in range(len(PREDICTORS))]
+
+    rng = np.random.default_rng(args.bootstrap_seed)
+    draws = np.empty((args.bootstrap_samples, len(PREDICTORS), 3))
+    n_init = stacked[0].shape[0]
+    for b in range(args.bootstrap_samples):
+        sample = np.concatenate(
+            [
+                cell[rng.integers(0, n_init, size=n_init)].reshape(-1, cell.shape[-1])
+                for cell in stacked
+            ]
+        )
+        for k in range(len(PREDICTORS)):
+            draws[b, k] = metrics(sample[:, 0], sample[:, 1 + k])
+    lo, hi = np.percentile(draws, [2.5, 97.5], axis=0)
+
+    lines = [
+        r"\begin{tabular}{@{}>{\raggedright\arraybackslash}p{2.4cm}ccc@{}}",
+        r"\toprule",
+        r"Predictor & MAE & Relative MAE & Correlation \\",
+        r"\midrule",
+    ]
+    for k, (_, label, _) in enumerate(PREDICTORS):
+        mae, rmae, corr = point[k]
+        lines.append(
+            f"{label} & {mae:.5f} [{lo[k,0]:.5f}, {hi[k,0]:.5f}] & "
+            f"{100*rmae:.2f}\\% [{100*lo[k,1]:.2f}, {100*hi[k,1]:.2f}] & "
+            f"{corr:.5f} [{lo[k,2]:.5f}, {hi[k,2]:.5f}] \\\\"
+        )
+    lines += [r"\bottomrule", r"\end{tabular}"]
+    args.table_outdir.mkdir(parents=True, exist_ok=True)
+    out = args.table_outdir / "finite_time_metrics_ci.tex"
+    out.write_text("\n".join(lines) + "\n")
+    print(f"wrote {out}")
+
+
+def supplementary_figure(rows: list[dict[str, str]], args: argparse.Namespace) -> None:
+    depths = sorted({int(r["depth"]) for r in rows})
+    widths = sorted({int(r["width"]) for r in rows})
+    depth_color = {d: plt.cm.viridis(i / (len(depths) - 1)) for i, d in enumerate(depths)}
+
+    fig, axes = plt.subplots(2, 3, figsize=(14.0, 8.2), dpi=220)
+
+    ax = axes[0, 0]
+    markers = ["o", "s", "^"]
+    for width, marker in zip(widths, markers):
+        means, sems = [], []
+        for depth in depths:
+            cell = [r for r in rows if int(r["width"]) == width and int(r["depth"]) == depth]
+            init_means = [
+                np.mean([float(r["empirical_gap"]) for r in cell if int(r["init_seed"]) == seed])
+                for seed in sorted({int(r["init_seed"]) for r in cell})
+            ]
+            means.append(np.mean(init_means))
+            sems.append(np.std(init_means, ddof=1) / np.sqrt(len(init_means)))
+        ax.errorbar(depths, means, yerr=sems, marker=marker, capsize=2, label=f"width {width}")
+    ax.set_xlabel("hidden-layer depth")
+    ax.set_ylabel("FA loss - BP loss")
+    ax.set_title("Measured cost by depth and width", fontsize=10)
+    ax.grid(alpha=0.16)
+    ax.legend(fontsize=7)
+
+    scatter_axes = [axes[0, 1], axes[0, 2], axes[1, 0]]
+    for (field, label, short), ax in zip(PREDICTORS, scatter_axes):
+        xs = np.asarray([float(r[field]) for r in rows])
+        ys = np.asarray([float(r["empirical_gap"]) for r in rows])
+        colors = [depth_color[int(r["depth"])] for r in rows]
+        ax.scatter(xs, ys, s=12, c=colors, alpha=0.6)
+        lo = min(xs.min(), ys.min())
+        hi = max(xs.max(), ys.max())
+        pad = 0.04 * (hi - lo)
+        ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], "k-", lw=1)
+        ax.set_xlabel("predicted gap")
+        ax.set_ylabel("measured gap")
+        ax.set_title(label, fontsize=10)
+        ax.grid(alpha=0.16)
+
+    residual_specs = [(axes[1, 1], "depth", "hidden-layer depth"), (axes[1, 2], "width", "hidden width")]
+    predictor_colors = ["#2f6f9f", "#c65f16", "#2b8a3e"]
+    offsets = [-0.18, 0.0, 0.18]
+    for ax, key, xlabel in residual_specs:
+        levels = depths if key == "depth" else widths
+        spacing = min(np.diff(levels)) if len(levels) > 1 else 1.0
+        for (field, label, short), color, off in zip(PREDICTORS, predictor_colors, offsets):
+            xs = np.asarray([float(r[key]) for r in rows]) + off * spacing
+            resid = np.asarray([float(r[field]) - float(r["empirical_gap"]) for r in rows])
+            ax.scatter(xs, resid, s=8, color=color, alpha=0.45, label=label)
+        ax.axhline(0.0, color="0.3", lw=1)
+        ax.set_xlabel(xlabel)
+        ax.set_ylabel("predicted minus measured")
+        ax.set_title(f"Signed residual by {key}", fontsize=10)
+        ax.grid(alpha=0.16)
+        if key == "depth":
+            ax.legend(fontsize=7)
+
+    for ax, letter in zip(axes.flat, "abcdef"):
+        ax.text(-0.14, 1.06, letter, transform=ax.transAxes, fontsize=12, fontweight="bold")
+
+    depth_handles = [
+        plt.Line2D([], [], marker="o", ls="", color=depth_color[d], label=f"depth {d}")
+        for d in depths
+    ]
+    axes[0, 2].legend(handles=depth_handles, fontsize=6, ncols=2, loc="upper left")
+
+    fig.tight_layout(w_pad=2.0, h_pad=2.4)
+    args.figure_outdir.mkdir(parents=True, exist_ok=True)
+    path = args.figure_outdir / "finite_time_supplement.png"
+    fig.savefig(path, bbox_inches="tight")
+    fig.savefig(path.with_suffix(".pdf"), bbox_inches="tight")
+    plt.close(fig)
+    print(f"figure: {path}")
+
+
+def main() -> None:
+    args = parse_args()
+    rows = load(args.rows)
+    bootstrap_table(rows, args)
+    supplementary_figure(rows, args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/generate_paper_tables.py b/scripts/generate_paper_tables.py
index ed66dae..95a67c8 100644
--- a/scripts/generate_paper_tables.py
+++ b/scripts/generate_paper_tables.py
@@ -26,6 +26,9 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--lr-dir", type=Path, default=Path("outputs/learning_rate_compensation")
     )
+    parser.add_argument(
+        "--teacher-dir", type=Path, default=Path("outputs/teacher_test_gap_sgd_T8000")
+    )
     parser.add_argument("--outdir", type=Path, default=Path("outputs/tables"))
     return parser.parse_args()
 
@@ -165,12 +168,39 @@ def learning_rate_table(lr_dir: Path, outdir: Path) -> None:
     write_fragment(outdir / "learning_rate_regimes.tex", lines)
 
 
+def teacher_table(teacher_dir: Path, outdir: Path) -> None:
+    rows = [r for r in read_rows(teacher_dir / "runs.csv") if r["run_type"] == "fa"]
+    by: dict[int, dict[int, dict[str, list[float]]]] = defaultdict(
+        lambda: defaultdict(lambda: defaultdict(list))
+    )
+    for row in rows:
+        cell = by[int(row["width"])][int(row["init_seed"])]
+        cell["train"].append(float(row["train_gap_to_bp"]))
+        cell["test"].append(float(row["test_gap_to_bp"]))
+    lines = [
+        r"\begin{tabular}{@{}rcc@{}}",
+        r"\toprule",
+        r"Width & Train difference & Test difference \\",
+        r"\midrule",
+    ]
+    for width, by_init in sorted(by.items()):
+        train_means = [statistics.mean(v["train"]) for _, v in sorted(by_init.items())]
+        test_means = [statistics.mean(v["test"]) for _, v in sorted(by_init.items())]
+        lines.append(
+            f"{width} & \\({statistics.mean(train_means):.4f}\\pm{sem(train_means):.4f}\\) & "
+            f"\\({statistics.mean(test_means):.4f}\\pm{sem(test_means):.4f}\\) \\\\"
+        )
+    lines += [r"\bottomrule", r"\end{tabular}"]
+    write_fragment(outdir / "teacher_nested_se.tex", lines)
+
+
 def main() -> None:
     args = parse_args()
     initialization_table(args.depth_dir, args.outdir)
     finite_time_table(args.depth_dir, args.outdir)
     cnn_table(args.cnn_dir, args.outdir)
     learning_rate_table(args.lr_dir, args.outdir)
+    teacher_table(args.teacher_dir, args.outdir)
 
 
 if __name__ == "__main__":