lars
2358a75ee1
CI / Sync project version with tag (pull_request) Has been skipped
CI / Format (ruff format) (pull_request) Successful in 43s
CI / Lint (ruff check) (pull_request) Successful in 54s
CI / Type check (ty) (pull_request) Successful in 53s
CI / Tests (pull_request) Successful in 2m33s
CI / Bump version, tag, and update changelog on merge to master (pull_request) Has been skipped
CI / Publish package to Gitea package registry (pull_request) Has been skipped
Closes the roadmap's long-standing "no eval-latency number exists for any configuration" gap. Instruments `giant rollout` to record per-physical-step wall-clock cost in its YAML sidecar, adds a measured Geant4/miniCaloSim per-step reference (giant/analysis/geant4_reference.py, from a 3-energy, 4-event-count-per-energy local benchmark), and wires both into a new eval_cost_per_step PlotSpec in the giant analyze gallery. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
559 lines
21 KiB
Python
559 lines
21 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.
|
||
|
||
Every rollout series gets a stable color via ``ps.get_color(i)``, ``i`` being
|
||
its position in ``payload["series"]`` — that position is fixed by the run's
|
||
YAML/``--label`` order (threaded unchanged from ``condor.RunMeta.rollouts``
|
||
through every ``PlotSpec``), so a given rollout keeps the same color across
|
||
every plot in a run. The reference, where a plot has one, always draws in one
|
||
fixed, distinct style (dark ink, dashed) instead of taking a slot in that
|
||
cycle.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import dataclasses
|
||
import subprocess
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import plotstyle as ps
|
||
from matplotlib.colors import LogNorm
|
||
import yaml
|
||
|
||
from giant.analysis.reduced import Reduced
|
||
|
||
_REFERENCE_LABEL = "reference (Geant4)"
|
||
|
||
_TEX_ESCAPE_MAP = {
|
||
"\\": r"\textbackslash{}",
|
||
"%": r"\%",
|
||
"&": r"\&",
|
||
"#": r"\#",
|
||
"$": r"\$",
|
||
"_": r"\_",
|
||
"{": r"\{",
|
||
"}": r"\}",
|
||
}
|
||
|
||
|
||
def _tex_escape(text: str) -> str:
|
||
"""Escape characters LaTeX treats specially in catalog-authored title/xlabel
|
||
text (e.g. a literal ``%`` in a "90% of deposited energy" title, which
|
||
``usetex`` otherwise reads as a comment marker and aborts the whole figure —
|
||
see gitea #81). A single pass over the *original* characters, so the
|
||
backslashes an escape itself introduces (e.g. ``\textbackslash{}``) are
|
||
never re-escaped."""
|
||
return "".join(_TEX_ESCAPE_MAP.get(c, c) for c in text)
|
||
|
||
|
||
def _ref_color() -> str:
|
||
return ps.colors.INK["primary"]
|
||
|
||
|
||
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, payload: dict, log_y: bool) -> None:
|
||
if "reference" in payload:
|
||
ax.stairs(
|
||
_density(payload["reference"], edges), edges, label=_REFERENCE_LABEL, color=_ref_color(), linestyle="--"
|
||
)
|
||
for i, (name, counts) in enumerate(payload.get("series", {}).items()):
|
||
ax.stairs(_density(counts, edges), edges, label=name, color=ps.get_color(i))
|
||
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, meta: dict) -> dict:
|
||
"""`_figure_params_single` 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 meta.get("training_epoch") is not None:
|
||
params["epoch"] = meta["training_epoch"]
|
||
if meta.get("best_val_loss") is not None:
|
||
params["best_val_loss"] = round(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 meta.get("steps") is not None:
|
||
params["steps"] = meta["steps"]
|
||
return params
|
||
|
||
|
||
def _figure_params_single(meta: dict) -> dict:
|
||
"""Curated run identity for the figure subtitle (``new_figure(params=...)``),
|
||
for exactly one rollout's ``plot_meta``.
|
||
|
||
``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 = meta.get("model_config") or {}
|
||
if "stage1_model" in mc:
|
||
return _figure_params_v2(mc, 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 meta.get("training_epoch") is not None:
|
||
params["epoch"] = meta["training_epoch"]
|
||
if meta.get("best_val_loss") is not None:
|
||
params["best_val_loss"] = round(meta["best_val_loss"], 4)
|
||
if mode == "wgan":
|
||
if mc.get("noise_dim") is not None:
|
||
params["noise_dim"] = mc["noise_dim"]
|
||
elif meta.get("steps") is not None:
|
||
params["steps"] = meta["steps"]
|
||
return params
|
||
|
||
|
||
def _figure_params(run_meta: dict) -> dict:
|
||
"""Curated run identity for the figure subtitle.
|
||
|
||
A single-rollout run reuses that rollout's ``plot_meta`` (same curated
|
||
model/training/rollout subset as always — see ``_figure_params_single``);
|
||
a multi-rollout run instead names the series being compared, since no
|
||
single ``model_config`` applies to the figure as a whole (each plot's own
|
||
gallery YAML still carries every rollout's full ``plot_meta`` for
|
||
after-the-fact lookup, via ``_plot_metadata``).
|
||
"""
|
||
rollouts = run_meta.get("rollouts") or {}
|
||
if len(rollouts) == 1:
|
||
((_, meta),) = rollouts.items()
|
||
return _figure_params_single(meta)
|
||
return {"rollouts": ", ".join(rollouts)} if rollouts else {}
|
||
|
||
|
||
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)
|
||
for i, (name, counts) in enumerate(r.payload.get("series", {}).items()):
|
||
ax.stairs(_density(counts, edges), edges, label=name, color=ps.get_color(i))
|
||
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)
|
||
if "reference" in r.payload:
|
||
ref = r.payload["reference"]
|
||
mean, std = np.asarray(ref["mean"]), np.asarray(ref["std"])
|
||
color = _ref_color()
|
||
ax.plot(centers, mean, label=_REFERENCE_LABEL, color=color, linestyle="--")
|
||
ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=color)
|
||
for i, (name, side) in enumerate(r.payload.get("series", {}).items()):
|
||
mean, std = np.asarray(side["mean"]), np.asarray(side["std"])
|
||
color = ps.get_color(i)
|
||
ax.plot(centers, mean, label=name, color=color)
|
||
ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=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))
|
||
series = r.payload.get("series", {})
|
||
has_ref = "reference" in r.payload
|
||
n_bars = len(series) + (1 if has_ref else 0)
|
||
width = 0.8 / max(n_bars, 1)
|
||
offsets = np.linspace(-0.4 + width / 2, 0.4 - width / 2, n_bars)
|
||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||
idx = 0
|
||
if has_ref:
|
||
ax.bar(x + offsets[idx], r.payload["reference"], width, label=_REFERENCE_LABEL, color=_ref_color())
|
||
idx += 1
|
||
for i, (name, vals) in enumerate(series.items()):
|
||
ax.bar(x + offsets[idx], vals, width, label=name, color=ps.get_color(i))
|
||
idx += 1
|
||
ax.set_xticks(x)
|
||
ax.set_xticklabels(labels, rotation=45, ha="right")
|
||
ax.set_ylabel(r.payload.get("ylabel", "value"))
|
||
if r.payload.get("log_y"):
|
||
ax.set_yscale("log")
|
||
ps.style_legend(ax, title="source")
|
||
return fig
|
||
|
||
|
||
def _render_router_gating(r: Reduced, params: dict):
|
||
series = r.payload.get("series", {})
|
||
names = list(series)
|
||
log_x = r.payload.get("log_x", False)
|
||
fig, axes = ps.new_figure("slide-16x9", title=r.title, params=params, nrows=len(names), ncols=2, squeeze=False)
|
||
for row, name in enumerate(names):
|
||
entry = series[name]
|
||
n_experts = entry["n_experts"]
|
||
for col, key in enumerate(("rollout", "reference")):
|
||
ax = axes[row, col]
|
||
side = entry.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)
|
||
panel_label = _REFERENCE_LABEL if key == "reference" else "rollout"
|
||
ax.set_title(f"{name} — {panel_label}", fontsize=8)
|
||
if row == len(names) - 1:
|
||
ax.set_xlabel(r.xlabel)
|
||
axes[row, 0].set_ylabel("mean gate weight")
|
||
if names:
|
||
ps.style_legend(axes[0, 0], title=f"{series[names[0]]['router_type']} router")
|
||
return fig
|
||
|
||
|
||
def _render_router_share(r: Reduced, params: dict):
|
||
series = r.payload.get("series", {})
|
||
names = list(series)
|
||
present: tuple[str, ...] = ("rollout", "reference")
|
||
if names:
|
||
present = tuple(k for k in ("rollout", "reference") if k in series[names[0]])
|
||
ncols = max(len(present), 1)
|
||
fig, axes = ps.new_figure("slide-16x9", title=r.title, params=params, nrows=len(names), ncols=ncols, squeeze=False)
|
||
for row, name in enumerate(names):
|
||
entry = series[name]
|
||
n_experts = entry["n_experts"]
|
||
cats = entry["categories"]
|
||
x = np.arange(len(cats))
|
||
for col, key in enumerate(present):
|
||
ax = axes[row, col]
|
||
side = entry.get(key)
|
||
if side is not None:
|
||
shares = np.array([side[c] for c in cats]) # (n_cat, n_experts)
|
||
bottom = np.zeros(len(cats))
|
||
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(cats, rotation=45, ha="right")
|
||
ax.set_ylim(0, 1)
|
||
panel_label = _REFERENCE_LABEL if key == "reference" else "rollout"
|
||
ax.set_title(f"{name} — {panel_label}", fontsize=8)
|
||
axes[row, 0].set_ylabel("share of rows dispatched to expert")
|
||
if names:
|
||
ps.style_legend(axes[0, 0], title=f"{series[names[0]]['router_type']} router")
|
||
return fig
|
||
|
||
|
||
def _render_router_specialization(r: Reduced, params: dict):
|
||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||
series = r.payload.get("series", {})
|
||
chance_levels: set[float] = set()
|
||
for i, (name, entry) in enumerate(series.items()):
|
||
color = ps.get_color(i)
|
||
if entry.get("chance_level") is not None:
|
||
chance_levels.add(entry["chance_level"])
|
||
for key, linestyle, label in (
|
||
("rollout", "-", name),
|
||
("reference", "--", f"{name} ({_REFERENCE_LABEL})"),
|
||
):
|
||
side = entry.get(key)
|
||
if side and side["centers"]:
|
||
ax.plot(
|
||
side["centers"],
|
||
side["score"],
|
||
label=label,
|
||
color=color,
|
||
linestyle=linestyle,
|
||
marker="o",
|
||
markersize=3,
|
||
)
|
||
for lvl in sorted(chance_levels):
|
||
ax.axhline(lvl, linestyle=":", color="gray")
|
||
if r.payload.get("log_x"):
|
||
ax.set_xscale("log")
|
||
ax.set_ylim(0, 1)
|
||
ax.set_xlabel(r.xlabel)
|
||
ax.set_ylabel("max gate weight")
|
||
ps.style_legend(ax, title="router")
|
||
return fig
|
||
|
||
|
||
def _render_heatmap(r: Reduced, params: dict):
|
||
series = dict(r.payload["series"])
|
||
row_labels = r.payload["row_labels"]
|
||
col_labels = r.payload["col_labels"]
|
||
# A heatmap-shaped plot is one matrix per rollout, so the reference (when the
|
||
# comparison has one — the distance scorecard doesn't) becomes one more panel
|
||
# rather than another line.
|
||
if r.payload.get("reference") is not None:
|
||
series["reference"] = r.payload["reference"]
|
||
names = list(series)
|
||
norm = LogNorm(vmin=1) if r.payload.get("log_color") else None
|
||
fig, axes = ps.new_figure(
|
||
"slide-16x9" if len(names) > 1 else "thesis-single",
|
||
title=r.title,
|
||
params=params,
|
||
nrows=1,
|
||
ncols=len(names),
|
||
squeeze=False,
|
||
)
|
||
flat = axes.ravel()
|
||
im = None
|
||
for ax, name in zip(flat, names):
|
||
mat = np.asarray(series[name], dtype=float)
|
||
im = ax.imshow(
|
||
mat,
|
||
origin="upper",
|
||
aspect="auto",
|
||
cmap=r.payload.get("cmap", "viridis"),
|
||
norm=norm,
|
||
vmin=None if norm else r.payload.get("vmin"),
|
||
vmax=None if norm else r.payload.get("vmax"),
|
||
)
|
||
ax.set_xticks(range(len(col_labels)))
|
||
ax.set_xticklabels(col_labels, rotation=45, ha="right")
|
||
ax.set_yticks(range(len(row_labels)))
|
||
ax.set_yticklabels(row_labels)
|
||
ax.set_xlabel(r.xlabel)
|
||
if len(names) > 1:
|
||
ax.set_title(name, fontsize=8)
|
||
flat[0].set_ylabel(r.payload.get("ylabel", ""))
|
||
fig.colorbar(im, ax=list(flat), label=r.payload.get("cbar_label", "value"))
|
||
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,
|
||
"router_specialization": _render_router_specialization,
|
||
"heatmap": _render_heatmap,
|
||
"unavailable": _render_unavailable,
|
||
}
|
||
|
||
|
||
def render(r: Reduced, run_meta: dict | None = None):
|
||
"""Build the matplotlib figure for one reduced artifact (dispatch on kind).
|
||
|
||
``title``/``xlabel`` are LaTeX-escaped here, at the one point every kind's
|
||
renderer draws them from — ``_plot_metadata`` deliberately keeps using the
|
||
unescaped ``r`` for the gallery YAML, which isn't LaTeX.
|
||
"""
|
||
escaped = dataclasses.replace(r, title=_tex_escape(r.title), xlabel=_tex_escape(r.xlabel))
|
||
return _RENDERERS[r.kind](escaped, _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(s) compared against a held-out Geant4 reference steps file.",
|
||
"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,
|
||
"reference": meta.reference,
|
||
"rollouts": {ro["name"]: ro["plot_meta"] for ro in meta.rollouts},
|
||
}
|
||
return render_all(run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery)
|