Add smooth capacity transition overlay plot

1 files changed, 287 insertions, 0 deletions

diff --git a/scripts/plot_capacity_transition_overlay.py b/scripts/plot_capacity_transition_overlay.py
new file mode 100644
index 0000000..e636725
--- /dev/null
+++ b/scripts/plot_capacity_transition_overlay.py
@@ -0,0 +1,287 @@
+#!/usr/bin/env python3
+"""Plot smooth capacity-transition prediction against FA trajectory ensembles."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import math
+from collections import defaultdict
+from dataclasses import dataclass
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+@dataclass(frozen=True)
+class SummaryRow:
+    width: float
+    parameter_count: float
+    task_dimension: float
+    fa_constraint_rank: float
+    bp_capacity_margin: float
+    fa_capacity_margin: float
+    train_gap_mean: float
+
+
+@dataclass(frozen=True)
+class RunRow:
+    width: float
+    feedback_seed: int
+    run_type: str
+    train_gap_to_bp: float
+    fa_capacity_margin: float
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--run-dir",
+        type=Path,
+        default=Path("outputs/downstream_capacity_random_main_fast"),
+    )
+    parser.add_argument("--samples", type=int, default=2048)
+    parser.add_argument(
+        "--transition-width",
+        type=float,
+        default=70.0,
+        help="Soft capacity-transition width measured in scalar output constraints.",
+    )
+    parser.add_argument("--xmin", type=float, default=-430.0)
+    parser.add_argument("--xmax", type=float, default=600.0)
+    parser.add_argument("--outname", default="capacity_transition_theory_vs_trajectories.png")
+    return parser.parse_args()
+
+
+def read_summary(path: Path) -> list[SummaryRow]:
+    rows: list[SummaryRow] = []
+    with path.open() as handle:
+        for row in csv.DictReader(handle):
+            rows.append(
+                SummaryRow(
+                    width=float(row["width"]),
+                    parameter_count=float(row["parameter_count"]),
+                    task_dimension=float(row["task_dimension"]),
+                    fa_constraint_rank=float(row["fa_constraint_rank"]),
+                    bp_capacity_margin=float(row["bp_capacity_margin"]),
+                    fa_capacity_margin=float(row["fa_capacity_margin"]),
+                    train_gap_mean=float(row["train_gap_mean"]),
+                )
+            )
+    return rows
+
+
+def read_runs(path: Path) -> list[RunRow]:
+    rows: list[RunRow] = []
+    with path.open() as handle:
+        for row in csv.DictReader(handle):
+            rows.append(
+                RunRow(
+                    width=float(row["width"]),
+                    feedback_seed=int(row["feedback_seed"]),
+                    run_type=row["run_type"],
+                    train_gap_to_bp=float(row["train_gap_to_bp"]),
+                    fa_capacity_margin=float(row["fa_capacity_margin"]),
+                )
+            )
+    return rows
+
+
+def soft_overlap_gap_counts(
+    width: np.ndarray,
+    input_dim: float,
+    output_dim: float,
+    task_dim: float,
+    samples: int,
+    transition_width: float,
+    rng: np.random.Generator,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    parameter_count = input_dim * width + width * width + width * output_dim
+    constraint_rank = np.maximum(width * width - 1.0, 0.0) + np.maximum(
+        width * output_dim - 1.0, 0.0
+    )
+    hard_fa_margin = parameter_count - constraint_rank - task_dim
+
+    raw_gaps = np.zeros((samples, width.size), dtype=np.float64)
+
+    def softplus_hinge(value: np.ndarray | float) -> np.ndarray:
+        scaled = np.asarray(value, dtype=np.float64) / transition_width
+        return transition_width * np.logaddexp(0.0, scaled)
+
+    for index, (p_count, k_rank) in enumerate(zip(parameter_count, constraint_rank)):
+        bp_missing = softplus_hinge(task_dim - p_count)
+        task_subspace_dim = min(task_dim, p_count)
+
+        if task_subspace_dim >= p_count or k_rank <= 0:
+            overlap = np.full(samples, min(k_rank, p_count), dtype=np.float64)
+        else:
+            mean = k_rank * task_subspace_dim / p_count
+            var = (
+                2.0
+                * k_rank
+                * task_subspace_dim
+                * (p_count - k_rank)
+                * (p_count - task_subspace_dim)
+                / (p_count * p_count * (p_count - 1.0) * (p_count + 2.0))
+            )
+            overlap = rng.normal(mean, math.sqrt(max(var, 0.0)), size=samples)
+            overlap = np.clip(overlap, 0.0, min(k_rank, task_subspace_dim))
+
+        finite_width_noise = rng.normal(0.0, transition_width, size=samples)
+        fa_missing = softplus_hinge(
+            task_dim - (parameter_count[index] - overlap) + finite_width_noise
+        )
+        raw_gaps[:, index] = np.maximum(fa_missing - bp_missing, 0.0)
+
+    mean = raw_gaps.mean(axis=0)
+    lo = np.quantile(raw_gaps, 0.1, axis=0)
+    hi = np.quantile(raw_gaps, 0.9, axis=0)
+    return hard_fa_margin, mean, lo, hi
+
+
+def infer_dims(run_dir: Path) -> tuple[float, float, float]:
+    payload = json.loads((run_dir / "summary.json").read_text())
+    config = payload["config"]
+    return (
+        float(config["input_dim"]),
+        float(config["output_dim"]),
+        float(config["train_samples"] * config["output_dim"]),
+    )
+
+
+def calibrate_scale(summary: list[SummaryRow], raw_mean_at_widths: np.ndarray) -> float:
+    observed = np.array([row.train_gap_mean for row in summary], dtype=np.float64)
+    denom = float(np.dot(raw_mean_at_widths, raw_mean_at_widths))
+    if denom <= 0:
+        return 1.0
+    return float(np.dot(raw_mean_at_widths, observed) / denom)
+
+
+def plot_transition(
+    run_dir: Path,
+    samples: int,
+    transition_width: float,
+    outname: str,
+    xmin: float,
+    xmax: float,
+) -> Path:
+    summary = read_summary(run_dir / "width_summary.csv")
+    runs = read_runs(run_dir / "runs.csv")
+    input_dim, output_dim, task_dim = infer_dims(run_dir)
+    rng = np.random.default_rng(0)
+
+    widths = np.array([row.width for row in summary], dtype=np.float64)
+    dense_width = np.linspace(widths.min(), widths.max(), 240)
+
+    _, raw_mean_at_widths, _, _ = soft_overlap_gap_counts(
+        widths, input_dim, output_dim, task_dim, samples, transition_width, rng
+    )
+    scale = calibrate_scale(summary, raw_mean_at_widths)
+
+    dense_margin, dense_mean_raw, dense_lo_raw, dense_hi_raw = soft_overlap_gap_counts(
+        dense_width,
+        input_dim,
+        output_dim,
+        task_dim,
+        samples,
+        transition_width,
+        np.random.default_rng(1),
+    )
+    dense_mean = scale * dense_mean_raw
+    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)
+    for row in runs:
+        if row.run_type == "fa":
+            by_seed[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)
+        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",
+            )
+
+    axes[0].set_title("P1: smooth capacity-transition prediction")
+    axes[1].set_title("P2: transparent trajectory ensemble 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,
+    )
+    fig.tight_layout()
+
+    outpath = run_dir / outname
+    fig.savefig(outpath, dpi=220, bbox_inches="tight")
+    plt.close(fig)
+    return outpath
+
+
+def main() -> None:
+    args = parse_args()
+    outpath = plot_transition(
+        args.run_dir,
+        args.samples,
+        args.transition_width,
+        args.outname,
+        args.xmin,
+        args.xmax,
+    )
+    print(outpath)
+
+
+if __name__ == "__main__":
+    main()