giant/giant/data/dataset.py

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,
    ) -> 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

    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        files = self.files
        if worker_info is not None:
            files = files[worker_info.id :: worker_info.num_workers]

        if self.shuffle:
            files = list(files)
            np.random.default_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 = np.random.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,
        )