lars
878e9ddca3
CI / Format (ruff format) (push) Failing after 28s
CI / Lint (ruff check) (push) Successful in 29s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 33s
CI / Type check (ty) (push) Successful in 37s
CI / Format (ruff format) (pull_request) Failing after 37s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 37s
CI / Tests (pull_request) Successful in 2m49s
CI / Tests (push) Successful in 2m55s
The design doc and its followups doc are no longer needed as a live reference now that the v0.3.0 redesign is implemented — comments and docstrings across the codebase cited it extensively (file path, "design doc §X.Y", "decision N", or bare "§X.Y" section numbers) as design rationale. Removed docs/ and edited every citing comment/docstring to drop the now-dangling reference while keeping the substantive explanation next to it. CLAUDE.md's v0.3.0 roadmap bullet loses its trailing pointer to the deleted file. Verified: no remaining "docs/v0.3.0", "design doc", "decision N", or "§N.N" references (repo-wide grep); ruff and ty clean; full test suite on the heaviest-touched modules (network, sample, rollout, migration, config, train) passes. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
406 lines
14 KiB
Python
406 lines
14 KiB
Python
"""Render reduced artifacts to styled PDFs + gallery metadata (the local step).
|
||
|
||
This is the *only* module that imports ``plotstyle`` (ETPlot's KIT matplotlib
|
||
theme), which renders through a real LaTeX toolchain — so it runs on the
|
||
submit/login node, never on a compute worker. It reads nothing but the small
|
||
``Reduced`` JSON files a run produced, so it is fully decoupled from the heavy
|
||
streaming compute.
|
||
|
||
For each reduced artifact it writes ``<out>/<family>/<id>.pdf`` plus a sibling
|
||
``<id>.yaml`` (per-plot gallery metadata) and a per-family ``metadata.yaml``.
|
||
Optionally runs ``gallery generate`` to build the static HTML site.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import subprocess
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import plotstyle as ps
|
||
import yaml
|
||
|
||
from giant.analysis.reduced import Reduced
|
||
|
||
_SERIES_LABELS = {"rollout": "rollout", "reference": "reference (Geant4)"}
|
||
|
||
|
||
def _density(counts: list[int] | np.ndarray, edges: np.ndarray) -> np.ndarray:
|
||
counts = np.asarray(counts, dtype=np.float64)
|
||
total = counts.sum()
|
||
if total == 0:
|
||
return counts
|
||
return counts / (total * (edges[1] - edges[0]))
|
||
|
||
|
||
def _overlay(ax, edges: np.ndarray, series: dict[str, list], log_y: bool) -> None:
|
||
for key in ("reference", "rollout"):
|
||
if key in series:
|
||
ax.stairs(_density(series[key], edges), edges, label=_SERIES_LABELS[key])
|
||
if log_y:
|
||
ax.set_yscale("log")
|
||
|
||
|
||
def _router_summary(router_cfg: dict) -> str:
|
||
if not router_cfg.get("enabled"):
|
||
return "off"
|
||
return f"{router_cfg.get('type', '?')}×{router_cfg.get('n_experts', '?')}"
|
||
|
||
|
||
def _figure_params_v2(mc: dict, run_meta: dict) -> dict:
|
||
"""`_figure_params` for a new-shape (nested) `model_config` — has a
|
||
`stage1_model` key. Reports stage 1's architecture (the headline
|
||
generator); stage 2's generator is only added (`mode_s2`) when it
|
||
differs from stage 1's, since a mixed run (the `stage1=flow` +
|
||
`stage2=wgan` case) is the interesting exception, not
|
||
the common case."""
|
||
s1 = mc["stage1_model"]
|
||
s2 = mc.get("stage2_model") or {}
|
||
mode = s1.get("generator")
|
||
params: dict = {}
|
||
if s1.get("hidden_dim") is not None:
|
||
params["hidden_dim"] = s1["hidden_dim"]
|
||
if s1.get("n_res_blocks") is not None:
|
||
params["n_res_blocks"] = s1["n_res_blocks"]
|
||
if mode is not None:
|
||
params["mode"] = mode
|
||
s2_mode = s2.get("generator")
|
||
if s2_mode is not None and s2_mode != mode:
|
||
params["mode_s2"] = s2_mode
|
||
particle_type = ((mc.get("conditioning") or {}).get("particle") or {}).get("type")
|
||
if particle_type is not None:
|
||
params["conditioning"] = particle_type
|
||
params["router"] = _router_summary(s1.get("router") or {})
|
||
if run_meta.get("training_epoch") is not None:
|
||
params["epoch"] = run_meta["training_epoch"]
|
||
if run_meta.get("best_val_loss") is not None:
|
||
params["best_val_loss"] = round(run_meta["best_val_loss"], 4)
|
||
if mode == "wgan":
|
||
noise_dim = (s1.get("wgan") or {}).get("noise_dim")
|
||
if noise_dim is not None:
|
||
params["noise_dim"] = noise_dim
|
||
elif run_meta.get("steps") is not None:
|
||
params["steps"] = run_meta["steps"]
|
||
return params
|
||
|
||
|
||
def _figure_params(run_meta: dict) -> dict:
|
||
"""Curated run identity for the figure subtitle (``new_figure(params=...)``).
|
||
|
||
``run_meta``/each plot's own ``<id>.yaml`` (see ``_plot_metadata``) already
|
||
carry every threaded model/training/rollout/dataset parameter for
|
||
after-the-fact lookup — this picks only the handful that matter for
|
||
telling figures apart at a glance while flipping through a gallery, since
|
||
the subtitle is one unwrapped line of text. The last slot is
|
||
architecture-conditional: flow/ddpm runs show the ODE ``steps`` used for
|
||
this rollout, wgan runs show ``noise_dim`` instead since wgan sampling is
|
||
single-pass and has no ODE step count.
|
||
|
||
Handles both a v0.2 checkpoint's flat ``model_config`` and a v0.3.0
|
||
nested one (has a ``stage1_model`` key — see ``_figure_params_v2``).
|
||
"""
|
||
mc = run_meta.get("model_config") or {}
|
||
if "stage1_model" in mc:
|
||
return _figure_params_v2(mc, run_meta)
|
||
|
||
mode = mc.get("mode")
|
||
params: dict = {}
|
||
if mc.get("hidden_dim") is not None:
|
||
params["hidden_dim"] = mc["hidden_dim"]
|
||
if mc.get("n_blocks") is not None:
|
||
params["n_blocks"] = mc["n_blocks"]
|
||
if mode is not None:
|
||
params["mode"] = mode
|
||
if mc.get("conditioning") is not None:
|
||
params["conditioning"] = mc["conditioning"]
|
||
params["router"] = _router_summary(mc.get("router") or {})
|
||
if run_meta.get("training_epoch") is not None:
|
||
params["epoch"] = run_meta["training_epoch"]
|
||
if run_meta.get("best_val_loss") is not None:
|
||
params["best_val_loss"] = round(run_meta["best_val_loss"], 4)
|
||
if mode == "wgan":
|
||
if mc.get("noise_dim") is not None:
|
||
params["noise_dim"] = mc["noise_dim"]
|
||
elif run_meta.get("steps") is not None:
|
||
params["steps"] = run_meta["steps"]
|
||
return params
|
||
|
||
|
||
def _render_overlay(r: Reduced, params: dict):
|
||
edges = np.asarray(r.payload["edges"])
|
||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||
_overlay(ax, edges, r.payload, r.payload.get("log_y", False))
|
||
ax.set_xlabel(r.xlabel)
|
||
ax.set_ylabel("density")
|
||
ps.style_legend(ax, title="source")
|
||
return fig
|
||
|
||
|
||
def _render_single(r: Reduced, params: dict):
|
||
edges = np.asarray(r.payload["edges"])
|
||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||
ax.stairs(
|
||
_density(r.payload["rollout"], edges), edges, label=_SERIES_LABELS["rollout"]
|
||
)
|
||
if r.payload.get("log_y"):
|
||
ax.set_yscale("log")
|
||
if r.payload.get("log_x"):
|
||
ax.set_xscale("log")
|
||
ax.set_xlabel(r.xlabel)
|
||
ax.set_ylabel("density")
|
||
ps.style_legend(ax, title="source")
|
||
return fig
|
||
|
||
|
||
def _render_grouped(r: Reduced, params: dict):
|
||
edges = np.asarray(r.payload["edges"])
|
||
groups = r.payload["groups"]
|
||
labels = list(groups)
|
||
n = len(labels)
|
||
ncols = min(3, n) or 1
|
||
nrows = (n + ncols - 1) // ncols
|
||
fig, axes = ps.new_figure(
|
||
"slide-16x9",
|
||
title=r.title,
|
||
params=params,
|
||
nrows=nrows,
|
||
ncols=ncols,
|
||
squeeze=False,
|
||
)
|
||
flat = axes.ravel()
|
||
for i, lbl in enumerate(labels):
|
||
ax = flat[i]
|
||
_overlay(ax, edges, groups[lbl], r.payload.get("log_y", False))
|
||
ax.set_title(lbl, fontsize=8)
|
||
ax.set_xlabel(r.xlabel)
|
||
for j in range(n, len(flat)):
|
||
flat[j].set_visible(False)
|
||
ps.style_legend(flat[0], title="source")
|
||
return fig
|
||
|
||
|
||
def _render_profile(r: Reduced, params: dict):
|
||
edges = np.asarray(r.payload["edges"])
|
||
centers = 0.5 * (edges[:-1] + edges[1:])
|
||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||
for key in ("reference", "rollout"):
|
||
mean = np.asarray(r.payload[f"{key}_mean"])
|
||
std = np.asarray(r.payload[f"{key}_std"])
|
||
(line,) = ax.plot(centers, mean, label=_SERIES_LABELS[key])
|
||
ax.fill_between(
|
||
centers, mean - std, mean + std, alpha=0.2, color=line.get_color()
|
||
)
|
||
ax.set_xlabel(r.xlabel)
|
||
ax.set_ylabel(r.payload.get("ylabel", "mean deposited energy [MeV]"))
|
||
ps.style_legend(ax, title="source")
|
||
return fig
|
||
|
||
|
||
def _render_bar(r: Reduced, params: dict):
|
||
labels = r.payload["labels"]
|
||
x = np.arange(len(labels))
|
||
width = 0.4
|
||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||
ax.bar(
|
||
x - width / 2, r.payload["reference"], width, label=_SERIES_LABELS["reference"]
|
||
)
|
||
ax.bar(x + width / 2, r.payload["rollout"], width, label=_SERIES_LABELS["rollout"])
|
||
ax.set_xticks(x)
|
||
ax.set_xticklabels(labels, rotation=45, ha="right")
|
||
ax.set_ylabel(r.payload.get("ylabel", "value"))
|
||
ps.style_legend(ax, title="source")
|
||
return fig
|
||
|
||
|
||
def _render_router_gating(r: Reduced, params: dict):
|
||
n_experts = r.payload["n_experts"]
|
||
log_x = r.payload.get("log_x", False)
|
||
fig, axes = ps.new_figure(
|
||
"slide-16x9", title=r.title, params=params, nrows=1, ncols=2, squeeze=False
|
||
)
|
||
flat = axes.ravel()
|
||
for ax, key in zip(flat, ("rollout", "reference")):
|
||
side = r.payload.get(key, {})
|
||
centers = np.asarray(side.get("centers", []))
|
||
means = np.asarray(side.get("means", []))
|
||
if len(centers) and means.size:
|
||
cum = np.zeros(len(centers))
|
||
for i in range(n_experts):
|
||
ax.fill_between(
|
||
centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}"
|
||
)
|
||
cum = cum + means[:, i]
|
||
if log_x:
|
||
ax.set_xscale("log")
|
||
ax.set_ylim(0, 1)
|
||
ax.set_title(_SERIES_LABELS[key], fontsize=8)
|
||
ax.set_xlabel(r.xlabel)
|
||
flat[0].set_ylabel("mean gate weight")
|
||
ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
|
||
return fig
|
||
|
||
|
||
def _render_router_share(r: Reduced, params: dict):
|
||
categories = r.payload["categories"]
|
||
n_experts = r.payload["n_experts"]
|
||
x = np.arange(len(categories))
|
||
present = [k for k in ("rollout", "reference") if k in r.payload]
|
||
fig, axes = ps.new_figure(
|
||
"slide-16x9",
|
||
title=r.title,
|
||
params=params,
|
||
nrows=1,
|
||
ncols=len(present),
|
||
squeeze=False,
|
||
)
|
||
flat = axes.ravel()
|
||
for ax, key in zip(flat, present):
|
||
side = r.payload[key]
|
||
shares = np.array([side[c] for c in categories]) # (n_cat, n_experts)
|
||
bottom = np.zeros(len(categories))
|
||
for i in range(n_experts):
|
||
ax.bar(x, shares[:, i], bottom=bottom, label=f"expert {i}")
|
||
bottom += shares[:, i]
|
||
ax.set_xticks(x)
|
||
ax.set_xticklabels(categories, rotation=45, ha="right")
|
||
ax.set_ylim(0, 1)
|
||
ax.set_title(_SERIES_LABELS[key], fontsize=8)
|
||
flat[0].set_ylabel("share of rows dispatched to expert")
|
||
ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
|
||
return fig
|
||
|
||
|
||
def _render_unavailable(r: Reduced, params: dict):
|
||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||
ax.axis("off")
|
||
ax.text(
|
||
0.5,
|
||
0.5,
|
||
r.payload.get("note", "not available"),
|
||
ha="center",
|
||
va="center",
|
||
wrap=True,
|
||
fontsize=10,
|
||
transform=ax.transAxes,
|
||
)
|
||
return fig
|
||
|
||
|
||
_RENDERERS = {
|
||
"overlay_hist": _render_overlay,
|
||
"single_hist": _render_single,
|
||
"grouped_hist": _render_grouped,
|
||
"profile": _render_profile,
|
||
"bar": _render_bar,
|
||
"router_gating": _render_router_gating,
|
||
"router_share": _render_router_share,
|
||
"unavailable": _render_unavailable,
|
||
}
|
||
|
||
|
||
def render(r: Reduced, run_meta: dict | None = None):
|
||
"""Build the matplotlib figure for one reduced artifact (dispatch on kind)."""
|
||
return _RENDERERS[r.kind](r, _figure_params(run_meta or {}))
|
||
|
||
|
||
def _plot_metadata(r: Reduced, run_meta: dict) -> dict:
|
||
meta = {
|
||
"title": r.title,
|
||
"description": f"Rollout vs reference: {r.title}.",
|
||
"plot_type": r.kind,
|
||
"family": r.family,
|
||
}
|
||
meta.update(r.meta)
|
||
if "note" in r.payload:
|
||
meta["note"] = r.payload["note"]
|
||
if run_meta:
|
||
# Every threaded model/training/rollout/dataset parameter, so a
|
||
# single plot's metadata is self-contained for later comparison
|
||
# without cross-referencing the run's root metadata.yaml.
|
||
meta["parameters"] = {k: v for k, v in run_meta.items() if k != "title"}
|
||
return meta
|
||
|
||
|
||
def render_all(
|
||
reduced_dir: str | Path,
|
||
out_dir: str | Path,
|
||
run_meta: dict | None = None,
|
||
*,
|
||
run_gallery: bool = False,
|
||
) -> list[Path]:
|
||
"""Render every reduced artifact under ``reduced_dir`` to a PDF tree.
|
||
|
||
Writes ``<out>/<family>/<id>.pdf`` + ``<id>.yaml`` and a per-family
|
||
``metadata.yaml`` (carrying the run's checkpoint/paths as gallery params).
|
||
Returns the list of PDF paths written.
|
||
"""
|
||
ps.use()
|
||
run_meta = run_meta or {}
|
||
reduced_dir, out_dir = Path(reduced_dir), Path(out_dir)
|
||
pdfs: list[Path] = []
|
||
families: set[str] = set()
|
||
|
||
for jf in sorted(reduced_dir.glob("*.json")):
|
||
r = Reduced.load(jf)
|
||
family_dir = out_dir / r.family
|
||
family_dir.mkdir(parents=True, exist_ok=True)
|
||
families.add(r.family)
|
||
fig = render(r, run_meta)
|
||
ps.savefig(fig, str(family_dir / r.id), formats=("pdf",))
|
||
(family_dir / f"{r.id}.yaml").write_text(
|
||
yaml.safe_dump(_plot_metadata(r, run_meta), sort_keys=False)
|
||
)
|
||
pdfs.append(family_dir / f"{r.id}.pdf")
|
||
import matplotlib.pyplot as plt
|
||
|
||
plt.close(fig)
|
||
|
||
# Root + per-family gallery metadata.
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
(out_dir / "metadata.yaml").write_text(
|
||
yaml.safe_dump(
|
||
{
|
||
"title": run_meta.get("title", "GIANT rollout analysis"),
|
||
"description": "Autoregressive rollout compared against held-out Geant4 reference steps.",
|
||
"experiment": "GIANT",
|
||
"parameters": {k: v for k, v in run_meta.items() if k != "title"},
|
||
},
|
||
sort_keys=False,
|
||
)
|
||
)
|
||
for fam in families:
|
||
(out_dir / fam / "metadata.yaml").write_text(
|
||
yaml.safe_dump(
|
||
{"title": fam, "description": f"{fam} plots."}, sort_keys=False
|
||
)
|
||
)
|
||
|
||
if run_gallery:
|
||
subprocess.run(["gallery", "generate", "--source", str(out_dir)], check=True)
|
||
return pdfs
|
||
|
||
|
||
def render_run(run_dir: str | Path, *, run_gallery: bool = False) -> list[Path]:
|
||
"""Render a prepped run directory: ``<run_dir>/reduced`` → ``<run_dir>/plots``.
|
||
|
||
First joins every plot's chunk partials (``reduced_partial/<id>__*.json``)
|
||
into ``reduced/<id>.json`` via ``merge_all`` — a no-op merge when the run
|
||
wasn't chunked (``n_chunks=1``) — then pulls the rollout provenance
|
||
(checkpoint, paths, cutoffs) from ``run_meta.json`` into every plot's
|
||
gallery metadata and renders.
|
||
"""
|
||
from giant.analysis.condor import RunMeta, merge_all
|
||
|
||
run_dir = Path(run_dir)
|
||
merge_all(run_dir)
|
||
meta = RunMeta.load(run_dir / "run_meta.json")
|
||
run_meta = {
|
||
"title": meta.title,
|
||
"rollout": meta.rollout,
|
||
"reference": meta.reference,
|
||
**meta.plot_meta,
|
||
}
|
||
return render_all(
|
||
run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery
|
||
)
|