From 09e4c765c7d5480aa7c2389a84478cbb2945b257 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 29 Jul 2026 13:34:32 +0200 Subject: [PATCH] Speed up giant train's setup stage Fits normalizers over multi-hundred-million-row datasets, so the setup pass's per-row Python overhead compounds fast: encode_secondaries recomputed an O(K) prefix sum from scratch on every one of its 15 stick-breaking iterations, np.isin re-sorted the full train-event-id array on every chunk, and pdg/material/process index lookups ran a Python dict lookup per row. The normalizer-fit pass also computed encode_secondaries's stick-logit and direction-rotation blocks in full even though it only ever reads the mass/charge columns. Replace the prefix-sum recompute with a single np.cumsum, add a sorted_membership helper (searchsorted-based) in place of np.isin at both the setup-pass and per-epoch call sites, vectorize the index lookups via _vectorized_map_lookup, and add an opt-in phys_only path so the setup pass skips the stick-breaking/rotation work it discards anyway. All four changes are output-identical performance refactors, backed by new unit tests plus the existing suite. Co-Authored-By: Claude Sonnet 5 --- giant/data/dataset.py | 4 +- giant/data/transforms.py | 124 +++++++++++++++++++++++++++++---------- giant/pipeline.py | 18 ++++-- tests/test_phase2.py | 98 +++++++++++++++++++++++++++++++ tests/test_transforms.py | 51 ++++++++++++++++ 5 files changed, 256 insertions(+), 39 deletions(-) diff --git a/giant/data/dataset.py b/giant/data/dataset.py index c1d07d2..da601d9 100644 --- a/giant/data/dataset.py +++ b/giant/data/dataset.py @@ -7,7 +7,7 @@ import torch from torch.utils.data import IterableDataset from giant.data.loader import iter_file_chunks -from giant.data.transforms import Normalizer, build_features +from giant.data.transforms import Normalizer, build_features, sorted_membership def make_event_split( @@ -98,7 +98,7 @@ class StreamingStepsDataset(IterableDataset): for path in files: for chunk in iter_file_chunks(path): - mask = np.isin(chunk["event_id"], self._events_arr) + mask = sorted_membership(chunk["event_id"], self._events_arr) if not mask.any(): continue chunk = {k: v[mask] for k, v in chunk.items()} diff --git a/giant/data/transforms.py b/giant/data/transforms.py index a4a5222..867b718 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -276,6 +276,44 @@ class _ReservoirSampler: return self._reservoir.astype(np.float32) +def sorted_membership(values: np.ndarray, sorted_arr: np.ndarray) -> np.ndarray: + """Boolean membership of `values` (any order) in `sorted_arr` (ascending, unique). + + Equivalent to `np.isin(values, sorted_arr)`, but `np.isin`'s default path + sorts both inputs on every call — costly when `sorted_arr` is a large, + already-sorted array (e.g. all train-split event ids) reused across many + chunks. This does one `searchsorted` per call instead. `values` need not + be sorted; `sorted_arr` must be ascending and duplicate-free. + """ + values = np.asarray(values) + if sorted_arr.size == 0: + return np.zeros(values.shape, dtype=bool) + idx = np.searchsorted(sorted_arr, values) + idx = np.clip(idx, 0, len(sorted_arr) - 1) + return sorted_arr[idx] == values + + +def _vectorized_map_lookup(values: np.ndarray, mapping: dict) -> np.ndarray: + """Vectorized equivalent of `np.array([mapping[v] for v in values], dtype=np.int64)`. + + Replaces a per-element Python dict lookup with one `searchsorted` call. + Raises `KeyError` if any value in `values` isn't a key of `mapping`, + matching the dict-comprehension it replaces (never silently misassigns). + """ + keys = np.asarray(list(mapping.keys())) + vals = np.asarray(list(mapping.values()), dtype=np.int64) + order = np.argsort(keys, kind="stable") + keys_sorted, vals_sorted = keys[order], vals[order] + values = np.asarray(values) + pos = np.searchsorted(keys_sorted, values) + pos = np.clip(pos, 0, len(keys_sorted) - 1) + found = keys_sorted[pos] == values + if not found.all(): + missing = np.unique(values[~found]) + raise KeyError(f"value(s) not in mapping: {missing[:10].tolist()}") + return vals_sorted[pos] + + def travel_direction(pre_pos: np.ndarray, post_pos: np.ndarray) -> np.ndarray: """World-frame unit vector pointing from pre_pos to post_pos. @@ -338,6 +376,7 @@ def encode_secondaries( e_sec: np.ndarray, pre_dir: np.ndarray, sec_pdg_list: np.ndarray | None = None, + phys_only: bool = False, ) -> np.ndarray: """Encode per-secondary attributes into continuous per-slot targets. @@ -359,38 +398,52 @@ def encode_secondaries( `sec_pdg_list` is optional so callers that only need the continuous stick/dir block (e.g. inference-time re-encoding) can omit it; omitting it zero-fills the last two columns, matching the padding-slot convention. + + `phys_only=True` skips the stick-breaking and direction-rotation blocks + (zero-filling them instead) and computes only log_mass/charge — for + callers (normalizer fitting) that discard the other four columns anyway, + so computing them would be wasted work repeated over the whole dataset. """ N, K = sec_E_list.shape e_sec = np.asarray(e_sec, dtype=np.float64) - stick_logits = np.zeros((N, K), dtype=np.float32) - for i in range(K): - if i == 0: - remaining = np.maximum(e_sec, _EPS) - else: - remaining = np.maximum(e_sec - sec_E_list[:, :i].sum(axis=1), _EPS) - f = np.clip(sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS) - logit = np.log(f / (1.0 - f)).astype(np.float32) - # Last valid slot: give it the full remaining budget - is_last = sec_valid[:, i] & ~( - sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool) - ) - logit = np.where(is_last, _STICK_LOGIT_CLIP, logit) - logit = np.where( - sec_valid[:, i], np.clip(logit, -_STICK_LOGIT_CLIP, _STICK_LOGIT_CLIP), 0.0 - ) - stick_logits[:, i] = logit.astype(np.float32) - - # 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) - for i in range(K): - # Only rotate valid slots; leave padded slots as (0,0,1) or whatever. - valid_mask = sec_valid[:, i] - if valid_mask.any(): - dir_local[valid_mask, i] = local_frame_rotation( - pre_dir[valid_mask], sec_dir_list[valid_mask, i] + if phys_only: + stick_logits = np.zeros((N, K), dtype=np.float32) + dir_local = np.zeros((N, K, 3), dtype=np.float32) + else: + cumsum = np.cumsum(sec_E_list.astype(np.float64), axis=1) + stick_logits = np.zeros((N, K), dtype=np.float32) + for i in range(K): + if i == 0: + remaining = np.maximum(e_sec, _EPS) + else: + remaining = np.maximum(e_sec - cumsum[:, i - 1], _EPS) + f = np.clip( + sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS ) + logit = np.log(f / (1.0 - f)).astype(np.float32) + # Last valid slot: give it the full remaining budget + is_last = sec_valid[:, i] & ~( + sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool) + ) + logit = np.where(is_last, _STICK_LOGIT_CLIP, logit) + logit = np.where( + sec_valid[:, i], + np.clip(logit, -_STICK_LOGIT_CLIP, _STICK_LOGIT_CLIP), + 0.0, + ) + stick_logits[:, i] = logit.astype(np.float32) + + # 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) + for i in range(K): + # Only rotate valid slots; leave padded slots as (0,0,1) or whatever. + valid_mask = sec_valid[:, i] + if valid_mask.any(): + dir_local[valid_mask, i] = local_frame_rotation( + pre_dir[valid_mask], sec_dir_list[valid_mask, i] + ) if sec_pdg_list is not None: from giant.particles import particle_phys_array @@ -573,8 +626,8 @@ def build_cond_features( [cond_cont, _physical_cond_columns(data, conditioning)] ).astype(np.float32) - pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64) - mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64) + pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map) + mat_idx = _vectorized_map_lookup(data["material"], mat_map) cond_cat = np.column_stack([pdg_idx, mat_idx]) if cond_normalizer is not None: @@ -627,6 +680,7 @@ def build_features( proc_map: dict[str, int] | None = None, require_secondaries: bool = False, conditioning: str = "embedding", + sec_phys_only: bool = False, ) -> tuple[ np.ndarray, np.ndarray, @@ -653,6 +707,11 @@ def build_features( per-secondary list columns are absent (a mis-converted file that would otherwise silently zero all Stage-2 targets). Training paths set this; Stage-1-only callers (e.g. `giant predict`) leave it False. + + sec_phys_only: passed straight through to `encode_secondaries` — skips + the stick-breaking/direction-rotation blocks of `sec_cont` (zero-filled + instead) for callers (normalizer fitting) that only read + `sec_cont[:, :, 4:6]` and would otherwise discard that work. """ from giant.constants import K_MAX @@ -687,8 +746,8 @@ def build_features( [cond_cont, _physical_cond_columns(data, conditioning)] ).astype(np.float32) # (N, COND_DIM=15) - pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64) - mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64) + pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map) + mat_idx = _vectorized_map_lookup(data["material"], mat_map) cond_cat = np.column_stack([pdg_idx, mat_idx]) # (N, 2) n_sec_raw = data["n_sec"].astype( @@ -715,6 +774,7 @@ def build_features( data["e_sec"], data["pre_dir"], sec_pdg_list=sec_pdg_list, + phys_only=sec_phys_only, ) # (N, K_MAX, 6) else: # Guard against silently training Stage 2 on zeroed targets: if any step @@ -755,7 +815,7 @@ def build_features( process = data.get("process") if proc_map is not None and process is not None: - proc_idx = np.array([proc_map[str(p)] for p in process], dtype=np.int64) + proc_idx = _vectorized_map_lookup(process, proc_map) else: proc_idx = np.zeros(len(cond_cat), dtype=np.int64) diff --git a/giant/pipeline.py b/giant/pipeline.py index 63a6f98..15452ed 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -20,7 +20,12 @@ from giant.data.loader import ( build_index_maps_from_files, build_process_map_from_files, ) -from giant.data.transforms import build_features, _WelfordAccumulator, _ReservoirSampler +from giant.data.transforms import ( + build_features, + _WelfordAccumulator, + _ReservoirSampler, + sorted_membership, +) from giant.data.dataset import make_event_split, StreamingStepsDataset from giant.model.network import build_models, build_critics from giant.train import train as run_training @@ -50,11 +55,9 @@ def run_train_job( all_event_ids, val_fraction=t["val_fraction"] ) events_arr = np.array(sorted(train_events)) - n_train_steps = int(np.isin(all_event_ids, events_arr).sum()) - total_train_batches = n_train_steps // t["batch_size"] echo( f" {len(all_event_ids):,} steps | " - f"{len(train_events)} train events (~{n_train_steps:,} steps, ~{total_train_batches:,} batches) | " + f"{len(train_events)} train events | " f"{len(val_events)} val events" ) @@ -95,11 +98,13 @@ def run_train_job( energy_sampler = ( _ReservoirSampler(capacity=100_000) if energy_router_active else None ) + n_train_steps = 0 for path in files: for chunk in iter_file_chunks(path): - mask = np.isin(chunk["event_id"], events_arr) + mask = sorted_membership(chunk["event_id"], events_arr) if not mask.any(): continue + n_train_steps += int(mask.sum()) chunk_tr = {k: v[mask] for k, v in chunk.items()} cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _ = build_features( chunk_tr, @@ -108,6 +113,7 @@ def run_train_job( proc_map=proc_map, require_secondaries=True, conditioning=conditioning, + sec_phys_only=True, ) cond_acc.update(cond_cont) tgt_acc.update(target_s1) @@ -120,6 +126,8 @@ def run_train_job( cond_norm = cond_acc.to_normalizer() tgt_norm = tgt_acc.to_normalizer() sec_phys_norm = sec_phys_acc.to_normalizer() + total_train_batches = n_train_steps // t["batch_size"] + echo(f" ~{n_train_steps:,} train steps, ~{total_train_batches:,} batches") if energy_sampler is not None and energy_sampler.n_seen > 0: assert cond_norm.mean is not None and cond_norm.std is not None diff --git a/tests/test_phase2.py b/tests/test_phase2.py index c2857d1..c8e7fd1 100644 --- a/tests/test_phase2.py +++ b/tests/test_phase2.py @@ -231,6 +231,53 @@ def test_encode_secondaries_energy_conservation(): assert np.isfinite(sec_cont).all() +def test_encode_secondaries_stick_logits_match_naive_reference(): + """Cumsum-based remaining-budget computation must match a naive + per-row, per-slot Python reference (no cumsum) within float tolerance.""" + from giant.data.transforms import encode_secondaries, _EPS, _STICK_LOGIT_CLIP + + rng = np.random.default_rng(11) + N = 25 + n_sec = rng.integers(1, K_MAX + 1, size=N) + e_sec = rng.uniform(0.1, 10.0, size=N).astype(np.float32) + + sec_E_list = np.zeros((N, K_MAX), dtype=np.float32) + sec_dir_list = np.zeros((N, K_MAX, 3), dtype=np.float32) + sec_dir_list[:, :, 2] = 1.0 + sec_valid = np.zeros((N, K_MAX), dtype=bool) + for i in range(N): + k = n_sec[i] + energies = rng.dirichlet(np.ones(k)) * e_sec[i] + energies = np.sort(energies)[::-1] + sec_E_list[i, :k] = energies.astype(np.float32) + sec_valid[i, :k] = True + + pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32) + + sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir) + stick_logits = sec_cont[:, :, 0] + + # Naive reference: recompute the remaining budget from scratch each slot, + # exactly what the pre-cumsum implementation did. + expected = np.zeros((N, K_MAX), dtype=np.float64) + for row in range(N): + for i in range(K_MAX): + if not sec_valid[row, i]: + continue + remaining = max(float(e_sec[row]) - float(sec_E_list[row, :i].sum()), _EPS) + f = min(max(float(sec_E_list[row, i]) / remaining, _EPS), 1.0 - _EPS) + logit = np.log(f / (1.0 - f)) + is_last = not (i + 1 < K_MAX and sec_valid[row, i + 1]) + if is_last: + logit = _STICK_LOGIT_CLIP + logit = np.clip(logit, -_STICK_LOGIT_CLIP, _STICK_LOGIT_CLIP) + expected[row, i] = logit + + np.testing.assert_allclose( + stick_logits[sec_valid], expected[sec_valid], rtol=1e-4, atol=1e-4 + ) + + def test_encode_secondaries_direction_encoding(): """Local-frame secondary directions should be unit vectors for valid slots.""" from giant.data.transforms import encode_secondaries @@ -272,6 +319,57 @@ def test_encode_secondaries_physical_columns_without_pdg_list(): np.testing.assert_allclose(sec_cont[:, :, 4:6], 0.0) +def test_encode_secondaries_phys_only_matches_full_and_zero_fills_rest(): + """phys_only=True must reproduce the mass/charge columns exactly and + zero-fill the stick-logit/direction columns it skips computing.""" + from giant.data.transforms import encode_secondaries + + rng = np.random.default_rng(3) + N = 30 + n_sec = rng.integers(1, 5, size=N) + e_sec = rng.uniform(0.1, 10.0, size=N).astype(np.float32) + + sec_E_list = np.zeros((N, K_MAX), dtype=np.float32) + sec_dir_list = rng.standard_normal((N, K_MAX, 3)).astype(np.float32) + norms = np.linalg.norm(sec_dir_list, axis=-1, keepdims=True) + sec_dir_list /= np.where(norms > 0, norms, 1.0) + sec_pdg_list = np.zeros((N, K_MAX), dtype=np.int64) + sec_valid = np.zeros((N, K_MAX), dtype=bool) + for i in range(N): + k = n_sec[i] + energies = rng.dirichlet(np.ones(k)) * e_sec[i] + energies = np.sort(energies)[::-1] + sec_E_list[i, :k] = energies.astype(np.float32) + sec_valid[i, :k] = True + sec_pdg_list[i, :k] = 11 # electron — resolvable by giant.particles + + pre_dir = rng.standard_normal((N, 3)).astype(np.float32) + pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True) + + full = encode_secondaries( + sec_E_list, + sec_dir_list, + sec_valid, + e_sec, + pre_dir, + sec_pdg_list=sec_pdg_list, + phys_only=False, + ) + phys_only = encode_secondaries( + sec_E_list, + sec_dir_list, + sec_valid, + e_sec, + pre_dir, + sec_pdg_list=sec_pdg_list, + phys_only=True, + ) + + np.testing.assert_array_equal(phys_only[:, :, 4:6], full[:, :, 4:6]) + np.testing.assert_array_equal(phys_only[:, :, 0], np.zeros((N, K_MAX))) + np.testing.assert_array_equal(phys_only[:, :, 1:4], np.zeros((N, K_MAX, 3))) + + def test_encode_secondaries_physical_columns_match_ground_truth_pdg(): """log_mass/charge for a valid slot match giant.particles for that PDG.""" from giant.data.transforms import encode_secondaries, log_transform diff --git a/tests/test_transforms.py b/tests/test_transforms.py index c9bc7ac..1dfaae1 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -12,7 +12,9 @@ from giant.data.transforms import ( log_transform, Normalizer, reconstruct_post_pos, + sorted_membership, travel_direction, + _vectorized_map_lookup, ) @@ -440,3 +442,52 @@ def test_build_cond_features_rejects_legacy_normalizer_in_physical_mode( cond_normalizer=legacy_norm, conditioning="physical", ) + + +# ── sorted_membership / _vectorized_map_lookup ────────────────────────────── + + +def test_sorted_membership_matches_np_isin(): + rng = np.random.default_rng(0) + sorted_arr = np.unique(rng.integers(0, 10_000, size=500)) + values = rng.integers(-100, 10_100, size=2_000) # some in, some out of range + # values deliberately not sorted + rng.shuffle(values) + + result = sorted_membership(values, sorted_arr) + expected = np.isin(values, sorted_arr) + np.testing.assert_array_equal(result, expected) + + +def test_sorted_membership_empty_sorted_arr(): + values = np.array([1, 2, 3]) + sorted_arr = np.array([], dtype=np.int64) + result = sorted_membership(values, sorted_arr) + np.testing.assert_array_equal(result, np.zeros(3, dtype=bool)) + + +def test_vectorized_map_lookup_matches_dict_comprehension_int_keys(): + rng = np.random.default_rng(1) + keys = np.unique(rng.integers(-1000, 1000, size=200)) + mapping = {int(k): i for i, k in enumerate(keys)} + values = rng.choice(keys, size=500) + + result = _vectorized_map_lookup(values, mapping) + expected = np.array([mapping[int(v)] for v in values], dtype=np.int64) + np.testing.assert_array_equal(result, expected) + + +def test_vectorized_map_lookup_matches_dict_comprehension_str_keys(): + mapping = {"PbWO4": 0, "G4_AIR": 1, "G4_Fe": 2} + values = np.array(["G4_Fe", "PbWO4", "G4_AIR", "PbWO4"], dtype=object) + + result = _vectorized_map_lookup(values, mapping) + expected = np.array([mapping[str(v)] for v in values], dtype=np.int64) + np.testing.assert_array_equal(result, expected) + + +def test_vectorized_map_lookup_raises_keyerror_on_missing_value(): + mapping = {1: 0, 2: 1} + values = np.array([1, 2, 3]) + with pytest.raises(KeyError): + _vectorized_map_lookup(values, mapping)