from __future__ import annotations from pathlib import Path from typing import NamedTuple import numpy as np import torch from torch.utils.data import IterableDataset from giant.constants import K_MAX from giant.data.loader import event_id_offset, iter_file_chunks from giant.data.transforms import Normalizer, build_features, sorted_membership class StepBatch(NamedTuple): """One training batch, as yielded by `StreamingStepsDataset`. Field order is load-bearing for existing positional unpacking elsewhere (`trainers.py`, `validate.py`, test fixtures) — append only, never insert or reorder. cond_cont: (B, COND_DIM) float32 cond_cat: (B, 2/3/4) int64 — width 2 unless conditioning="onehot" target_s1: (B, 9) float32 — normalised Stage-1 primary target n_sec: (B,) int64 — true secondary count per step sec_cont: (B, k_max, SEC_SLOT_DIM) float32 — [stick_logit, local_dir, log_mass, charge] per slot (mass/charge normalised iff `sec_phys_normalizer` was given); always computed the same way regardless of stage2_model.particle_type.target, only actually used downstream under target="physical" proc_idx: (B,) int64 — process-class label (ProcessRouter supervision only; zeros when `proc_map` is None) sec_type_idx: (B, k_max) int64 — per-slot class index into `sec_type_class_map`, for particle_type.target in ("onehot", "embedding"); zeros (unused) otherwise """ cond_cont: torch.Tensor cond_cat: torch.Tensor target_s1: torch.Tensor n_sec: torch.Tensor sec_cont: torch.Tensor proc_idx: torch.Tensor sec_type_idx: torch.Tensor def make_event_split( all_event_ids: np.ndarray, val_fraction: float = 0.1, seed: int = 42, ) -> tuple[set, set]: """Assign unique event_ids to train/val sets by event_id, not by row.""" rng = np.random.default_rng(seed) unique = np.unique(all_event_ids) rng.shuffle(unique) # max(1, ...) only applies when a validation split was actually # requested — val_fraction=0.0 is an explicit "train on everything" # request and must not be silently overridden into holding out 1 event. n_val = max(1, int(len(unique) * val_fraction)) if val_fraction > 0 else 0 val_set = set(unique[:n_val].tolist()) train_set = set(unique[n_val:].tolist()) return train_set, val_set class StreamingStepsDataset(IterableDataset): """Streams parquet files one row-group at a time. Never loads more than `shuffle_buffer` rows into RAM simultaneously. Files are split evenly across DataLoader workers via worker_info. Yields whole batches (use with `DataLoader(..., batch_size=None)`) rather than single rows, so the batch is assembled with vectorized numpy slicing instead of a per-row Python loop in the default collate. Each batch is a `StepBatch` — see its docstring for field meanings. `k_max` (constructor arg, default the module constant) should match `stage2_model.k_max` — it sets the padded width of `sec_cont`/`sec_type_idx` above. """ def __init__( self, files: list[Path], split_events: set, pdg_map: dict[int, int], mat_map: dict[str, int], cond_normalizer: Normalizer, target_normalizer: Normalizer, batch_size: int, shuffle_buffer: int = 65536, shuffle: bool = True, proc_map: dict[str, int] | None = None, particle_conditioning: str = "embedding", material_conditioning: str = "embedding", sec_phys_normalizer: Normalizer | None = None, pdg_topn_map: dict[int, int] | None = None, mat_topn_map: dict[str, int] | None = None, sec_type_class_map: dict | None = None, k_max: int = K_MAX, seed: int = 0, ) -> None: self.files = list(files) self._offsets = {path: event_id_offset(i) for i, path in enumerate(self.files)} self.split_events = split_events self._events_arr = np.array(sorted(split_events)) self.pdg_map = pdg_map self.mat_map = mat_map self.cond_normalizer = cond_normalizer self.target_normalizer = target_normalizer self.batch_size = batch_size self.shuffle_buffer = max(shuffle_buffer, batch_size) self.shuffle = shuffle self.proc_map = proc_map self.particle_conditioning = particle_conditioning self.material_conditioning = material_conditioning self.sec_phys_normalizer = sec_phys_normalizer self.pdg_topn_map = pdg_topn_map self.mat_topn_map = mat_topn_map self.sec_type_class_map = sec_type_class_map self.k_max = k_max self.seed = seed self.epoch = 0 self._rng = np.random.default_rng() def set_epoch(self, epoch: int) -> None: """Select the shuffle stream for `epoch` (the DistributedSampler convention). The training loop calls this at the top of every epoch. Shuffling is seeded from `(seed, epoch, worker_id)` rather than the global numpy state so epoch *k*'s batch order is the same whether it runs as epoch *k* of one long `giant train`, or as its own resumed job in a per-epoch workflow chain (`giant/workflow/tasks.py:TrainEpochTask`). Workers are re-forked from this object each epoch (no `persistent_workers`), so setting it here reaches them. """ self.epoch = int(epoch) def __iter__(self): worker_info = torch.utils.data.get_worker_info() files = self.files worker_id = worker_info.id if worker_info is not None else 0 if worker_info is not None: files = files[worker_id :: worker_info.num_workers] self._rng = np.random.default_rng([self.seed, self.epoch, worker_id]) if self.shuffle: files = list(files) self._rng.shuffle(files) buf_cont: list[np.ndarray] = [] buf_cat: list[np.ndarray] = [] buf_tgt: list[np.ndarray] = [] buf_nsec: list[np.ndarray] = [] buf_sec: list[np.ndarray] = [] buf_proc: list[np.ndarray] = [] buf_type: list[np.ndarray] = [] buf_n = 0 for path in files: for chunk in iter_file_chunks(path, offset=self._offsets[path], k_max=self.k_max): mask = sorted_membership(chunk["event_id"], self._events_arr) if not mask.any(): continue chunk = {k: v[mask] for k, v in chunk.items()} feats = build_features( chunk, self.pdg_map, self.mat_map, cond_normalizer=self.cond_normalizer, target_normalizer=self.target_normalizer, sec_phys_normalizer=self.sec_phys_normalizer, proc_map=self.proc_map, require_secondaries=True, particle_conditioning=self.particle_conditioning, material_conditioning=self.material_conditioning, pdg_topn_map=self.pdg_topn_map, mat_topn_map=self.mat_topn_map, sec_type_class_map=self.sec_type_class_map, k_max=self.k_max, ) buf_cont.append(feats.cond_cont) buf_cat.append(feats.cond_cat) buf_tgt.append(feats.target_s1) buf_nsec.append(feats.n_sec) buf_sec.append(feats.sec_cont) buf_proc.append(feats.proc_idx) buf_type.append(feats.sec_type_idx) buf_n += len(feats.cond_cont) if buf_n >= self.shuffle_buffer: ( buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_proc, buf_type, buf_n, ) = yield from self._flush( buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_proc, buf_type, final=False, ) if buf_n > 0: yield from self._flush( buf_cont, buf_cat, buf_tgt, buf_nsec, buf_sec, buf_proc, buf_type, final=True, ) def _flush( self, buf_cont: list[np.ndarray], buf_cat: list[np.ndarray], buf_tgt: list[np.ndarray], buf_nsec: list[np.ndarray], buf_sec: list[np.ndarray], buf_proc: list[np.ndarray], buf_type: list[np.ndarray], final: bool, ): cont = np.concatenate(buf_cont) cat = np.concatenate(buf_cat) tgt = np.concatenate(buf_tgt) nsec = np.concatenate(buf_nsec) sec = np.concatenate(buf_sec) proc = np.concatenate(buf_proc) styp = np.concatenate(buf_type) if self.shuffle: idx = self._rng.permutation(len(cont)) cont, cat, tgt = cont[idx], cat[idx], tgt[idx] nsec, sec, proc, styp = nsec[idx], sec[idx], proc[idx], styp[idx] bs = self.batch_size n = len(cont) n_full = n // bs if not final else (n + bs - 1) // bs for start in range(0, n_full * bs, bs): end = min(start + bs, n) yield StepBatch( cond_cont=torch.from_numpy(cont[start:end]).float(), cond_cat=torch.from_numpy(cat[start:end]).long(), target_s1=torch.from_numpy(tgt[start:end]).float(), n_sec=torch.from_numpy(nsec[start:end]).long(), sec_cont=torch.from_numpy(sec[start:end]).float(), proc_idx=torch.from_numpy(proc[start:end]).long(), sec_type_idx=torch.from_numpy(styp[start:end]).long(), ) if final: return [], [], [], [], [], [], [], 0 rem = n_full * bs return ( [cont[rem:]], [cat[rem:]], [tgt[rem:]], [nsec[rem:]], [sec[rem:]], [proc[rem:]], [styp[rem:]], n - rem, )