giant/scripts/profile_analysis_costs.py

"""Benchmark `giant analyze compute-one`'s per-job cost against synthetic data.

Generates mock rollout+reference parquet files at a few row counts, times
`compute_reduced` for every chunkable catalog spec at each size (a single
chunk covering the whole mock file), fits a straight line (intercept, seconds
per row) through the timings, and prints the result as a Python dict literal
ready to paste into `giant/analysis/runtime_estimate.py::_COST_MODEL`.

The three `chunkable=False` router specs (`router_gating`,
`router_share_by_pdg`, `router_share_by_process`) need a live MoE checkpoint
to do any real work; without one (this machine has no `/ceph` access, so no
real checkpoint) they short-circuit almost instantly and are excluded here —
see `runtime_estimate.py`'s `_ROUTER_FIXED_S` for how those are handled
instead.

Usage: ``uv run python scripts/profile_analysis_costs.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.analysis.catalog import catalog_ids, get_spec
from giant.analysis.condor import compute_reduced
from giant.analysis.context import build_context

# Row counts (per side) to benchmark at. Kept in local memory/CPU range so the
# whole sweep finishes in about a minute; the fit is linear so it extrapolates
# fine to real multi-GB rollouts.
SIDE_ROW_COUNTS = [20_000, 100_000, 500_000, 2_000_000]

_MATERIALS = ["G4_PbWO4", "G4_Pb", "G4_lAr", "G4_Si"]
_PDGS = [11, -11, 22, 2112, 2212, 211, -211, 13]
_ROUTER_IDS = {"router_gating", "router_share_by_pdg", "router_share_by_process"}


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):
    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_rollout(n: int, n_events: int, seed: int) -> pl.DataFrame:
    rng = np.random.default_rng(seed)
    event_id = rng.integers(0, n_events, size=n)
    is_secondary = rng.random(n) < 0.15  # generation>0, step_no==0 birth rows
    is_synthetic = rng.random(n) < 0.05  # bookkeeping termination rows

    pre_E = rng.lognormal(mean=3.0, sigma=1.5, size=n)
    edep = rng.uniform(0, 1, size=n) * pre_E * 0.3
    post_E = np.clip(pre_E - edep, 0.0, None)
    pre_dir = _unit_vectors(n, rng)
    post_dir = _unit_vectors(n, rng)
    pos = rng.uniform(-50, 300, size=(n, 3))
    step_length = rng.uniform(0.1, 10.0, size=n)
    post_pos = pos + pre_dir * step_length[:, None]

    reasons = np.where(
        is_synthetic,
        rng.choice(["escaped", "energy_cutoff", "max_steps", "unknown_pdg"], size=n),
        "natural_end",
    )

    return pl.DataFrame(
        {
            "event_id": event_id,
            "track_id": rng.integers(0, 5, size=n),
            "parent_id": np.where(is_secondary, 0, -1),
            "generation": is_secondary.astype(np.int64),
            "step_no": np.where(is_secondary, 0, rng.integers(0, 20, size=n)),
            "pdg": rng.choice(_PDGS, size=n),
            "pre_x": pos[:, 0],
            "pre_y": pos[:, 1],
            "pre_z": pos[:, 2],
            "pre_E": pre_E,
            "pre_dx": pre_dir[:, 0],
            "pre_dy": pre_dir[:, 1],
            "pre_dz": pre_dir[:, 2],
            "post_x": post_pos[:, 0],
            "post_y": post_pos[:, 1],
            "post_z": post_pos[:, 2],
            "post_E": post_E,
            "post_dx": post_dir[:, 0],
            "post_dy": post_dir[:, 1],
            "post_dz": post_dir[:, 2],
            "edep": np.where(
                is_synthetic, np.where(reasons == "escaped", 0.0, pre_E), edep
            ),
            "step_length": np.where(is_synthetic, 0.0, step_length),
            "material": rng.choice(_MATERIALS, size=n),
            "layer_id": rng.integers(0, 30, size=n),
            "n_sec_pred": rng.integers(0, 4, size=n),
            "termination_reason": reasons,
        }
    )


def _make_reference(n: int, n_events: int, seed: int) -> pl.DataFrame:
    rng = np.random.default_rng(seed + 1)
    event_id = rng.integers(0, n_events, size=n)
    pre_E = rng.lognormal(mean=3.0, sigma=1.5, size=n)
    edep = rng.uniform(0, 1, size=n) * pre_E * 0.3
    post_E = np.clip(pre_E - edep, 0.0, None)
    pre_dir = _unit_vectors(n, rng)
    post_dir = _unit_vectors(n, rng)
    pos = rng.uniform(-50, 300, size=(n, 3))
    step_length = rng.uniform(0.1, 10.0, size=n)
    post_pos = pos + pre_dir * step_length[:, None]

    k = rng.poisson(0.3, size=n).clip(max=5).astype(np.int64)
    sec_pdg = _ragged_lists(k, rng, 0, 1)  # placeholder, overwritten below
    sec_E = _ragged_lists(k, rng, 0.1, 50.0)
    sec_dx = _ragged_lists(k, rng, -1.0, 1.0)
    sec_dy = _ragged_lists(k, rng, -1.0, 1.0)
    sec_dz = _ragged_lists(k, rng, -1.0, 1.0)
    total = int(k.sum())
    flat_pdg = rng.choice(_PDGS, size=total).tolist()
    idx = np.cumsum(k)[:-1]
    sec_pdg = [list(x) for x in np.split(np.array(flat_pdg), idx)]

    return pl.DataFrame(
        {
            "event_id": event_id,
            "track_id": rng.integers(0, 5, size=n),
            "step_no": rng.integers(0, 20, size=n),
            "pdg": rng.choice(_PDGS, size=n),
            "pre_x": pos[:, 0],
            "pre_y": pos[:, 1],
            "pre_z": pos[:, 2],
            "pre_E": pre_E,
            "pre_dx": pre_dir[:, 0],
            "pre_dy": pre_dir[:, 1],
            "pre_dz": pre_dir[:, 2],
            "post_x": post_pos[:, 0],
            "post_y": post_pos[:, 1],
            "post_z": post_pos[:, 2],
            "post_E": post_E,
            "post_dx": post_dir[:, 0],
            "post_dy": post_dir[:, 1],
            "post_dz": post_dir[:, 2],
            "edep": edep,
            "step_length": step_length,
            "material": rng.choice(_MATERIALS, size=n),
            "layer_id": rng.integers(0, 30, size=n),
            "process": rng.choice(["compt", "phot", "eBrem", "eIoni", "conv"], size=n),
            "sec_E_list": sec_E,
            "sec_pdg_list": sec_pdg,
            "sec_dx_list": sec_dx,
            "sec_dy_list": sec_dy,
            "sec_dz_list": sec_dz,
        }
    )


def _time(
    spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path
) -> float:
    t0 = time.perf_counter()
    compute_reduced(
        spec_id,
        rollout,
        reference,
        shared,
        out,
        checkpoint=None,
        chunk_index=0,
        n_chunks=1,
    )
    return time.perf_counter() - t0


def main() -> None:
    ids = [i for i in catalog_ids() if get_spec(i).chunkable]
    timings: dict[str, list[tuple[int, float]]] = {i: [] for i in ids}

    with TemporaryDirectory(prefix="giant-profile-") as tmp:
        tmp_path = Path(tmp)
        for n_side in SIDE_ROW_COUNTS:
            n_events = max(n_side // 20, 10)
            rollout = tmp_path / f"rollout_{n_side}.parquet"
            reference = tmp_path / f"reference_{n_side}.parquet"
            _make_rollout(n_side, n_events, seed=0).write_parquet(rollout)
            _make_reference(n_side, n_events, seed=0).write_parquet(reference)

            shared = tmp_path / f"shared_{n_side}.json"
            ctx = build_context(
                rollout,
                reference,
                n_energy_bins=4,
                n_marginal_bins=50,
                top_k_pdg=6,
                sample_rows=min(n_side, 200_000),
            )
            ctx.save(shared)

            # 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(rollout).select(pl.len()).collect()
            pl.scan_parquet(reference).select(pl.len()).collect()

            n_rows = 2 * n_side  # rollout + reference rows in this "chunk"
            for spec_id in ids:
                out = tmp_path / f"{spec_id}_{n_side}.json"
                dt = _time(spec_id, rollout, reference, shared, out)
                timings[spec_id].append((n_rows, dt))
                print(f"{spec_id:35s} n_rows={n_rows:>9d} time={dt:7.3f}s")

            rollout.unlink()
            reference.unlink()
            shared.unlink()

    print("\n# spec_id -> (intercept_s, seconds_per_row), fit by least squares")
    print("_COST_MODEL: dict[str, tuple[float, float]] = {")
    for spec_id in ids:
        xs = np.array([n for n, _ in timings[spec_id]], dtype=float)
        ys = np.array([t for _, t in timings[spec_id]], dtype=float)
        slope, intercept = np.polyfit(xs, ys, 1)
        intercept = max(intercept, 0.0)
        slope = max(slope, 0.0)
        print(f'    "{spec_id}": ({intercept:.6f}, {slope:.9f}),')
    print("}")

    if _ROUTER_IDS:
        print(
            "\n# router_* specs excluded: need a live MoE checkpoint to do real\n"
            "# work, none available on this machine — see _ROUTER_FIXED_S instead."
        )


if __name__ == "__main__":
    main()