1 files changed, 86 insertions, 0 deletions
diff --git a/scripts/plot_phase_transition_training_time.py b/scripts/plot_phase_transition_training_time.py
new file mode 100644
index 0000000..e51eaea
--- /dev/null
+++ b/scripts/plot_phase_transition_training_time.py
@@ -0,0 +1,86 @@
+#!/usr/bin/env python3
+"""Compare capacity-transition curves at different SGD horizons."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pandas as pd
+
+
+RUNS = [
+    ("T=3000, 256 traj", Path("outputs/phase_transition_sgd_256_lr001_T3000")),
+    ("T=10000 probe", Path("outputs/phase_transition_sgd_long_lr001_T10000_probe")),
+    ("T=30000 probe", Path("outputs/phase_transition_sgd_verylong_lr001_T30000_probe")),
+]
+
+
+def main() -> None:
+    outdir = Path("outputs/phase_transition_training_time_comparison")
+    outdir.mkdir(parents=True, exist_ok=True)
+
+    fig, ax = plt.subplots(figsize=(9.2, 5.7), dpi=180)
+    for label, run_dir in RUNS:
+        df = pd.read_csv(run_dir / "width_summary.csv").sort_values(
+            "fa_capacity_margin", ascending=False
+        )
+        ax.plot(
+            df["fa_capacity_margin"],
+            df["train_gap_mean"],
+            marker="o",
+            linewidth=2.0,
+            label=label,
+        )
+    ax.axvline(0, color="black", linestyle="--", linewidth=1.1, label="hard FA margin=0")
+    ax.axhline(0, color="black", linewidth=1.0)
+    ax.set_xlim(1100, -430)
+    ax.set_xlabel("hard FA capacity margin P - K_FA - N*out (capacity decreases ->)")
+    ax.set_ylabel("FA train MSE - BP train MSE")
+    ax.set_title("Training time shifts the apparent FA capacity transition")
+    ax.legend(frameon=True)
+    ax.grid(alpha=0.18)
+    fig.tight_layout()
+    path = outdir / "train_gap_vs_margin_by_training_time.png"
+    fig.savefig(path)
+    plt.close(fig)
+
+    fig2, ax2 = plt.subplots(figsize=(9.2, 5.7), dpi=180)
+    for label, run_dir in RUNS:
+        df = pd.read_csv(run_dir / "width_summary.csv").sort_values(
+            "fa_capacity_margin", ascending=False
+        )
+        ax2.plot(
+            df["fa_capacity_margin"],
+            df["bp_train_mean"],
+            marker="o",
+            linewidth=1.8,
+            linestyle="--",
+            label=f"BP {label}",
+        )
+        ax2.plot(
+            df["fa_capacity_margin"],
+            df["fa_train_mean"],
+            marker="o",
+            linewidth=2.0,
+            label=f"FA {label}",
+        )
+    ax2.axvline(0, color="black", linestyle="--", linewidth=1.1)
+    ax2.set_xlim(1100, -430)
+    ax2.set_yscale("log")
+    ax2.set_xlabel("hard FA capacity margin P - K_FA - N*out (capacity decreases ->)")
+    ax2.set_ylabel("train MSE")
+    ax2.set_title("BP/FA train losses across training time")
+    ax2.legend(frameon=True, fontsize=8, ncol=2)
+    ax2.grid(alpha=0.18)
+    fig2.tight_layout()
+    path2 = outdir / "bp_fa_train_loss_by_training_time.png"
+    fig2.savefig(path2)
+    plt.close(fig2)
+
+    print(f"plot: {path}")
+    print(f"plot: {path2}")
+
+
+if __name__ == "__main__":
+    main()