Compare commits
10 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| b8f8965338 | |||
| 81d22c1964 | |||
| 417b741484 | |||
| 37d73e6578 | |||
| ff204732d7 | |||
| 02ed4e531c | |||
| 1b6c8b33b7 | |||
| 7560e2bff0 | |||
| ffb7c0cc2a | |||
| bdebd83c8b |
+1
-1
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.3.6"
|
||||
current_version = "0.3.8"
|
||||
parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
|
||||
serialize = ["{major}.{minor}.{patch}"]
|
||||
search = "{current_version}"
|
||||
|
||||
@@ -1,5 +1,22 @@
|
||||
# Changelog
|
||||
|
||||
## [0.3.8] - 2026-08-24
|
||||
|
||||
### Added
|
||||
|
||||
- Add giant analyze metrics plots for training progress [gitea #75](https://git.larsbogner.de/lars/giant/issues/75)
|
||||
|
||||
|
||||
### Fixed
|
||||
|
||||
- Fix LaTeX-unavailable skip check in analyze metrics smoke test
|
||||
|
||||
## [0.3.7] - 2026-08-24
|
||||
|
||||
### Added
|
||||
|
||||
- Add rollout-quality distance, confusion, containment and router plots [gitea #76](https://git.larsbogner.de/lars/giant/issues/76)
|
||||
|
||||
## [0.3.6] - 2026-08-24
|
||||
|
||||
### Changed
|
||||
|
||||
@@ -47,6 +47,7 @@ from giant.analysis.reduce import (
|
||||
leakage_fraction,
|
||||
profile_finalize,
|
||||
profile_partial,
|
||||
sec_count_by_event,
|
||||
species_share,
|
||||
sum_merge,
|
||||
transverse_expr,
|
||||
@@ -56,6 +57,7 @@ from giant.analysis.router_gating import (
|
||||
compute_router_gating,
|
||||
compute_router_share_by_pdg,
|
||||
compute_router_share_by_process,
|
||||
compute_router_specialization,
|
||||
)
|
||||
from giant.analysis.sources import Side, open_side, physical_steps, secondaries
|
||||
from giant.analysis.type_embedding_distance import compute_type_embedding_l1_distance
|
||||
@@ -176,6 +178,68 @@ def _np_hist_pair(r: np.ndarray, t: np.ndarray, nbins: int) -> tuple[np.ndarray,
|
||||
return edges, np.histogram(r, edges)[0], np.histogram(t, edges)[0]
|
||||
|
||||
|
||||
def _ks_statistic(r_counts, t_counts) -> float:
|
||||
"""KS statistic (max |CDF diff|) between two same-edge binned histograms.
|
||||
|
||||
``nan`` when neither side has any mass (nothing to compare); 1.0 (maximal
|
||||
mismatch) when exactly one side is entirely empty and the other isn't —
|
||||
correctly the worst score rather than an undefined one.
|
||||
"""
|
||||
r_counts = np.asarray(r_counts, dtype=np.float64)
|
||||
t_counts = np.asarray(t_counts, dtype=np.float64)
|
||||
r_tot, t_tot = r_counts.sum(), t_counts.sum()
|
||||
if r_tot == 0 and t_tot == 0:
|
||||
return float("nan")
|
||||
if r_tot == 0 or t_tot == 0:
|
||||
return 1.0
|
||||
r_cdf = np.cumsum(r_counts) / r_tot
|
||||
t_cdf = np.cumsum(t_counts) / t_tot
|
||||
return float(np.max(np.abs(r_cdf - t_cdf)))
|
||||
|
||||
|
||||
def _integer_confusion(t: np.ndarray, r: np.ndarray, max_bins: int = 21) -> tuple[list[str], np.ndarray]:
|
||||
"""Confusion matrix of two paired small-integer arrays (e.g. secondary counts).
|
||||
|
||||
Bins are consecutive integers ``0..cap``, with the last bin an overflow
|
||||
``"cap+"`` bucket, so an occasional pathological count doesn't blow up the
|
||||
heatmap. Returns ``(labels, matrix)`` with ``matrix[i, j]`` counting pairs
|
||||
with ``t == i`` and ``r == j`` (both clipped into ``[0, cap]``).
|
||||
"""
|
||||
cap = min(max(int(t.max()) if len(t) else 0, int(r.max()) if len(r) else 0, 1), max_bins - 1)
|
||||
t_c = np.clip(t.astype(np.int64), 0, cap)
|
||||
r_c = np.clip(r.astype(np.int64), 0, cap)
|
||||
n = cap + 1
|
||||
mat = np.zeros((n, n), dtype=np.int64)
|
||||
np.add.at(mat, (t_c, r_c), 1)
|
||||
labels = [str(i) for i in range(cap)] + [f"{cap}+"]
|
||||
return labels, mat
|
||||
|
||||
|
||||
def _containment_depths(mat: np.ndarray, edges: np.ndarray, quantile: float) -> np.ndarray:
|
||||
"""Per-event depth containing ``quantile`` of that event's deposited energy.
|
||||
|
||||
``mat`` is a ``(n_events, n_bins)`` edep-per-depth-bin sum matrix (see
|
||||
``reduce.profile_partial``); bins are ordered by increasing depth (matching
|
||||
``edges``, monotonic). Zero-energy events are dropped — containment depth is
|
||||
undefined for them.
|
||||
"""
|
||||
totals = mat.sum(axis=1)
|
||||
valid = totals > 0
|
||||
mat, totals = mat[valid], totals[valid]
|
||||
cum = np.cumsum(mat, axis=1) / totals[:, None]
|
||||
idx = (cum >= quantile).argmax(axis=1) # first bin whose cumulative fraction reaches quantile
|
||||
return edges[1:][idx]
|
||||
|
||||
|
||||
def _group_keys(ctx: Context, axis: str) -> list:
|
||||
"""The group keys ``_marginal_grouped_finalize`` iterates for ``axis``."""
|
||||
if axis == "pdg":
|
||||
return list(ctx.top_pdgs)
|
||||
if axis == "material":
|
||||
return list(ctx.materials)
|
||||
return list(range(len(ctx.energy_edges) - 1)) # energy
|
||||
|
||||
|
||||
# Human-readable figure titles per marginal variable (the axis labels carry units;
|
||||
# these read cleanly as a title without them).
|
||||
_TITLE_NAMES = {
|
||||
@@ -299,6 +363,64 @@ def _marginal_grouped_finalize(parts: list[dict], ctx: Context, var: str, axis:
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# distance summary: a var x group-axis scorecard, reusing the marginal hists
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _distance_summary_partial(b: Bundle) -> dict:
|
||||
out: dict[str, dict] = {}
|
||||
for var in MARGINAL_VARS:
|
||||
out[var] = {"overall": _marginal_overall_partial(b, var)}
|
||||
for axis in GROUPING_AXES:
|
||||
out[var][axis] = _marginal_grouped_partial(b, var, axis)
|
||||
return out
|
||||
|
||||
|
||||
def _distance_summary_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
col_labels = ["overall", *GROUPING_AXES]
|
||||
matrix: list[list[float]] = []
|
||||
for var in MARGINAL_VARS:
|
||||
edges = _marginal_edges(ctx, var)
|
||||
nb = len(edges) - 1
|
||||
row: list[float] = []
|
||||
|
||||
r = sum_merge([p[var]["overall"]["r"] for p in parts])
|
||||
t = sum_merge([p[var]["overall"]["t"] for p in parts])
|
||||
row.append(_ks_statistic(_finalize_counts(r, 0, nb), _finalize_counts(t, 0, nb)))
|
||||
|
||||
for axis in GROUPING_AXES:
|
||||
r = sum_merge([p[var][axis]["r"] for p in parts])
|
||||
t = sum_merge([p[var][axis]["t"] for p in parts])
|
||||
dists, weights = [], []
|
||||
for k in _group_keys(ctx, axis):
|
||||
rc, tc = _finalize_counts(r, k, nb), _finalize_counts(t, k, nb)
|
||||
w = sum(rc) + sum(tc)
|
||||
if w == 0:
|
||||
continue
|
||||
dists.append(_ks_statistic(rc, tc))
|
||||
weights.append(w)
|
||||
row.append(float(np.average(dists, weights=weights)) if dists else float("nan"))
|
||||
matrix.append(row)
|
||||
|
||||
return Reduced(
|
||||
id="marginal_distance_summary",
|
||||
family="quality",
|
||||
kind="heatmap",
|
||||
title="Marginal distance summary (KS statistic, rollout vs reference)",
|
||||
xlabel="grouping axis",
|
||||
payload={
|
||||
"matrix": matrix,
|
||||
"row_labels": [_TITLE_NAMES[v] for v in MARGINAL_VARS],
|
||||
"col_labels": col_labels,
|
||||
"ylabel": "marginal variable",
|
||||
"cbar_label": "KS statistic (0 = identical, 1 = maximal mismatch)",
|
||||
"vmin": 0.0,
|
||||
"vmax": 1.0,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# per-event scalar observables
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -438,6 +560,41 @@ def _profile_finalize(
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# shower containment depth (reuses the longitudinal profile's per-event matrix)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_CONTAINMENT_QUANTILES: list[tuple[float, str]] = [
|
||||
(0.90, "shower_containment_depth_90"),
|
||||
(0.95, "shower_containment_depth_95"),
|
||||
]
|
||||
|
||||
|
||||
def _containment_finalize(parts: list[dict], ctx: Context, spec_id: str, quantile: float) -> Reduced:
|
||||
edges = np.asarray(ctx.depth_edges)
|
||||
nb = len(edges) - 1
|
||||
_assert_event_disjoint([p["r_ids"] for p in parts], spec_id, "rollout")
|
||||
_assert_event_disjoint([p["t_ids"] for p in parts], spec_id, "reference")
|
||||
r_full = np.concatenate([np.asarray(p["r_mat"], dtype=float).reshape(-1, nb) for p in parts], axis=0)
|
||||
t_full = np.concatenate([np.asarray(p["t_mat"], dtype=float).reshape(-1, nb) for p in parts], axis=0)
|
||||
r_depth = _containment_depths(r_full, edges, quantile)
|
||||
t_depth = _containment_depths(t_full, edges, quantile)
|
||||
hedges, rc, tc = _np_hist_pair(r_depth, t_depth, ctx.n_marginal_bins)
|
||||
return Reduced(
|
||||
id=spec_id,
|
||||
family="shower",
|
||||
kind="overlay_hist",
|
||||
title=f"Shower containment depth ({quantile:.0%} of deposited energy)",
|
||||
xlabel=f"depth containing {quantile:.0%} of deposited energy [mm]",
|
||||
payload={
|
||||
"edges": hedges.tolist(),
|
||||
_ROLL: rc.astype(np.int64).tolist(),
|
||||
_REF: tc.astype(np.int64).tolist(),
|
||||
"log_y": False,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# species share + leakage
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -627,6 +784,43 @@ def _sec_cos_angle_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
)
|
||||
|
||||
|
||||
def _n_sec_confusion_partial(b: Bundle) -> dict:
|
||||
r_sec, t_sec = _sec_frames(b)
|
||||
r_ids, r_n = sec_count_by_event(b.r_phys, r_sec)
|
||||
t_ids, t_n = sec_count_by_event(b.t_all, t_sec)
|
||||
return {"r_ids": r_ids.tolist(), "r_n": r_n.tolist(), "t_ids": t_ids.tolist(), "t_n": t_n.tolist()}
|
||||
|
||||
|
||||
def _n_sec_confusion_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
r_ids = np.concatenate([np.asarray(p["r_ids"], dtype=np.int64) for p in parts])
|
||||
r_n = np.concatenate([np.asarray(p["r_n"], dtype=np.int64) for p in parts])
|
||||
t_ids = np.concatenate([np.asarray(p["t_ids"], dtype=np.int64) for p in parts])
|
||||
t_n = np.concatenate([np.asarray(p["t_n"], dtype=np.int64) for p in parts])
|
||||
# event-disjoint chunking (see Bundle.open) means each event_id appears in
|
||||
# exactly one part on each side, so a plain dict build is a safe merge.
|
||||
r_map = dict(zip(r_ids.tolist(), r_n.tolist()))
|
||||
t_map = dict(zip(t_ids.tolist(), t_n.tolist()))
|
||||
common = sorted(set(r_map) & set(t_map))
|
||||
true_n = np.array([t_map[e] for e in common], dtype=np.int64)
|
||||
pred_n = np.array([r_map[e] for e in common], dtype=np.int64)
|
||||
labels, mat = _integer_confusion(true_n, pred_n)
|
||||
return Reduced(
|
||||
id="n_sec_confusion",
|
||||
family="secondaries",
|
||||
kind="heatmap",
|
||||
title="Predicted vs true secondary count per event",
|
||||
xlabel="predicted secondaries (rollout)",
|
||||
payload={
|
||||
"matrix": mat.tolist(),
|
||||
"row_labels": labels,
|
||||
"col_labels": labels,
|
||||
"ylabel": "true secondaries (reference)",
|
||||
"cbar_label": "event count",
|
||||
"vmin": 0.0,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# router diagnostics (not chunked — already bounded/subsampled)
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -640,6 +834,9 @@ _router_share_pdg_partial, _router_share_pdg_finalize = _unchunkable(
|
||||
_router_share_process_partial, _router_share_process_finalize = _unchunkable(
|
||||
lambda b: compute_router_share_by_process(b.checkpoint, b.t_phys)
|
||||
)
|
||||
_router_specialization_partial, _router_specialization_finalize = _unchunkable(
|
||||
lambda b: compute_router_specialization(b.checkpoint, b.r_phys, b.t_phys)
|
||||
)
|
||||
_type_embedding_l1_distance_partial, _type_embedding_l1_distance_finalize = _unchunkable(
|
||||
lambda b: compute_type_embedding_l1_distance(b.type_embedding_l1_dist)
|
||||
)
|
||||
@@ -676,6 +873,15 @@ def build_catalog() -> list[PlotSpec]:
|
||||
)
|
||||
)
|
||||
|
||||
specs.append(
|
||||
PlotSpec(
|
||||
"marginal_distance_summary",
|
||||
"quality",
|
||||
compute_partial=_distance_summary_partial,
|
||||
finalize=_distance_summary_finalize,
|
||||
)
|
||||
)
|
||||
|
||||
specs += [
|
||||
PlotSpec(
|
||||
"event_total_edep",
|
||||
@@ -745,6 +951,17 @@ def build_catalog() -> list[PlotSpec]:
|
||||
"transverse_edges",
|
||||
),
|
||||
),
|
||||
]
|
||||
for quantile, spec_id in _CONTAINMENT_QUANTILES:
|
||||
specs.append(
|
||||
PlotSpec(
|
||||
spec_id,
|
||||
"shower",
|
||||
compute_partial=lambda b: _profile_partial(b, depth_expr, "depth_edges"),
|
||||
finalize=lambda parts, ctx, q=quantile, sid=spec_id: _containment_finalize(parts, ctx, sid, q),
|
||||
)
|
||||
)
|
||||
specs += [
|
||||
PlotSpec(
|
||||
"species_edep_share",
|
||||
"species",
|
||||
@@ -781,6 +998,12 @@ def build_catalog() -> list[PlotSpec]:
|
||||
compute_partial=_sec_cos_angle_partial,
|
||||
finalize=_sec_cos_angle_finalize,
|
||||
),
|
||||
PlotSpec(
|
||||
"n_sec_confusion",
|
||||
"secondaries",
|
||||
compute_partial=_n_sec_confusion_partial,
|
||||
finalize=_n_sec_confusion_finalize,
|
||||
),
|
||||
PlotSpec(
|
||||
"router_gating",
|
||||
"model",
|
||||
@@ -802,6 +1025,13 @@ def build_catalog() -> list[PlotSpec]:
|
||||
finalize=_router_share_process_finalize,
|
||||
chunkable=False,
|
||||
),
|
||||
PlotSpec(
|
||||
"router_specialization",
|
||||
"model",
|
||||
compute_partial=_router_specialization_partial,
|
||||
finalize=_router_specialization_finalize,
|
||||
chunkable=False,
|
||||
),
|
||||
PlotSpec(
|
||||
"type_embedding_l1_distance",
|
||||
"model",
|
||||
|
||||
@@ -271,3 +271,20 @@ def leakage_fraction(lf: pl.LazyFrame) -> np.ndarray:
|
||||
escaped = per_event["escaped"].fill_null(0.0).to_numpy()
|
||||
total = deposited + escaped
|
||||
return np.where(total > 0, escaped / total, 0.0)
|
||||
|
||||
|
||||
def sec_count_by_event(lf_all: pl.LazyFrame, sec_lf: pl.LazyFrame) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Per-event secondary count, zero-filled for events that produced none.
|
||||
|
||||
Two bounded per-event ``group_by``s — the full event set (from ``lf_all``)
|
||||
and the secondary counts (from ``sec_lf``, see ``sources.secondaries``) —
|
||||
merged in Python via a dict. Both results are event-granularity (not
|
||||
per-row), so this stays in the same bounded-memory budget as
|
||||
``event_scalars``; a plain ``group_by`` on ``sec_lf`` alone would silently
|
||||
drop zero-secondary events instead of zero-filling them.
|
||||
"""
|
||||
ev = lf_all.select("event_id").unique().collect(engine="streaming")["event_id"].to_numpy()
|
||||
cnt_df = sec_lf.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")
|
||||
cnt = dict(zip(cnt_df["event_id"].to_list(), cnt_df["n"].to_list()))
|
||||
counts = np.array([cnt.get(int(e), 0) for e in ev], dtype=np.int64)
|
||||
return ev, counts
|
||||
|
||||
@@ -19,6 +19,10 @@ from pathlib import Path
|
||||
# "single_hist" one series only (e.g. rollout leakage; reference has none)
|
||||
# "router_gating" stacked mean MoE gate weight vs energy, rollout + reference
|
||||
# "router_share" stacked bar of MoE top-1 dispatch share by category
|
||||
# "router_specialization" max gate weight vs energy, rollout + reference (one
|
||||
# scalar trend line summarizing "router_gating")
|
||||
# "heatmap" row x col matrix + colorbar (distance scorecard or a
|
||||
# predicted-vs-true confusion matrix)
|
||||
# "unavailable" plot not applicable to this run (e.g. non-MoE checkpoint)
|
||||
|
||||
|
||||
|
||||
@@ -260,6 +260,47 @@ def _render_router_share(r: Reduced, params: dict):
|
||||
return fig
|
||||
|
||||
|
||||
def _render_router_specialization(r: Reduced, params: dict):
|
||||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||||
for key in ("reference", "rollout"):
|
||||
side = r.payload.get(key)
|
||||
if side and side["centers"]:
|
||||
ax.plot(side["centers"], side["score"], label=_SERIES_LABELS[key], marker="o", markersize=3)
|
||||
chance = r.payload.get("chance_level")
|
||||
if chance is not None:
|
||||
ax.axhline(chance, linestyle="--", color="gray", label="chance level (1/n_experts)")
|
||||
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=f"{r.payload.get('router_type', '')} router")
|
||||
return fig
|
||||
|
||||
|
||||
def _render_heatmap(r: Reduced, params: dict):
|
||||
mat = np.asarray(r.payload["matrix"], dtype=float)
|
||||
row_labels = r.payload["row_labels"]
|
||||
col_labels = r.payload["col_labels"]
|
||||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||||
im = ax.imshow(
|
||||
mat,
|
||||
origin="upper",
|
||||
aspect="auto",
|
||||
cmap=r.payload.get("cmap", "viridis"),
|
||||
vmin=r.payload.get("vmin"),
|
||||
vmax=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)
|
||||
ax.set_ylabel(r.payload.get("ylabel", ""))
|
||||
fig.colorbar(im, ax=ax, 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")
|
||||
@@ -284,6 +325,8 @@ _RENDERERS = {
|
||||
"bar": _render_bar,
|
||||
"router_gating": _render_router_gating,
|
||||
"router_share": _render_router_share,
|
||||
"router_specialization": _render_router_specialization,
|
||||
"heatmap": _render_heatmap,
|
||||
"unavailable": _render_unavailable,
|
||||
}
|
||||
|
||||
|
||||
@@ -203,6 +203,7 @@ _TITLES = {
|
||||
"router_gating": "Router gating (mixture-of-experts decision boundaries)",
|
||||
"router_share_by_pdg": "Router expert share by particle species",
|
||||
"router_share_by_process": "Router expert share by physics process",
|
||||
"router_specialization": "Router specialization score vs energy (max gate weight)",
|
||||
}
|
||||
|
||||
|
||||
@@ -250,6 +251,54 @@ def compute_router_gating(
|
||||
)
|
||||
|
||||
|
||||
def compute_router_specialization(
|
||||
checkpoint: str | Path | None,
|
||||
r_phys: pl.LazyFrame,
|
||||
t_phys: pl.LazyFrame,
|
||||
seed: int = 0,
|
||||
) -> Reduced:
|
||||
"""Scalar specialization trend: max gate weight vs energy, per side.
|
||||
|
||||
Summarizes `router_gating`'s full per-expert stacked area into one curve —
|
||||
the routing plan's own "how sharp is the boundary here" number (1/n_experts
|
||||
= uniform/no specialization, 1.0 = one expert fully owns that energy). Same
|
||||
quantile energy bins as `router_gating` (`_quantile_bins`), so this is
|
||||
directly comparable to that plot's ceiling described in the roadmap's MoE
|
||||
writeup.
|
||||
"""
|
||||
handle = load_router(checkpoint) if checkpoint else None
|
||||
if handle is None:
|
||||
return _unavailable("router_specialization")
|
||||
|
||||
sides: dict[str, dict] = {}
|
||||
for name, lf in (("rollout", r_phys), ("reference", t_phys)):
|
||||
df = _subsample(lf, _SAMPLE_ROWS, seed)
|
||||
df, gate = _gate_for_df(handle, df)
|
||||
x = df["pre_E"].to_numpy()
|
||||
if len(x):
|
||||
binned = _quantile_bins(x, gate, _N_BINS)
|
||||
means = np.asarray(binned["means"])
|
||||
score = means.max(axis=1).tolist() if means.size else []
|
||||
sides[name] = {"centers": binned["centers"], "score": score}
|
||||
else:
|
||||
sides[name] = {"centers": [], "score": []}
|
||||
|
||||
return Reduced(
|
||||
id="router_specialization",
|
||||
family="model",
|
||||
kind="router_specialization",
|
||||
title=_TITLES["router_specialization"],
|
||||
xlabel="pre-step energy [MeV]",
|
||||
payload={
|
||||
"router_type": handle.router_type,
|
||||
"n_experts": handle.router.n_experts,
|
||||
"log_x": True,
|
||||
"chance_level": 1.0 / handle.router.n_experts,
|
||||
**sides,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def compute_router_share_by_pdg(
|
||||
checkpoint: str | Path | None,
|
||||
r_phys: pl.LazyFrame,
|
||||
|
||||
@@ -6,8 +6,8 @@ streaming `group_by` pass(es) over the chunk (see `catalog.py`/`reduce.py`).
|
||||
`_COST_MODEL` below is ``spec_id -> (intercept_s, seconds_per_row)``.
|
||||
``n_rows`` is the combined rollout+reference row count of the job's input:
|
||||
the chunk's row count for `chunkable=True` specs, the whole dataset's for the
|
||||
three `chunkable=False` router specs (they always run as a single job
|
||||
regardless of chunk count).
|
||||
`chunkable=False` router specs in `_ROUTER_IDS` (they always run as a single
|
||||
job regardless of chunk count).
|
||||
|
||||
Calibrated 2026-07-27 from real HTCondor timings (`condor_history`
|
||||
``RemoteWallClockTime``) of a production run: prediction ``563f5ee3``
|
||||
@@ -54,7 +54,7 @@ _FIXED_OVERHEAD_S = 60.0
|
||||
# scan. Calibrated from the 3 real router jobs' observed wall times (119, 66,
|
||||
# 124s) — max minus _FIXED_OVERHEAD_S, on top of it.
|
||||
_ROUTER_FIXED_S = 64.0
|
||||
_ROUTER_IDS = frozenset({"router_gating", "router_share_by_pdg", "router_share_by_process"})
|
||||
_ROUTER_IDS = frozenset({"router_gating", "router_share_by_pdg", "router_share_by_process", "router_specialization"})
|
||||
|
||||
# Conservative fallback for any catalog id not in _COST_MODEL (e.g. a plot
|
||||
# added after the last calibration run) — the most expensive fitted per-row
|
||||
|
||||
@@ -1644,6 +1644,25 @@ def analyze_render(
|
||||
typer.echo(f"rendered {len(pdfs)} plots → {Path(run_dir) / 'plots'}")
|
||||
|
||||
|
||||
@analyze_app.command("metrics")
|
||||
def analyze_metrics(
|
||||
run_dir: Annotated[Path, typer.Argument(help="Run directory containing metrics.csv (from `giant train`)")],
|
||||
out_dir: Annotated[
|
||||
Optional[Path],
|
||||
typer.Option(
|
||||
"--out",
|
||||
"-o",
|
||||
help="Override the output directory (default: <cwd>/analysis_runs/metrics_<run_dir name>)",
|
||||
),
|
||||
] = None,
|
||||
) -> None:
|
||||
"""Render training-progress plots (loss/lr/accuracy/grad-norm/router/wgan/throughput) from <run_dir>/metrics.csv."""
|
||||
from giant.training.plots import render_metrics
|
||||
|
||||
paths = render_metrics(run_dir, out_dir, default_base=Path.cwd() / "analysis_runs")
|
||||
typer.echo(f"rendered {len(paths)} plots -> {paths[0].parent if paths else '(nothing to render)'}")
|
||||
|
||||
|
||||
@analyze_app.command("submit")
|
||||
def analyze_submit(
|
||||
rollout_yaml: Annotated[Path, typer.Argument(help="giant rollout YAML sidecar")],
|
||||
|
||||
@@ -0,0 +1,356 @@
|
||||
"""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.condor.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
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "giant"
|
||||
version = "0.3.6"
|
||||
version = "0.3.8"
|
||||
description = "Geant4 step-function surrogate via conditional flow matching"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
|
||||
@@ -159,6 +159,20 @@ def test_secondaries_rollout_vs_reference_align():
|
||||
assert t["pdg"].to_list() == [22, 22]
|
||||
|
||||
|
||||
def test_sec_count_by_event_zero_fills_events_with_no_secondaries():
|
||||
r_phys = physical_steps(_rollout_frame(), Side.rollout)
|
||||
r_sec = secondaries(_rollout_frame(), Side.rollout)
|
||||
ev, n = R.sec_count_by_event(r_phys, r_sec)
|
||||
# event 1 has one secondary track; event 2 has none and must still appear (as 0),
|
||||
# not silently drop out of a plain group_by on the secondaries frame alone.
|
||||
assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 0}
|
||||
|
||||
t_all = _reference_frame()
|
||||
t_sec = secondaries(t_all, Side.reference)
|
||||
ev, n = R.sec_count_by_event(t_all, t_sec)
|
||||
assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 1}
|
||||
|
||||
|
||||
def test_leakage_fraction():
|
||||
frac = R.leakage_fraction(_rollout_frame())
|
||||
# event 1: escaped pre_E=30, deposited=90 -> 30/120 = 0.25; event 2: 0
|
||||
|
||||
+66
-2
@@ -6,7 +6,13 @@ import numpy as np
|
||||
import pytest
|
||||
|
||||
from giant.analysis import build_catalog, catalog_ids, get_spec
|
||||
from giant.analysis.catalog import Bundle, PlotSpec
|
||||
from giant.analysis.catalog import (
|
||||
Bundle,
|
||||
PlotSpec,
|
||||
_containment_depths,
|
||||
_integer_confusion,
|
||||
_ks_statistic,
|
||||
)
|
||||
from giant.analysis.context import Context, build_context
|
||||
from tests.test_analysis_reduce import _reference_frame, _rollout_frame
|
||||
|
||||
@@ -53,6 +59,8 @@ def test_every_spec_computes_valid_reduced(bundle: Bundle):
|
||||
"single_hist",
|
||||
"router_gating",
|
||||
"router_share",
|
||||
"router_specialization",
|
||||
"heatmap",
|
||||
"unavailable",
|
||||
}
|
||||
assert r.title and r.xlabel
|
||||
@@ -88,6 +96,14 @@ def _validate_payload(r) -> None:
|
||||
for side in ("rollout", "reference"):
|
||||
if side in p:
|
||||
assert cat in p[side]
|
||||
elif r.kind == "router_specialization":
|
||||
for side in ("rollout", "reference"):
|
||||
if side in p:
|
||||
assert len(p[side]["centers"]) == len(p[side]["score"])
|
||||
elif r.kind == "heatmap":
|
||||
assert len(p["matrix"]) == len(p["row_labels"])
|
||||
for row in p["matrix"]:
|
||||
assert len(row) == len(p["col_labels"])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -98,7 +114,10 @@ def _validate_payload(r) -> None:
|
||||
# sec_count_per_species via pdg-keyed sums), concat-then-finalize with
|
||||
# data-dependent edges (event_total_edep), concat-then-mean/std (shower_
|
||||
# longitudinal), concat-then-max-edge (leakage_fraction), pdg-keyed sum with a
|
||||
# ratio (species_edep_share), and a chunkable=False passthrough (router_gating).
|
||||
# ratio (species_edep_share), a chunkable=False passthrough (router_gating),
|
||||
# nested sum-merge into a scorecard (marginal_distance_summary), concat-then-
|
||||
# event-id-join (n_sec_confusion), and concat-then-per-event-derived-quantity
|
||||
# (shower_containment_depth_90, reusing the profile matrix's own merge shape).
|
||||
_CHUNK_EQUIVALENCE_IDS = [
|
||||
"marginal_edep",
|
||||
"species_edep_share",
|
||||
@@ -107,6 +126,9 @@ _CHUNK_EQUIVALENCE_IDS = [
|
||||
"leakage_fraction",
|
||||
"sec_count_per_species",
|
||||
"router_gating",
|
||||
"marginal_distance_summary",
|
||||
"n_sec_confusion",
|
||||
"shower_containment_depth_90",
|
||||
]
|
||||
|
||||
|
||||
@@ -146,3 +168,45 @@ def test_chunked_matches_unchunked(ctx: Context, spec_id: str):
|
||||
assert chunked.id == unchunked.id
|
||||
assert chunked.kind == unchunked.kind
|
||||
_assert_payload_close(unchunked.payload, chunked.payload)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# new (gitea #76) reductions: KS distance, confusion matrix, containment depth
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_ks_statistic():
|
||||
assert _ks_statistic([10, 10], [10, 10]) == 0.0 # identical shape -> 0
|
||||
assert _ks_statistic([10, 0], [0, 10]) == 1.0 # fully disjoint -> 1
|
||||
assert _ks_statistic([0, 0], [0, 0]) != _ks_statistic([0, 0], [0, 0]) # nan (no data either side)
|
||||
assert _ks_statistic([10, 0], [0, 0]) == 1.0 # one side empty, other isn't -> maximal mismatch
|
||||
|
||||
|
||||
def test_integer_confusion_matches_event_pairing():
|
||||
# true (reference) n_sec = [1, 1]; predicted (rollout) n_sec = [1, 0]
|
||||
labels, mat = _integer_confusion(np.array([1, 1]), np.array([1, 0]))
|
||||
assert labels == ["0", "1+"]
|
||||
assert mat.tolist() == [[0, 0], [1, 1]] # row=true, col=pred
|
||||
|
||||
|
||||
def test_integer_confusion_caps_pathological_outliers():
|
||||
labels, mat = _integer_confusion(np.array([0, 500]), np.array([0, 0]), max_bins=5)
|
||||
assert labels[-1] == "4+"
|
||||
assert mat.shape == (5, 5)
|
||||
assert mat.sum() == 2
|
||||
|
||||
|
||||
def test_containment_depths_simple_ramp():
|
||||
# one event, edep concentrated in the first bin -> 90%/95% containment
|
||||
# depth is the first bin's right edge; a zero-energy event is dropped.
|
||||
mat = np.array([[9.0, 1.0, 0.0], [0.0, 0.0, 0.0]])
|
||||
edges = np.array([0.0, 1.0, 2.0, 3.0])
|
||||
depths = _containment_depths(mat, edges, 0.90)
|
||||
assert depths.tolist() == [1.0]
|
||||
|
||||
|
||||
def test_n_sec_confusion_spec(bundle):
|
||||
spec = get_spec("n_sec_confusion")
|
||||
r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
|
||||
assert r.payload["row_labels"] == r.payload["col_labels"] == ["0", "1+"]
|
||||
assert r.payload["matrix"] == [[0, 0], [1, 1]]
|
||||
|
||||
@@ -0,0 +1,333 @@
|
||||
"""Tests for giant.training.plots (gitea #75) — render smoke tests skipped
|
||||
where plotstyle/LaTeX is unavailable, plus pure-function column-classification
|
||||
coverage that needs neither."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("plotstyle")
|
||||
|
||||
from giant.training import plots as plots_mod # noqa: E402
|
||||
from giant.training.plots import MetricsTable, derive_metrics_dir, render_metrics # noqa: E402
|
||||
|
||||
# --- fixtures ----------------------------------------------------------
|
||||
|
||||
|
||||
_RICH_HEADER = [
|
||||
"epoch",
|
||||
"stage1/train/loss",
|
||||
"stage1/train/loss_gen",
|
||||
"stage1/train/nsec_acc",
|
||||
"stage1/train/grad_norm",
|
||||
"stage1/val/loss",
|
||||
"stage1/val/loss_gen",
|
||||
"stage1/val/nsec_acc",
|
||||
"stage1/lr",
|
||||
"stage1/router/entropy",
|
||||
"stage1/router/util_min",
|
||||
"stage1/router/util_max",
|
||||
"stage1/router/util_std",
|
||||
"stage2/train/d_loss",
|
||||
"stage2/train/g_loss",
|
||||
"stage2/train/wasserstein",
|
||||
"stage2/train/gp_loss",
|
||||
"stage2/train/loss_nsec",
|
||||
"stage2/train/nsec_acc",
|
||||
"stage2/train/grad_norm_d",
|
||||
"stage2/train/grad_norm_g",
|
||||
"stage2/lr",
|
||||
"stage2/critic_lr",
|
||||
"val/loss",
|
||||
"val/marginal_kl",
|
||||
"grad_norm",
|
||||
"gpu_mem_mb",
|
||||
"samples_per_sec",
|
||||
"is_best",
|
||||
"epoch_time_s",
|
||||
]
|
||||
|
||||
_RICH_ROWS = [
|
||||
[
|
||||
1,
|
||||
1.0,
|
||||
0.8,
|
||||
0.5,
|
||||
1.2,
|
||||
0.9,
|
||||
0.7,
|
||||
0.6,
|
||||
3e-4,
|
||||
1.5,
|
||||
0.05,
|
||||
0.3,
|
||||
0.1,
|
||||
-0.2,
|
||||
0.3,
|
||||
0.5,
|
||||
0.1,
|
||||
0.4,
|
||||
0.4,
|
||||
0.9,
|
||||
1.1,
|
||||
3e-4,
|
||||
1e-4,
|
||||
0.85,
|
||||
0.4,
|
||||
2.1,
|
||||
512.0,
|
||||
100.0,
|
||||
1,
|
||||
5.0,
|
||||
],
|
||||
[
|
||||
2,
|
||||
0.8,
|
||||
0.6,
|
||||
0.6,
|
||||
1.0,
|
||||
0.7,
|
||||
0.5,
|
||||
0.7,
|
||||
2e-4,
|
||||
1.6,
|
||||
0.06,
|
||||
0.28,
|
||||
0.09,
|
||||
-0.1,
|
||||
0.25,
|
||||
0.4,
|
||||
0.09,
|
||||
0.3,
|
||||
0.5,
|
||||
0.8,
|
||||
1.0,
|
||||
2e-4,
|
||||
8e-5,
|
||||
0.7,
|
||||
0.35,
|
||||
1.9,
|
||||
520.0,
|
||||
105.0,
|
||||
0,
|
||||
5.1,
|
||||
],
|
||||
]
|
||||
|
||||
_MINIMAL_HEADER = [
|
||||
"epoch",
|
||||
"stage1/train/loss",
|
||||
"stage1/train/loss_gen",
|
||||
"stage1/val/loss",
|
||||
"stage1/val/loss_gen",
|
||||
"stage1/lr",
|
||||
"val/loss",
|
||||
"grad_norm",
|
||||
"gpu_mem_mb",
|
||||
"samples_per_sec",
|
||||
"is_best",
|
||||
"epoch_time_s",
|
||||
]
|
||||
|
||||
_MINIMAL_ROWS = [
|
||||
[1, 1.0, 0.8, 0.9, 0.7, 3e-4, 0.85, 0.4, 0.0, 100.0, 0, 5.0],
|
||||
[2, 0.8, 0.6, 0.7, 0.5, 2e-4, 0.7, 0.35, 0.0, 105.0, 1, 5.1],
|
||||
]
|
||||
|
||||
|
||||
def _write_csv(path: Path, header: list[str], rows: list[list]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(path, "w", newline="") as f:
|
||||
writer = csv.writer(f)
|
||||
writer.writerow(header)
|
||||
writer.writerows(rows)
|
||||
|
||||
|
||||
# --- MetricsTable --------------------------------------------------------
|
||||
|
||||
|
||||
def test_metrics_table_load_round_trips(tmp_path: Path):
|
||||
csv_path = tmp_path / "metrics.csv"
|
||||
_write_csv(csv_path, _MINIMAL_HEADER, _MINIMAL_ROWS)
|
||||
table = MetricsTable.load(csv_path)
|
||||
assert table.epochs == [1, 2]
|
||||
assert table.columns["stage1/train/loss"] == [1.0, 0.8]
|
||||
assert "epoch" not in table.columns
|
||||
assert table.best_epochs() == [2]
|
||||
|
||||
|
||||
def test_metrics_table_best_epochs_empty_without_is_best_column():
|
||||
table = MetricsTable(epochs=[1, 2], columns={"stage1/train/loss": [1.0, 0.5]})
|
||||
assert table.best_epochs() == []
|
||||
|
||||
|
||||
# --- column classification (pure functions, no matplotlib) --------------
|
||||
|
||||
|
||||
def _rich_columns() -> dict[str, list]:
|
||||
return {name: [0.0] for name in _RICH_HEADER if name != "epoch"}
|
||||
|
||||
|
||||
def test_stages_detects_only_stages_present():
|
||||
assert plots_mod._stages(_rich_columns()) == ["stage1", "stage2"]
|
||||
assert plots_mod._stages({"stage2/train/loss": [0.0]}) == ["stage2"]
|
||||
assert plots_mod._stages({"val/loss": [0.0]}) == []
|
||||
|
||||
|
||||
def test_split_matches_stage_and_split_prefix_only():
|
||||
cols = _rich_columns()
|
||||
train = plots_mod._split(cols, "stage1", "train")
|
||||
assert train == {
|
||||
"loss": "stage1/train/loss",
|
||||
"loss_gen": "stage1/train/loss_gen",
|
||||
"nsec_acc": "stage1/train/nsec_acc",
|
||||
"grad_norm": "stage1/train/grad_norm",
|
||||
}
|
||||
assert plots_mod._split(cols, "stage2", "val") == {}
|
||||
|
||||
|
||||
def test_point_in_time_excludes_train_val_router():
|
||||
cols = _rich_columns()
|
||||
pit = plots_mod._point_in_time(cols, "stage1")
|
||||
assert pit == {"lr": "stage1/lr"}
|
||||
pit2 = plots_mod._point_in_time(cols, "stage2")
|
||||
assert pit2 == {"lr": "stage2/lr", "critic_lr": "stage2/critic_lr"}
|
||||
|
||||
|
||||
def test_router_columns():
|
||||
cols = _rich_columns()
|
||||
assert plots_mod._router(cols, "stage1") == {
|
||||
"entropy": "stage1/router/entropy",
|
||||
"util_min": "stage1/router/util_min",
|
||||
"util_max": "stage1/router/util_max",
|
||||
"util_std": "stage1/router/util_std",
|
||||
}
|
||||
assert plots_mod._router(cols, "stage2") == {}
|
||||
|
||||
|
||||
def test_run_level_excludes_stage_prefixed_columns_including_val_loss_lookalike():
|
||||
cols = _rich_columns()
|
||||
run_level = plots_mod._run_level(cols)
|
||||
assert set(run_level) == {
|
||||
"val/loss",
|
||||
"val/marginal_kl",
|
||||
"grad_norm",
|
||||
"gpu_mem_mb",
|
||||
"samples_per_sec",
|
||||
"is_best",
|
||||
"epoch_time_s",
|
||||
}
|
||||
# stage-prefixed "val/loss" lookalike (stage1/val/loss) must not leak in
|
||||
assert "stage1/val/loss" not in run_level
|
||||
|
||||
|
||||
def test_loss_keys_excludes_acc_and_wgan_and_grad_norm():
|
||||
train = {"loss": "x", "loss_gen": "x", "nsec_acc": "x", "grad_norm": "x", "d_loss": "x"}
|
||||
val = {"loss": "x", "loss_gen": "x"}
|
||||
assert plots_mod._loss_keys(train, val) == ["loss", "loss_gen"]
|
||||
|
||||
|
||||
# --- derive_metrics_dir ---------------------------------------------------
|
||||
|
||||
|
||||
def test_derive_metrics_dir_explicit_out_dir_wins():
|
||||
assert derive_metrics_dir("runs/my-run", out_dir="/somewhere") == Path("/somewhere")
|
||||
|
||||
|
||||
def test_derive_metrics_dir_default_base():
|
||||
assert derive_metrics_dir("runs/my-run", default_base="/data/analysis_runs") == Path(
|
||||
"/data/analysis_runs/metrics_my-run"
|
||||
)
|
||||
|
||||
|
||||
def test_derive_metrics_dir_falls_back_to_cwd_analysis_runs(monkeypatch, tmp_path):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
assert derive_metrics_dir("runs/my-run") == tmp_path / "analysis_runs" / "metrics_my-run"
|
||||
|
||||
|
||||
# --- render_metrics end to end -------------------------------------------
|
||||
|
||||
|
||||
def _try_render(run_dir: Path, out_dir: Path) -> list[Path]:
|
||||
try:
|
||||
return render_metrics(run_dir, out_dir)
|
||||
except RuntimeError as e: # LaTeX missing at render time
|
||||
pytest.skip(f"LaTeX rendering unavailable: {e}")
|
||||
|
||||
|
||||
def test_render_metrics_rich_run_produces_expected_plots_outside_run_dir(tmp_path: Path):
|
||||
run_dir = tmp_path / "run"
|
||||
out_dir = tmp_path / "out"
|
||||
_write_csv(run_dir / "metrics.csv", _RICH_HEADER, _RICH_ROWS)
|
||||
|
||||
paths = _try_render(run_dir, out_dir)
|
||||
|
||||
names = {p.stem for p in paths}
|
||||
assert names == {
|
||||
"overview",
|
||||
"stage1_loss",
|
||||
"stage2_loss",
|
||||
"lr",
|
||||
"stage1_accuracy",
|
||||
"stage2_accuracy",
|
||||
"grad_norm",
|
||||
"stage1_router",
|
||||
"stage2_wgan_balance",
|
||||
"throughput",
|
||||
}
|
||||
assert all(p.exists() for p in paths)
|
||||
assert all(p.is_relative_to(out_dir) for p in paths)
|
||||
# nothing written into the training run directory itself
|
||||
assert not any(run_dir.rglob("*.pdf"))
|
||||
|
||||
|
||||
def test_render_metrics_minimal_run_omits_router_wgan_accuracy(tmp_path: Path):
|
||||
run_dir = tmp_path / "run"
|
||||
out_dir = tmp_path / "out"
|
||||
_write_csv(run_dir / "metrics.csv", _MINIMAL_HEADER, _MINIMAL_ROWS)
|
||||
|
||||
paths = _try_render(run_dir, out_dir)
|
||||
|
||||
names = {p.stem for p in paths}
|
||||
assert names == {"overview", "stage1_loss", "lr", "grad_norm", "throughput"}
|
||||
assert "stage1_accuracy" not in names
|
||||
assert "stage1_router" not in names
|
||||
assert "stage1_wgan_balance" not in names
|
||||
|
||||
|
||||
def test_render_metrics_default_out_dir_uses_analysis_runs_convention(tmp_path: Path):
|
||||
run_dir = tmp_path / "runs" / "my-run"
|
||||
_write_csv(run_dir / "metrics.csv", _MINIMAL_HEADER, _MINIMAL_ROWS)
|
||||
default_base = tmp_path / "analysis_runs"
|
||||
|
||||
try:
|
||||
paths = render_metrics(run_dir, default_base=default_base)
|
||||
except RuntimeError as e:
|
||||
pytest.skip(f"LaTeX rendering unavailable: {e}")
|
||||
|
||||
assert paths
|
||||
assert all(p.is_relative_to(default_base / "metrics_my-run") for p in paths)
|
||||
|
||||
|
||||
# --- CLI -------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_cli_analyze_metrics_smoke(tmp_path: Path):
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from giant.cli import app
|
||||
|
||||
run_dir = tmp_path / "run"
|
||||
out_dir = tmp_path / "out"
|
||||
_write_csv(run_dir / "metrics.csv", _MINIMAL_HEADER, _MINIMAL_ROWS)
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(app, ["analyze", "metrics", str(run_dir), "--out", str(out_dir)])
|
||||
|
||||
if result.exit_code != 0 and "latex" in (str(result.output) + str(result.exception)).lower():
|
||||
pytest.skip("LaTeX rendering unavailable")
|
||||
assert result.exit_code == 0, result.output or result.exception
|
||||
assert any(out_dir.glob("*.pdf"))
|
||||
Reference in New Issue
Block a user