lars
74343d3e48
- 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>
140 lines
4.9 KiB
Python
140 lines
4.9 KiB
Python
import numpy as np
|
|
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():
|
|
rng = np.random.default_rng(42)
|
|
event_ids = rng.integers(0, 50, size=1000)
|
|
train_set, val_set = make_event_split(event_ids, val_fraction=0.2)
|
|
unique = np.unique(event_ids)
|
|
assert len(train_set) + len(val_set) == len(unique)
|
|
|
|
|
|
def test_make_event_split_no_overlap():
|
|
rng = np.random.default_rng(7)
|
|
event_ids = rng.integers(0, 50, size=1000)
|
|
train_set, val_set = make_event_split(event_ids, val_fraction=0.2)
|
|
assert train_set.isdisjoint(val_set)
|
|
|
|
|
|
def test_make_event_split_no_empty_sets():
|
|
rng = np.random.default_rng(0)
|
|
event_ids = rng.integers(0, 20, size=500)
|
|
train_set, val_set = make_event_split(event_ids, val_fraction=0.2)
|
|
assert len(train_set) > 0
|
|
assert len(val_set) > 0
|
|
|
|
|
|
def test_make_event_split_val_fraction_zero_holds_out_nothing():
|
|
"""val_fraction=0.0 is an explicit "train on everything" request and
|
|
must not be silently overridden into holding out 1 event."""
|
|
rng = np.random.default_rng(3)
|
|
event_ids = rng.integers(0, 50, size=1000)
|
|
train_set, val_set = make_event_split(event_ids, val_fraction=0.0)
|
|
assert val_set == set()
|
|
assert train_set == set(np.unique(event_ids).tolist())
|
|
|
|
|
|
def test_make_event_split_reproducible():
|
|
event_ids = np.arange(100)
|
|
a_tr, a_val = make_event_split(event_ids, val_fraction=0.1, seed=42)
|
|
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())
|