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 <noreply@anthropic.com>

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()}

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)
 

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

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

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)