lars/giant
1
0
Fork 0
Code Issues 2 Pull Requests 1 Actions Packages Projects Releases Wiki Activity
Files
a5683517be235ad1d5e1dbc86cb12d37bfa91fe3
giant/giant/rollout.py
T
lars 670f57c309 Rescale secondary energies to exactly consume the e_sec budget 
decode_secondaries's stick-breaking only guarantees valid secondary slots
sum to <= e_sec, leaving a shortfall that rollout.py silently dumped into
that step's edep. Rescale the valid slots by one common per-row factor
instead, so they sum to exactly e_sec whenever n_sec > 0: this spreads any
shortfall proportionally across all secondaries rather than concentrating
it in whichever slot is last by energy rank (which would let that one
low-energy secondary balloon and distort the shower's topology). Rows
where every valid slot decodes to ~zero fall back to an even split.

n_sec == 0 rows are unchanged (still nothing to carry the budget, so
rollout.py's edep top-up still applies there) — narrowed the related
caveat in load_rollout_vs_truth's docstring to just that case.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-13 13:45:13 +02:00

513 lines
16 KiB
Python
Raw Blame History

"""Autoregressive shower rollout driver.
Steps the two-stage GIANT surrogate forward into a full particle shower: each
primary post-step becomes the next pre-step, secondaries are pushed as new tracks,
and the material/layer_id conditioning at every step comes from a `GeometryOracle`
(the surrogate does not predict them).
Tracks are advanced breadth-first: every sweep steps all currently-active tracks
once (in `batch_size` chunks), so many tracks share each model forward pass. A
track terminates on one of the recorded `termination_reason`s in constants.py.
Energy accounting: on every terminal stop except escape, the track's remaining
energy is deposited locally so the shower conserves energy; escaped energy is
treated as detector leakage and not deposited.
"""
from __future__ import annotations
import numpy as np
import torch
from giant.constants import (
TERM_ENERGY_CUTOFF,
TERM_ESCAPED,
TERM_MAX_STEPS,
TERM_NATURAL_END,
TERM_UNKNOWN_PDG,
)
from giant.data.transforms import (
Normalizer,
build_cond_features,
decode_secondaries,
energy_simplex_decode,
inv_local_frame_rotation,
inv_log_transform,
reconstruct_post_pos,
)
from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx
# Record columns produced per step / per terminal marker.
_RECORD_KEYS = [
"event_id",
"track_id",
"parent_id",
"generation",
"step_no",
"pdg",
"pre_x",
"pre_y",
"pre_z",
"pre_E",
"pre_dx",
"pre_dy",
"pre_dz",
"post_x",
"post_y",
"post_z",
"post_E",
"post_dx",
"post_dy",
"post_dz",
"edep",
"step_length",
"material",
"layer_id",
"n_sec_pred",
"termination_reason",
]
def _empty_frontier() -> dict[str, np.ndarray]:
return {
"event_id": np.empty(0, dtype=np.int64),
"track_id": np.empty(0, dtype=np.int64),
"parent_id": np.empty(0, dtype=np.int64),
"generation": np.empty(0, dtype=np.int64),
"step_in_track": np.empty(0, dtype=np.int64),
"pdg": np.empty(0, dtype=np.int64),
"pre_pos": np.empty((0, 3), dtype=np.float64),
"pre_E": np.empty(0, dtype=np.float64),
"pre_dir": np.empty((0, 3), dtype=np.float64),
}
def _concat_frontiers(parts: list[dict[str, np.ndarray]]) -> dict[str, np.ndarray]:
parts = [p for p in parts if len(p["event_id"]) > 0]
if not parts:
return _empty_frontier()
return {k: np.concatenate([p[k] for p in parts], axis=0) for k in parts[0]}
class _Recorder:
"""Accumulates per-step rows into column lists, materialised at the end."""
def __init__(self) -> None:
self._cols: dict[str, list] = {k: [] for k in _RECORD_KEYS}
def add(self, **cols) -> None:
n = len(cols["event_id"])
if n == 0:
return
for k in _RECORD_KEYS:
v = cols[k]
self._cols[k].append(np.asarray(v).reshape(n))
def to_dict(self) -> dict[str, np.ndarray]:
out = {}
for k, chunks in self._cols.items():
if chunks:
out[k] = np.concatenate(chunks, axis=0)
else:
out[k] = np.empty(
0,
dtype=object
if k in ("material", "termination_reason")
else np.float64,
)
return out
def make_seed_frontier(
event_id: np.ndarray,
pdg: np.ndarray,
pre_pos: np.ndarray,
pre_E: np.ndarray,
pre_dir: np.ndarray,
) -> tuple[dict[str, np.ndarray], dict[int, int]]:
"""Build the initial frontier from primary entry states.
Returns (frontier, event_track_count) where the latter tracks the next
unused track_id per event (each primary gets a fresh id starting from 0).
"""
event_id = np.asarray(event_id, dtype=np.int64)
n = len(event_id)
track_id = np.empty(n, dtype=np.int64)
counts: dict[int, int] = {}
for i, ev in enumerate(event_id.tolist()):
c = counts.get(ev, 0)
track_id[i] = c
counts[ev] = c + 1
dir_ = np.asarray(pre_dir, dtype=np.float64)
dir_ = dir_ / np.clip(np.linalg.norm(dir_, axis=1, keepdims=True), 1e-12, None)
frontier = {
"event_id": event_id,
"track_id": track_id,
"parent_id": np.full(n, -1, dtype=np.int64),
"generation": np.zeros(n, dtype=np.int64),
"step_in_track": np.zeros(n, dtype=np.int64),
"pdg": np.asarray(pdg, dtype=np.int64),
"pre_pos": np.asarray(pre_pos, dtype=np.float64),
"pre_E": np.asarray(pre_E, dtype=np.float64),
"pre_dir": dir_,
}
return frontier, counts
def _terminal_rows(tr: dict[str, np.ndarray], sel: np.ndarray, reason: str, edep):
"""Assemble terminal-marker record columns for the selected tracks."""
pos = tr["pre_pos"][sel]
dir_ = tr["pre_dir"][sel]
n = int(sel.sum())
return dict(
event_id=tr["event_id"][sel],
track_id=tr["track_id"][sel],
parent_id=tr["parent_id"][sel],
generation=tr["generation"][sel],
step_no=tr["step_in_track"][sel],
pdg=tr["pdg"][sel],
pre_x=pos[:, 0],
pre_y=pos[:, 1],
pre_z=pos[:, 2],
pre_E=tr["pre_E"][sel],
pre_dx=dir_[:, 0],
pre_dy=dir_[:, 1],
pre_dz=dir_[:, 2],
post_x=pos[:, 0],
post_y=pos[:, 1],
post_z=pos[:, 2],
post_E=np.zeros(n),
post_dx=dir_[:, 0],
post_dy=dir_[:, 1],
post_dz=dir_[:, 2],
edep=np.asarray(edep, dtype=np.float64).reshape(n),
step_length=np.zeros(n),
material=tr.get("_material", np.full(len(sel), "", dtype=object))[sel],
layer_id=tr.get("_layer_id", np.zeros(len(sel), dtype=np.int64))[sel],
n_sec_pred=np.zeros(n, dtype=np.int64),
termination_reason=np.full(n, reason, dtype=object),
)
@torch.no_grad()
def rollout(
stage1_model: torch.nn.Module,
sec_decoder: torch.nn.Module,
oracle,
seeds: dict[str, np.ndarray],
cond_norm: Normalizer,
tgt_norm: Normalizer,
pdg_map: dict[int, int],
mat_map: dict[str, int],
*,
energy_cutoff: float,
max_steps: int,
steps: int = 10,
batch_size: int = 4096,
device: torch.device | None = None,
max_tracks_per_event: int | None = None,
escape_threshold: float | None = None,
) -> dict[str, np.ndarray]:
"""Run showers to completion; return a step-record dict (see _RECORD_KEYS)."""
device = device or torch.device("cpu")
stage1_model.eval()
sec_decoder.eval()
if escape_threshold is not None:
oracle.escape_threshold = float(escape_threshold)
pdg_map_inv = {v: k for k, v in pdg_map.items()}
pdg_emb_weight = stage1_model.pdg_embedding_weight()
frontier, counts = make_seed_frontier(
seeds["event_id"],
seeds["pdg"],
seeds["pre_pos"],
seeds["pre_E"],
seeds["pre_dir"],
)
rec = _Recorder()
while len(frontier["event_id"]) > 0:
next_parts: list[dict[str, np.ndarray]] = []
n_total = len(frontier["event_id"])
for start in range(0, n_total, batch_size):
chunk = {k: v[start : start + batch_size] for k, v in frontier.items()}
next_parts.append(
_step_chunk(
chunk,
stage1_model,
sec_decoder,
oracle,
cond_norm,
tgt_norm,
pdg_map,
mat_map,
pdg_map_inv,
pdg_emb_weight,
rec,
counts,
energy_cutoff,
max_steps,
steps,
device,
max_tracks_per_event,
)
)
frontier = _concat_frontiers(next_parts)
return rec.to_dict()
def _step_chunk(
tr,
stage1_model,
sec_decoder,
oracle,
cond_norm,
tgt_norm,
pdg_map,
mat_map,
pdg_map_inv,
pdg_emb_weight,
rec,
counts,
energy_cutoff,
max_steps,
steps,
device,
max_tracks_per_event,
) -> dict[str, np.ndarray]:
"""Advance one chunk of tracks by a single step; return the next frontier."""
n = len(tr["event_id"])
# --- Geometry lookup + material/layer conditioning ---
material, layer_id, escaped = oracle.query(tr["pre_pos"])
tr = dict(tr)
tr["_material"] = material
tr["_layer_id"] = layer_id
known_pdg = np.array([int(p) in pdg_map for p in tr["pdg"]], dtype=bool)
# --- Pre-step termination gates (in priority order; each track picks one) ---
stop = np.zeros(n, dtype=bool)
escaped_sel = escaped & ~stop
rec.add(
**_terminal_rows(
tr, escaped_sel, TERM_ESCAPED, edep=np.zeros(int(escaped_sel.sum()))
)
)
stop |= escaped_sel
unknown_sel = ~known_pdg & ~stop
rec.add(
**_terminal_rows(
tr, unknown_sel, TERM_UNKNOWN_PDG, edep=tr["pre_E"][unknown_sel]
)
)
stop |= unknown_sel
cutoff_sel = (tr["pre_E"] < energy_cutoff) & ~stop
rec.add(
**_terminal_rows(
tr, cutoff_sel, TERM_ENERGY_CUTOFF, edep=tr["pre_E"][cutoff_sel]
)
)
stop |= cutoff_sel
maxstep_sel = (tr["step_in_track"] >= max_steps) & ~stop
rec.add(
**_terminal_rows(tr, maxstep_sel, TERM_MAX_STEPS, edep=tr["pre_E"][maxstep_sel])
)
stop |= maxstep_sel
active = ~stop
if not active.any():
return _empty_frontier()
tr = {k: v[active] for k, v in tr.items()}
material = tr["_material"]
layer_id = tr["_layer_id"]
# --- Build conditioning and run the two stages ---
cond_dict = {
"pre_pos": tr["pre_pos"],
"pre_E": tr["pre_E"],
"pre_dir": tr["pre_dir"],
"layer_id": layer_id,
"material": material,
"pdg": tr["pdg"],
}
cond_cont, cond_cat = build_cond_features(cond_dict, pdg_map, mat_map, cond_norm)
cc = torch.from_numpy(cond_cont).float().to(device)
ck = torch.from_numpy(cond_cat).long().to(device)
stage1_norm, n_sec_pred = sample_flow(stage1_model, cc, ck, steps=steps)
raw = tgt_norm.inverse_transform(stage1_norm.cpu().numpy())
step_length = inv_log_transform(raw[:, 0])
edep, e_sec, post_E, _delta = energy_simplex_decode(raw[:, 1:3], tr["pre_E"])
post_dir_local = raw[:, 3:6].copy()
post_dir_local /= np.clip(
np.linalg.norm(post_dir_local, axis=1, keepdims=True), 1e-8, None
)
post_dir_world = inv_local_frame_rotation(tr["pre_dir"], post_dir_local)
travel_dir_local = raw[:, 6:9].copy()
travel_dir_local /= np.clip(
np.linalg.norm(travel_dir_local, axis=1, keepdims=True), 1e-8, None
)
post_pos = reconstruct_post_pos(
tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local
)
n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64)
# --- Secondaries ---
sec_cont, sec_type_emb, _valid = sample_secondaries(
sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
)
sec_pdg_idx = snap_type_to_pdg_idx(sec_type_emb, pdg_emb_weight)
sec_E, sec_dir_world, sec_pdg_code, sec_valid = decode_secondaries(
sec_cont.cpu().numpy(),
sec_pdg_idx.cpu().numpy(),
n_sec_np,
e_sec,
tr["pre_dir"],
pdg_map_inv,
)
edep = edep.astype(np.float64)
post_E = post_E.astype(np.float64)
# --- Spawn secondaries (with per-event track cap) ---
new_tracks, dropped_edep = _spawn_secondaries(
tr,
post_pos,
sec_valid,
sec_E,
sec_dir_world,
sec_pdg_code,
counts,
max_tracks_per_event,
)
# Energy bookkeeping so each step conserves exactly (edep + carried + post_E
# == pre_E): `decode_secondaries` already rescales valid slots to sum to
# exactly `e_sec` whenever n_sec > 0, so `residual` here is ~0 except when
# n_sec == 0 (no secondary to carry the budget at all — the whole `e_sec`
# becomes residual). Also deposit the energy of any sub-cap secondaries
# we dropped for hitting `max_tracks_per_event`.
sec_E_valid_sum = (sec_E * sec_valid).sum(axis=1)
residual = np.maximum(e_sec - sec_E_valid_sum, 0.0)
edep = edep + residual + dropped_edep
# --- Record the stepped rows; mark natural_end where the primary died ---
natural = post_E <= 0.0
reason = np.where(natural, TERM_NATURAL_END, "").astype(object)
rec.add(
event_id=tr["event_id"],
track_id=tr["track_id"],
parent_id=tr["parent_id"],
generation=tr["generation"],
step_no=tr["step_in_track"],
pdg=tr["pdg"],
pre_x=tr["pre_pos"][:, 0],
pre_y=tr["pre_pos"][:, 1],
pre_z=tr["pre_pos"][:, 2],
pre_E=tr["pre_E"],
pre_dx=tr["pre_dir"][:, 0],
pre_dy=tr["pre_dir"][:, 1],
pre_dz=tr["pre_dir"][:, 2],
post_x=post_pos[:, 0],
post_y=post_pos[:, 1],
post_z=post_pos[:, 2],
post_E=post_E,
post_dx=post_dir_world[:, 0],
post_dy=post_dir_world[:, 1],
post_dz=post_dir_world[:, 2],
edep=edep,
step_length=step_length,
material=material,
layer_id=layer_id,
n_sec_pred=n_sec_np,
termination_reason=reason,
)
# --- Continue surviving primaries ---
cont = ~natural
cont_frontier = {
"event_id": tr["event_id"][cont],
"track_id": tr["track_id"][cont],
"parent_id": tr["parent_id"][cont],
"generation": tr["generation"][cont],
"step_in_track": tr["step_in_track"][cont] + 1,
"pdg": tr["pdg"][cont],
"pre_pos": post_pos[cont],
"pre_E": post_E[cont],
"pre_dir": post_dir_world[cont],
}
return _concat_frontiers([cont_frontier, new_tracks])
def _spawn_secondaries(
tr,
post_pos,
sec_valid,
sec_E,
sec_dir_world,
sec_pdg_code,
counts,
max_tracks_per_event,
) -> tuple[dict[str, np.ndarray], np.ndarray]:
"""Turn valid secondaries into new tracks; return (frontier, per-parent dropped edep).
Secondaries are born at their parent's post_pos. When `max_tracks_per_event`
is set and an event is at its cap, further secondaries are not spawned; their
energy is returned as `dropped_edep` (indexed by parent row) so it is
deposited into the parent step instead of vanishing.
"""
B = len(tr["event_id"])
dropped_edep = np.zeros(B, dtype=np.float64)
pr, sl = np.nonzero(sec_valid) # parent-row idx, slot idx
if len(pr) == 0:
return _empty_frontier(), dropped_edep
# Assign a fresh per-event track_id to each candidate in stable parent order,
# applying the per-event cap. The candidate count per chunk is small
# (<= batch_size * K_MAX), so a plain loop is clear and fast enough.
order = np.lexsort((sl, pr)) # group by parent row, slot ascending
kept_pr, kept_sl, kept_tid = [], [], []
for j in order:
ev = int(tr["event_id"][pr[j]])
cur = counts.get(ev, 0)
if max_tracks_per_event is not None and cur >= max_tracks_per_event:
dropped_edep[pr[j]] += float(sec_E[pr[j], sl[j]])
continue
kept_pr.append(pr[j])
kept_sl.append(sl[j])
kept_tid.append(cur)
counts[ev] = cur + 1
if not kept_pr:
return _empty_frontier(), dropped_edep
pr_k = np.array(kept_pr, dtype=np.int64)
sl_k = np.array(kept_sl, dtype=np.int64)
tid_k = np.array(kept_tid, dtype=np.int64)
frontier = {
"event_id": tr["event_id"][pr_k],
"track_id": tid_k,
"parent_id": tr["track_id"][pr_k],
"generation": tr["generation"][pr_k] + 1,
"step_in_track": np.zeros(len(pr_k), dtype=np.int64),
"pdg": sec_pdg_code[pr_k, sl_k].astype(np.int64),
"pre_pos": post_pos[pr_k],
"pre_E": sec_E[pr_k, sl_k].astype(np.float64),
"pre_dir": sec_dir_world[pr_k, sl_k].astype(np.float64),
}
return frontier, dropped_edep
Reference in New Issue View Git Blame Copy Permalink