Add data-integrity guards against silent NaN/Inf propagation and races
- log_transform / _validate_unit_pre_dir now raise on non-finite input instead of letting a NaN row silently poison the persisted normalizer cache (norm < 1e-6 was always False for NaN, so the existing guard never caught it). - encode_secondaries warns when a row's secondary energies cumulatively exceed e_sec, instead of silently saturating the overflowing slot's stick-breaking logit via the _EPS floor. - EVENT_ID_FILE_STRIDE overflow now raises instead of silently colliding two files' event ids together (reintroducing train/val leakage). - make_event_split(val_fraction=0.0) now actually holds out nothing, instead of always forcing at least 1 validation event. - setup_cache.save() is now serialized with a flock, since two concurrent writers (a real scenario on this repo's shared portal/condor machines) could otherwise race and silently drop one writer's freshly-computed cache section. - Documented (no behavior change) the pre_dir ≈ -ẑ antipodal rotation singularity in _rodrigues_axis, which is real but inherent to any single-valued local-frame convention. Each fix has a regression test. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
@@ -19,7 +19,10 @@ def make_event_split(
|
||||
rng = np.random.default_rng(seed)
|
||||
unique = np.unique(all_event_ids)
|
||||
rng.shuffle(unique)
|
||||
n_val = max(1, int(len(unique) * val_fraction))
|
||||
# max(1, ...) only applies when a validation split was actually
|
||||
# requested — val_fraction=0.0 is an explicit "train on everything"
|
||||
# request and must not be silently overridden into holding out 1 event.
|
||||
n_val = max(1, int(len(unique) * val_fraction)) if val_fraction > 0 else 0
|
||||
val_set = set(unique[:n_val].tolist())
|
||||
train_set = set(unique[n_val:].tolist())
|
||||
return train_set, val_set
|
||||
|
||||
+23
-3
@@ -27,6 +27,26 @@ def event_id_offset(file_index: int) -> int:
|
||||
return file_index * EVENT_ID_FILE_STRIDE
|
||||
|
||||
|
||||
def _offset_event_id(raw_ids: np.ndarray, offset: int) -> np.ndarray:
|
||||
"""Add this file's `event_id_offset`, after checking the raw ids fit in one stride block.
|
||||
|
||||
Without this check, a file whose own raw event_id numbering reaches
|
||||
`EVENT_ID_FILE_STRIDE` (an unusually large job, or non-contiguous
|
||||
numbering) would silently collide into the next file's offset block,
|
||||
merging unrelated events across files — reintroducing exactly the
|
||||
train/val event leakage this offset scheme exists to prevent.
|
||||
"""
|
||||
raw_ids = np.asarray(raw_ids, dtype=np.int64)
|
||||
if raw_ids.size and int(raw_ids.max()) >= EVENT_ID_FILE_STRIDE:
|
||||
raise ValueError(
|
||||
f"event_id {int(raw_ids.max())} >= EVENT_ID_FILE_STRIDE "
|
||||
f"({EVENT_ID_FILE_STRIDE}) — this file has a larger event_id "
|
||||
"than the per-file offset scheme can support without colliding "
|
||||
"with the next file's id block."
|
||||
)
|
||||
return raw_ids + offset
|
||||
|
||||
|
||||
def _read_manifest(path: Path) -> list[Path]:
|
||||
files = []
|
||||
for line in path.read_text().splitlines():
|
||||
@@ -98,7 +118,7 @@ def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
|
||||
has_sec_lists = "sec_E_list" in df.columns
|
||||
|
||||
d: dict[str, np.ndarray] = {
|
||||
"event_id": df["event_id"].to_numpy().astype(np.int64) + offset,
|
||||
"event_id": _offset_event_id(df["event_id"].to_numpy(), 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),
|
||||
@@ -142,7 +162,7 @@ def load_steps(path: str | Path, offset: int = 0) -> dict[str, 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."""
|
||||
ids = pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
|
||||
return ids.astype(np.int64) + offset
|
||||
return _offset_event_id(ids, offset)
|
||||
|
||||
|
||||
def iter_file_chunks(
|
||||
@@ -173,7 +193,7 @@ _COND_COLS = [
|
||||
|
||||
def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
|
||||
return {
|
||||
"event_id": df["event_id"].to_numpy().astype(np.int64) + offset,
|
||||
"event_id": _offset_event_id(df["event_id"].to_numpy(), 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),
|
||||
|
||||
@@ -13,6 +13,7 @@ before reuse — see `load`/`save`.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import fcntl
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
@@ -270,15 +271,32 @@ def save(
|
||||
Best-effort: any OSError (permission denied on a read-only mount, disk
|
||||
full, ...) is caught, echoed as a warning, and swallowed — a failure to
|
||||
cache must never fail training.
|
||||
|
||||
The load-merge-write is serialized with an exclusive flock on a sidecar
|
||||
lockfile: `os.replace` alone only guarantees the *file* is never
|
||||
corrupt, not that concurrent writers don't race. Without the lock, two
|
||||
concurrent `giant train`/condor jobs against the same `data` path (this
|
||||
repo's shared-portal/condor usage makes that a real scenario, not just
|
||||
theoretical) could both `load()` the same base state, merge their own
|
||||
`sections` in independently, and whichever `os.replace()` lands last
|
||||
silently discards the other's freshly-computed section.
|
||||
"""
|
||||
path = sidecar_path(data)
|
||||
lock_path = path.parent / f".{path.name}.lock"
|
||||
tmp = path.parent / f".{path.name}.tmp.{os.getpid()}"
|
||||
try:
|
||||
base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty(files)
|
||||
merged = base.merge(sections)
|
||||
payload = json.dumps(merged.to_json(), separators=(",", ":"))
|
||||
tmp.write_text(payload)
|
||||
os.replace(tmp, path)
|
||||
with open(lock_path, "a") as lock_file:
|
||||
fcntl.flock(lock_file, fcntl.LOCK_EX)
|
||||
try:
|
||||
base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty(
|
||||
files
|
||||
)
|
||||
merged = base.merge(sections)
|
||||
payload = json.dumps(merged.to_json(), separators=(",", ":"))
|
||||
tmp.write_text(payload)
|
||||
os.replace(tmp, path)
|
||||
finally:
|
||||
fcntl.flock(lock_file, fcntl.LOCK_UN)
|
||||
except OSError as exc:
|
||||
echo(
|
||||
f"setup cache: could not write {path} ({exc}) — continuing without caching"
|
||||
|
||||
@@ -12,7 +12,17 @@ _SIMPLEX_FLOOR = 1e-5
|
||||
|
||||
|
||||
def log_transform(x: np.ndarray, eps: float = _EPS) -> np.ndarray:
|
||||
return np.log(np.asarray(x, dtype=np.float32) + eps)
|
||||
x = np.asarray(x, dtype=np.float32)
|
||||
y = np.log(x + eps)
|
||||
if not np.all(np.isfinite(y)):
|
||||
bad = int(np.sum(~np.isfinite(y)))
|
||||
raise ValueError(
|
||||
f"log_transform: {bad} value(s) produced non-finite output (input "
|
||||
f"< -eps={eps:g}, or already NaN/Inf); every quantity this is "
|
||||
"applied to should be non-negative, so this indicates upstream "
|
||||
"data corruption rather than expected float noise."
|
||||
)
|
||||
return y
|
||||
|
||||
|
||||
def inv_log_transform(y: np.ndarray, eps: float = _EPS) -> np.ndarray:
|
||||
@@ -110,8 +120,22 @@ def _rodrigues_axis(pre_dir: np.ndarray) -> np.ndarray:
|
||||
[pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1
|
||||
)
|
||||
axis_norm = np.linalg.norm(axis, axis=1, keepdims=True) # (N,1)
|
||||
# Replace zero-norm axes with x̂ (the Rodrigues terms that involve the axis
|
||||
# are multiplied by sin_t≈0 and (1-cos_t)≈0, so the choice is irrelevant).
|
||||
# axis_norm ~ 0 happens at BOTH poles: pre_dir ~ +ẑ (forward) and
|
||||
# pre_dir ~ -ẑ (near-exact backscatter) — ‖pre_dir × ẑ‖ = sin(angle to
|
||||
# ẑ) vanishes at both. The "choice is irrelevant" claim below only holds
|
||||
# at +ẑ, where sin_t~0 AND (1-cos_t)~0 so every axis-dependent Rodrigues
|
||||
# term vanishes. At -ẑ, sin_t~0 but (1-cos_t)~2 — not negligible — so
|
||||
# snapping to a fixed x̂ there is a genuine (if physically rare)
|
||||
# modeling choice, not a no-op: it picks one representative out of an
|
||||
# inherently ambiguous family of 180°-about-any-transverse-axis
|
||||
# rotations (no single-valued frame convention can be continuous through
|
||||
# this antipode — same obstruction as a sphere's tangent frame having no
|
||||
# continuous choice at a pole). x̂ is still fine to use — it's a fixed,
|
||||
# self-consistent convention that `local_frame_rotation`/
|
||||
# `inv_local_frame_rotation` (same threshold) round-trip correctly
|
||||
# through — but steps whose pre_dir falls in this tiny near-backscatter
|
||||
# cone get a discontinuous "roll" relative to their non-degenerate
|
||||
# neighbors, injecting a small amount of label noise there.
|
||||
safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm)
|
||||
return np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm)
|
||||
|
||||
@@ -134,8 +158,20 @@ def _validate_unit_pre_dir(pre_dir: np.ndarray) -> np.ndarray:
|
||||
drift is corrected silently; a near-zero-norm row has no well-defined
|
||||
direction, so it's raised loudly instead of producing a meaningless
|
||||
rotation (previously it fell through to an arbitrary axis with no error).
|
||||
|
||||
NaN/Inf rows are also raised on explicitly: `norm < 1e-6` is False for a
|
||||
NaN norm, so without this check a non-finite row would silently pass
|
||||
through and poison everything downstream (e.g. the persisted normalizer
|
||||
stats in `setup_cache`, if the row is swept into a Welford accumulator).
|
||||
"""
|
||||
pre_dir = np.asarray(pre_dir, dtype=np.float32)
|
||||
if not np.all(np.isfinite(pre_dir)):
|
||||
bad = int(np.sum(~np.all(np.isfinite(pre_dir), axis=1)))
|
||||
raise ValueError(
|
||||
f"pre_dir has {bad} row(s) with non-finite (NaN/Inf) components; "
|
||||
"local/inv_local_frame_rotation require a well-defined incoming "
|
||||
"direction for every row."
|
||||
)
|
||||
norm = np.linalg.norm(pre_dir, axis=1, keepdims=True)
|
||||
if np.any(norm < 1e-6):
|
||||
raise ValueError(
|
||||
@@ -423,11 +459,21 @@ def encode_secondaries(
|
||||
else:
|
||||
cumsum = np.cumsum(sec_E_list.astype(np.float64), axis=1)
|
||||
stick_logits = np.zeros((N, K), dtype=np.float32)
|
||||
# A valid slot whose cumulative secondary energy so far exceeds
|
||||
# e_sec by more than float noise means sec_E_list sums to more than
|
||||
# e_sec — a real upstream data mismatch, not something to paper
|
||||
# over. Flagged once after the loop rather than let `remaining`'s
|
||||
# np.maximum(..., _EPS) floor silently absorb it by saturating that
|
||||
# slot's stick-breaking logit with no signal that anything was off.
|
||||
_SHORTFALL_TOL = 1e-3
|
||||
shortfall_flagged = np.zeros(N, dtype=bool)
|
||||
for i in range(K):
|
||||
if i == 0:
|
||||
remaining = np.maximum(e_sec, _EPS)
|
||||
remaining_raw = e_sec
|
||||
else:
|
||||
remaining = np.maximum(e_sec - cumsum[:, i - 1], _EPS)
|
||||
remaining_raw = e_sec - cumsum[:, i - 1]
|
||||
shortfall_flagged |= sec_valid[:, i] & (remaining_raw < -_SHORTFALL_TOL)
|
||||
remaining = np.maximum(remaining_raw, _EPS)
|
||||
f = np.clip(
|
||||
sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS
|
||||
)
|
||||
@@ -444,6 +490,17 @@ def encode_secondaries(
|
||||
)
|
||||
stick_logits[:, i] = logit.astype(np.float32)
|
||||
|
||||
if shortfall_flagged.any():
|
||||
n = int(shortfall_flagged.sum())
|
||||
warnings.warn(
|
||||
f"encode_secondaries: {n}/{N} row(s) have sec_E_list summing "
|
||||
"to more than e_sec (beyond float noise) — the overflowing "
|
||||
"slot(s)' stick-breaking logit was saturated instead of "
|
||||
"reflecting a real fraction; check upstream secondary "
|
||||
"energy accounting for these rows.",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
# Rotate each slot's direction into the local frame of the primary.
|
||||
# pre_dir is broadcast across all K slots.
|
||||
dir_local = np.zeros((N, K, 3), dtype=np.float32)
|
||||
|
||||
Reference in New Issue
Block a user