diff --git a/giant/cli.py b/giant/cli.py index 39eace5..f36e074 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -804,7 +804,29 @@ def rollout( seeds = _seed_from_data(files, n_events) typer.echo(f"seeded {len(seeds['event_id']):,} shower(s)") - records = run_rollout( + out, dataset_path, pred_uuid = _resolve_prediction_output(data, out) + out.parent.mkdir(parents=True, exist_ok=True) + + # Written incrementally as each batch of steps is produced, rather than + # buffering the whole run (which scales with n_events * max_steps * + # avg_tracks_per_event) — mirrors the row-group streaming `giant predict` + # already does on its input side. + writer: pq.ParquetWriter | None = None + + def _write_chunk(row: dict[str, np.ndarray]) -> None: + nonlocal writer + table = pa.table(row) + if writer is None: + table = table.replace_schema_metadata( + { + PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE, + PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION, + } + ) + writer = pq.ParquetWriter(out, table.schema) + writer.write_table(table) + + summary = run_rollout( model, sec_decoder, oracle, @@ -820,18 +842,10 @@ def rollout( device=_device, max_tracks_per_event=max_tracks_per_event, escape_threshold=escape_threshold, + on_chunk=_write_chunk, ) - - out, dataset_path, pred_uuid = _resolve_prediction_output(data, out) - out.parent.mkdir(parents=True, exist_ok=True) - - table = pa.table(records).replace_schema_metadata( - { - PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE, - PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION, - } - ) - pq.write_table(table, out) + if writer is not None: + writer.close() ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path) ref = yaml.safe_load(ref_path.read_text()) @@ -848,10 +862,8 @@ def rollout( ) ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False)) - n_rows = len(records["event_id"]) - reasons = Counter(r for r in records["termination_reason"].tolist() if r) - typer.echo(f"wrote {n_rows:,} step rows → {out}") - typer.echo(f"terminations: {dict(reasons)}") + typer.echo(f"wrote {summary['n_rows']:,} step rows → {out}") + typer.echo(f"terminations: {summary['termination_reason_counts']}") typer.echo(f"reference: {ref_path}") diff --git a/giant/rollout.py b/giant/rollout.py index 109faa0..228c3c5 100644 --- a/giant/rollout.py +++ b/giant/rollout.py @@ -16,6 +16,9 @@ treated as detector leakage and not deposited. from __future__ import annotations +from collections import Counter +from typing import Callable + import numpy as np import torch @@ -89,32 +92,95 @@ def _concat_frontiers(parts: list[dict[str, np.ndarray]]) -> dict[str, np.ndarra return {k: np.concatenate([p[k] for p in parts], axis=0) for k in parts[0]} -class _Recorder: - """Accumulates per-step rows into column lists, materialised at the end.""" +# Fixed per-key dtype, so every chunk table has an identical schema — needed +# for `giant rollout --on_chunk` to stream chunks straight into one +# pq.ParquetWriter (which requires matching schemas across writes), and a +# side benefit even in the buffered path since np.concatenate would otherwise +# silently upcast any stray int32/float32 chunk to the majority dtype. +_RECORD_DTYPES: dict[str, type] = { + "event_id": np.int64, + "track_id": np.int64, + "parent_id": np.int64, + "generation": np.int64, + "step_no": np.int64, + "pdg": np.int64, + "pre_x": np.float64, + "pre_y": np.float64, + "pre_z": np.float64, + "pre_E": np.float64, + "pre_dx": np.float64, + "pre_dy": np.float64, + "pre_dz": np.float64, + "post_x": np.float64, + "post_y": np.float64, + "post_z": np.float64, + "post_E": np.float64, + "post_dx": np.float64, + "post_dy": np.float64, + "post_dz": np.float64, + "edep": np.float64, + "step_length": np.float64, + "material": object, + "layer_id": np.int64, + "n_sec_pred": np.int64, + "termination_reason": object, +} - def __init__(self) -> None: - self._cols: dict[str, list] = {k: [] for k in _RECORD_KEYS} + +class _Recorder: + """Accumulates per-step rows into column lists, materialised at the end — + or, when `sink` is given, streams each non-empty chunk to it immediately + instead, keeping only row-count / termination-reason summaries in memory. + + The streaming path is what lets `giant rollout` write output incrementally + (see `rollout`'s `on_chunk` parameter): without it, a whole run's steps — + scaling with `n_events * max_steps * avg_tracks_per_event` — would sit in + RAM until the very end. + """ + + def __init__( + self, sink: Callable[[dict[str, np.ndarray]], None] | None = None + ) -> None: + self._sink = sink + self._cols: dict[str, list] | None = ( + None if sink is not None else {k: [] for k in _RECORD_KEYS} + ) + self.n_rows = 0 + self.termination_reason_counts: Counter[str] = Counter() def add(self, **cols) -> None: n = len(cols["event_id"]) if n == 0: return - for k in _RECORD_KEYS: - v = cols[k] - self._cols[k].append(np.asarray(v).reshape(n)) + row = { + k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n) + for k in _RECORD_KEYS + } + self.n_rows += n + reasons = row["termination_reason"] + nonempty = reasons[reasons != ""] + if len(nonempty): + for r, c in zip(*np.unique(nonempty, return_counts=True)): + self.termination_reason_counts[str(r)] += int(c) + + if self._sink is not None: + self._sink(row) + else: + assert self._cols is not None + for k in _RECORD_KEYS: + self._cols[k].append(row[k]) def to_dict(self) -> dict[str, np.ndarray]: + assert self._cols is not None, ( + "to_dict() is unavailable when streaming to a sink — use " + "n_rows/termination_reason_counts instead" + ) out = {} for k, chunks in self._cols.items(): if chunks: out[k] = np.concatenate(chunks, axis=0) else: - out[k] = np.empty( - 0, - dtype=object - if k in ("material", "termination_reason") - else np.float64, - ) + out[k] = np.empty(0, dtype=_RECORD_DTYPES[k]) return out @@ -209,8 +275,20 @@ def rollout( device: torch.device | None = None, max_tracks_per_event: int | None = None, escape_threshold: float | None = None, + on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None, ) -> dict[str, np.ndarray]: - """Run showers to completion; return a step-record dict (see _RECORD_KEYS).""" + """Run showers to completion. + + By default, returns a step-record dict (see _RECORD_KEYS) with the whole + run's rows materialised in memory. + + If `on_chunk` is given, every non-empty batch of rows is streamed to it as + soon as it's produced instead — no per-run buffering — and this returns a + small summary dict instead: `{"n_rows": int, "termination_reason_counts": + dict[str, int]}`. Use this for large `--n-events`/`--max-steps` runs, + where the full record set would otherwise scale with + `n_events * max_steps * avg_tracks_per_event`. + """ device = device or torch.device("cpu") stage1_model.eval() sec_decoder.eval() @@ -227,7 +305,7 @@ def rollout( seeds["pre_E"], seeds["pre_dir"], ) - rec = _Recorder() + rec = _Recorder(sink=on_chunk) while len(frontier["event_id"]) > 0: next_parts: list[dict[str, np.ndarray]] = [] @@ -257,6 +335,11 @@ def rollout( ) frontier = _concat_frontiers(next_parts) + if on_chunk is not None: + return { + "n_rows": rec.n_rows, + "termination_reason_counts": dict(rec.termination_reason_counts), + } return rec.to_dict() diff --git a/tests/test_rollout.py b/tests/test_rollout.py index f744755..4d7c747 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -1,5 +1,6 @@ """Tests for the autoregressive shower rollout driver.""" +from collections import Counter from pathlib import Path from unittest.mock import patch @@ -161,3 +162,94 @@ def test_max_tracks_cap_conserves_energy(): leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum() assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4) assert len(np.unique(rec["track_id"][m])) <= 3 + + +# ── Streaming output (on_chunk) ────────────────────────────────────────────── + + +def _run_streaming(on_chunk, **kwargs): + torch.manual_seed(0) + np.random.seed(0) + s1, s2 = _models() + cond, tgt = _norms() + seeds = kwargs.pop("seeds", None) or _seeds() + return rollout( + s1, + s2, + _oracle(), + seeds, + cond, + tgt, + PDG_MAP, + MAT_MAP, + energy_cutoff=kwargs.pop("energy_cutoff", 1.0), + max_steps=kwargs.pop("max_steps", 30), + steps=4, + batch_size=128, + max_tracks_per_event=kwargs.pop("max_tracks_per_event", 300), + escape_threshold=kwargs.pop("escape_threshold", 1e9), + on_chunk=on_chunk, + ) + + +def test_on_chunk_receives_every_row_exactly_once(): + """Concatenating the streamed chunks must reproduce the buffered result.""" + from giant.rollout import _RECORD_KEYS + + buffered = _run() + + chunks: list[dict[str, np.ndarray]] = [] + summary = _run_streaming(chunks.append) + + streamed = {k: np.concatenate([c[k] for c in chunks]) for k in _RECORD_KEYS} + assert summary["n_rows"] == len(buffered["event_id"]) + assert len(streamed["event_id"]) == len(buffered["event_id"]) + for k in _RECORD_KEYS: + np.testing.assert_array_equal(streamed[k], buffered[k]) + + +def test_on_chunk_summary_termination_reason_counts_match_buffered(): + buffered = _run() + summary = _run_streaming(lambda row: None) + + expected = Counter(r for r in buffered["termination_reason"].tolist() if r) + assert summary["termination_reason_counts"] == dict(expected) + + +def test_on_chunk_never_buffers_full_records(): + """Streaming mode must not accumulate rows for later to_dict() retrieval.""" + from giant.rollout import _Recorder + + rec = _Recorder(sink=lambda row: None) + rec.add( + event_id=np.array([0]), + track_id=np.array([0]), + parent_id=np.array([-1]), + generation=np.array([0]), + step_no=np.array([0]), + pdg=np.array([11]), + pre_x=np.array([0.0]), + pre_y=np.array([0.0]), + pre_z=np.array([0.0]), + pre_E=np.array([1.0]), + pre_dx=np.array([0.0]), + pre_dy=np.array([0.0]), + pre_dz=np.array([1.0]), + post_x=np.array([0.0]), + post_y=np.array([0.0]), + post_z=np.array([1.0]), + post_E=np.array([0.0]), + post_dx=np.array([0.0]), + post_dy=np.array([0.0]), + post_dz=np.array([1.0]), + edep=np.array([1.0]), + step_length=np.array([1.0]), + material=np.array(["G4_AIR"], dtype=object), + layer_id=np.array([0]), + n_sec_pred=np.array([0]), + termination_reason=np.array(["natural_end"], dtype=object), + ) + assert rec.n_rows == 1 + assert rec.termination_reason_counts == {"natural_end": 1} + with pytest.raises(AssertionError): + rec.to_dict()