Expand empirical transition overlays

1 files changed, 56 insertions, 53 deletions

diff --git a/scripts/plot_capacity_transition_overlay.py b/scripts/plot_capacity_transition_overlay.py
index e636725..cc53012 100644
--- a/scripts/plot_capacity_transition_overlay.py
+++ b/scripts/plot_capacity_transition_overlay.py
@@ -28,6 +28,7 @@ class SummaryRow:
 
 @dataclass(frozen=True)
 class RunRow:
+    source_id: int
     width: float
     feedback_seed: int
     run_type: str
@@ -42,6 +43,7 @@ def parse_args() -> argparse.Namespace:
         type=Path,
         default=Path("outputs/downstream_capacity_random_main_fast"),
     )
+    parser.add_argument("--extra-run-dir", type=Path, action="append", default=[])
     parser.add_argument("--samples", type=int, default=2048)
     parser.add_argument(
         "--transition-width",
@@ -73,12 +75,13 @@ def read_summary(path: Path) -> list[SummaryRow]:
     return rows
 
 
-def read_runs(path: Path) -> list[RunRow]:
+def read_runs(path: Path, source_id: int) -> list[RunRow]:
     rows: list[RunRow] = []
     with path.open() as handle:
         for row in csv.DictReader(handle):
             rows.append(
                 RunRow(
+                    source_id=source_id,
                     width=float(row["width"]),
                     feedback_seed=int(row["feedback_seed"]),
                     run_type=row["run_type"],
@@ -161,6 +164,7 @@ def calibrate_scale(summary: list[SummaryRow], raw_mean_at_widths: np.ndarray) -
 
 def plot_transition(
     run_dir: Path,
+    extra_run_dirs: list[Path],
     samples: int,
     transition_width: float,
     outname: str,
@@ -168,7 +172,9 @@ def plot_transition(
     xmax: float,
 ) -> Path:
     summary = read_summary(run_dir / "width_summary.csv")
-    runs = read_runs(run_dir / "runs.csv")
+    runs = read_runs(run_dir / "runs.csv", 0)
+    for source_id, extra_run_dir in enumerate(extra_run_dirs, start=1):
+        runs.extend(read_runs(extra_run_dir / "runs.csv", source_id))
     input_dim, output_dim, task_dim = infer_dims(run_dir)
     rng = np.random.default_rng(0)
 
@@ -193,71 +199,67 @@ def plot_transition(
     dense_lo = scale * dense_lo_raw
     dense_hi = scale * dense_hi_raw
 
-    summary_margin = np.array([row.fa_capacity_margin for row in summary], dtype=np.float64)
-    summary_gap = np.array([row.train_gap_mean for row in summary], dtype=np.float64)
-
-    by_seed: dict[int, list[RunRow]] = defaultdict(list)
+    by_seed: dict[tuple[int, int], list[RunRow]] = defaultdict(list)
     for row in runs:
         if row.run_type == "fa":
-            by_seed[row.feedback_seed].append(row)
+            by_seed[(row.source_id, row.feedback_seed)].append(row)
 
     fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2), sharex=True, sharey=True)
 
-    for axis, show_trajectories in zip(axes, [False, True]):
-        axis.fill_between(
-            dense_margin,
-            dense_lo,
-            dense_hi,
-            color="tab:orange",
-            alpha=0.22,
-            linewidth=0,
-            label="theory 10-90% band",
-        )
-        axis.plot(
-            dense_margin,
-            dense_mean,
-            color="black",
-            linewidth=2.4,
-            label="theory mean",
-        )
-        axis.axvline(0.0, color="black", linestyle="--", linewidth=1.1, alpha=0.8)
+    axes[0].fill_between(
+        dense_margin,
+        dense_lo,
+        dense_hi,
+        color="tab:orange",
+        alpha=0.22,
+        linewidth=0,
+        label="theory 10-90% band",
+    )
+    axes[0].plot(
+        dense_margin,
+        dense_mean,
+        color="black",
+        linewidth=2.4,
+        label="theory mean",
+    )
+    axes[0].axvline(0.0, color="black", linestyle="--", linewidth=1.1, alpha=0.8)
+
+    for axis in axes:
         axis.axhline(0.0, color="black", linewidth=0.9, alpha=0.75)
         axis.set_xlim(xmin, xmax)
         axis.set_xlabel("hard FA capacity margin  P - K_FA - N*out")
         axis.grid(True, color="#dddddd", linewidth=0.8, alpha=0.7)
 
-        if show_trajectories:
-            for seed, rows in sorted(by_seed.items()):
-                ordered = sorted(rows, key=lambda item: item.fa_capacity_margin)
-                axis.plot(
-                    [row.fa_capacity_margin for row in ordered],
-                    [row.train_gap_to_bp for row in ordered],
-                    color="tab:blue",
-                    alpha=0.12,
-                    linewidth=0.9,
-                )
-            axis.plot(
-                summary_margin,
-                summary_gap,
-                color="tab:blue",
-                linewidth=2.2,
-                alpha=0.95,
-                label="empirical mean",
-            )
-            axis.scatter(
-                summary_margin,
-                summary_gap,
-                color="black",
-                s=34,
-                zorder=4,
-                label="empirical points",
-            )
+    for _seed, rows in sorted(by_seed.items()):
+        ordered = sorted(rows, key=lambda item: item.fa_capacity_margin)
+        axes[1].plot(
+            [row.fa_capacity_margin for row in ordered],
+            [row.train_gap_to_bp for row in ordered],
+            color="tab:blue",
+            alpha=0.07,
+            linewidth=0.85,
+        )
+
+    axes[1].text(
+        0.98,
+        0.96,
+        f"{len(by_seed)} empirical FA trajectories",
+        transform=axes[1].transAxes,
+        ha="right",
+        va="top",
+        fontsize=10,
+        bbox={
+            "boxstyle": "round,pad=0.25",
+            "facecolor": "white",
+            "edgecolor": "#cccccc",
+            "alpha": 0.85,
+        },
+    )
 
     axes[0].set_title("P1: smooth capacity-transition prediction")
-    axes[1].set_title("P2: transparent trajectory ensemble overlay")
+    axes[1].set_title("P2: empirical trajectory overlay")
     axes[0].set_ylabel("FA train MSE - BP train MSE")
     axes[0].legend(loc="upper right")
-    axes[1].legend(loc="upper right")
     fig.suptitle(
         f"Random-label memorization transition, scale={scale:.4g}, transition width={transition_width:g}",
         y=1.02,
@@ -274,6 +276,7 @@ def main() -> None:
     args = parse_args()
     outpath = plot_transition(
         args.run_dir,
+        args.extra_run_dir,
         args.samples,
         args.transition_width,
         args.outname,