giant/giant/tools/profile_setup_scan.py

"""Benchmark `giant.pipeline.run_setup_stage`'s cold-cache scan against synthetic data.

Generates a schema-complete synthetic steps parquet (matching
`tests/test_pipeline.py`'s `_make_synthetic_steps`, but built with vectorized
numpy instead of a per-row Python loop so it scales to millions of rows) at a
few row counts, times `run_setup_stage` with `cache_setup=False` (so every
call is a genuine cold scan, never served from the sidecar), and prints a
before/after-style table. Run this on `master` before a change and again
after to see what a step actually bought — see the "speed up dwarf
warm-cache" plan for the pass-by-pass breakdown this benchmark is meant to
attribute (giant/data/loader.py, giant/data/scan.py, giant/pipeline.py).

Usage: ``uv run python giant/tools/profile_setup_scan.py``
"""

from __future__ import annotations

import time
from pathlib import Path
from tempfile import TemporaryDirectory

import numpy as np
import polars as pl

from giant import config as gconfig
from giant.pipeline import run_setup_stage

ROW_COUNTS = [20_000, 100_000, 500_000, 2_000_000]

_MATERIALS = ["G4_AIR", "G4_Fe"]
_PDGS = [11, 22]
_PROCESSES = ["eIoni", "phot", "compt"]


def _unit_vectors(n: int, rng: np.random.Generator) -> np.ndarray:
    v = rng.normal(size=(n, 3))
    return v / np.linalg.norm(v, axis=1, keepdims=True)


def _ragged_lists(k: np.ndarray, rng: np.random.Generator, lo: float, hi: float) -> list[list[float]]:
    total = int(k.sum())
    flat = rng.uniform(lo, hi, size=total)
    idx = np.cumsum(k)[:-1]
    return [arr.tolist() for arr in np.split(flat, idx)]


def _make_synthetic_steps(n: int, seed: int = 0) -> pl.DataFrame:
    """Vectorized equivalent of tests/test_pipeline.py's `_make_synthetic_steps`.

    event_id is assigned so each event gets 2-3 steps (matching that
    fixture's structure), and pdg/material/process cycle deterministically
    by row index rather than being drawn at random, same as the original.
    """
    rng = np.random.default_rng(seed)
    n_events = max(n // 3, 1)

    pre_E = rng.uniform(50.0, 500.0, size=n)
    n_sec = rng.integers(0, 3, size=n)
    frac_dep = rng.uniform(0.05, 0.3, size=n)
    frac_sec = np.where(n_sec > 0, rng.uniform(0.05, 0.2, size=n), 0.0)
    frac_post = 1.0 - frac_dep - frac_sec
    edep = pre_E * frac_dep
    e_sec = pre_E * frac_sec
    post_E = pre_E * frac_post
    pre_pos = rng.uniform(-10, 10, size=(n, 3))
    step_length = rng.uniform(0.1, 5.0, size=n)
    pre_dir = np.zeros((n, 3))
    pre_dir[:, 2] = 1.0
    post_dir = _unit_vectors(n, rng)
    post_pos = pre_pos + step_length[:, None] * pre_dir

    row_idx = np.arange(n)
    event_id = row_idx % n_events

    sec_E = _ragged_lists(n_sec, rng, 0.1, 1.0)  # placeholder magnitude, rescaled below
    sec_dx = _ragged_lists(n_sec, rng, -1.0, 1.0)
    sec_dy = _ragged_lists(n_sec, rng, -1.0, 1.0)
    sec_dz = _ragged_lists(n_sec, rng, -1.0, 1.0)
    total_sec = int(n_sec.sum())
    flat_pdg = [_PDGS[(row_idx[i] + j) % 2] for i in range(n) for j in range(n_sec[i])]
    idx = np.cumsum(n_sec)[:-1]
    sec_pdg = (
        [list(x) for x in np.split(np.array(flat_pdg, dtype=np.int64), idx)] if total_sec else [[] for _ in range(n)]
    )
    # Rescale each row's secondary energies to sum to that row's e_sec (a
    # Dirichlet split, like the original fixture) rather than the raw
    # uniform placeholder.
    sec_E_scaled = []
    for i in range(n):
        vals = np.array(sec_E[i])
        if vals.size:
            sec_E_scaled.append((vals / vals.sum() * e_sec[i]).tolist())
        else:
            sec_E_scaled.append([])

    return pl.DataFrame(
        {
            "event_id": event_id,
            "pdg": np.array(_PDGS)[row_idx % 2],
            "pre_x": pre_pos[:, 0],
            "pre_y": pre_pos[:, 1],
            "pre_z": pre_pos[:, 2],
            "pre_E": pre_E,
            "pre_dx": pre_dir[:, 0],
            "pre_dy": pre_dir[:, 1],
            "pre_dz": pre_dir[:, 2],
            "material": np.array(_MATERIALS)[row_idx % 2],
            "layer_id": row_idx % 5,
            "child_track_ids": [list(range(int(k))) for k in n_sec],
            "e_sec": e_sec,
            "process": np.array(_PROCESSES)[row_idx % 3],
            "step_length": step_length,
            "post_E": post_E,
            "edep": edep,
            "post_dx": post_dir[:, 0],
            "post_dy": post_dir[:, 1],
            "post_dz": post_dir[:, 2],
            "post_x": post_pos[:, 0],
            "post_y": post_pos[:, 1],
            "post_z": post_pos[:, 2],
            "sec_E_list": sec_E_scaled,
            "sec_pdg_list": sec_pdg,
            "sec_dx_list": sec_dx,
            "sec_dy_list": sec_dy,
            "sec_dz_list": sec_dz,
        }
    )


def _time_setup_stage(data: Path) -> float:
    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {})
    gconfig.validate_config(cfg)
    t0 = time.perf_counter()
    run_setup_stage(
        data,
        val_fraction=cfg["train"]["val_fraction"],
        seed=cfg["train"]["seed"],
        cfg=cfg,
        cache_setup=False,
        echo=lambda *a, **k: None,
    )
    return time.perf_counter() - t0


def main() -> None:
    with TemporaryDirectory(prefix="giant-setup-scan-profile-") as tmp:
        tmp_path = Path(tmp)
        print(f"{'n_rows':>10s}  {'time (s)':>10s}  {'rows/s':>12s}")
        for n in ROW_COUNTS:
            path = tmp_path / f"steps_{n}.parquet"
            _make_synthetic_steps(n).write_parquet(path)
            # warm the OS page cache so the timed pass measures compute, not
            # the one-time cold read of a freshly-written file.
            pl.scan_parquet(path).select(pl.len()).collect()

            dt = _time_setup_stage(path)
            print(f"{n:>10,d}  {dt:>10.3f}  {n / dt:>12,.0f}")
            path.unlink()


if __name__ == "__main__":
    main()