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>
293 lines
11 KiB
Python
293 lines
11 KiB
Python
"""Tests for the plot catalog: id uniqueness + every spec computes a valid Reduced."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from giant.analysis import build_catalog, catalog_ids, get_spec
|
|
from giant.analysis.catalog import (
|
|
Bundle,
|
|
PlotSpec,
|
|
_containment_depths,
|
|
_ks_statistic,
|
|
)
|
|
from giant.analysis.context import Context, build_context
|
|
from giant.analysis.grouping import pdg_label
|
|
from giant.analysis.sources import RolloutSpec
|
|
from tests.test_analysis_reduce import _reference_frame, _rollout_frame
|
|
|
|
|
|
def _build_ctx() -> Context:
|
|
r, t = _rollout_frame(), _reference_frame()
|
|
return build_context(
|
|
[RolloutSpec("rollout", r)], t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000
|
|
)
|
|
|
|
|
|
def _two_rollout_specs() -> list[RolloutSpec]:
|
|
# Two distinct rollout sources so multi-series merging/finalize code is
|
|
# exercised even though the underlying frame is the same fixture.
|
|
return [RolloutSpec("flow", _rollout_frame()), RolloutSpec("wgan", _rollout_frame())]
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def ctx() -> Context:
|
|
return _build_ctx()
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def two_ctx() -> Context:
|
|
t = _reference_frame()
|
|
return build_context(_two_rollout_specs(), t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def bundle(ctx: Context) -> Bundle:
|
|
return Bundle.open([RolloutSpec("rollout", _rollout_frame())], _reference_frame(), ctx)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def two_bundle(two_ctx: Context) -> Bundle:
|
|
return Bundle.open(_two_rollout_specs(), _reference_frame(), two_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.finalize([spec.compute_partial(bundle)], bundle.ctx)
|
|
assert r.id == spec.id
|
|
assert r.kind in {
|
|
"overlay_hist",
|
|
"grouped_hist",
|
|
"profile",
|
|
"bar",
|
|
"single_hist",
|
|
"router_gating",
|
|
"router_share",
|
|
"router_specialization",
|
|
"heatmap",
|
|
"unavailable",
|
|
}
|
|
assert r.title and r.xlabel
|
|
_validate_payload(r, ["rollout"])
|
|
|
|
|
|
def test_every_spec_computes_valid_reduced_with_two_rollouts(two_bundle: Bundle):
|
|
for spec in build_catalog():
|
|
r = spec.finalize([spec.compute_partial(two_bundle)], two_bundle.ctx)
|
|
assert r.id == spec.id
|
|
_validate_payload(r, ["flow", "wgan"])
|
|
|
|
|
|
def _validate_payload(r, names: list[str]) -> None:
|
|
p = r.payload
|
|
if r.kind == "overlay_hist":
|
|
n = len(p["edges"]) - 1
|
|
assert list(p["series"]) == names
|
|
for v in p["series"].values():
|
|
assert len(v) == n
|
|
assert len(p["reference"]) == n
|
|
elif r.kind == "single_hist":
|
|
assert list(p["series"]) == names
|
|
for v in p["series"].values():
|
|
assert len(v) == 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 list(g["series"]) == names
|
|
for v in g["series"].values():
|
|
assert len(v) == n
|
|
assert len(g["reference"]) == n
|
|
elif r.kind == "profile":
|
|
n = len(p["edges"]) - 1
|
|
assert list(p["series"]) == names
|
|
for side in p["series"].values():
|
|
assert len(side["mean"]) == n and len(side["std"]) == n
|
|
assert len(p["reference"]["mean"]) == n and len(p["reference"]["std"]) == n
|
|
elif r.kind == "bar":
|
|
assert list(p["series"]) == names
|
|
for v in p["series"].values():
|
|
assert len(p["labels"]) == len(v)
|
|
assert len(p["labels"]) == len(p["reference"])
|
|
elif r.kind == "unavailable":
|
|
assert p["note"]
|
|
elif r.kind == "router_gating":
|
|
for entry in p["series"].values():
|
|
for side in ("rollout", "reference"):
|
|
if side in entry:
|
|
assert len(entry[side]["centers"]) == len(entry[side]["means"])
|
|
elif r.kind == "router_share":
|
|
for entry in p["series"].values():
|
|
for cat in entry["categories"]:
|
|
for side in ("rollout", "reference"):
|
|
if side in entry:
|
|
assert cat in entry[side]
|
|
elif r.kind == "router_specialization":
|
|
for entry in p["series"].values():
|
|
for side in ("rollout", "reference"):
|
|
if side in entry:
|
|
assert len(entry[side]["centers"]) == len(entry[side]["score"])
|
|
elif r.kind == "heatmap":
|
|
assert list(p["series"]) == names
|
|
for mat in p["series"].values():
|
|
assert len(mat) == len(p["row_labels"])
|
|
for row in mat:
|
|
assert len(row) == len(p["col_labels"])
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# chunked (compute_partial x N -> finalize) must match the unchunked (N=1) result
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# One representative id per merge shape: sum-mergeable (marginal_edep,
|
|
# 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), a chunkable=False passthrough (router_gating),
|
|
# sum-mergeable-with-a-zero-fill-denominator (sec_count_per_step{,_by_species}),
|
|
# nested sum-merge into a scorecard (marginal_distance_summary), 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",
|
|
"event_total_edep",
|
|
"shower_longitudinal",
|
|
"leakage_fraction",
|
|
"sec_count_per_species",
|
|
"sec_count_per_step",
|
|
"sec_count_per_step_by_species",
|
|
"router_gating",
|
|
"marginal_distance_summary",
|
|
"shower_containment_depth_90",
|
|
]
|
|
|
|
|
|
def _assert_payload_close(a, b, path: str = "payload") -> None:
|
|
"""Recursively compare two JSON-shaped payloads (float-tolerant)."""
|
|
assert type(a) is type(b), f"{path}: {type(a)} != {type(b)}"
|
|
if isinstance(a, dict):
|
|
assert set(a) == set(b), f"{path}: key mismatch {set(a)} != {set(b)}"
|
|
for k in a:
|
|
_assert_payload_close(a[k], b[k], f"{path}.{k}")
|
|
elif isinstance(a, list):
|
|
assert len(a) == len(b), f"{path}: length mismatch"
|
|
for i, (x, y) in enumerate(zip(a, b)):
|
|
_assert_payload_close(x, y, f"{path}[{i}]")
|
|
elif isinstance(a, float):
|
|
assert np.isclose(a, b, atol=1e-9), f"{path}: {a} != {b}"
|
|
else:
|
|
assert a == b, f"{path}: {a} != {b}"
|
|
|
|
|
|
@pytest.mark.parametrize("spec_id", _CHUNK_EQUIVALENCE_IDS)
|
|
def test_chunked_matches_unchunked(two_ctx: Context, spec_id: str):
|
|
"""A plot computed over N event-disjoint chunks then merged must equal the
|
|
same plot computed in one unchunked pass — the core chunking correctness
|
|
guarantee (see the analysis-rollout-plots chunking plan). Exercised with
|
|
two rollout series so the per-rollout merge path is covered too."""
|
|
spec: PlotSpec = get_spec(spec_id)
|
|
rollouts, t = _two_rollout_specs(), _reference_frame()
|
|
|
|
unchunked_bundle = Bundle.open(rollouts, t, two_ctx)
|
|
unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], two_ctx)
|
|
|
|
# 4 chunks over only 2 distinct event_ids also exercises empty chunks.
|
|
n_chunks = 4 if spec.chunkable else 1
|
|
parts = [spec.compute_partial(Bundle.open(rollouts, t, two_ctx, chunk=(k, n_chunks))) for k in range(n_chunks)]
|
|
chunked = spec.finalize(parts, two_ctx)
|
|
|
|
assert chunked.id == unchunked.id
|
|
assert chunked.kind == unchunked.kind
|
|
_assert_payload_close(unchunked.payload, chunked.payload)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# new (gitea #76) reductions: KS distance and 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_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_sec_count_per_step_counts_empty_steps(bundle):
|
|
spec = get_spec("sec_count_per_step")
|
|
r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
|
|
# reference: 3 steps, two of which emit exactly one secondary
|
|
assert r.payload["reference"][:2] == [1, 2]
|
|
# rollout: 4 physical steps, one of which emits a single secondary
|
|
assert r.payload["series"]["rollout"][:2] == [3, 1]
|
|
assert sum(r.payload["reference"]) == 3
|
|
|
|
|
|
def test_sec_count_per_step_by_species_zero_row_is_per_species(bundle):
|
|
spec = get_spec("sec_count_per_step_by_species")
|
|
r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
|
|
cols = r.payload["col_labels"]
|
|
ref = r.payload["reference"]
|
|
g = cols.index(pdg_label(22))
|
|
# two reference steps emit one photon each; the third emits none
|
|
assert [row[g] for row in ref][:2] == [1, 2]
|
|
# every other species column is "no such secondary" on all 3 steps
|
|
for j, _ in enumerate(cols):
|
|
if j != g:
|
|
assert ref[0][j] == 3 and sum(row[j] for row in ref[1:]) == 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# eval_cost_per_step
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_eval_cost_per_step_unavailable_without_timing(bundle: Bundle):
|
|
# `bundle`'s RolloutSpec carries no `timing` -> no rollout to compare.
|
|
spec = get_spec("eval_cost_per_step")
|
|
r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
|
|
assert r.kind == "unavailable"
|
|
assert r.payload["note"]
|
|
|
|
|
|
def test_eval_cost_per_step_bar_with_timing(ctx: Context):
|
|
spec = get_spec("eval_cost_per_step")
|
|
rs = RolloutSpec(
|
|
"rollout",
|
|
_rollout_frame(),
|
|
timing={"us_per_step": 12.5, "write_us_per_step": 2.5},
|
|
)
|
|
b = Bundle.open([rs], _reference_frame(), ctx)
|
|
r = spec.finalize([spec.compute_partial(b)], ctx)
|
|
assert r.kind == "bar"
|
|
assert r.payload["series"]["rollout"] == [12.5, 2.5, 15.0]
|
|
assert len(r.payload["reference"]) == 3
|
|
assert r.payload["log_y"] is True
|
|
assert "rollout" in r.meta["speedup_vs_geant4_total"]
|