lars
f4c2545e8b
Replace the monolithic giant/analysis.py (predict-local + RolloutVsTruth
diagnostics) with a lean giant/analysis/ package that compares one
autoregressive `giant rollout` for a checkpoint against a held-out
miniCaloSim reference file, and generates publication-styled plots in
parallel on HTCondor.
Rollout output and a raw reference file share a world-frame physical
column subset under identical names, so the old ALR/local-frame decode
machinery is gone — everything is world-frame mm/MeV.
- sources.py: canonical LazyFrames, synthetic-termination-row filtering,
the secondary view (rollout generation>0 tracks vs reference sec_*_list).
- reduce.py: streaming primitives — a single hist1d group_by pass, per-event
scalars, edep-weighted depth/transverse profiles, species share, leakage.
- context.py/grouping.py: prep resolves fixed bin edges + energy/pdg/material
group sets once into shared.json, so each compute job is one pass, no range
scan (histogram efficiency).
- catalog.py: declarative PlotSpec registry — marginals x {overall,energy,pdg,
material}, per-event totals, shower profiles, species/leakage, secondaries.
- render.py: the only plotstyle/LaTeX importer; PDFs + gallery metadata.
- condor.py + `giant analyze` CLI (prep/compute-one/list/render/submit):
one job per plot, compute/render split (workers polars-only, no LaTeX).
Styling via ETPlot's plotstyle (added to the analysis extra). New tests cover
the reduce primitives, catalog id uniqueness + compute, condor submit, and a
guarded render smoke test. Delete the two predict-diagnostics notebooks.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
172 lines
6.2 KiB
Python
172 lines
6.2 KiB
Python
"""Tests for the streaming compute primitives (giant.analysis.reduce/sources/grouping)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import polars as pl
|
|
|
|
from giant.analysis import grouping as G
|
|
from giant.analysis import reduce as R
|
|
from giant.analysis.sources import (
|
|
SYNTHETIC_TERMINATION_REASONS,
|
|
Side,
|
|
physical_steps,
|
|
secondaries,
|
|
)
|
|
|
|
|
|
def _rollout_frame() -> pl.LazyFrame:
|
|
# event 1: primary (2 steps) + 1 secondary track + 1 escaped bookkeeping row
|
|
# event 2: primary (1 step)
|
|
return pl.DataFrame(
|
|
{
|
|
"event_id": [1, 1, 1, 1, 2],
|
|
"track_id": [0, 0, 1, 0, 0],
|
|
"parent_id": [-1, -1, 0, -1, -1],
|
|
"generation": [0, 0, 1, 0, 0],
|
|
"step_no": [0, 1, 0, 99, 0],
|
|
"pdg": [11, 11, 22, 11, 11],
|
|
"pre_x": [0.0, 0.0, 0.0, 0.0, 0.0],
|
|
"pre_y": [0.0, 0.0, 0.0, 0.0, 0.0],
|
|
"pre_z": [0.0, 1.0, 1.0, 2.0, 0.0],
|
|
"pre_E": [100.0, 60.0, 20.0, 30.0, 50.0],
|
|
"pre_dx": [0.0, 0.0, 1.0, 0.0, 0.0],
|
|
"pre_dy": [0.0, 0.0, 0.0, 0.0, 0.0],
|
|
"pre_dz": [1.0, 1.0, 0.0, 1.0, 1.0],
|
|
"post_x": [0.0, 0.0, 1.0, 0.0, 0.0],
|
|
"post_y": [0.0, 0.0, 0.0, 0.0, 0.0],
|
|
"post_z": [1.0, 2.0, 1.0, 2.0, 1.0],
|
|
"post_E": [60.0, 30.0, 0.0, 0.0, 20.0],
|
|
"post_dx": [0.0, 0.0, 1.0, 0.0, 0.0],
|
|
"post_dy": [0.0, 0.0, 0.0, 0.0, 0.0],
|
|
"post_dz": [1.0, 1.0, 0.0, 1.0, 1.0],
|
|
"edep": [40.0, 30.0, 20.0, 0.0, 30.0],
|
|
"step_length": [1.0, 1.0, 1.0, 0.0, 1.0],
|
|
"material": ["G4_PbWO4"] * 5,
|
|
"layer_id": [0, 1, 1, -1, 0],
|
|
"n_sec_pred": [1, 0, 0, 0, 0],
|
|
"termination_reason": [
|
|
"",
|
|
"natural_end",
|
|
"natural_end",
|
|
"escaped",
|
|
"natural_end",
|
|
],
|
|
}
|
|
).lazy()
|
|
|
|
|
|
def _reference_frame() -> pl.LazyFrame:
|
|
return pl.DataFrame(
|
|
{
|
|
"event_id": [1, 1, 2],
|
|
"track_id": [0, 0, 0],
|
|
"step_no": [0, 1, 0],
|
|
"pdg": [11, 11, 11],
|
|
"pre_x": [0.0, 0.0, 0.0],
|
|
"pre_y": [0.0, 0.0, 0.0],
|
|
"pre_z": [0.0, 1.0, 0.0],
|
|
"pre_E": [100.0, 60.0, 50.0],
|
|
"pre_dx": [0.0, 0.0, 0.0],
|
|
"pre_dy": [0.0, 0.0, 0.0],
|
|
"pre_dz": [1.0, 1.0, 1.0],
|
|
"post_x": [0.0, 0.0, 0.0],
|
|
"post_y": [0.0, 0.0, 0.0],
|
|
"post_z": [1.0, 2.0, 1.0],
|
|
"post_E": [60.0, 30.0, 20.0],
|
|
"post_dx": [0.0, 0.0, 0.0],
|
|
"post_dy": [0.0, 0.0, 0.0],
|
|
"post_dz": [1.0, 1.0, 1.0],
|
|
"edep": [40.0, 30.0, 30.0],
|
|
"step_length": [1.0, 1.0, 1.0],
|
|
"material": ["G4_PbWO4", "G4_PbWO4", "G4_Pb"],
|
|
"layer_id": [0, 1, 0],
|
|
"sec_E_list": [[20.0], [], [10.0]],
|
|
"sec_pdg_list": [[22], [], [22]],
|
|
"sec_dx_list": [[1.0], [], [0.0]],
|
|
"sec_dy_list": [[0.0], [], [0.0]],
|
|
"sec_dz_list": [[0.0], [], [1.0]],
|
|
}
|
|
).lazy()
|
|
|
|
|
|
def test_hist1d_overall_and_grouped():
|
|
lf = _rollout_frame()
|
|
edges = np.linspace(0.0, 50.0, 6) # width 10
|
|
h = R.hist1d(lf, pl.col("edep"), edges)
|
|
# edep values: 40,30,20,0,30 -> bins [0),[10),[20),[30),[40)
|
|
assert h[0].tolist() == [1, 0, 1, 2, 1]
|
|
# grouped by pdg: pdg 22 has a single edep=20
|
|
hg = R.hist1d(lf, pl.col("edep"), edges, group=pl.col("pdg"))
|
|
assert hg[22].tolist() == [0, 0, 1, 0, 0]
|
|
assert hg[11].sum() == 4
|
|
|
|
|
|
def test_physical_steps_drops_synthetic_rollout_rows_only():
|
|
lf = _rollout_frame()
|
|
phys = physical_steps(lf, Side.rollout).collect()
|
|
assert phys.height == 4 # dropped the escaped bookkeeping row
|
|
assert "escaped" not in phys["termination_reason"].to_list()
|
|
assert SYNTHETIC_TERMINATION_REASONS # non-empty guard
|
|
# reference passes through unchanged
|
|
ref = _reference_frame()
|
|
assert physical_steps(ref, Side.reference).collect().height == ref.collect().height
|
|
|
|
|
|
def test_event_scalars_totals_include_all_rows():
|
|
lf = _rollout_frame()
|
|
es = R.event_scalars(lf).sort("event_id")
|
|
row1 = es.filter(pl.col("event_id") == 1).to_dicts()[0]
|
|
assert row1["total_edep"] == 90.0 # 40+30+20+0
|
|
assert row1["incident_E"] == 100.0
|
|
assert row1["n_steps"] == 4
|
|
|
|
|
|
def test_secondaries_rollout_vs_reference_align():
|
|
r = secondaries(_rollout_frame(), Side.rollout).collect().sort("event_id")
|
|
assert r["energy"].to_list() == [20.0] # only the generation>0, step_no==0 row
|
|
assert r["pdg"].to_list() == [22]
|
|
t = secondaries(_reference_frame(), Side.reference).collect().sort("event_id")
|
|
# two secondaries (event 1 and event 2); empty list dropped
|
|
assert sorted(t["energy"].to_list()) == [10.0, 20.0]
|
|
assert t["pdg"].to_list() == [22, 22]
|
|
|
|
|
|
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
|
|
assert sorted(round(f, 6) for f in frac) == [0.0, 0.25]
|
|
|
|
|
|
def test_weighted_profile_matches_manual_bincount():
|
|
lf = _rollout_frame()
|
|
ea = R.entry_axis(lf)
|
|
lf2 = R.attach_entry_axis(lf, ea)
|
|
edges = np.linspace(0.0, 3.0, 4) # depth bins along +z
|
|
mean, std = R.weighted_profile(lf2, R.depth_expr(), edges, pl.col("edep"))
|
|
assert mean.shape == (3,)
|
|
# totals conserved: sum over bins == mean total edep per event
|
|
assert np.isclose(mean.sum() * 1, (90.0 + 30.0) / 2) # 2 events
|
|
|
|
|
|
def test_energy_bins_edges_and_event_map():
|
|
incident = np.array([100.0, 100.0, 1000.0, 1000.0])
|
|
edges = G.energy_bin_edges(incident, n_bins=2)
|
|
assert len(edges) == 3 and edges[0] <= 100.0 < edges[-1]
|
|
ids, bins = G.event_energy_bins(_rollout_frame(), edges)
|
|
assert set(bins.tolist()) <= {0, 1}
|
|
assert len(ids) == 2
|
|
|
|
|
|
def test_digitize_expr_matches_numpy():
|
|
edges = np.array([0.0, 10.0, 100.0, 1000.0])
|
|
df = pl.DataFrame({"v": [5.0, 50.0, 500.0, 2000.0]})
|
|
got = df.select(G.digitize_expr(pl.col("v"), edges).alias("b"))["b"].to_list()
|
|
assert got == np.digitize([5.0, 50.0, 500.0, 2000.0], edges[1:-1]).tolist()
|
|
|
|
|
|
def test_pdg_and_material_labels():
|
|
assert G.pdg_label(22) == "gamma"
|
|
assert G.pdg_label(999999) == "999999"
|
|
assert G.material_label("G4_PbWO4") == "PbWO4"
|