giant/giant/data/dataset.py

from __future__ import annotations

from pathlib import Path

import numpy as np
import torch
from torch.utils.data import IterableDataset

from giant.data.loader import iter_file_chunks
from giant.data.transforms import Normalizer, build_features


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)
    n_val = max(1, int(len(unique) * val_fraction))
    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 tuple:
        (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx)
    where:
        cond_cont:   (B, COND_DIM) float32
        cond_cat:    (B, 2) int64
        target_s1:   (B, 9) float32 — normalised Stage-1 primary target
        n_sec:       (B,) int64 — true secondary count per step
        sec_cont:    (B, K_MAX, 4) float32 — [stick_logit, local_dir] per slot
        sec_pdg_idx: (B, K_MAX) int64 — PDG model-index per secondary slot
        proc_idx:    (B,) int64 — process-class label (ProcessRouter supervision
                     only; zeros when `proc_map` is None)
    """

    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,
    ) -> None:
        self.files = list(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

    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_spdg: list[np.ndarray] = []
        buf_proc: list[np.ndarray] = []
        buf_n = 0

        for path in files:
            for chunk in iter_file_chunks(path):
                mask = np.isin(chunk["event_id"], self._events_arr)
                if not mask.any():
                    continue
                chunk = {k: v[mask] for k, v in chunk.items()}

                (
                    cond_cont,
                    cond_cat,
                    target_s1,
                    n_sec,
                    sec_cont,
                    sec_pdg_idx,
                    proc_idx,
                    _,
                    _,
                ) = build_features(
                    chunk,
                    self.pdg_map,
                    self.mat_map,
                    cond_normalizer=self.cond_normalizer,
                    target_normalizer=self.target_normalizer,
                    proc_map=self.proc_map,
                    require_secondaries=True,
                )
                buf_cont.append(cond_cont)
                buf_cat.append(cond_cat)
                buf_tgt.append(target_s1)
                buf_nsec.append(n_sec)
                buf_sec.append(sec_cont)
                buf_spdg.append(sec_pdg_idx)
                buf_proc.append(proc_idx)
                buf_n += len(cond_cont)

                if buf_n >= self.shuffle_buffer:
                    (
                        buf_cont,
                        buf_cat,
                        buf_tgt,
                        buf_nsec,
                        buf_sec,
                        buf_spdg,
                        buf_proc,
                        buf_n,
                    ) = yield from self._flush(
                        buf_cont,
                        buf_cat,
                        buf_tgt,
                        buf_nsec,
                        buf_sec,
                        buf_spdg,
                        buf_proc,
                        final=False,
                    )

        if buf_n > 0:
            yield from self._flush(
                buf_cont,
                buf_cat,
                buf_tgt,
                buf_nsec,
                buf_sec,
                buf_spdg,
                buf_proc,
                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_spdg: list[np.ndarray],
        buf_proc: 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)
        spdg = np.concatenate(buf_spdg)
        proc = np.concatenate(buf_proc)

        if self.shuffle:
            idx = np.random.permutation(len(cont))
            cont, cat, tgt = cont[idx], cat[idx], tgt[idx]
            nsec, sec, spdg, proc = nsec[idx], sec[idx], spdg[idx], proc[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 (
                torch.from_numpy(cont[start:end]).float(),
                torch.from_numpy(cat[start:end]).long(),
                torch.from_numpy(tgt[start:end]).float(),
                torch.from_numpy(nsec[start:end]).long(),
                torch.from_numpy(sec[start:end]).float(),
                torch.from_numpy(spdg[start:end]).long(),
                torch.from_numpy(proc[start:end]).long(),
            )

        if final:
            return [], [], [], [], [], [], [], 0
        rem = n_full * bs
        return (
            [cont[rem:]],
            [cat[rem:]],
            [tgt[rem:]],
            [nsec[rem:]],
            [sec[rem:]],
            [spdg[rem:]],
            [proc[rem:]],
            n - rem,
        )