giant/giant/data/loader.py

from pathlib import Path
from typing import Iterator

import numpy as np
import pandas as pd
import pyarrow.parquet as pq

# A manifest is a plain text file listing one parquet path per line, used to
# name a curated subset of files (e.g. a train/holdout pool) without copying
# or symlinking the underlying parquet files. Lines are resolved relative to
# the manifest's own directory, so the manifest stays valid if the whole
# dataset tree is moved or copied elsewhere intact.
MANIFEST_SUFFIX = ".manifest"

# Each input parquet file is a separate Geant4 job converted 1:1 from its own
# ROOT file (scripts/steps_to_parquet.py), and a job's event_id numbering
# always restarts from 0 — so when multiple files are loaded together (a
# directory or .manifest), raw event_id values collide across files even
# though they refer to unrelated events. Every per-file event_id column gets
# offset by its file's index in the (deterministically ordered) files list
# so ids stay globally unique across a multi-file load; the stride is far
# larger than any realistic per-file event count.
EVENT_ID_FILE_STRIDE = 1_000_000


def event_id_offset(file_index: int) -> int:
    return file_index * EVENT_ID_FILE_STRIDE


def _read_manifest(path: Path) -> list[Path]:
    files = []
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        resolved = (path.parent / line).resolve()
        if not resolved.is_file():
            raise FileNotFoundError(f"{path} lists missing file: {resolved}")
        files.append(resolved)
    if not files:
        raise FileNotFoundError(f"manifest {path} lists no files")
    return files


def find_parquet_files(path: str | Path) -> list[Path]:
    p = Path(path)
    if p.suffix == MANIFEST_SUFFIX:
        return _read_manifest(p)
    if p.is_dir():
        files = sorted(p.glob("*.parquet"))
        if not files:
            raise FileNotFoundError(f"no .parquet files found in {p}")
        return files
    return [p]


def _pad_list_col(series: pd.Series, K: int, fill: float = 0.0) -> np.ndarray:
    """Pad / truncate a list-valued Series to fixed width K → (N, K) float32."""
    out = np.full((len(series), K), fill, dtype=np.float32)
    for i, lst in enumerate(series):
        if lst is not None and len(lst) > 0:
            n = min(len(lst), K)
            out[i, :n] = lst[:n]
    return out


def _pad_list_col_int(series: pd.Series, K: int, fill: int = 0) -> np.ndarray:
    """Pad / truncate a list-valued integer Series to fixed width K → (N, K) int64."""
    out = np.full((len(series), K), fill, dtype=np.int64)
    for i, lst in enumerate(series):
        if lst is not None and len(lst) > 0:
            n = min(len(lst), K)
            out[i, :n] = lst[:n]
    return out


def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndarray:
    """Pad three list-valued direction columns → (N, K, 3) float32.

    Padding direction defaults to (0,0,1) (forward) so it is a valid unit vector.
    """
    N = len(dx)
    out = np.zeros((N, K, 3), dtype=np.float32)
    out[:, :, 2] = 1.0
    for i in range(N):
        lx, ly, lz = dx.iloc[i], dy.iloc[i], dz.iloc[i]
        if lx is not None and len(lx) > 0:
            n = min(len(lx), K)
            out[i, :n, 0] = lx[:n]
            out[i, :n, 1] = ly[:n]
            out[i, :n, 2] = lz[:n]
    return out


def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
    from giant.constants import K_MAX

    has_sec_lists = "sec_E_list" in df.columns

    d: dict[str, np.ndarray] = {
        "event_id": df["event_id"].to_numpy().astype(np.int64) + offset,
        "pdg": df["pdg"].to_numpy(dtype=np.int32),
        "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
        "pre_E": df["pre_E"].to_numpy(dtype=np.float32),
        "pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32),
        "material": df["material"].to_numpy(dtype=object),
        "layer_id": df["layer_id"].to_numpy(dtype=np.int32),
        "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
        "e_sec": df["e_sec"].to_numpy(dtype=np.float32),
        # The physics process that ended the step (e.g. "compt", "phot",
        # "eBrem") — a post-step outcome, so it's a router/classifier
        # supervision label only, never conditioning (see build_process_map*
        # / ProcessRouter). Guarded like has_sec_lists: older parquet
        # conversions predating this column still load fine.
        "process": (
            df["process"].to_numpy(dtype=object)
            if "process" in df.columns
            else np.full(len(df), "", dtype=object)
        ),
        "step_length": df["step_length"].to_numpy(dtype=np.float32),
        "post_E": df["post_E"].to_numpy(dtype=np.float32),
        "delta_e": (df["pre_E"] - df["post_E"]).to_numpy(dtype=np.float32),
        "edep": df["edep"].to_numpy(dtype=np.float32),
        "post_dir": df[["post_dx", "post_dy", "post_dz"]].to_numpy(dtype=np.float32),
        "post_pos": df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32),
    }

    if has_sec_lists:
        d["sec_E_list"] = _pad_list_col(df["sec_E_list"], K_MAX)
        d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], K_MAX)
        d["sec_dir_list"] = _pad_dir_col(
            df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], K_MAX
        )

    return d


def load_steps(path: str | Path, offset: int = 0) -> dict[str, np.ndarray]:
    return _df_to_dict(pd.read_parquet(path), offset=offset)


def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
    """Read only the event_id column — cheap scan for split assignment."""
    ids = pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
    return ids.astype(np.int64) + offset


def iter_file_chunks(
    path: str | Path, offset: int = 0
) -> Iterator[dict[str, np.ndarray]]:
    """Yield one parquet row-group at a time so a large file never fully loads."""
    pf = pq.ParquetFile(path)
    for i in range(pf.num_row_groups):
        yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset)


_COND_COLS = [
    "event_id",
    "pdg",
    "pre_x",
    "pre_y",
    "pre_z",
    "pre_E",
    "pre_dx",
    "pre_dy",
    "pre_dz",
    "material",
    "layer_id",
    "child_track_ids",
    "e_sec",
]


def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
    return {
        "event_id": df["event_id"].to_numpy().astype(np.int64) + offset,
        "pdg": df["pdg"].to_numpy(dtype=np.int32),
        "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
        "pre_E": df["pre_E"].to_numpy(dtype=np.float32),
        "pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32),
        "material": df["material"].to_numpy(dtype=object),
        "layer_id": df["layer_id"].to_numpy(dtype=np.int32),
        "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
        "e_sec": df["e_sec"].to_numpy(dtype=np.float32),
    }


def iter_cond_chunks(
    path: str | Path, offset: int = 0
) -> Iterator[dict[str, np.ndarray]]:
    """Yield conditioning-only row-groups (no post-step columns read from disk)."""
    pf = pq.ParquetFile(path)
    for i in range(pf.num_row_groups):
        yield _cond_df_to_dict(
            pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset
        )


def build_index_maps(
    data: dict[str, np.ndarray],
) -> tuple[dict[int, int], dict[str, int]]:
    pdg_vals = sorted(int(v) for v in np.unique(data["pdg"]))
    mat_vals = sorted(str(v) for v in np.unique(data["material"]))
    return (
        {v: i for i, v in enumerate(pdg_vals)},
        {v: i for i, v in enumerate(mat_vals)},
    )


def build_index_maps_from_files(
    files: list[Path],
) -> tuple[dict[int, int], dict[str, int]]:
    """Scan only pdg and material columns across all files (2-column read)."""
    pdg_vals: set[int] = set()
    mat_vals: set[str] = set()
    for path in files:
        df = pd.read_parquet(path, columns=["pdg", "material"])
        pdg_vals.update(int(v) for v in df["pdg"].unique())
        mat_vals.update(str(v) for v in df["material"].unique())
    return (
        {v: i for i, v in enumerate(sorted(pdg_vals))},
        {v: i for i, v in enumerate(sorted(mat_vals))},
    )


def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str, int]:
    """Scan the `process` column and build a frequency-capped process->index map.

    Physics processes have a long tail (rare nuclear captures, decays, ...)
    while `ProcessRouter` needs a fixed number of expert slots, so only the
    `n_experts - 1` most frequent processes get their own index; every rarer
    process is bucketed into a shared "other" index (`n_experts - 1`). This
    mirrors how `build_features` clamps the n_sec label to K_MAX for the
    fixed-width n_sec_head classifier.
    """
    counts: dict[str, int] = {}
    for path in files:
        df = pd.read_parquet(path, columns=["process"])
        for name, count in df["process"].value_counts().items():
            name = str(name)
            counts[name] = counts.get(name, 0) + int(count)
    ranked = sorted(counts, key=lambda name: counts[name], reverse=True)
    keep = ranked[: max(n_experts - 1, 0)]
    proc_map = {name: i for i, name in enumerate(keep)}
    other_idx = n_experts - 1
    for name in ranked[len(keep) :]:
        proc_map[name] = other_idx
    return proc_map