lars/giant
1
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity
Files
ff204732d72cf4f3d73cc152276fd37e045c16b4
giant/giant/rollout.py
T
lars c1c4957e2f Implement n_sec.mode = "stop_token" for the AR secondary decoder (gitea #40) 
Stage 2's autoregressive decoder still predicted multiplicity the v0.2 way:
a one-shot n_sec_head classifier over conditioning alone, run before any
secondary token existed, with the AR loop then always executing k_max slots
and discarding the tail. This adds a real per-slot EOS mechanism instead:

- Stage2Autoregressive gains a stop_head (build_stop_head=True) that predicts
  P(n_sec == k | prefix) at each slot, mutually exclusive with n_sec_head
  (n_sec.mode = "stop_token" builds no n_sec_head at all).
- sample_secondaries_ar accepts n_sec_pred=None to drive generation off the
  stop head instead of a pre-resolved count: each row stops the first slot
  its stop logit fires (stage2_model.n_sec.stop_sampling = "greedy" — the
  default, threshold at 0 — or "sample", a Bernoulli draw), and the whole
  batch loop breaks once every row has stopped, so cost scales with the
  realized n_sec instead of a fixed k_max. Passing n_sec_pred explicitly
  (the scheduled-sampling self-sample path) is unchanged.
- resolve_n_sec returns None for a stop-token decoder instead of raising;
  rollout.py/cli.py/validate.py now derive the realized count from
  sample_stage2's returned sec_valid (sec_valid.sum(-1)) after sampling,
  rather than resolving it up front — a no-op reordering under every other
  n_sec.mode, where sec_valid was already built from n_sec_pred.
- Training: _stop_target_and_mask (giant/training/stage2_inputs.py) builds
  the per-slot target/mask (one slot wider than the existing token-content
  sec_mask, since the stop slot itself needs supervision) and
  StageTrainer._stop_loss trains it with masked BCE, gated on stop_head
  exactly like _n_sec_loss gates on n_sec_head. Wired into both the
  flow/ddpm trainer and the WGAN trainer (whose skip_g_step now also checks
  stop_head), weighted by the existing stage2_model.n_sec.lambda — the stop
  head replaces n_sec_head under this mode, so no new weight key.
- validate_config now accepts stop_token (requires decoder="autoregressive"
  and n_sec.owner="stage2") instead of always rejecting it.

Decisions made during planning: stop_sampling defaults to "greedy" for
deterministic rollouts; the stop head reuses stage2_model.heads.n_sec's
HeadConfig shape and stage2_model.n_sec.lambda's weight rather than adding
new config keys, since the two heads never coexist.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-08-17 10:44:14 +02:00

859 lines
32 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
from collections import Counter
from typing import TYPE_CHECKING, Callable, TypedDict
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,
decode_secondary_cont,
energy_simplex_decode,
inv_local_frame_rotation,
inv_log_transform,
reconstruct_post_pos,
)
from giant.particles import (
decode_embedding_nearest,
decode_topn_class,
invert_dense_map,
nearest_known_pdg,
particle_phys_array,
)
from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
if TYPE_CHECKING:
from giant.data.loader import TopNMap
class L1DistCollector:
"""Accumulates the L1-distance diagnostic across a whole rollout
run: the L1 distance between each emitted secondary's raw predicted
embedding vector and the nearest table row it snapped to (only
meaningful under `particle_type.target = "embedding"` —
`giant.particles.decode_embedding_nearest`). A heavy tail means the
decoder is emitting vectors off the embedding manifold — the direct
analogue of the species-collapse symptom the v0.3.0 redesign exists to
fix.
Not folded into `rollout()`'s own return value (which is shape-typed as
step records, see `_RECORD_KEYS`/`RolloutSummary`) — passed in and read
back by the caller instead, mirroring the existing `on_chunk` pattern.
O(1) memory via a fixed log-spaced histogram rather than raw samples,
since a heavy right tail is exactly what this diagnostic watches for.
"""
def __init__(self, n_bins: int = 50, lo: float = 1e-3, hi: float = 1e3) -> None:
self.n = 0
self.total = 0.0
self.total_sq = 0.0
self.minimum = float("inf")
self.maximum = 0.0
self.hist_edges = np.geomspace(lo, hi, n_bins + 1)
self.hist_counts = np.zeros(n_bins, dtype=np.int64)
def add(self, dist: np.ndarray, valid: np.ndarray) -> None:
vals = np.asarray(dist)[np.asarray(valid)]
if vals.size == 0:
return
self.n += int(vals.size)
self.total += float(vals.sum())
self.total_sq += float(np.square(vals).sum())
self.minimum = min(self.minimum, float(vals.min()))
self.maximum = max(self.maximum, float(vals.max()))
self.hist_counts += np.histogram(vals, bins=self.hist_edges)[0]
def summary(self) -> dict | None:
"""`None` if nothing was ever added (target != "embedding", or a
run with zero secondaries) — the caller should omit the diagnostic
entirely rather than write a degenerate summary."""
if self.n == 0:
return None
mean = self.total / self.n
variance = max(self.total_sq / self.n - mean**2, 0.0)
return {
"n": self.n,
"mean": mean,
"std": variance**0.5,
"min": self.minimum,
"max": self.maximum,
"hist_edges": self.hist_edges.tolist(),
"hist_counts": self.hist_counts.tolist(),
}
def decode_secondary_identity(
sec_decoder: torch.nn.Module,
sec_cont: torch.Tensor,
sec_type: torch.Tensor,
n_sec_np: np.ndarray,
e_sec: np.ndarray,
pre_dir: np.ndarray,
sec_phys_norm: Normalizer,
pdg_map: dict[int, int],
sec_type_topn_map: "TopNMap | None",
other_policy: str,
rng: np.random.Generator | None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None]:
"""Decode Stage 2's raw (sec_cont, sec_type) output into physical
secondary attributes, branching on `sec_decoder.particle_type_cfg`:
- `"physical"`: unchanged v0.2 path — `sec_type` already *is* (log_mass,
charge), used as the secondary's identity as-is (no snapping).
- `"onehot"`: `sec_type` is per-slot class logits — argmax, then
`giant.particles.decode_topn_class` (+ `other_policy`) resolves a
concrete PDG, whose real physics (log_mass, charge) then come from
`giant.particles.particle_phys_array` — unlike "physical", the PDG
resolution IS the secondary's identity here, not just a reporting
label.
- `"embedding"`: `sec_type` is a raw vector in the conditioning's own
embedding space — `giant.particles.decode_embedding_nearest` L1-snaps
it to the nearest table row for the PDG (+ physics via
`particle_phys_array`), and also returns the L1 distance (see this
module's `L1DistCollector`).
Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg,
sec_type_l1_dist) — the last is `None` except under `"embedding"`.
"""
target = sec_decoder.particle_type_cfg.target
if target == "physical":
sec_full = torch.cat([sec_cont, sec_type], dim=-1).cpu().numpy()
sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid = decode_secondaries(
sec_full, n_sec_np, e_sec, pre_dir, sec_phys_normalizer=sec_phys_norm
)
sec_pdg = nearest_known_pdg(sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()).reshape(
sec_mass.shape
)
return sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, None
sec_E, sec_dir_world, sec_valid = decode_secondary_cont(sec_cont.cpu().numpy(), n_sec_np, e_sec, pre_dir)
sec_type_np = sec_type.cpu().numpy()
l1_dist = None
if target == "onehot":
if sec_type_topn_map is None:
raise RuntimeError(
"particle_type.target='onehot' rollout needs sec_type_topn_map "
"(the checkpoint's saved top-N map) — see ckpt['sec_type_topn_map']"
)
class_idx = sec_type_np.argmax(axis=-1)
sec_pdg = decode_topn_class(
class_idx,
sec_type_topn_map,
n_classes=sec_decoder.type_dim,
other_policy=other_policy,
rng=rng,
)
else: # "embedding"
idx_to_pdg = invert_dense_map(pdg_map)
emb_weight = sec_decoder.cond_enc.pdg_emb.weight.detach().cpu().numpy()
sec_pdg, l1_dist = decode_embedding_nearest(sec_type_np, emb_weight, idx_to_pdg)
l1_dist = np.where(sec_valid, l1_dist, 0.0).astype(np.float32)
sec_mass, sec_charge = particle_phys_array(sec_pdg.reshape(-1)).T
sec_mass = np.where(sec_valid, sec_mass.reshape(sec_pdg.shape), 0.0).astype(np.float32)
sec_charge = np.where(sec_valid, sec_charge.reshape(sec_pdg.shape), 0.0).astype(np.float32)
sec_pdg = np.where(sec_valid, sec_pdg, 0).astype(np.int64)
return sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, l1_dist
# 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),
"mass": np.empty(0, dtype=np.float64),
"charge": np.empty(0, 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]}
# Fixed per-key dtype, so every chunk table has an identical schema — needed
# for `giant rollout --on_chunk` to stream chunks straight into one
# pq.ParquetWriter (which requires matching schemas across writes), and a
# side benefit even in the buffered path since np.concatenate would otherwise
# silently upcast any stray int32/float32 chunk to the majority dtype.
_RECORD_DTYPES: dict[str, type] = {
"event_id": np.int64,
"track_id": np.int64,
"parent_id": np.int64,
"generation": np.int64,
"step_no": np.int64,
"pdg": np.int64,
"pre_x": np.float64,
"pre_y": np.float64,
"pre_z": np.float64,
"pre_E": np.float64,
"pre_dx": np.float64,
"pre_dy": np.float64,
"pre_dz": np.float64,
"post_x": np.float64,
"post_y": np.float64,
"post_z": np.float64,
"post_E": np.float64,
"post_dx": np.float64,
"post_dy": np.float64,
"post_dz": np.float64,
"edep": np.float64,
"step_length": np.float64,
"material": object,
"layer_id": np.int64,
"n_sec_pred": np.int64,
"termination_reason": object,
}
class RolloutSummary(TypedDict):
"""`rollout()`'s return shape when streaming to `on_chunk` instead of materializing rows."""
n_rows: int
termination_reason_counts: dict[str, int]
class _Recorder:
"""Accumulates per-step rows into column lists, materialised at the end —
or, when `sink` is given, streams each non-empty chunk to it immediately
instead, keeping only row-count / termination-reason summaries in memory.
The streaming path is what lets `giant rollout` write output incrementally
(see `rollout`'s `on_chunk` parameter): without it, a whole run's steps —
scaling with `n_events * max_steps * avg_tracks_per_event` — would sit in
RAM until the very end.
"""
def __init__(self, sink: Callable[[dict[str, np.ndarray]], None] | None = None) -> None:
self._sink = sink
self._cols: dict[str, list] | None = None if sink is not None else {k: [] for k in _RECORD_KEYS}
self.n_rows = 0
self.termination_reason_counts: Counter[str] = Counter()
def add(self, **cols) -> None:
n = len(cols["event_id"])
if n == 0:
return
row = {k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n) for k in _RECORD_KEYS}
self.n_rows += n
reasons = row["termination_reason"]
nonempty = reasons[reasons != ""]
if len(nonempty):
for r, c in zip(*np.unique(nonempty, return_counts=True)):
self.termination_reason_counts[str(r)] += int(c)
if self._sink is not None:
self._sink(row)
else:
assert self._cols is not None
for k in _RECORD_KEYS:
self._cols[k].append(row[k])
def to_dict(self) -> dict[str, np.ndarray]:
assert self._cols is not None, (
"to_dict() is unavailable when streaming to a sink — use n_rows/termination_reason_counts instead"
)
out = {}
for k, chunks in self._cols.items():
if chunks:
out[k] = np.concatenate(chunks, axis=0)
else:
out[k] = np.empty(0, dtype=_RECORD_DTYPES[k])
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,
particle_conditioning: str = "embedding",
) -> 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)
pdg_arr = np.asarray(pdg, dtype=np.int64)
if particle_conditioning == "physical":
# Real primaries always have a genuine ground-truth PDG code, looked
# up once here and carried forward unchanged for the track's lifetime
# (its species never changes mid-track) — same lifecycle as "pdg"
# itself.
mass, charge = particle_phys_array(pdg_arr).T
else:
# "embedding"/"onehot" never read mass/charge (see
# _physical_cond_columns), so resolving them here would only risk
# crashing a rollout on a PDG code giant.particles can't resolve, for
# a value that's never used.
mass = np.zeros(n, dtype=np.float64)
charge = np.zeros(n, dtype=np.float64)
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": pdg_arr,
"pre_pos": np.asarray(pre_pos, dtype=np.float64),
"pre_E": np.asarray(pre_E, dtype=np.float64),
"pre_dir": dir_,
"mass": mass.astype(np.float64),
"charge": charge.astype(np.float64),
}
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,
sec_phys_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,
on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None,
particle_conditioning: str = "embedding",
material_conditioning: str = "embedding",
pdg_topn_map: "TopNMap | None" = None,
mat_topn_map: "TopNMap | None" = None,
sec_type_topn_map: "TopNMap | None" = None,
other_policy: str = "sample",
seed: int | None = None,
stage1_ddpm_steps: int = 1000,
stage2_ddpm_steps: int = 1000,
l1_dist_collector: "L1DistCollector | None" = None,
) -> dict[str, np.ndarray] | RolloutSummary:
"""Run showers to completion.
By default, returns a step-record dict (see _RECORD_KEYS) with the whole
run's rows materialised in memory.
If `on_chunk` is given, every non-empty batch of rows is streamed to it as
soon as it's produced instead — no per-run buffering — and this returns a
small summary dict instead: `{"n_rows": int, "termination_reason_counts":
dict[str, int]}`. Use this for large `--n-events`/`--max-steps` runs,
where the full record set would otherwise scale with
`n_events * max_steps * avg_tracks_per_event`.
There is no `mode` parameter — each stage's generative objective is read
directly off the model instance's own `generator_kind` (stage 1 and
stage 2 objectives are independent, e.g. `stage1_model.generator="flow"`
+ `stage2_model.generator="wgan"`), and the decoder (one-shot vs
autoregressive) is inferred from `sec_decoder`'s own class — see
`sample_stage1`/`sample_stage2` (giant.sample).
`pdg_topn_map`/`mat_topn_map`/`sec_type_topn_map` serve three independent
purposes, no longer required to share one map (see gitea #29):
`pdg_topn_map`/`mat_topn_map` are required whenever
`particle_conditioning`/`material_conditioning` is `"onehot"` (feeds
`build_cond_features`'s extra `cond_cat` top-N columns); `sec_type_topn_map`/
`other_policy` are required instead under
`stage2_model.particle_type.target = "onehot"` (secondary-species
decode) — its class count (`stage2_model.particle_type.n_classes`) may
differ from `pdg_topn_map`'s. `seed` seeds the `other_policy = "sample"`
draw only (torch/numpy sampling itself is seeded by the caller, same as
today).
`l1_dist_collector`, if given, accumulates the embedding-distance
diagnostic across the whole run — see `L1DistCollector`. Only populated
under `particle_type.target = "embedding"`; a no-op otherwise.
"""
if particle_conditioning == "onehot" and pdg_topn_map is None:
raise RuntimeError(
"conditioning.particle.type='onehot' rollout needs pdg_topn_map "
"(the checkpoint's saved top-N map) — see ckpt['pdg_topn_map']"
)
if sec_decoder.particle_type_cfg.target == "onehot" and sec_type_topn_map is None:
raise RuntimeError(
"stage2_model.particle_type.target='onehot' rollout needs sec_type_topn_map "
"(the checkpoint's saved top-N map) — see ckpt['sec_type_topn_map']"
)
if material_conditioning == "onehot" and mat_topn_map is None:
raise RuntimeError(
"conditioning.material.type='onehot' rollout needs mat_topn_map "
"(the checkpoint's saved top-N map) — see ckpt['mat_topn_map']"
)
device = device or torch.device("cpu")
stage1_model.eval()
sec_decoder.eval()
if escape_threshold is not None:
oracle.escape_threshold = float(escape_threshold)
rng = np.random.default_rng(seed)
frontier, counts = make_seed_frontier(
seeds["event_id"],
seeds["pdg"],
seeds["pre_pos"],
seeds["pre_E"],
seeds["pre_dir"],
particle_conditioning=particle_conditioning,
)
rec = _Recorder(sink=on_chunk)
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,
sec_phys_norm,
pdg_map,
mat_map,
rec,
counts,
energy_cutoff,
max_steps,
steps,
device,
max_tracks_per_event,
particle_conditioning,
material_conditioning,
pdg_topn_map,
mat_topn_map,
sec_type_topn_map,
other_policy,
rng,
stage1_ddpm_steps,
stage2_ddpm_steps,
l1_dist_collector,
)
)
frontier = _concat_frontiers(next_parts)
if on_chunk is not None:
return {
"n_rows": rec.n_rows,
"termination_reason_counts": dict(rec.termination_reason_counts),
}
return rec.to_dict()
def _step_chunk(
tr,
stage1_model,
sec_decoder,
oracle,
cond_norm,
tgt_norm,
sec_phys_norm,
pdg_map,
mat_map,
rec,
counts,
energy_cutoff,
max_steps,
steps,
device,
max_tracks_per_event,
particle_conditioning,
material_conditioning,
pdg_topn_map,
mat_topn_map,
sec_type_topn_map,
other_policy,
rng,
stage1_ddpm_steps,
stage2_ddpm_steps,
l1_dist_collector,
) -> 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
if particle_conditioning == "physical":
# Under physical-property conditioning, mass/charge (already resolved
# on every track — see the cond_dict comment below) drive the model,
# not a training-vocab PDG embedding — build_cond_features passes
# strict=False for exactly this mode, so an out-of-vocab species no
# longer raises. Terminating on it here would defeat the entire
# point of physical conditioning: generalizing to a held-out species.
known_pdg = np.ones(n, dtype=bool)
else:
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 ---
# "mass"/"charge" are the track's own already-resolved physical identity
# (real for a primary, the model's raw predicted values with no snapping
# for a track descended from a secondary — see _spawn_secondaries), used
# directly instead of re-deriving via a pdg lookup. "pdg" still flows
# through for cond_cat's embedding-mode index and the known_pdg gate.
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"],
"mass": tr["mass"],
"charge": tr["charge"],
}
cond_cont, cond_cat = build_cond_features(
cond_dict,
pdg_map,
mat_map,
cond_norm,
particle_conditioning=particle_conditioning,
material_conditioning=material_conditioning,
pdg_topn_map=pdg_topn_map.class_map if particle_conditioning == "onehot" else None,
mat_topn_map=mat_topn_map.class_map if material_conditioning == "onehot" else None,
)
cc = torch.from_numpy(cond_cont).float().to(device)
ck = torch.from_numpy(cond_cat).long().to(device)
stage1_norm, n_sec_pred_stage1 = sample_stage1(stage1_model, cc, ck, steps, stage1_ddpm_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_pred = resolve_n_sec(stage1_model, sec_decoder, cc, ck, stage1_norm, n_sec_pred_stage1)
# --- Secondaries ---
# No snapping for "physical"/history-facing state elsewhere in the
# pipeline: sec_mass/sec_charge (or, for "onehot"/"embedding", the
# resolved sec_pdg -> real physics) are the secondary's identity, used
# as-is for the spawned track's own future conditioning — see
# decode_secondary_identity's docstring for how each
# particle_type.target differs on whether PDG resolution is a real
# identity decision or just a reporting label.
sec_cont, sec_type, sec_valid = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
# A stop-token decoder resolves n_sec_pred=None above — the real count
# only exists once sample_stage2 has actually generated (or stopped
# generating) tokens, so read it back off sec_valid here. Under every
# other n_sec.mode sec_valid was built FROM n_sec_pred, so this is a
# no-op round trip in those cases.
n_sec_np = sec_valid.sum(dim=-1).cpu().numpy().astype(np.int64)
sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, sec_type_l1_dist = decode_secondary_identity(
sec_decoder,
sec_cont,
sec_type,
n_sec_np,
e_sec,
tr["pre_dir"],
sec_phys_norm,
pdg_map,
sec_type_topn_map,
other_policy,
rng,
)
sec_valid = np.arange(sec_E.shape[1])[None, :] < n_sec_np[:, None]
if l1_dist_collector is not None and sec_type_l1_dist is not None:
l1_dist_collector.add(sec_type_l1_dist, sec_valid)
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,
sec_mass,
sec_charge,
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],
"mass": tr["mass"][cont],
"charge": tr["charge"][cont],
}
return _concat_frontiers([cont_frontier, new_tracks])
def _spawn_secondaries(
tr,
post_pos,
sec_valid,
sec_E,
sec_dir_world,
sec_pdg_code,
sec_mass,
sec_charge,
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),
# Reporting-only nominal PDG (nearest-known-PDG label, never fed back
# into the model) — the track's actual physical identity going
# forward is "mass"/"charge" below, the model's raw prediction.
"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),
"mass": sec_mass[pr_k, sl_k].astype(np.float64),
"charge": sec_charge[pr_k, sl_k].astype(np.float64),
}
return frontier, dropped_edep
Reference in New Issue View Git Blame Copy Permalink