Add AAAI depth experiments and diagnostic figures

1 files changed, 311 insertions, 0 deletions

diff --git a/scripts/plot_aaai_main_figures.py b/scripts/plot_aaai_main_figures.py
new file mode 100644
index 0000000..cd12011
--- /dev/null
+++ b/scripts/plot_aaai_main_figures.py
@@ -0,0 +1,311 @@
+#!/usr/bin/env python3
+"""Render the two three-panel composite figures for the focused AAAI story."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import math
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--depth-dir", type=Path, default=Path("outputs/aaai_depth_experiments"))
+    parser.add_argument(
+        "--cnn-dir", type=Path, default=Path("outputs/cnn_initialization_validation")
+    )
+    parser.add_argument(
+        "--mnist-dir", type=Path, default=Path("outputs/real_data_validation_mnist")
+    )
+    parser.add_argument(
+        "--null-dir", type=Path, default=Path("outputs/soft_erosion_theory_vs_empirical")
+    )
+    parser.add_argument("--outdir", type=Path, default=Path("outputs/aaai_main_figures"))
+    return parser.parse_args()
+
+
+def read_csv(path: Path) -> list[dict[str, str]]:
+    if not path.exists():
+        raise FileNotFoundError(path)
+    with path.open() as handle:
+        return list(csv.DictReader(handle))
+
+
+def panel_label(ax: plt.Axes, label: str) -> None:
+    ax.text(-0.13, 1.06, label, transform=ax.transAxes, fontsize=12, fontweight="bold")
+
+
+def draw_schematic(ax: plt.Axes) -> None:
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, 1)
+    ax.axis("off")
+    rows = [(0.82, "BP", "#2f6f9f"), (0.50, "FA", "#c65f16"), (0.18, "DFA", "#2b8a3e")]
+    x_positions = [0.18, 0.42, 0.66, 0.90]
+    labels = ["$h_1$", "$h_2$", "$h_3$", "error"]
+    for y, name, color in rows:
+        ax.text(0.01, y, name, va="center", ha="left", fontweight="bold", color=color)
+        for x, node in zip(x_positions, labels):
+            ax.add_patch(plt.Circle((x, y), 0.035, facecolor="white", edgecolor=color, lw=1.5))
+            ax.text(x, y - 0.075, node, ha="center", va="top", fontsize=8)
+        if name in {"BP", "FA"}:
+            arrow_label = r"$W^\top$" if name == "BP" else "$B$"
+            for right, left in zip(x_positions[:0:-1], x_positions[-2::-1]):
+                ax.annotate(
+                    "",
+                    xy=(left + 0.04, y),
+                    xytext=(right - 0.04, y),
+                    arrowprops=dict(arrowstyle="->", color=color, lw=1.5),
+                )
+                ax.text((left + right) / 2, y + 0.035, arrow_label, ha="center", fontsize=8)
+        else:
+            for left in x_positions[:-1]:
+                ax.annotate(
+                    "",
+                    xy=(left + 0.04, y),
+                    xytext=(x_positions[-1] - 0.04, y),
+                    arrowprops=dict(arrowstyle="->", color=color, lw=1.2, alpha=0.9),
+                )
+            ax.text(0.56, y + 0.055, "direct random maps", ha="center", fontsize=8)
+    ax.set_title("Fixed random feedback rules", fontsize=10)
+
+
+def figure_one(args: argparse.Namespace) -> Path:
+    init_rows = read_csv(args.depth_dir / "initialization_rows.csv")
+    cnn_rows = read_csv(args.cnn_dir / "cnn_initialization_rows.csv")
+    mnist_rows = read_csv(args.mnist_dir / "e0_mnist_rows.csv")
+    theory_rows = read_csv(args.null_dir / "theory_gap_summary.csv")
+    empirical_rows = read_csv(args.null_dir / "empirical_gap_summary.csv")
+
+    fig, axes = plt.subplots(1, 3, figsize=(14.2, 4.1), dpi=220)
+    draw_schematic(axes[0])
+    panel_label(axes[0], "a")
+
+    ax = axes[1]
+    theory_by_width = {int(row["width"]): float(row["mean"]) for row in theory_rows}
+    empirical_by_width = {int(row["width"]): float(row["mean"]) for row in empirical_rows}
+    widths = sorted(set(theory_by_width) & set(empirical_by_width))
+    ax.plot(
+        widths,
+        [theory_by_width[width] for width in widths],
+        "o-",
+        color="#8c6d31",
+        label="random-direction deletion",
+    )
+    ax.plot(
+        widths,
+        [empirical_by_width[width] for width in widths],
+        "s-",
+        color="#c65f16",
+        label="measured FA cost",
+    )
+    ax.set_yscale("log")
+    ax.set_xlabel("width")
+    ax.set_ylabel("finite-time loss gap")
+    ax.grid(alpha=0.16, which="both")
+    ax2 = ax.twinx()
+    ax2.plot(widths, [1.0 / (width * width) for width in widths], "k--", lw=1.4, label="$1/D$")
+    ax2.set_yscale("log")
+    ax2.set_ylabel("best guaranteed squared alignment")
+    handles1, labels1 = ax.get_legend_handles_labels()
+    handles2, labels2 = ax2.get_legend_handles_labels()
+    ax.legend(handles1 + handles2, labels1 + labels2, fontsize=7, loc="upper right")
+    worst_ratio = max(theory_by_width[width] / empirical_by_width[width] for width in widths)
+    ax.text(0.04, 0.05, f"null overestimate: up to {worst_ratio:.0f}×", transform=ax.transAxes, fontsize=8)
+    ax.set_title("Matrix mismatch is not optimization cost", fontsize=10)
+    panel_label(ax, "b")
+
+    ax = axes[2]
+    depths = sorted({int(row["depth"]) for row in init_rows})
+    cmap = plt.cm.viridis
+    depth_color = {depth: cmap(index / max(len(depths) - 1, 1)) for index, depth in enumerate(depths)}
+    for row in init_rows:
+        depth = int(row["depth"])
+        marker = "o" if row["rule"] == "FA" else "s"
+        ax.errorbar(
+            float(row["predicted_deficit"]),
+            float(row["empirical_deficit"]),
+            yerr=2 * float(row["empirical_stderr"]),
+            fmt=marker,
+            ms=3.8,
+            color=depth_color[depth],
+            alpha=0.68,
+            lw=0.7,
+            capsize=1,
+        )
+    for row in mnist_rows:
+        ax.plot(float(row["hidden_share"]), float(row["empirical_mean"]), "^", color="#7b6fa6", ms=4, alpha=0.65)
+    for row in cnn_rows:
+        marker = "P" if row["rule"] == "FA" else "X"
+        ax.errorbar(
+            float(row["predicted_deficit"]),
+            float(row["empirical_deficit"]),
+            yerr=2 * float(row["empirical_stderr"]),
+            fmt=marker,
+            color="#b23a48",
+            ms=5,
+            capsize=1,
+        )
+    values = [
+        float(row[key])
+        for row in init_rows
+        for key in ("predicted_deficit", "empirical_deficit")
+    ]
+    values += [
+        float(row[key])
+        for row in cnn_rows
+        for key in ("predicted_deficit", "empirical_deficit")
+    ]
+    values += [float(row[key]) for row in mnist_rows for key in ("hidden_share", "empirical_mean")]
+    lo, hi = min(values), max(values)
+    pad = 0.04 * (hi - lo + 1e-12)
+    ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], "k-", lw=1)
+    legend_handles = [
+        plt.Line2D([], [], marker="o", ls="", color="0.35", label="synthetic FA"),
+        plt.Line2D([], [], marker="s", ls="", color="0.35", label="synthetic DFA"),
+        plt.Line2D([], [], marker="^", ls="", color="#7b6fa6", label="MNIST MLP"),
+        plt.Line2D([], [], marker="P", ls="", color="#b23a48", label="CE CNN FA/DFA"),
+    ]
+    shape_legend = ax.legend(handles=legend_handles, fontsize=7, loc="lower right")
+    ax.add_artist(shape_legend)
+    depth_handles = [
+        plt.Line2D([], [], marker="o", ls="", color=depth_color[depth], label=f"depth {depth}")
+        for depth in depths
+    ]
+    ax.legend(handles=depth_handles, fontsize=6, ncols=2, loc="upper left")
+    ax.set_xlabel("exact expected initial deficit")
+    ax.set_ylabel("measured mean deficit")
+    ax.set_title("One calibration across rule, depth, and data", fontsize=10)
+    ax.grid(alpha=0.16)
+    panel_label(ax, "c")
+
+    fig.tight_layout(w_pad=2.0)
+    path = args.outdir / "figure1_mismatch_vs_initial_cost.png"
+    fig.savefig(path, bbox_inches="tight")
+    fig.savefig(path.with_suffix(".pdf"), bbox_inches="tight")
+    plt.close(fig)
+    return path
+
+
+def scatter_predictor(
+    ax: plt.Axes,
+    rows: list[dict[str, str]],
+    fields: list[tuple[str, str, str]],
+    title: str,
+) -> None:
+    depths = sorted({int(row["depth"]) for row in rows})
+    colors = {depth: plt.cm.viridis(i / max(len(depths) - 1, 1)) for i, depth in enumerate(depths)}
+    values: list[float] = []
+    for field, marker, label in fields:
+        for row in rows:
+            x = float(row[field])
+            y = float(row["empirical_gap"])
+            values.extend([x, y])
+            ax.scatter(x, y, s=13, marker=marker, color=colors[int(row["depth"])], alpha=0.60)
+        ax.scatter([], [], marker=marker, color="0.35", label=label)
+    lo, hi = min(values), max(values)
+    pad = 0.04 * (hi - lo + 1e-12)
+    ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], "k-", lw=1)
+    ax.set_xlabel("predicted gap")
+    ax.set_ylabel("measured gap")
+    ax.set_title(title, fontsize=10)
+    ax.grid(alpha=0.16)
+    ax.legend(fontsize=7)
+
+
+def figure_two(args: argparse.Namespace) -> Path:
+    rows = read_csv(args.depth_dir / "finite_time_rows.csv")
+    fig, axes = plt.subplots(1, 3, figsize=(14.0, 4.1), dpi=220)
+
+    ax = axes[0]
+    widths = sorted({int(row["width"]) for row in rows})
+    markers = ["o", "s", "^"]
+    for width, marker in zip(widths, markers):
+        depths = sorted({int(row["depth"]) for row in rows if int(row["width"]) == width})
+        means = []
+        sems = []
+        for depth in depths:
+            cell = [
+                row
+                for row in rows
+                if int(row["width"]) == width and int(row["depth"]) == depth
+            ]
+            values = np.asarray([float(row["empirical_gap"]) for row in cell])
+            init_means = np.asarray(
+                [
+                    np.mean(
+                        [
+                            float(row["empirical_gap"])
+                            for row in cell
+                            if int(row["init_seed"]) == init_seed
+                        ]
+                    )
+                    for init_seed in sorted({int(row["init_seed"]) for row in cell})
+                ]
+            )
+            means.append(float(values.mean()))
+            sems.append(float(init_means.std(ddof=1) / math.sqrt(len(init_means))))
+        ax.errorbar(depths, means, yerr=sems, marker=marker, capsize=2, label=f"width {width}")
+    ax.set_xlabel("hidden-layer depth")
+    ax.set_ylabel("FA loss - BP loss")
+    ax.set_title("Later cost changes with training dynamics", fontsize=10)
+    ax.grid(alpha=0.16)
+    ax.legend(fontsize=7)
+    panel_label(ax, "a")
+
+    scatter_predictor(axes[1], rows, [("fixed_gap", "^", "frozen $K(0)$")], "Frozen initialization misses drift")
+    panel_label(axes[1], "b")
+
+    scatter_predictor(
+        axes[2],
+        rows,
+        [
+            ("velocity_gap", "o", "linear velocity"),
+            ("retangent_gap", "s", "early retangent"),
+        ],
+        "Two snapshots recover the later cost",
+    )
+    panel_label(axes[2], "c")
+
+    predictor_legend = axes[2].get_legend()
+    if predictor_legend is not None:
+        axes[2].add_artist(predictor_legend)
+
+    depth_handles = [
+        plt.Line2D(
+            [],
+            [],
+            marker="o",
+            ls="",
+            color=plt.cm.viridis(i / max(len(sorted({int(row['depth']) for row in rows})) - 1, 1)),
+            label=f"depth {depth}",
+        )
+        for i, depth in enumerate(sorted({int(row["depth"]) for row in rows}))
+    ]
+    axes[2].legend(handles=depth_handles, fontsize=6, ncols=2, loc="lower right")
+
+    fig.tight_layout(w_pad=2.0)
+    path = args.outdir / "figure2_training_dynamics.png"
+    fig.savefig(path, bbox_inches="tight")
+    fig.savefig(path.with_suffix(".pdf"), bbox_inches="tight")
+    plt.close(fig)
+    return path
+
+
+def main() -> None:
+    args = parse_args()
+    args.outdir.mkdir(parents=True, exist_ok=True)
+    path1 = figure_one(args)
+    path2 = figure_two(args)
+    print(f"figure: {path1}")
+    print(f"figure: {path2}")
+
+
+if __name__ == "__main__":
+    main()