Merge pull request 'Offset event_id per file to avoid cross-file collisions' (#23) from fix/event-id-file-offset into master

Reviewed-on: #23

giant/cli.py

@@ -26,6 +26,7 @@ from giant.constants import (
     ROLLOUT_COORD_VALUE,
 )
 from giant.data.loader import (
+    event_id_offset,
     find_parquet_files,
     iter_file_chunks,
     iter_cond_chunks,
@@ -911,8 +912,8 @@ def predict(
     buffer: dict[str, np.ndarray] | None = None
 
     bar = tqdm(total=total_rows, desc="predict", unit="row", dynamic_ncols=True)
-    for path in files:
-        for chunk in chunk_iter(path):
+    for i, path in enumerate(files):
+        for chunk in chunk_iter(path, offset=event_id_offset(i)):
             N_in = len(chunk["event_id"])
 
             pdg_mask = np.array([int(p) in pdg_map for p in chunk["pdg"]])
@@ -961,8 +962,8 @@ def _seed_from_data(files: list[Path], n_events: int | None) -> dict[str, np.nda
     """
     best_E: dict[int, float] = {}
     best: dict[int, tuple] = {}
-    for path in files:
-        for chunk in iter_cond_chunks(path):
+    for file_idx, path in enumerate(files):
+        for chunk in iter_cond_chunks(path, offset=event_id_offset(file_idx)):
             ev = chunk["event_id"]
             pe = chunk["pre_E"]
             for i in range(len(ev)):

giant/data/dataset.py

@@ -6,7 +6,7 @@ import numpy as np
 import torch
 from torch.utils.data import IterableDataset
 
-from giant.data.loader import iter_file_chunks
+from giant.data.loader import event_id_offset, iter_file_chunks
 from giant.data.transforms import Normalizer, build_features, sorted_membership
 
 
@@ -65,6 +65,7 @@ class StreamingStepsDataset(IterableDataset):
         sec_phys_normalizer: Normalizer | None = None,
     ) -> None:
         self.files = list(files)
+        self._offsets = {path: event_id_offset(i) for i, path in enumerate(self.files)}
         self.split_events = split_events
         self._events_arr = np.array(sorted(split_events))
         self.pdg_map = pdg_map
@@ -97,7 +98,7 @@ class StreamingStepsDataset(IterableDataset):
         buf_n = 0
 
         for path in files:
-            for chunk in iter_file_chunks(path):
+            for chunk in iter_file_chunks(path, offset=self._offsets[path]):
                 mask = sorted_membership(chunk["event_id"], self._events_arr)
                 if not mask.any():
                     continue

giant/data/loader.py

@@ -12,6 +12,20 @@ import pyarrow.parquet as pq
 # dataset tree is moved or copied elsewhere intact.
 MANIFEST_SUFFIX = ".manifest"
 
+# Each input parquet file is a separate Geant4 job converted 1:1 from its own
+# ROOT file (scripts/steps_to_parquet.py), and a job's event_id numbering
+# always restarts from 0 — so when multiple files are loaded together (a
+# directory or .manifest), raw event_id values collide across files even
+# though they refer to unrelated events. Every per-file event_id column gets
+# offset by its file's index in the (deterministically ordered) files list
+# so ids stay globally unique across a multi-file load; the stride is far
+# larger than any realistic per-file event count.
+EVENT_ID_FILE_STRIDE = 1_000_000
+
+
+def event_id_offset(file_index: int) -> int:
+    return file_index * EVENT_ID_FILE_STRIDE
+
 
 def _read_manifest(path: Path) -> list[Path]:
     files = []
@@ -78,13 +92,13 @@ def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndar
     return out
 
 
-def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
+def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
     from giant.constants import K_MAX
 
     has_sec_lists = "sec_E_list" in df.columns
 
     d: dict[str, np.ndarray] = {
-        "event_id": df["event_id"].to_numpy(),
+        "event_id": df["event_id"].to_numpy().astype(np.int64) + offset,
         "pdg": df["pdg"].to_numpy(dtype=np.int32),
         "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
         "pre_E": df["pre_E"].to_numpy(dtype=np.float32),
@@ -121,20 +135,23 @@ def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
     return d
 
 
-def load_steps(path: str | Path) -> dict[str, np.ndarray]:
-    return _df_to_dict(pd.read_parquet(path))
+def load_steps(path: str | Path, offset: int = 0) -> dict[str, np.ndarray]:
+    return _df_to_dict(pd.read_parquet(path), offset=offset)
 
 
-def load_event_ids(path: str | Path) -> np.ndarray:
+def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
     """Read only the event_id column — cheap scan for split assignment."""
-    return pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
+    ids = pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
+    return ids.astype(np.int64) + offset
 
 
-def iter_file_chunks(path: str | Path) -> Iterator[dict[str, np.ndarray]]:
+def iter_file_chunks(
+    path: str | Path, offset: int = 0
+) -> Iterator[dict[str, np.ndarray]]:
     """Yield one parquet row-group at a time so a large file never fully loads."""
     pf = pq.ParquetFile(path)
     for i in range(pf.num_row_groups):
-        yield _df_to_dict(pf.read_row_group(i).to_pandas())
+        yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset)
 
 
 _COND_COLS = [
@@ -154,9 +171,9 @@ _COND_COLS = [
 ]
 
 
-def _cond_df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
+def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
     return {
-        "event_id": df["event_id"].to_numpy(),
+        "event_id": df["event_id"].to_numpy().astype(np.int64) + offset,
         "pdg": df["pdg"].to_numpy(dtype=np.int32),
         "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
         "pre_E": df["pre_E"].to_numpy(dtype=np.float32),
@@ -168,11 +185,15 @@ def _cond_df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
     }
 
 
-def iter_cond_chunks(path: str | Path) -> Iterator[dict[str, np.ndarray]]:
+def iter_cond_chunks(
+    path: str | Path, offset: int = 0
+) -> Iterator[dict[str, np.ndarray]]:
     """Yield conditioning-only row-groups (no post-step columns read from disk)."""
     pf = pq.ParquetFile(path)
     for i in range(pf.num_row_groups):
-        yield _cond_df_to_dict(pf.read_row_group(i, columns=_COND_COLS).to_pandas())
+        yield _cond_df_to_dict(
+            pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset
+        )
 
 
 def build_index_maps(

giant/data/setup_cache.py

@@ -22,13 +22,16 @@ import numpy as np
 
 from giant import config
 from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM
-from giant.data.loader import load_event_ids
+from giant.data.loader import event_id_offset, load_event_ids
 from giant.data.transforms import Normalizer, sorted_membership
 
 # Bump manually on a change to the data-encoding semantics (e.g. a future
 # energy_simplex_encode bugfix) that doesn't also move one of _DIMS below —
 # a dims change already hard-invalidates on its own.
-_CACHE_FORMAT_VERSION = 1
+# v2: event_id is now offset per-file (see loader.event_id_offset) to avoid
+# cross-file collisions, so a v1 sidecar's event_index/normalizers were
+# computed against collided ids and must not be reused.
+_CACHE_FORMAT_VERSION = 2
 
 _DIMS = {
     "COND_DIM": COND_DIM,
@@ -263,7 +266,9 @@ def compute_event_index_from_files(files: list[Path]) -> tuple[np.ndarray, np.nd
     """Unique event ids + per-event row (step) counts, across all `files`."""
     if not files:
         return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64)
-    all_ids = np.concatenate([load_event_ids(f) for f in files])
+    all_ids = np.concatenate(
+        [load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)]
+    )
     unique_ids, counts = np.unique(all_ids, return_counts=True)
     return unique_ids, counts
 

giant/pipeline.py

@@ -16,6 +16,7 @@ from giant.constants import (
 )
 from giant.data import setup_cache
 from giant.data.loader import (
+    event_id_offset,
     find_parquet_files,
     iter_file_chunks,
     build_index_maps_from_files,
@@ -167,8 +168,8 @@ def run_setup_stage(
         energy_sampler = (
             _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
         )
-        for path in files:
-            for chunk in iter_file_chunks(path):
+        for i, path in enumerate(files):
+            for chunk in iter_file_chunks(path, offset=event_id_offset(i)):
                 mask = sorted_membership(chunk["event_id"], events_arr)
                 if not mask.any():
                     continue

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])