diff options
Diffstat (limited to 'scripts/plot_aaai_main_figures.py')
| -rw-r--r-- | scripts/plot_aaai_main_figures.py | 311 |
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() |