Add early-kernel operator predictor

1 files changed, 311 insertions, 0 deletions

diff --git a/scripts/plot_compressed_operator_results.py b/scripts/plot_compressed_operator_results.py
new file mode 100644
index 0000000..c23fa14
--- /dev/null
+++ b/scripts/plot_compressed_operator_results.py
@@ -0,0 +1,311 @@
+#!/usr/bin/env python3
+"""Plot early-kernel finite-time FA/BP gap predictions."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+from dataclasses import dataclass
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+
+
+@dataclass(frozen=True)
+class MetricsRow:
+    predictor: str
+    rows: int
+    mae: float
+    bias: float
+    corr: float
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--glob",
+        default="outputs/compressed_operator_s20_256traj_T50_width64_N*/compressed_operator_rows.csv",
+    )
+    parser.add_argument("--input-dim", type=int, default=16)
+    parser.add_argument("--output-dim", type=int, default=4)
+    parser.add_argument("--width", type=int, default=64)
+    parser.add_argument(
+        "--outdir",
+        type=Path,
+        default=Path("outputs/compressed_operator_s20_256traj_T50_width64_plots"),
+    )
+    return parser.parse_args()
+
+
+def read_rows(pattern: str, input_dim: int, output_dim: int, width: int) -> pd.DataFrame:
+    frames = []
+    for path in sorted(Path().glob(pattern)):
+        dirname = path.parent.name
+        try:
+            train_samples = int(dirname.rsplit("N", 1)[1])
+        except (IndexError, ValueError) as exc:
+            raise ValueError(f"Cannot infer train sample count from {dirname}") from exc
+        frame = pd.read_csv(path)
+        frame["train_samples"] = train_samples
+        frames.append(frame)
+    if not frames:
+        raise FileNotFoundError(f"No rows found for pattern: {pattern}")
+    df = pd.concat(frames, ignore_index=True)
+    parameter_count = input_dim * width + width * width + width * output_dim
+    constraint_rank = max(width * width - 1, 0) + max(width * output_dim - 1, 0)
+    df["hard_fa_capacity_margin"] = (
+        parameter_count - constraint_rank - output_dim * df["train_samples"]
+    )
+    df["trajectory_id"] = (
+        df["init_seed"].astype(str) + ":" + df["feedback_seed"].astype(str)
+    )
+    return df.sort_values(["hard_fa_capacity_margin", "trajectory_id"]).reset_index(drop=True)
+
+
+def predictor_metrics(df: pd.DataFrame) -> list[MetricsRow]:
+    rows: list[MetricsRow] = []
+    for predictor, column in [
+        ("fixed K(0)", "fixed_gap"),
+        ("FA compressed", "compressed_gap"),
+        ("linear velocity", "linear_gap"),
+        ("early re-tangent", "retangent_gap"),
+    ]:
+        error = df[column].to_numpy() - df["empirical_gap"].to_numpy()
+        corr = float(np.corrcoef(df[column], df["empirical_gap"])[0, 1])
+        rows.append(
+            MetricsRow(
+                predictor=predictor,
+                rows=len(df),
+                mae=float(np.mean(np.abs(error))),
+                bias=float(np.mean(error)),
+                corr=corr,
+            )
+        )
+    return rows
+
+
+def write_metrics(path: Path, rows: list[MetricsRow]) -> None:
+    with path.open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(MetricsRow.__annotations__.keys()))
+        writer.writeheader()
+        for row in rows:
+            writer.writerow(row.__dict__)
+
+
+def summarize_by_margin(df: pd.DataFrame, column: str) -> pd.DataFrame:
+    records = []
+    for margin, group in df.groupby("hard_fa_capacity_margin"):
+        values = group[column].to_numpy()
+        records.append(
+            {
+                "margin": margin,
+                "mean": float(values.mean()),
+                "q05": float(np.quantile(values, 0.05)),
+                "q25": float(np.quantile(values, 0.25)),
+                "q75": float(np.quantile(values, 0.75)),
+                "q95": float(np.quantile(values, 0.95)),
+            }
+        )
+    return pd.DataFrame(records).sort_values("margin")
+
+
+def plot_theory_band_with_empirical(df: pd.DataFrame, outdir: Path) -> Path:
+    theory = summarize_by_margin(df, "linear_gap")
+    empirical = summarize_by_margin(df, "empirical_gap")
+    fig, ax = plt.subplots(figsize=(11.5, 6.4), dpi=180)
+
+    for _tid, group in df.groupby("trajectory_id"):
+        group = group.sort_values("hard_fa_capacity_margin")
+        ax.plot(
+            group["hard_fa_capacity_margin"],
+            group["empirical_gap"],
+            color="#c65f16",
+            alpha=0.16,
+            linewidth=1.0,
+        )
+
+    ax.fill_between(
+        theory["margin"],
+        theory["q05"],
+        theory["q95"],
+        color="#1f5a9d",
+        alpha=0.16,
+        linewidth=0,
+        label="velocity theory 5-95%",
+    )
+    ax.fill_between(
+        theory["margin"],
+        theory["q25"],
+        theory["q75"],
+        color="#1f5a9d",
+        alpha=0.28,
+        linewidth=0,
+        label="velocity theory 25-75%",
+    )
+    ax.plot(
+        theory["margin"],
+        theory["mean"],
+        color="#174f91",
+        linewidth=2.3,
+        marker="o",
+        markersize=4.5,
+        label="velocity theory mean",
+    )
+    ax.plot(
+        empirical["margin"],
+        empirical["mean"],
+        color="#a84708",
+        linewidth=2.1,
+        marker="o",
+        markersize=4.2,
+        label="empirical mean",
+    )
+    ax.axvline(0, color="black", linestyle="--", linewidth=1.1)
+    ax.set_title("Early-kernel velocity prediction vs empirical FA/BP gap")
+    ax.set_xlabel("hard FA capacity margin")
+    ax.set_ylabel("FA train MSE - BP train MSE")
+    ax.legend(loc="upper left", frameon=True)
+    ax.grid(alpha=0.18)
+    fig.tight_layout()
+    outpath = outdir / "T50_s20_velocity_theory_band_empirical_trajectories.png"
+    fig.savefig(outpath)
+    plt.close(fig)
+    return outpath
+
+
+def plot_empirical_only(df: pd.DataFrame, outdir: Path) -> Path:
+    fig, ax = plt.subplots(figsize=(11.5, 6.4), dpi=180)
+    for _tid, group in df.groupby("trajectory_id"):
+        group = group.sort_values("hard_fa_capacity_margin")
+        ax.plot(
+            group["hard_fa_capacity_margin"],
+            group["empirical_gap"],
+            color="#c65f16",
+            alpha=0.20,
+            linewidth=1.05,
+        )
+    ax.axvline(0, color="black", linestyle="--", linewidth=1.1)
+    ax.set_title("Empirical FA/BP gap trajectories")
+    ax.set_xlabel("hard FA capacity margin")
+    ax.set_ylabel("FA train MSE - BP train MSE")
+    ax.grid(alpha=0.18)
+    fig.tight_layout()
+    outpath = outdir / "T50_s20_empirical_trajectories_only.png"
+    fig.savefig(outpath)
+    plt.close(fig)
+    return outpath
+
+
+def plot_prediction_scatter(df: pd.DataFrame, outdir: Path) -> Path:
+    fig, ax = plt.subplots(figsize=(7.0, 6.2), dpi=180)
+    empirical = df["empirical_gap"].to_numpy()
+    predictors = [
+        ("fixed K(0)", "fixed_gap", "#7f7f7f", 0.38),
+        ("linear velocity", "linear_gap", "#1f5a9d", 0.58),
+        ("early re-tangent", "retangent_gap", "#2b8a3e", 0.45),
+    ]
+    max_value = float(
+        max(
+            empirical.max(),
+            *(df[column].max() for _label, column, _color, _alpha in predictors),
+        )
+    )
+    min_value = float(
+        min(
+            empirical.min(),
+            *(df[column].min() for _label, column, _color, _alpha in predictors),
+        )
+    )
+    for label, column, color, alpha in predictors:
+        ax.scatter(
+            df[column],
+            empirical,
+            s=19,
+            alpha=alpha,
+            color=color,
+            label=label,
+            edgecolors="none",
+        )
+    ax.plot([min_value, max_value], [min_value, max_value], color="black", linewidth=1.0)
+    ax.set_title("Predicted vs empirical FA/BP train-gap")
+    ax.set_xlabel("predicted gap")
+    ax.set_ylabel("empirical gap")
+    ax.legend(loc="upper left", frameon=True)
+    ax.grid(alpha=0.18)
+    fig.tight_layout()
+    outpath = outdir / "T50_s20_prediction_scatter.png"
+    fig.savefig(outpath)
+    plt.close(fig)
+    return outpath
+
+
+def plot_error_by_margin(df: pd.DataFrame, outdir: Path) -> Path:
+    fig, ax = plt.subplots(figsize=(10.5, 5.8), dpi=180)
+    for label, column, color in [
+        ("fixed K(0)", "fixed_gap", "#7f7f7f"),
+        ("linear velocity", "linear_gap", "#1f5a9d"),
+        ("early re-tangent", "retangent_gap", "#2b8a3e"),
+    ]:
+        records = []
+        for margin, group in df.groupby("hard_fa_capacity_margin"):
+            error = group[column].to_numpy() - group["empirical_gap"].to_numpy()
+            records.append(
+                {
+                    "margin": margin,
+                    "mean": float(error.mean()),
+                    "q25": float(np.quantile(error, 0.25)),
+                    "q75": float(np.quantile(error, 0.75)),
+                }
+            )
+        summary = pd.DataFrame(records).sort_values("margin")
+        ax.plot(summary["margin"], summary["mean"], color=color, marker="o", label=label)
+        ax.fill_between(
+            summary["margin"],
+            summary["q25"],
+            summary["q75"],
+            color=color,
+            alpha=0.15,
+            linewidth=0,
+        )
+    ax.axhline(0, color="black", linewidth=1.0)
+    ax.axvline(0, color="black", linestyle="--", linewidth=1.1)
+    ax.set_title("Prediction error by capacity margin")
+    ax.set_xlabel("hard FA capacity margin")
+    ax.set_ylabel("predicted gap - empirical gap")
+    ax.legend(loc="upper left", frameon=True)
+    ax.grid(alpha=0.18)
+    fig.tight_layout()
+    outpath = outdir / "T50_s20_prediction_error_by_margin.png"
+    fig.savefig(outpath)
+    plt.close(fig)
+    return outpath
+
+
+def main() -> None:
+    args = parse_args()
+    args.outdir.mkdir(parents=True, exist_ok=True)
+    df = read_rows(args.glob, args.input_dim, args.output_dim, args.width)
+    df.to_csv(args.outdir / "combined_rows.csv", index=False)
+    metrics = predictor_metrics(df)
+    write_metrics(args.outdir / "predictor_metrics.csv", metrics)
+    paths = [
+        plot_theory_band_with_empirical(df, args.outdir),
+        plot_empirical_only(df, args.outdir),
+        plot_prediction_scatter(df, args.outdir),
+        plot_error_by_margin(df, args.outdir),
+    ]
+    print(f"combined rows: {args.outdir / 'combined_rows.csv'}")
+    print(f"metrics: {args.outdir / 'predictor_metrics.csv'}")
+    for row in metrics:
+        print(
+            f"{row.predictor}: rows={row.rows}, MAE={row.mae:.6g}, "
+            f"bias={row.bias:.6g}, corr={row.corr:.6g}"
+        )
+    for path in paths:
+        print(f"plot: {path}")
+
+
+if __name__ == "__main__":
+    main()