Offset event_id per file to avoid cross-file collisions #23

Merged
lars merged 1 commits from fix/event-id-file-offset into master 2026-07-30 11:44:31 +02:00
8 changed files with 268 additions and 24 deletions
+5 -4
View File
@@ -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)):
+3 -2
View File
@@ -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
+33 -12
View File
@@ -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(
+8 -3
View File
@@ -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
+3 -2
View File
@@ -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
+97 -1
View File
@@ -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())
+87
View File
@@ -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])
+32
View File
@@ -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])