Offset event_id per file to avoid cross-file collisions
CI / Format (ruff format) (push) Successful in 25s
CI / Lint (ruff check) (push) Successful in 27s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 23s
CI / Lint (ruff check) (pull_request) Successful in 27s
CI / Tests (push) Successful in 1m14s
CI / Format (ruff format) (pull_request) Successful in 27s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 31s
CI / Tests (pull_request) Successful in 59s
CI / Format (ruff format) (push) Successful in 25s
CI / Lint (ruff check) (push) Successful in 27s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 23s
CI / Lint (ruff check) (pull_request) Successful in 27s
CI / Tests (push) Successful in 1m14s
CI / Format (ruff format) (pull_request) Successful in 27s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 31s
CI / Tests (pull_request) Successful in 59s
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>
This commit is contained in:
+33
-12
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user