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>
109 lines
3.7 KiB
Python
109 lines
3.7 KiB
Python
"""Grouping axes (overall / energy / pdg / material) and their labels.
|
|
|
|
Energy grouping is by the event's **incident (primary) energy** — the largest
|
|
``pre_E`` in the event — so every step of a shower lands in one bin, the physically
|
|
meaningful stratification for a calorimeter surrogate. The quantile bin *edges* are
|
|
sized once in ``prep`` (from a subsample) and shipped in ``shared.json``; a compute
|
|
job that needs them re-derives the small per-event ``event_id -> bin`` map itself
|
|
(one bounded streaming ``group_by`` over ``pre_E``), so ``shared.json`` stays tiny.
|
|
|
|
Pure/plotstyle-free so it can run on the compute workers.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import polars as pl
|
|
|
|
# Common electromagnetic/hadronic species; anything else falls back to its code.
|
|
PDG_NAMES: dict[int, str] = {
|
|
11: "e-",
|
|
-11: "e+",
|
|
22: "gamma",
|
|
2112: "n",
|
|
2212: "p",
|
|
-2212: "pbar",
|
|
111: "pi0",
|
|
211: "pi+",
|
|
-211: "pi-",
|
|
13: "mu-",
|
|
-13: "mu+",
|
|
321: "K+",
|
|
-321: "K-",
|
|
130: "K0L",
|
|
}
|
|
|
|
|
|
def pdg_label(code: int) -> str:
|
|
"""Human-readable species label for a PDG code (falls back to the code)."""
|
|
code = int(code)
|
|
if code in PDG_NAMES:
|
|
return PDG_NAMES[code]
|
|
if abs(code) > 1_000_000_000:
|
|
return f"ion {code}"
|
|
return str(code)
|
|
|
|
|
|
def material_label(name: str) -> str:
|
|
"""Display label for a Geant4 material, dropping the ``G4_`` prefix."""
|
|
return name[3:] if name.startswith("G4_") else name
|
|
|
|
|
|
def energy_bin_edges(incident_E: np.ndarray, n_bins: int = 4) -> np.ndarray:
|
|
"""Equal-population (quantile) bin edges over per-event incident energies.
|
|
|
|
Returns ``n_bins + 1`` monotonically non-decreasing edges. The top edge is
|
|
nudged up so the largest value falls inside the last bin under a
|
|
right-open convention. Degenerate (single-value) input widens by +/-0.5.
|
|
"""
|
|
incident_E = np.asarray(incident_E, dtype=np.float64)
|
|
edges = np.quantile(incident_E, np.linspace(0.0, 1.0, n_bins + 1))
|
|
edges = np.unique(edges)
|
|
if edges.size < 2:
|
|
v = edges[0] if edges.size else 0.0
|
|
edges = np.array([v - 0.5, v + 0.5])
|
|
edges[-1] = np.nextafter(edges[-1], np.inf)
|
|
return edges
|
|
|
|
|
|
def energy_bin_labels(edges: np.ndarray) -> list[str]:
|
|
"""``E in [lo, hi)`` labels for each bin defined by ``edges`` (MeV)."""
|
|
return [
|
|
f"E in [{edges[i]:.3g}, {edges[i + 1]:.3g}) MeV" for i in range(len(edges) - 1)
|
|
]
|
|
|
|
|
|
def digitize_expr(value: pl.Expr, edges: np.ndarray) -> pl.Expr:
|
|
"""Bin index of ``value`` under arbitrary (possibly non-uniform) ``edges``.
|
|
|
|
``bin = (#interior edges <= value)``, clipped to ``[0, n_bins-1]`` — matches
|
|
``np.digitize(value, edges[1:-1])`` and works for the quantile energy edges.
|
|
Vectorized as a sum of boolean comparisons; no per-row Python.
|
|
"""
|
|
interior = [float(e) for e in edges[1:-1]]
|
|
n_bins = len(edges) - 1
|
|
idx = pl.lit(0, dtype=pl.Int32)
|
|
for e in interior:
|
|
idx = idx + (value >= e).cast(pl.Int32)
|
|
return idx.clip(0, n_bins - 1)
|
|
|
|
|
|
def event_energy_bins(
|
|
lf: pl.LazyFrame, edges: np.ndarray
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Per-event incident-energy bin: ``(event_ids, bin_idx)`` numpy arrays.
|
|
|
|
Incident energy is ``max(pre_E)`` per event (the primary). One bounded
|
|
streaming ``group_by``; the tiny per-event result is digitized in numpy.
|
|
"""
|
|
per_event = (
|
|
lf.group_by("event_id")
|
|
.agg(pl.col("pre_E").max().alias("incident_E"))
|
|
.collect(engine="streaming")
|
|
.sort("event_id")
|
|
)
|
|
event_ids = per_event["event_id"].to_numpy()
|
|
incident = per_event["incident_E"].to_numpy()
|
|
bin_idx = np.clip(np.digitize(incident, edges[1:-1]), 0, len(edges) - 2)
|
|
return event_ids, bin_idx.astype(np.int64)
|