From b944bba8fb492b9c622694d22bdce15fc61cef8f Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 30 Jul 2026 11:36:48 +0200 Subject: [PATCH] Offset event_id per file to avoid cross-file collisions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- giant/cli.py | 9 ++-- giant/data/dataset.py | 5 +- giant/data/loader.py | 45 +++++++++++++----- giant/data/setup_cache.py | 11 +++-- giant/pipeline.py | 5 +- tests/test_dataset.py | 98 ++++++++++++++++++++++++++++++++++++++- tests/test_loader.py | 87 ++++++++++++++++++++++++++++++++++ tests/test_setup_cache.py | 32 +++++++++++++ 8 files changed, 268 insertions(+), 24 deletions(-) diff --git a/giant/cli.py b/giant/cli.py index 462355b..a759861 100644 --- a/giant/cli.py +++ b/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)): diff --git a/giant/data/dataset.py b/giant/data/dataset.py index da601d9..94eff71 100644 --- a/giant/data/dataset.py +++ b/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 diff --git a/giant/data/loader.py b/giant/data/loader.py index 4ddbf7f..febf885 100644 --- a/giant/data/loader.py +++ b/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( diff --git a/giant/data/setup_cache.py b/giant/data/setup_cache.py index cf210dd..f0a707f 100644 --- a/giant/data/setup_cache.py +++ b/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 diff --git a/giant/pipeline.py b/giant/pipeline.py index b6bdfa5..c2de00e 100644 --- a/giant/pipeline.py +++ b/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 diff --git a/tests/test_dataset.py b/tests/test_dataset.py index ad7f677..23bf20f 100644 --- a/tests/test_dataset.py +++ b/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()) diff --git a/tests/test_loader.py b/tests/test_loader.py index b6af7bb..80fab28 100644 --- a/tests/test_loader.py +++ b/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]) diff --git a/tests/test_setup_cache.py b/tests/test_setup_cache.py index c0b1e48..04f5c00 100644 --- a/tests/test_setup_cache.py +++ b/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])