Offset event_id per file to avoid cross-file collisions

Each input parquet file is one Geant4 job (scripts/steps_to_parquet.py),
and a job's event_id numbering always restarts from 0 — so loading
multiple files together (a directory or .manifest) let same-numbered
events from different files collapse into one during the event index
scan and train/val split, corrupting both. Every per-file event_id now
gets offset by file index * EVENT_ID_FILE_STRIDE (giant/data/loader.py),
threaded through the setup-cache event index, the streaming dataset,
and predict/rollout seeding. Bumps the setup-cache format version so
stale sidecars computed pre-fix are invalidated.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

tests/test_dataset.py

@@ -1,5 +1,10 @@
 import numpy as np
-from giant.data.dataset import make_event_split
+import pandas as pd
+
+from giant.constants import COND_DIM, X_DIM
+from giant.data import setup_cache
+from giant.data.dataset import StreamingStepsDataset, make_event_split
+from giant.data.transforms import Normalizer
 
 
 def test_make_event_split_sizes():
@@ -31,3 +36,94 @@ def test_make_event_split_reproducible():
     b_tr, b_val = make_event_split(event_ids, val_fraction=0.1, seed=42)
     assert a_tr == b_tr
     assert a_val == b_val
+
+
+# ── StreamingStepsDataset: cross-file event_id offsetting ──────────────────
+
+
+def _steps_df(event_ids, n_per_event=3, pre_E=100.0):
+    """A schema-complete but minimal steps DataFrame — no secondaries, so
+    `require_secondaries=True` never needs the per-secondary list columns."""
+    rows = []
+    for eid in event_ids:
+        for s in range(n_per_event):
+            rows.append(
+                {
+                    "event_id": eid,
+                    "pdg": 11,
+                    "pre_x": 0.0,
+                    "pre_y": 0.0,
+                    "pre_z": 0.0,
+                    "pre_E": pre_E,
+                    "pre_dx": 0.0,
+                    "pre_dy": 0.0,
+                    "pre_dz": 1.0,
+                    "material": "G4_AIR",
+                    "layer_id": s,
+                    "child_track_ids": [],
+                    "e_sec": 0.0,
+                    "step_length": 1.0,
+                    "post_E": pre_E * 0.9,
+                    "edep": pre_E * 0.1,
+                    "post_dx": 0.0,
+                    "post_dy": 0.0,
+                    "post_dz": 1.0,
+                    "post_x": 0.0,
+                    "post_y": 0.0,
+                    "post_z": 1.0,
+                }
+            )
+    return pd.DataFrame(rows)
+
+
+def _dummy_normalizer(width):
+    norm = Normalizer()
+    norm.mean = np.zeros(width, dtype=np.float32)
+    norm.std = np.ones(width, dtype=np.float32)
+    return norm
+
+
+def test_streaming_dataset_offsets_colliding_event_ids_across_files(tmp_path):
+    """Two files that each restart event_id from 0 (one Geant4 job per file,
+    see scripts/steps_to_parquet.py) must not have their same-numbered events
+    collapsed together: every row from every file must show up in exactly one
+    of train/val, and the number of distinct events must be the sum across
+    files, not the union of raw ids."""
+    n_events, n_per_event = 5, 3
+    path_a = tmp_path / "a.parquet"
+    path_b = tmp_path / "b.parquet"
+    _steps_df(range(n_events), n_per_event=n_per_event).to_parquet(path_a)
+    _steps_df(range(n_events), n_per_event=n_per_event).to_parquet(path_b)
+    files = [path_a, path_b]
+
+    unique_ids, counts = setup_cache.compute_event_index_from_files(files)
+    assert len(unique_ids) == 2 * n_events
+
+    train_events, val_events = make_event_split(unique_ids, val_fraction=0.4, seed=0)
+    assert train_events.isdisjoint(val_events)
+
+    pdg_map, mat_map = {11: 0}, {"G4_AIR": 0}
+    cond_norm = _dummy_normalizer(COND_DIM)
+    tgt_norm = _dummy_normalizer(X_DIM)
+
+    def _count_rows(split_events):
+        ds = StreamingStepsDataset(
+            files=files,
+            split_events=split_events,
+            pdg_map=pdg_map,
+            mat_map=mat_map,
+            cond_normalizer=cond_norm,
+            target_normalizer=tgt_norm,
+            batch_size=4,
+            shuffle=False,
+            conditioning="embedding",
+        )
+        return sum(len(batch[0]) for batch in ds)
+
+    n_train = _count_rows(train_events)
+    n_val = _count_rows(val_events)
+
+    total_rows = 2 * n_events * n_per_event
+    assert n_train + n_val == total_rows
+    assert n_train == int(counts[np.isin(unique_ids, list(train_events))].sum())
+    assert n_val == int(counts[np.isin(unique_ids, list(val_events))].sum())

tests/test_loader.py

@@ -3,10 +3,16 @@ import pandas as pd
 import pytest
 
 from giant.data.loader import (
+    EVENT_ID_FILE_STRIDE,
     build_index_maps,
     build_index_maps_from_files,
     build_process_map_from_files,
+    event_id_offset,
     find_parquet_files,
+    iter_cond_chunks,
+    iter_file_chunks,
+    load_event_ids,
+    load_steps,
 )
 
 
@@ -279,3 +285,84 @@ def test_build_index_maps_from_files_matches_build_index_maps(tmp_path):
     from_files = build_index_maps_from_files([path])
     from_memory = build_index_maps({"pdg": pdg, "material": material})
     assert from_files == from_memory
+
+
+# ── event_id_offset / per-file event_id offsetting ─────────────────────────
+
+
+def _steps_df(event_ids):
+    """Minimal schema-complete steps rows (no secondaries) for _df_to_dict."""
+    return pd.DataFrame(
+        [
+            {
+                "event_id": eid,
+                "pdg": 11,
+                "pre_x": 0.0,
+                "pre_y": 0.0,
+                "pre_z": 0.0,
+                "pre_E": 100.0,
+                "pre_dx": 0.0,
+                "pre_dy": 0.0,
+                "pre_dz": 1.0,
+                "material": "G4_AIR",
+                "layer_id": 0,
+                "child_track_ids": [],
+                "e_sec": 0.0,
+                "step_length": 1.0,
+                "post_E": 90.0,
+                "edep": 10.0,
+                "post_dx": 0.0,
+                "post_dy": 0.0,
+                "post_dz": 1.0,
+                "post_x": 0.0,
+                "post_y": 0.0,
+                "post_z": 1.0,
+            }
+            for eid in event_ids
+        ]
+    )
+
+
+def test_event_id_offset_scales_by_file_index():
+    assert event_id_offset(0) == 0
+    assert event_id_offset(1) == EVENT_ID_FILE_STRIDE
+    assert event_id_offset(3) == 3 * EVENT_ID_FILE_STRIDE
+
+
+def test_load_event_ids_default_offset_is_zero(tmp_path):
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"event_id": [5, 6, 7]}).to_parquet(path)
+    np.testing.assert_array_equal(load_event_ids(path), [5, 6, 7])
+
+
+def test_load_event_ids_applies_offset(tmp_path):
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"event_id": [0, 1, 2]}).to_parquet(path)
+    offset = event_id_offset(1)
+    np.testing.assert_array_equal(
+        load_event_ids(path, offset=offset), [offset, offset + 1, offset + 2]
+    )
+
+
+def test_load_steps_applies_offset_to_event_id(tmp_path):
+    path = tmp_path / "a.parquet"
+    _steps_df([0, 1]).to_parquet(path)
+    offset = event_id_offset(2)
+    d = load_steps(path, offset=offset)
+    np.testing.assert_array_equal(d["event_id"], [offset, offset + 1])
+
+
+def test_iter_file_chunks_applies_offset(tmp_path):
+    path = tmp_path / "a.parquet"
+    _steps_df([0, 1, 2]).to_parquet(path)
+    offset = event_id_offset(1)
+    ids = np.concatenate([c["event_id"] for c in iter_file_chunks(path, offset=offset)])
+    np.testing.assert_array_equal(sorted(ids), [offset, offset + 1, offset + 2])
+
+
+def test_iter_cond_chunks_applies_offset(tmp_path):
+    path = tmp_path / "a.parquet"
+    _steps_df([0, 1]).to_parquet(path)
+    offset = event_id_offset(5)
+    ids = np.concatenate([c["event_id"] for c in iter_cond_chunks(path, offset=offset)])
+    np.testing.assert_array_equal(sorted(ids), [offset, offset + 1])

tests/test_setup_cache.py

@@ -225,3 +225,35 @@ def test_n_train_steps_for_split_matches_full_scan():
     result = setup_cache.n_train_steps_for_split(unique_ids, counts, train_events_arr)
 
     assert result == 20 + 40 + 50
+
+
+# ── compute_event_index_from_files: cross-file event_id offsetting ─────────
+
+
+def test_compute_event_index_from_files_offsets_colliding_ids(tmp_path):
+    """Two files that each restart event_id from 0 (one Geant4 job per file)
+    must not have their same-numbered events collapsed into one by
+    np.unique — each file's ids get shifted by a distinct offset first (see
+    giant.data.loader.event_id_offset)."""
+    path_a = tmp_path / "a.parquet"
+    path_b = tmp_path / "b.parquet"
+    pd.DataFrame({"event_id": [0, 1, 2]}).to_parquet(path_a)
+    pd.DataFrame({"event_id": [0, 1, 2]}).to_parquet(path_b)
+
+    unique_ids, counts = setup_cache.compute_event_index_from_files([path_a, path_b])
+
+    assert len(unique_ids) == 6
+    assert int(counts.sum()) == 6
+    assert np.all(counts == 1)
+
+
+def test_compute_event_index_from_files_single_file_unaffected(tmp_path):
+    """A single file's ids are offset by 0 (event_id_offset(0) == 0), so a
+    single-file load's unique ids/counts are unchanged by the offsetting."""
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"event_id": [5, 5, 7]}).to_parquet(path)
+
+    unique_ids, counts = setup_cache.compute_event_index_from_files([path])
+
+    np.testing.assert_array_equal(unique_ids, [5, 7])
+    np.testing.assert_array_equal(counts, [2, 1])