giant/giant/training/plots.py

"""Training-progress plots from `<run_dir>/metrics.csv` (gitea #75).

`MetricsCollector` (`giant.training.metrics`) writes one row per epoch with a
column set that varies by run — flow/ddpm vs wgan, routed vs not (see the
`MetricSpec` declarations in `giant.training.trainers`). This module reads
that header dynamically rather than hardcoding a column list, buckets columns
by the fixed naming convention `MetricsCollector` itself documents
(`<stage>/train/<key>`, `<stage>/val/<key>`, `<stage>/router/<key>`,
`<stage>/<key>` for point-in-time values, and an unprefixed run-level tail —
see `giant.training.metrics`'s module docstring), and renders one PDF per
applicable figure with the same `plotstyle` conventions
`giant.analysis.render` uses, for visual consistency with the
rollout-vs-reference plots.

Unlike `giant.analysis`, there is no reduce/chunk/condor split here — the CSV
is tiny and this always runs as one local pass — but the CLI entry point
still lives under `giant analyze` (`analyze metrics`) as the shared home for
plotstyle-rendered diagnostics, and shares its `analysis_runs/` output
convention (see `derive_metrics_dir`) so training-progress plots don't get
written into the training run directory itself.
"""

from __future__ import annotations

import csv
import math
from dataclasses import dataclass
from pathlib import Path

# Stage names are always exactly these two — hardcoded in
# `giant.training.trainers.build_stage_trainers` — so a column belongs to a
# stage iff it's prefixed by one of these, and everything else (bar `epoch`)
# is run-level. This is what makes dynamic header parsing tractable without
# needing to know the per-run metric keys themselves.
_STAGE_NAMES = ("stage1", "stage2")

_ACC_KEYS = {"nsec_acc", "stop_acc", "type_acc"}
_WGAN_BALANCE_KEYS = {"d_loss", "g_loss", "wasserstein", "gp_loss"}
_ROUTER_KEYS = ("entropy", "util_min", "util_max", "util_std")


@dataclass
class MetricsTable:
    """`<run_dir>/metrics.csv`, parsed with no hardcoded column list."""

    epochs: list[int]
    columns: dict[str, list[float]]

    @classmethod
    def load(cls, path: str | Path) -> "MetricsTable":
        with open(path, newline="") as f:
            rows = list(csv.DictReader(f))
        epochs = [int(float(r["epoch"])) for r in rows]
        fieldnames = rows[0].keys() if rows else []
        columns = {name: [float(r[name]) for r in rows] for name in fieldnames if name != "epoch"}
        return cls(epochs=epochs, columns=columns)

    def best_epochs(self) -> list[int]:
        is_best = self.columns.get("is_best")
        if not is_best:
            return []
        return [epoch for epoch, flag in zip(self.epochs, is_best) if flag]


# --- column classification --------------------------------------------------


def _stages(columns: dict) -> list[str]:
    return [s for s in _STAGE_NAMES if any(name.startswith(f"{s}/") for name in columns)]


def _split(columns: dict, stage: str, split: str) -> dict[str, str]:
    prefix = f"{stage}/{split}/"
    return {name[len(prefix) :]: name for name in columns if name.startswith(prefix)}


def _point_in_time(columns: dict, stage: str) -> dict[str, str]:
    prefix = f"{stage}/"
    out = {}
    for name in columns:
        if not name.startswith(prefix):
            continue
        rest = name[len(prefix) :]
        head = rest.split("/", 1)[0]
        if head not in ("train", "val", "router"):
            out[rest] = name
    return out


def _router(columns: dict, stage: str) -> dict[str, str]:
    prefix = f"{stage}/router/"
    return {name[len(prefix) :]: name for name in columns if name.startswith(prefix)}


def _run_level(columns: dict) -> dict[str, str]:
    known_prefixes = tuple(f"{s}/" for s in _STAGE_NAMES)
    return {name: name for name in columns if not name.startswith(known_prefixes)}


def _loss_keys(train: dict[str, str], val: dict[str, str]) -> list[str]:
    keys = {k for k in train if k not in _ACC_KEYS and k not in _WGAN_BALANCE_KEYS and k != "grad_norm"}
    keys |= {k for k in val if k not in _ACC_KEYS and k not in _WGAN_BALANCE_KEYS and k != "grad_norm"}
    return sorted(keys)


# --- output location ---------------------------------------------------------


def derive_metrics_dir(
    run_dir: str | Path,
    out_dir: str | Path | None = None,
    default_base: str | Path | None = None,
) -> Path:
    """Plots output directory.

    Precedence: an explicit `out_dir` always wins. Otherwise
    `default_base / f"metrics_{run_dir.name}"` (the CLI passes the repo's
    gitignored `analysis_runs/`, matching `giant.analysis.run.derive_run_dir`'s
    convention) — training-progress plots live alongside rollout-vs-reference
    analysis runs, not inside the training run directory itself.
    """
    if out_dir is not None:
        return Path(out_dir)
    base = Path(default_base) if default_base is not None else Path.cwd() / "analysis_runs"
    return base / f"metrics_{Path(run_dir).name}"


# --- figures ------------------------------------------------------------------


def _mark_best(ax, table: MetricsTable) -> None:
    for epoch in table.best_epochs():
        ax.axvline(epoch, color="grey", linestyle="--", linewidth=0.8, alpha=0.7)


def _overview_figure(table: MetricsTable):
    import plotstyle as ps

    run_level = _run_level(table.columns)
    if "val/loss" not in run_level:
        return None
    fig, ax = ps.new_figure("thesis-single", title="training overview")
    ax.plot(table.epochs, table.columns["val/loss"], label="val/loss")
    if "val/marginal_kl" in run_level:
        kl = table.columns["val/marginal_kl"]
        if any(math.isfinite(v) for v in kl):
            ax.plot(table.epochs, kl, label="val/marginal_kl")
    _mark_best(ax, table)
    best = table.best_epochs()
    if best:
        idx = table.epochs.index(best[-1])
        ax.annotate(
            f"best: epoch {best[-1]}\nval/loss={table.columns['val/loss'][idx]:.4g}",
            xy=(best[-1], table.columns["val/loss"][idx]),
            xytext=(0.98, 0.95),
            textcoords="axes fraction",
            ha="right",
            va="top",
            fontsize=8,
        )
    ax.set_xlabel("epoch")
    ax.set_ylabel("loss")
    ps.style_legend(ax, title="series")
    return fig


def _loss_figure(table: MetricsTable, stage: str):
    import plotstyle as ps

    train = _split(table.columns, stage, "train")
    val = _split(table.columns, stage, "val")
    keys = _loss_keys(train, val)
    if not keys:
        return None
    n = len(keys)
    ncols = min(3, n)
    nrows = (n + ncols - 1) // ncols
    fig, axes = ps.new_figure(
        "slide-16x9",
        title=f"{stage} loss",
        nrows=nrows,
        ncols=ncols,
        squeeze=False,
    )
    flat = axes.ravel()
    for ax, key in zip(flat, keys):
        if key in train:
            ax.plot(table.epochs, table.columns[train[key]], label="train")
        if key in val:
            ax.plot(table.epochs, table.columns[val[key]], label="val")
        ax.set_yscale("log")
        ax.set_title(key, fontsize=8)
        ax.set_xlabel("epoch")
    for j in range(n, len(flat)):
        flat[j].set_visible(False)
    ps.style_legend(flat[0], title="series")
    return fig


def _lr_figure(table: MetricsTable):
    import plotstyle as ps

    series: dict[str, str] = {}
    for stage in _stages(table.columns):
        for key, col in _point_in_time(table.columns, stage).items():
            series[f"{stage}/{key}"] = col
    if not series:
        return None
    fig, ax = ps.new_figure("thesis-single", title="learning rate schedule")
    for label, col in series.items():
        ax.plot(table.epochs, table.columns[col], label=label)
    ax.set_xlabel("epoch")
    ax.set_ylabel("learning rate")
    ps.style_legend(ax, title="series")
    return fig


def _accuracy_figure(table: MetricsTable, stage: str):
    import plotstyle as ps

    train = _split(table.columns, stage, "train")
    val = _split(table.columns, stage, "val")
    keys = sorted((set(train) | set(val)) & _ACC_KEYS)
    if not keys:
        return None
    n = len(keys)
    fig, axes = ps.new_figure("slide-16x9", title=f"{stage} accuracy", nrows=1, ncols=n, squeeze=False)
    flat = axes.ravel()
    for ax, key in zip(flat, keys):
        if key in train:
            ax.plot(table.epochs, table.columns[train[key]], label="train")
        if key in val:
            ax.plot(table.epochs, table.columns[val[key]], label="val")
        ax.set_title(key, fontsize=8)
        ax.set_xlabel("epoch")
        ax.set_ylim(0, 1)
    ps.style_legend(flat[0], title="series")
    return fig


def _grad_norm_figure(table: MetricsTable):
    import plotstyle as ps

    run_level = _run_level(table.columns)
    if "grad_norm" not in run_level:
        return None
    fig, ax = ps.new_figure("thesis-single", title="gradient norm")
    ax.plot(table.epochs, table.columns["grad_norm"], label="grad_norm")
    for stage in _stages(table.columns):
        train = _split(table.columns, stage, "train")
        for key in ("grad_norm_d", "grad_norm_g", "grad_norm_type_slice", "grad_norm_cont_slice"):
            if key in train:
                ax.plot(table.epochs, table.columns[train[key]], label=f"{stage}/{key}")
    ax.set_yscale("log")
    ax.set_xlabel("epoch")
    ax.set_ylabel("grad norm")
    ps.style_legend(ax, title="series")
    return fig


def _router_figure(table: MetricsTable, stage: str):
    import plotstyle as ps

    router = _router(table.columns, stage)
    if "entropy" not in router:
        return None
    fig, ax = ps.new_figure("thesis-single", title=f"{stage} router health")
    ax.plot(table.epochs, table.columns[router["entropy"]], label="entropy", color="black")
    ax.set_xlabel("epoch")
    ax.set_ylabel("entropy [bits]")
    ax2 = ax.twinx()
    for key in ("util_min", "util_max", "util_std"):
        if key in router:
            ax2.plot(table.epochs, table.columns[router[key]], label=key, linestyle="--")
    ax2.set_ylabel("expert utilization")
    ax2.set_ylim(0, 1)
    lines1, labels1 = ax.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax.legend(lines1 + lines2, labels1 + labels2, loc="upper right", frameon=False, fontsize=7)
    return fig


def _wgan_balance_figure(table: MetricsTable, stage: str):
    import plotstyle as ps

    train = _split(table.columns, stage, "train")
    keys = [k for k in _WGAN_BALANCE_KEYS if k in train]
    if not keys:
        return None
    fig, ax = ps.new_figure("thesis-single", title=f"{stage} WGAN critic/generator balance")
    for key in sorted(keys):
        ax.plot(table.epochs, table.columns[train[key]], label=key)
    ax.set_xlabel("epoch")
    ax.set_ylabel("value")
    ps.style_legend(ax, title="series")
    return fig


def _throughput_figure(table: MetricsTable):
    import plotstyle as ps

    run_level = _run_level(table.columns)
    keys = [k for k in ("samples_per_sec", "gpu_mem_mb", "epoch_time_s") if k in run_level]
    if not keys:
        return None
    fig, axes = ps.new_figure("slide-16x9", title="throughput / resources", nrows=1, ncols=len(keys), squeeze=False)
    flat = axes.ravel()
    for ax, key in zip(flat, keys):
        ax.plot(table.epochs, table.columns[key])
        _mark_best(ax, table)
        ax.set_title(key, fontsize=8)
        ax.set_xlabel("epoch")
    return fig


# --- entry point ---------------------------------------------------------


def render_metrics(
    run_dir: str | Path,
    out_dir: str | Path | None = None,
    default_base: str | Path | None = None,
) -> list[Path]:
    """`<run_dir>/metrics.csv` -> `<plots dir>/<name>.pdf`.

    See `derive_metrics_dir` for how the plots directory is resolved.
    """
    import matplotlib.pyplot as plt
    import plotstyle as ps

    ps.use()
    table = MetricsTable.load(Path(run_dir) / "metrics.csv")
    plots_dir = derive_metrics_dir(run_dir, out_dir, default_base)
    plots_dir.mkdir(parents=True, exist_ok=True)

    figures = [("overview", _overview_figure(table))]
    for stage in _stages(table.columns):
        figures.append((f"{stage}_loss", _loss_figure(table, stage)))
    figures.append(("lr", _lr_figure(table)))
    for stage in _stages(table.columns):
        figures.append((f"{stage}_accuracy", _accuracy_figure(table, stage)))
    figures.append(("grad_norm", _grad_norm_figure(table)))
    for stage in _stages(table.columns):
        figures.append((f"{stage}_router", _router_figure(table, stage)))
        figures.append((f"{stage}_wgan_balance", _wgan_balance_figure(table, stage)))
    figures.append(("throughput", _throughput_figure(table)))

    paths: list[Path] = []
    for name, fig in figures:
        if fig is None:
            continue
        path = plots_dir / name
        ps.savefig(fig, str(path), formats=("pdf",))
        plt.close(fig)
        paths.append(path.with_suffix(".pdf"))
    return paths