lars
60c2ca1985
New "model" family in the gallery: router_gating (mean soft gate weight vs. pre-step energy, showing the router's soft decision boundaries) and router_share_by_pdg/router_share_by_process (stacked top-1 dispatch share by species / true physics process). Needs a live checkpoint's Router, so it's a documented exception to the rest of the package's polars/numpy-only contract; gracefully degrades to a placeholder for non-MoE checkpoints. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
"""Tests for the plot catalog: id uniqueness + every spec computes a valid Reduced."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from giant.analysis import build_catalog, catalog_ids, get_spec
|
|
from giant.analysis.catalog import Bundle
|
|
from giant.analysis.context import build_context
|
|
from tests.test_analysis_reduce import _reference_frame, _rollout_frame
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def bundle() -> Bundle:
|
|
r, t = _rollout_frame(), _reference_frame()
|
|
ctx = build_context(
|
|
r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000
|
|
)
|
|
return Bundle.open(r, t, ctx)
|
|
|
|
|
|
def test_catalog_ids_unique_and_nonempty():
|
|
ids = catalog_ids()
|
|
assert ids and len(ids) == len(set(ids))
|
|
# the required families are all present
|
|
fams = {s.family for s in build_catalog()}
|
|
assert {"marginals", "event", "shower", "species", "secondaries"} <= fams
|
|
|
|
|
|
def test_get_spec_roundtrip_and_unknown():
|
|
spec = get_spec("marginal_edep")
|
|
assert spec.id == "marginal_edep" and spec.family == "marginals"
|
|
with pytest.raises(KeyError):
|
|
get_spec("does_not_exist")
|
|
|
|
|
|
def test_every_spec_computes_valid_reduced(bundle: Bundle):
|
|
for spec in build_catalog():
|
|
r = spec.compute(bundle)
|
|
assert r.id == spec.id
|
|
assert r.kind in {
|
|
"overlay_hist",
|
|
"grouped_hist",
|
|
"profile",
|
|
"bar",
|
|
"single_hist",
|
|
"router_gating",
|
|
"router_share",
|
|
"unavailable",
|
|
}
|
|
assert r.title and r.xlabel
|
|
_validate_payload(r)
|
|
|
|
|
|
def _validate_payload(r) -> None:
|
|
p = r.payload
|
|
if r.kind == "overlay_hist":
|
|
n = len(p["edges"]) - 1
|
|
assert len(p["rollout"]) == n and len(p["reference"]) == n
|
|
elif r.kind == "single_hist":
|
|
assert len(p["rollout"]) == len(p["edges"]) - 1
|
|
elif r.kind == "grouped_hist":
|
|
n = len(p["edges"]) - 1
|
|
assert p["groups"], "grouped hist must have at least one group"
|
|
for g in p["groups"].values():
|
|
assert len(g["rollout"]) == n and len(g["reference"]) == n
|
|
elif r.kind == "profile":
|
|
n = len(p["edges"]) - 1
|
|
for k in ("rollout_mean", "rollout_std", "reference_mean", "reference_std"):
|
|
assert len(p[k]) == n
|
|
elif r.kind == "bar":
|
|
assert len(p["labels"]) == len(p["rollout"]) == len(p["reference"])
|
|
elif r.kind == "unavailable":
|
|
assert p["note"]
|
|
elif r.kind == "router_gating":
|
|
for side in ("rollout", "reference"):
|
|
if side in p:
|
|
assert len(p[side]["centers"]) == len(p[side]["means"])
|
|
elif r.kind == "router_share":
|
|
for cat in p["categories"]:
|
|
for side in ("rollout", "reference"):
|
|
if side in p:
|
|
assert cat in p[side]
|