Plot operator overlap stress panels

1 files changed, 192 insertions, 0 deletions

diff --git a/scripts/plot_operator_overlap_grid.py b/scripts/plot_operator_overlap_grid.py
new file mode 100644
index 0000000..e1e599f
--- /dev/null
+++ b/scripts/plot_operator_overlap_grid.py
@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+"""Plot theory/empirical overlap panels for operator stress settings."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+
+
+SETTING_ORDER = [
+    "d1_w64_T50_s20",
+    "d2_w32_T50_s20",
+    "d2_w64_T25_s10",
+    "d2_w64_T50_s20_256traj",
+    "d2_w64_T100_s40",
+    "d2_w96_T50_s20",
+    "d3_w64_T50_s20",
+]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--rows",
+        type=Path,
+        default=Path("outputs/operator_stress_grid_summary/stress_grid_rows.csv"),
+    )
+    parser.add_argument(
+        "--metrics",
+        type=Path,
+        default=Path("outputs/operator_stress_grid_summary/stress_grid_metrics.csv"),
+    )
+    parser.add_argument(
+        "--outdir",
+        type=Path,
+        default=Path("outputs/operator_stress_grid_summary"),
+    )
+    return parser.parse_args()
+
+
+def summarize_by_margin(group: pd.DataFrame, column: str) -> pd.DataFrame:
+    records = []
+    for margin, by_margin in group.groupby("hard_fa_capacity_margin"):
+        values = by_margin[column].to_numpy(dtype=np.float64)
+        records.append(
+            {
+                "margin": float(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 trajectory_id(frame: pd.DataFrame) -> pd.Series:
+    return frame["init_seed"].astype(str) + ":" + frame["feedback_seed"].astype(str)
+
+
+def pretty_setting(setting: str) -> str:
+    if setting == "d2_w64_T50_s20_256traj":
+        return "d2 w64 T50 s20 (256 traj)"
+    return setting.replace("_", " ")
+
+
+def panel_ylim(group: pd.DataFrame) -> tuple[float, float]:
+    values = np.concatenate(
+        [
+            group["empirical_gap"].to_numpy(dtype=np.float64),
+            group["linear_gap"].to_numpy(dtype=np.float64),
+        ]
+    )
+    lo = float(values.min())
+    hi = float(values.max())
+    pad = max((hi - lo) * 0.12, 0.004)
+    return max(0.0, lo - pad), hi + pad
+
+
+def plot_panel(
+    ax: plt.Axes,
+    group: pd.DataFrame,
+    metrics: pd.Series,
+    show_legend: bool,
+) -> None:
+    group = group.sort_values(["hard_fa_capacity_margin", "init_seed", "feedback_seed"])
+    theory = summarize_by_margin(group, "linear_gap")
+    empirical = summarize_by_margin(group, "empirical_gap")
+    group = group.assign(_trajectory_id=trajectory_id(group))
+
+    for _tid, traj in group.groupby("_trajectory_id"):
+        traj = traj.sort_values("hard_fa_capacity_margin")
+        ax.plot(
+            traj["hard_fa_capacity_margin"],
+            traj["empirical_gap"],
+            color="#c65f16",
+            alpha=0.16 if len(group) > 80 else 0.26,
+            linewidth=1.05,
+        )
+
+    ax.fill_between(
+        theory["margin"],
+        theory["q05"],
+        theory["q95"],
+        color="#1f5a9d",
+        alpha=0.14,
+        linewidth=0,
+        label="theory 5-95%",
+    )
+    ax.fill_between(
+        theory["margin"],
+        theory["q25"],
+        theory["q75"],
+        color="#1f5a9d",
+        alpha=0.26,
+        linewidth=0,
+        label="theory 25-75%",
+    )
+    ax.plot(
+        theory["margin"],
+        theory["mean"],
+        color="#174f91",
+        marker="o",
+        markersize=3.3,
+        linewidth=2.0,
+        label="theory mean",
+    )
+    ax.plot(
+        empirical["margin"],
+        empirical["mean"],
+        color="#a84708",
+        marker="o",
+        markersize=3.0,
+        linewidth=1.8,
+        label="empirical mean",
+    )
+    ax.axvline(0.0, color="black", linestyle="--", linewidth=0.9)
+    ax.set_ylim(*panel_ylim(group))
+    ax.grid(alpha=0.17)
+    ax.set_title(
+        f"{pretty_setting(str(metrics.name))}\n"
+        f"MAE={metrics['linear_mae']:.4g}, "
+        f"bias={metrics['linear_bias']:.4g}, "
+        f"corr={metrics['linear_corr']:.4f}",
+        fontsize=10,
+    )
+    if show_legend:
+        ax.legend(loc="upper left", fontsize=8, frameon=True)
+
+
+def main() -> None:
+    args = parse_args()
+    args.outdir.mkdir(parents=True, exist_ok=True)
+    rows = pd.read_csv(args.rows)
+    metrics = pd.read_csv(args.metrics).set_index("setting")
+    settings = [setting for setting in SETTING_ORDER if setting in set(rows["setting"])]
+
+    fig, axes = plt.subplots(4, 2, figsize=(15.5, 18.5), dpi=180)
+    axes_flat = axes.ravel()
+    for index, setting in enumerate(settings):
+        group = rows[rows["setting"] == setting]
+        plot_panel(
+            axes_flat[index],
+            group,
+            metrics.loc[setting],
+            show_legend=index == 0,
+        )
+        axes_flat[index].set_xlabel("hard FA capacity margin")
+        axes_flat[index].set_ylabel("FA train MSE - BP train MSE")
+
+    for index in range(len(settings), len(axes_flat)):
+        axes_flat[index].axis("off")
+
+    fig.suptitle(
+        "Early-kernel velocity theory vs empirical trajectory distributions",
+        fontsize=18,
+        y=0.995,
+    )
+    fig.tight_layout(rect=(0, 0, 1, 0.985))
+    outpath = args.outdir / "stress_grid_overlap_panels.png"
+    fig.savefig(outpath)
+    plt.close(fig)
+    print(f"plot: {outpath}")
+
+
+if __name__ == "__main__":
+    main()