1 files changed, 143 insertions, 0 deletions
diff --git a/scripts/plot_phase_transition_zoom_comparison.py b/scripts/plot_phase_transition_zoom_comparison.py
new file mode 100644
index 0000000..8e0dc03
--- /dev/null
+++ b/scripts/plot_phase_transition_zoom_comparison.py
@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+"""Plot trajectory-level capacity transition at two training horizons."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+
+
+RUNS = [
+    (
+        "T=10000, 1024 FA trajectories",
+        Path("outputs/phase_transition_zoom_T10000_1024traj"),
+        "#c65f16",
+    ),
+    (
+        "T=30000, 256 FA trajectories",
+        Path("outputs/phase_transition_zoom_T30000_256traj"),
+        "#174f91",
+    ),
+]
+
+
+def summarize_gaps(runs: pd.DataFrame) -> pd.DataFrame:
+    fa = runs[runs["run_type"] == "fa"].copy()
+    records = []
+    for margin, group in fa.groupby("fa_capacity_margin"):
+        gaps = group["train_gap_to_bp"].to_numpy(dtype=np.float64)
+        records.append(
+            {
+                "margin": float(margin),
+                "mean": float(gaps.mean()),
+                "q05": float(np.quantile(gaps, 0.05)),
+                "q25": float(np.quantile(gaps, 0.25)),
+                "q75": float(np.quantile(gaps, 0.75)),
+                "q95": float(np.quantile(gaps, 0.95)),
+            }
+        )
+    return pd.DataFrame(records).sort_values("margin", ascending=False)
+
+
+def add_summary(ax: plt.Axes, runs: pd.DataFrame, label: str, color: str) -> pd.DataFrame:
+    summary = summarize_gaps(runs)
+    x = summary["margin"].to_numpy(dtype=np.float64)
+    ax.fill_between(
+        x,
+        summary["q05"],
+        summary["q95"],
+        color=color,
+        alpha=0.13,
+        linewidth=0,
+    )
+    ax.fill_between(
+        x,
+        summary["q25"],
+        summary["q75"],
+        color=color,
+        alpha=0.24,
+        linewidth=0,
+    )
+    ax.plot(
+        x,
+        summary["mean"],
+        color=color,
+        marker="o",
+        linewidth=2.3,
+        label=label,
+    )
+    return summary
+
+
+def add_trajectories(ax: plt.Axes, runs: pd.DataFrame, color: str) -> None:
+    fa = runs[runs["run_type"] == "fa"].copy()
+    fa["_tid"] = fa["init_seed"].astype(str) + ":" + fa["feedback_seed"].astype(str)
+    for _tid, group in fa.groupby("_tid"):
+        group = group.sort_values("fa_capacity_margin", ascending=False)
+        ax.plot(
+            group["fa_capacity_margin"],
+            group["train_gap_to_bp"],
+            color=color,
+            alpha=0.11,
+            linewidth=0.85,
+        )
+
+
+def main() -> None:
+    outdir = Path("outputs/phase_transition_zoom_comparison")
+    outdir.mkdir(parents=True, exist_ok=True)
+
+    fig, axes = plt.subplots(1, 2, figsize=(13.2, 5.2), dpi=180)
+    for ax, (label, run_dir, color) in zip(axes, RUNS, strict=True):
+        runs = pd.read_csv(run_dir / "runs.csv")
+        add_trajectories(ax, runs, color)
+        summary = add_summary(ax, runs, label, color)
+        ax.axvline(0, color="black", linestyle="--", linewidth=1.1)
+        ax.axhline(0, color="black", linewidth=1.0)
+        ax.set_xlabel("hard FA capacity margin P - K_FA - N*out")
+        ax.set_title(label)
+        ax.grid(alpha=0.18)
+        ax.legend(frameon=True)
+        fa = runs[runs["run_type"] == "fa"]
+        ymax = max(
+            0.01,
+            float(summary["q95"].max()) * 1.18,
+            float(fa["train_gap_to_bp"].max()) * 1.03,
+        )
+        ax.set_ylim(-0.004 * ymax, ymax)
+    axes[0].set_xlim(285, -280)
+    axes[1].set_xlim(90, -150)
+    for ax in axes:
+        ax.set_ylabel("FA train MSE - BP train MSE")
+    fig.suptitle("More trajectories separate finite-time ramp from capacity-limited gap")
+    fig.tight_layout()
+    path = outdir / "trajectory_gap_clouds_T10000_vs_T30000.png"
+    fig.savefig(path)
+    plt.close(fig)
+
+    fig2, ax = plt.subplots(figsize=(8.8, 5.3), dpi=180)
+    for label, run_dir, color in RUNS:
+        runs = pd.read_csv(run_dir / "runs.csv")
+        add_summary(ax, runs, label, color)
+    ax.axvline(0, color="black", linestyle="--", linewidth=1.1, label="hard margin=0")
+    ax.axhline(0, color="black", linewidth=1.0)
+    ax.set_xlim(285, -280)
+    ax.set_xlabel("hard FA capacity margin P - K_FA - N*out")
+    ax.set_ylabel("FA train MSE - BP train MSE")
+    ax.set_title("Training horizon compresses the apparent transition")
+    ax.grid(alpha=0.18)
+    ax.legend(frameon=True)
+    fig2.tight_layout()
+    path2 = outdir / "trajectory_gap_quantiles_T10000_vs_T30000.png"
+    fig2.savefig(path2)
+    plt.close(fig2)
+
+    print(f"plot: {path}")
+    print(f"plot: {path2}")
+
+
+if __name__ == "__main__":
+    main()