"""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()