plot_gradient_alignment_results.py

#!/usr/bin/env python3
"""Render plots from a gradient-alignment metrics.json file."""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
from typing import Dict, List, Sequence

import matplotlib.pyplot as plt
import numpy as np

METRIC_SPECS = [
    {
        "key": "train_loss",
        "title": "Mean train loss on fixed train batches",
        "ylabel": "Cross-entropy loss",
        "filename": "train_loss.png",
    },
    {
        "key": "eval_loss",
        "title": "Mean eval loss on held-out validation batches",
        "ylabel": "Cross-entropy loss",
        "filename": "eval_loss.png",
    },
    {
        "key": "parameter_norm",
        "title": "Mean parameter vector norm across seeds",
        "ylabel": "Norm",
        "filename": "parameter_norm.png",
    },
    {
        "key": "whole_model_cosine",
        "title": "Whole-model gradient cosine vs training time",
        "ylabel": "Mean pairwise cosine",
        "filename": "whole_model_cosine.png",
    },
    {
        "key": "function_update_cosine",
        "title": "Function-space update alignment (same batch) vs training time",
        "ylabel": "Mean pairwise cosine",
        "filename": "function_update_cosine.png",
    },
    {
        "key": "function_update_cross_batch_cosine",
        "title": "Function-space update alignment (held-out readout batch) vs training time",
        "ylabel": "Mean pairwise cosine",
        "filename": "function_update_cross_batch_cosine.png",
    },
    {
        "key": "function_update_alignment_to_prev",
        "title": "Function-space update alignment to previous checkpoint (same batch)",
        "ylabel": "Mean cosine",
        "filename": "function_update_alignment_to_prev.png",
    },
    {
        "key": "function_update_cross_batch_alignment_to_prev",
        "title": "Function-space update alignment to previous checkpoint (held-out batch)",
        "ylabel": "Mean cosine",
        "filename": "function_update_cross_batch_alignment_to_prev.png",
    },
    {
        "key": "hessian_grad_proxy_cosine",
        "title": "Cross-seed alignment of Hessian proxy Hg_hat",
        "ylabel": "Mean pairwise cosine",
        "filename": "hessian_grad_proxy_cosine.png",
    },
    {
        "key": "gradient_hessian_proxy_alignment",
        "title": "Alignment between gradient and Hessian proxy Hg_hat",
        "ylabel": "Mean cosine",
        "filename": "gradient_hessian_proxy_alignment.png",
    },
    {
        "key": "local_learning_coefficient",
        "title": "Sparse local learning coefficient on fixed train subset",
        "ylabel": "LLC estimate",
        "filename": "local_learning_coefficient.png",
        "error_key": "local_learning_coefficient_ci95_halfwidth",
    },
    {
        "key": "local_learning_coefficient_ci95_halfwidth",
        "title": "Sparse LLC 95% confidence half-width",
        "ylabel": "95% CI half-width",
        "filename": "local_learning_coefficient_ci95_halfwidth.png",
    },
    {
        "key": "local_learning_coefficient_seed_std",
        "title": "Across-seed dispersion of sparse LLC estimates",
        "ylabel": "Std across seeds",
        "filename": "local_learning_coefficient_seed_std.png",
    },
    {
        "key": "empirical_fisher_top1_eigenvalue",
        "title": "Top empirical Fisher eigenvalue estimate",
        "ylabel": "Estimated eigenvalue",
        "filename": "empirical_fisher_top1_eigenvalue.png",
    },
    {
        "key": "empirical_fisher_hessian_top1_alignment",
        "title": "Alignment between top empirical Fisher and top Hessian vectors",
        "ylabel": "Mean absolute cosine",
        "filename": "empirical_fisher_hessian_top1_alignment.png",
    },
    {
        "key": "gradient_to_hessian_top1_alignment",
        "title": "Alignment between gradient and top Hessian vector",
        "ylabel": "Mean absolute cosine",
        "filename": "gradient_to_hessian_top1_alignment.png",
    },
    {
        "key": "hessian_top1_pairwise_alignment",
        "title": "Top Hessian vector alignment across seeds",
        "ylabel": "Mean absolute cosine",
        "filename": "hessian_top1_pairwise_alignment.png",
    },
    {
        "key": "hessian_top1_alignment_to_init",
        "title": "Top Hessian vector alignment to initialization",
        "ylabel": "Mean absolute cosine",
        "filename": "hessian_top1_alignment_to_init.png",
    },
    {
        "key": "hessian_top1_alignment_to_prev",
        "title": "Top Hessian vector alignment to previous sparse checkpoint",
        "ylabel": "Mean absolute cosine",
        "filename": "hessian_top1_alignment_to_prev.png",
    },
    {
        "key": "hessian_top1_eigenvalue",
        "title": "Leading Hessian Rayleigh quotient estimate",
        "ylabel": "Estimated eigenvalue",
        "filename": "hessian_top1_eigenvalue.png",
    },
    {
        "key": "hessian_trace_estimate",
        "title": "Hutchinson estimate of Hessian trace",
        "ylabel": "Trace estimate",
        "filename": "hessian_trace_estimate.png",
    },
    {
        "key": "jacobian_frobenius_sq_estimate",
        "title": "Jacobian Frobenius norm squared estimate",
        "ylabel": "Norm squared estimate",
        "filename": "jacobian_frobenius_sq_estimate.png",
    },
    {
        "key": "metric_time_total_s",
        "title": "Total metric compute time per checkpoint",
        "ylabel": "Seconds",
        "filename": "metric_time_total_s.png",
    },
    {
        "key": "gradient_covariance_entropy_rank",
        "title": "Gradient covariance entropy effective rank",
        "ylabel": "Effective rank",
        "filename": "gradient_covariance_entropy_rank.png",
    },
    {
        "key": "gradient_covariance_participation_ratio",
        "title": "Gradient covariance participation ratio",
        "ylabel": "Participation ratio",
        "filename": "gradient_covariance_participation_ratio.png",
    },
    {
        "key": "topk_explained_variance",
        "title": "Top-k explained variance of gradient covariance",
        "ylabel": "Variance fraction",
        "filename": "topk_explained_variance.png",
    },
    {
        "key": "top1_explained_variance",
        "title": "Top-1 explained variance of gradient covariance",
        "ylabel": "Variance fraction",
        "filename": "top1_explained_variance.png",
    },
    {
        "key": "projected_gradient_energy_fraction",
        "title": "Gradient energy in top-k eigenspace",
        "ylabel": "Energy fraction",
        "filename": "projected_gradient_energy_fraction.png",
    },
    {
        "key": "residual_gradient_energy_fraction",
        "title": "Gradient energy outside top-k eigenspace",
        "ylabel": "Energy fraction",
        "filename": "residual_gradient_energy_fraction.png",
    },
    {
        "key": "subspace_similarity",
        "title": "Gradient subspace similarity vs training time",
        "ylabel": "Mean principal-overlap score",
        "filename": "subspace_similarity.png",
    },
    {
        "key": "subspace_overlap_to_init",
        "title": "Top-k gradient subspace overlap to initialization",
        "ylabel": "Mean principal-overlap score",
        "filename": "subspace_overlap_to_init.png",
    },
    {
        "key": "subspace_overlap_to_prev",
        "title": "Top-k gradient subspace overlap to previous checkpoint",
        "ylabel": "Mean principal-overlap score",
        "filename": "subspace_overlap_to_prev.png",
    },
    {
        "key": "top1_alignment_to_init",
        "title": "Top eigenvector alignment to initialization",
        "ylabel": "Mean absolute cosine",
        "filename": "top1_alignment_to_init.png",
    },
    {
        "key": "top1_alignment_to_prev",
        "title": "Top eigenvector alignment to previous checkpoint",
        "ylabel": "Mean absolute cosine",
        "filename": "top1_alignment_to_prev.png",
    },
]


_TIMING_COMPONENTS = [
    ("training_time_s", "Training between checkpoints"),
    ("metric_time_losses_s", "Loss probes"),
    ("metric_time_grad_snapshots_s", "Gradient snapshots"),
    ("metric_time_function_updates_s", "Function updates"),
    ("metric_time_covariance_s", "Covariance geometry"),
    ("metric_time_hessian_proxy_s", "Hessian proxy"),
    ("metric_time_sparse_hessian_top1_s", "Top Hessian eigenpair"),
    ("metric_time_sparse_hessian_trace_s", "Hessian trace"),
    ("metric_time_sparse_jacobian_s", "Jacobian probes"),
    ("metric_time_local_learning_coefficient_s", "LLC estimation"),
]


def _mean_finite(values: Sequence[float]) -> float:
    finite = [v for v in values if np.isfinite(v)]
    return float(np.mean(finite)) if finite else 0.0


def _timing_breakdown(
    records: Sequence[Dict[str, object]],
) -> tuple[List[str], List[float]]:
    """Return (labels, mean_seconds) sorted ascending by time, filtering zero/absent."""
    labels: List[str] = []
    means: List[float] = []
    for key, label in _TIMING_COMPONENTS:
        if key not in records[0]:
            continue
        mean = _mean_finite([float(r[key]) for r in records])
        if mean > 0:
            labels.append(label)
            means.append(mean)
    if not labels:
        return [], []
    order = list(np.argsort(means))
    return [labels[i] for i in order], [means[i] for i in order]


def plot_metric_time_breakdown(
    records: Sequence[Dict[str, object]],
    output_path: Path,
) -> None:
    labels, means = _timing_breakdown(records)
    if not labels:
        return
    total = sum(means)
    cmap = plt.get_cmap("tab10")
    colors = [cmap(i / max(len(labels), 1)) for i in range(len(labels))]

    fig, ax = plt.subplots(figsize=(8, max(3, 0.5 * len(labels) + 1.5)))
    y_pos = range(len(labels))
    bars = ax.barh(y_pos, means, color=colors, edgecolor="white", linewidth=0.5)
    for bar, mean in zip(bars, means):
        pct = 100 * mean / total if total > 0 else 0
        ax.text(
            bar.get_width() + total * 0.01,
            bar.get_y() + bar.get_height() / 2,
            f"{mean:.2f}s ({pct:.0f}%)",
            va="center",
            fontsize=9,
        )
    ax.set_yticks(list(y_pos))
    ax.set_yticklabels(labels)
    ax.set_xlabel("Mean seconds per invocation")
    ax.set_title("Metric computation time breakdown")
    ax.set_xlim(0, max(means) * 1.35)
    ax.grid(True, axis="x", alpha=0.3)
    fig.tight_layout()
    fig.savefig(output_path, dpi=180)
    plt.close(fig)


def draw_metric_time_breakdown_on_axis(
    ax,
    records: Sequence[Dict[str, object]],
) -> None:
    labels, means = _timing_breakdown(records)
    if not labels:
        ax.axis("off")
        return
    total = sum(means)
    cmap = plt.get_cmap("tab10")
    colors = [cmap(i / max(len(labels), 1)) for i in range(len(labels))]

    y_pos = range(len(labels))
    bars = ax.barh(y_pos, means, color=colors, edgecolor="white", linewidth=0.5)
    for bar, mean in zip(bars, means):
        pct = 100 * mean / total if total > 0 else 0
        ax.text(
            bar.get_width() + total * 0.008,
            bar.get_y() + bar.get_height() / 2,
            f"{mean:.1f}s ({pct:.0f}%)",
            va="center",
            fontsize=7,
        )
    ax.set_yticks(list(y_pos))
    ax.set_yticklabels(labels, fontsize=7)
    ax.set_xlabel("Mean seconds per invocation")
    ax.set_title("Metric time breakdown", fontsize=10)
    ax.set_xlim(0, max(means) * 1.4)
    ax.grid(True, axis="x", alpha=0.3)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--input-dir",
        type=Path,
        help="Directory containing metrics.json and where plots will be written. Defaults to the most recently modified results_* directory.",
    )
    return parser.parse_args()


def resolve_input_dir(input_dir: Path | None) -> Path:
    if input_dir is not None:
        return input_dir

    repo_root = Path(__file__).resolve().parent
    candidates = [
        path
        for path in repo_root.glob("results_*")
        if path.is_dir() and (path / "metrics.json").exists()
    ]
    if not candidates:
        raise FileNotFoundError(
            f"Could not find any results_* directory with metrics.json under {repo_root}"
        )
    latest = max(candidates, key=lambda path: path.stat().st_mtime)
    print(f"Using latest results directory: {latest}")
    return latest


def load_records(input_dir: Path) -> List[Dict[str, object]]:
    metrics_path = input_dir / "metrics.json"
    if not metrics_path.exists():
        raise FileNotFoundError(f"Could not find metrics file at {metrics_path}")
    return json.loads(metrics_path.read_text(encoding="utf-8"))


def load_config(input_dir: Path) -> Dict[str, object]:
    config_path = input_dir / "config.json"
    if not config_path.exists():
        return {}
    return json.loads(config_path.read_text(encoding="utf-8"))


def available_metric_specs(records: Sequence[Dict[str, object]]) -> List[Dict[str, str]]:
    if not records:
        return []
    return [spec for spec in METRIC_SPECS if spec["key"] in records[0]]


def plot_metric(
    steps: Sequence[int],
    values: Sequence[float],
    title: str,
    ylabel: str,
    output_path: Path,
    use_markers: bool = False,
    errors: Sequence[float] | None = None,
) -> None:
    plt.figure(figsize=(7, 4))
    plot_kwargs = {"linewidth": 2.0}
    contains_gaps = any(not np.isfinite(value) for value in values)
    if use_markers or contains_gaps:
        plot_kwargs.update({"marker": "o", "markersize": 4.0})
    line = plt.plot(steps, values, **plot_kwargs)[0]
    if errors is not None:
        finite_points = [
            (step, value, error)
            for step, value, error in zip(steps, values, errors)
            if np.isfinite(value) and np.isfinite(error)
        ]
        if finite_points:
            err_steps, err_values, err_halfwidths = zip(*finite_points)
            plt.errorbar(
                err_steps,
                err_values,
                yerr=err_halfwidths,
                fmt="none",
                ecolor=line.get_color(),
                elinewidth=1.0,
                alpha=0.35,
                capsize=2.0,
            )
    plt.xlabel("Training step")
    plt.ylabel(ylabel)
    plt.title(title)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(output_path, dpi=180)
    plt.close()


def draw_metric_on_axis(
    ax,
    steps: Sequence[int],
    values: Sequence[float],
    title: str,
    ylabel: str,
    use_markers: bool = False,
    errors: Sequence[float] | None = None,
) -> None:
    plot_kwargs = {"linewidth": 1.8}
    contains_gaps = any(not np.isfinite(value) for value in values)
    if use_markers or contains_gaps:
        plot_kwargs.update({"marker": "o", "markersize": 3.5})
    line = ax.plot(steps, values, **plot_kwargs)[0]
    if errors is not None:
        finite_points = [
            (step, value, error)
            for step, value, error in zip(steps, values, errors)
            if np.isfinite(value) and np.isfinite(error)
        ]
        if finite_points:
            err_steps, err_values, err_halfwidths = zip(*finite_points)
            ax.errorbar(
                err_steps,
                err_values,
                yerr=err_halfwidths,
                fmt="none",
                ecolor=line.get_color(),
                elinewidth=0.9,
                alpha=0.3,
                capsize=2.0,
            )
    ax.set_title(title, fontsize=10)
    ax.set_xlabel("Training step")
    ax.set_ylabel(ylabel)
    ax.grid(True, alpha=0.3)


def gradient_trace_ratio(records: Sequence[Dict[str, object]]) -> List[float]:
    ratios: List[float] = []
    for record in records:
        empirical_fisher = float(record["empirical_fisher_trace"])
        gradient_covariance = float(record["gradient_covariance_trace"])
        if not np.isfinite(empirical_fisher) or not np.isfinite(gradient_covariance):
            ratios.append(np.nan)
        elif abs(gradient_covariance) <= 1e-12:
            ratios.append(np.nan)
        else:
            ratios.append(empirical_fisher / gradient_covariance)
    return ratios


def mean_gradient_norm_series(records: Sequence[Dict[str, object]]) -> tuple[List[float], str]:
    if records and "mean_gradient_norm" in records[0]:
        return (
            [float(record["mean_gradient_norm"]) for record in records],
            "Mean gradient norm across covariance batches",
        )

    norms: List[float] = []
    for record in records:
        empirical_fisher = float(record["empirical_fisher_trace"])
        gradient_covariance = float(record["gradient_covariance_trace"])
        if not np.isfinite(empirical_fisher) or not np.isfinite(gradient_covariance):
            norms.append(np.nan)
        else:
            mean_grad_sq = max(empirical_fisher - gradient_covariance, 0.0)
            norms.append(float(np.sqrt(mean_grad_sq)))
    return norms, "Mean gradient norm across covariance batches (derived from trace gap)"


def plot_trace_comparison(
    records: Sequence[Dict[str, object]],
    output_path: Path,
    use_markers: bool = False,
) -> None:
    steps = [int(record["step"]) for record in records]
    empirical_fisher = [float(record["empirical_fisher_trace"]) for record in records]
    gradient_covariance = [float(record["gradient_covariance_trace"]) for record in records]

    plt.figure(figsize=(7, 4))
    plot_kwargs = {"linewidth": 2.0}
    contains_gaps = any(
        not np.isfinite(value) for value in empirical_fisher + gradient_covariance
    )
    if use_markers or contains_gaps:
        plot_kwargs.update({"marker": "o", "markersize": 4.0})
    plt.plot(steps, empirical_fisher, label="Empirical Fisher trace", **plot_kwargs)
    plt.plot(steps, gradient_covariance, label="Gradient covariance trace", **plot_kwargs)
    plt.xlabel("Training step")
    plt.ylabel("Trace estimate")
    plt.title("Centered vs uncentered gradient trace")
    plt.grid(True, alpha=0.3)
    plt.legend(frameon=False)
    plt.tight_layout()
    plt.savefig(output_path, dpi=180)
    plt.close()


def draw_trace_comparison_on_axis(
    ax,
    records: Sequence[Dict[str, object]],
    use_markers: bool = False,
) -> None:
    steps = [int(record["step"]) for record in records]
    empirical_fisher = [float(record["empirical_fisher_trace"]) for record in records]
    gradient_covariance = [float(record["gradient_covariance_trace"]) for record in records]

    plot_kwargs = {"linewidth": 1.8}
    contains_gaps = any(
        not np.isfinite(value) for value in empirical_fisher + gradient_covariance
    )
    if use_markers or contains_gaps:
        plot_kwargs.update({"marker": "o", "markersize": 3.5})
    ax.plot(steps, empirical_fisher, label="Empirical Fisher", **plot_kwargs)
    ax.plot(steps, gradient_covariance, label="Covariance", **plot_kwargs)
    ax.set_title("Centered vs uncentered gradient trace", fontsize=10)
    ax.set_xlabel("Training step")
    ax.set_ylabel("Trace estimate")
    ax.grid(True, alpha=0.3)
    ax.legend(fontsize=7, frameon=False, loc="best")


def plot_trace_ratio(
    records: Sequence[Dict[str, object]],
    output_path: Path,
    use_markers: bool = False,
) -> None:
    steps = [int(record["step"]) for record in records]
    ratios = gradient_trace_ratio(records)

    plt.figure(figsize=(7, 4))
    plot_kwargs = {"linewidth": 2.0}
    contains_gaps = any(not np.isfinite(value) for value in ratios)
    if use_markers or contains_gaps:
        plot_kwargs.update({"marker": "o", "markersize": 4.0})
    plt.plot(steps, ratios, **plot_kwargs)
    plt.axhline(1.0, color="black", linewidth=0.8, alpha=0.35)
    plt.xlabel("Training step")
    plt.ylabel("Uncentered / centered")
    plt.title("Gradient trace ratio")
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(output_path, dpi=180)
    plt.close()


def draw_trace_ratio_on_axis(
    ax,
    records: Sequence[Dict[str, object]],
    use_markers: bool = False,
) -> None:
    steps = [int(record["step"]) for record in records]
    ratios = gradient_trace_ratio(records)

    plot_kwargs = {"linewidth": 1.8}
    contains_gaps = any(not np.isfinite(value) for value in ratios)
    if use_markers or contains_gaps:
        plot_kwargs.update({"marker": "o", "markersize": 3.5})
    ax.plot(steps, ratios, **plot_kwargs)
    ax.axhline(1.0, color="black", linewidth=0.7, alpha=0.35)
    ax.set_title("Gradient trace ratio", fontsize=10)
    ax.set_xlabel("Training step")
    ax.set_ylabel("Uncentered / centered")
    ax.grid(True, alpha=0.3)


def plot_mean_gradient_norm(
    records: Sequence[Dict[str, object]],
    output_path: Path,
    use_markers: bool = False,
) -> None:
    steps = [int(record["step"]) for record in records]
    norms, title = mean_gradient_norm_series(records)

    plt.figure(figsize=(7, 4))
    plot_kwargs = {"linewidth": 2.0}
    contains_gaps = any(not np.isfinite(value) for value in norms)
    if use_markers or contains_gaps:
        plot_kwargs.update({"marker": "o", "markersize": 4.0})
    plt.plot(steps, norms, **plot_kwargs)
    plt.xlabel("Training step")
    plt.ylabel("Norm")
    plt.title(title)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(output_path, dpi=180)
    plt.close()


def draw_mean_gradient_norm_on_axis(
    ax,
    records: Sequence[Dict[str, object]],
    use_markers: bool = False,
) -> None:
    steps = [int(record["step"]) for record in records]
    norms, title = mean_gradient_norm_series(records)

    plot_kwargs = {"linewidth": 1.8}
    contains_gaps = any(not np.isfinite(value) for value in norms)
    if use_markers or contains_gaps:
        plot_kwargs.update({"marker": "o", "markersize": 3.5})
    ax.plot(steps, norms, **plot_kwargs)
    ax.set_title(title, fontsize=10)
    ax.set_xlabel("Training step")
    ax.set_ylabel("Norm")
    ax.grid(True, alpha=0.3)


def plot_layer_lines(
    records: Sequence[Dict[str, object]],
    output_path: Path,
    use_markers: bool = False,
) -> None:
    steps = [int(record["step"]) for record in records]
    layer_names = sorted(set().union(*(record["per_layer_cosine"].keys() for record in records)))

    plt.figure(figsize=(max(8, len(records) * 0.5), 5))
    for layer_name in layer_names:
        values = [record["per_layer_cosine"].get(layer_name, np.nan) for record in records]
        plot_kwargs = {"linewidth": 1.8, "label": layer_name}
        if use_markers:
            plot_kwargs.update({"marker": "o", "markersize": 3.5})
        plt.plot(steps, values, **plot_kwargs)

    plt.xlabel("Training step")
    plt.ylabel("Mean pairwise cosine")
    plt.title("Per-layer gradient alignment across seeds")
    plt.axhline(0.0, color="black", linewidth=0.8, alpha=0.35)
    plt.grid(True, alpha=0.3)
    plt.legend(loc="center left", bbox_to_anchor=(1.02, 0.5), frameon=False)
    plt.tight_layout()
    plt.savefig(output_path, dpi=180)
    plt.close()


def draw_layer_lines_on_axis(
    ax,
    records: Sequence[Dict[str, object]],
    use_markers: bool = False,
) -> None:
    steps = [int(record["step"]) for record in records]
    layer_names = sorted(set().union(*(record["per_layer_cosine"].keys() for record in records)))
    for layer_name in layer_names:
        values = [record["per_layer_cosine"].get(layer_name, np.nan) for record in records]
        plot_kwargs = {"linewidth": 1.2, "label": layer_name}
        if use_markers:
            plot_kwargs.update({"marker": "o", "markersize": 2.5})
        ax.plot(steps, values, **plot_kwargs)
    ax.set_title("Per-layer gradient alignment across seeds", fontsize=10)
    ax.set_xlabel("Training step")
    ax.set_ylabel("Mean pairwise cosine")
    ax.axhline(0.0, color="black", linewidth=0.7, alpha=0.35)
    ax.grid(True, alpha=0.3)
    ax.legend(fontsize=6, frameon=False, ncol=1, loc="best")


def plot_dashboard(
    input_dir: Path,
    records: Sequence[Dict[str, object]],
    metric_specs: Sequence[Dict[str, str]],
    use_markers: bool = False,
) -> None:
    steps = [int(r["step"]) for r in records]
    has_trace_plots = all(
        key in records[0] for key in ("empirical_fisher_trace", "gradient_covariance_trace")
    )
    extra_panels = 2 + (3 if has_trace_plots else 0)
    total_panels = len(metric_specs) + extra_panels
    ncols = 4
    nrows = math.ceil(total_panels / ncols)
    fig, axes = plt.subplots(nrows, ncols, figsize=(ncols * 5.2, nrows * 3.7))
    axes = np.atleast_1d(axes).reshape(-1)

    for ax, spec in zip(axes, metric_specs):
        values = [float(record[spec["key"]]) for record in records]
        errors = None
        if "error_key" in spec:
            errors = [float(record.get(spec["error_key"], np.nan)) for record in records]
        draw_metric_on_axis(
            ax,
            steps,
            values,
            title=spec["title"],
            ylabel=spec["ylabel"],
            use_markers=use_markers,
            errors=errors,
        )

    next_idx = len(metric_specs)
    if has_trace_plots and next_idx < len(axes):
        draw_trace_comparison_on_axis(axes[next_idx], records, use_markers=use_markers)
        next_idx += 1
    if has_trace_plots and next_idx < len(axes):
        draw_trace_ratio_on_axis(axes[next_idx], records, use_markers=use_markers)
        next_idx += 1
    if has_trace_plots and next_idx < len(axes):
        draw_mean_gradient_norm_on_axis(axes[next_idx], records, use_markers=use_markers)
        next_idx += 1
    if next_idx < len(axes):
        draw_layer_lines_on_axis(axes[next_idx], records, use_markers=use_markers)
        next_idx += 1
    if next_idx < len(axes):
        draw_metric_time_breakdown_on_axis(axes[next_idx], records)

    for ax in axes[total_panels:]:
        ax.axis("off")

    fig.suptitle("Training, Alignment, And Curvature Dashboard", fontsize=16)
    fig.tight_layout(rect=(0, 0, 1, 0.98))
    fig.savefig(input_dir / "dashboard.png", dpi=180)
    plt.close(fig)


def render_all_plots(
    input_dir: Path,
    records: Sequence[Dict[str, object]],
    config: Dict[str, object],
) -> None:
    if not records:
        return

    steps = [int(r["step"]) for r in records]
    use_markers = bool(config.get("checkpoint_powers_of_two", False))
    metric_specs = available_metric_specs(records)

    for spec in metric_specs:
        errors = None
        if "error_key" in spec:
            errors = [float(r.get(spec["error_key"], np.nan)) for r in records]
        plot_metric(
            steps,
            [float(r[spec["key"]]) for r in records],
            title=spec["title"],
            ylabel=spec["ylabel"],
            output_path=input_dir / spec["filename"],
            use_markers=use_markers,
            errors=errors,
        )
    if all(key in records[0] for key in ("empirical_fisher_trace", "gradient_covariance_trace")):
        (input_dir / "empirical_fisher_trace.png").unlink(missing_ok=True)
        (input_dir / "gradient_covariance_trace.png").unlink(missing_ok=True)
        plot_trace_comparison(
            records,
            input_dir / "gradient_trace_comparison.png",
            use_markers=use_markers,
        )
        plot_trace_ratio(
            records,
            input_dir / "gradient_trace_ratio.png",
            use_markers=use_markers,
        )
        plot_mean_gradient_norm(
            records,
            input_dir / "mean_gradient_norm.png",
            use_markers=use_markers,
        )
    plot_layer_lines(records, input_dir / "per_layer_cosine.png", use_markers=use_markers)
    plot_metric_time_breakdown(records, input_dir / "metric_time_breakdown.png")
    plot_dashboard(input_dir, records, metric_specs, use_markers=use_markers)


def main() -> None:
    args = parse_args()
    input_dir = resolve_input_dir(args.input_dir)
    records = load_records(input_dir)
    config = load_config(input_dir)
    render_all_plots(input_dir, records, config)


if __name__ == "__main__":
    main()