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

@@ -196,7 +196,7 @@ class Normalizer:
 
 
 class _WelfordAccumulator:
-    """Streaming mean/variance (Welford's online algorithm, batch update).
+    """Streaming mean/variance (Chan/Golub/LeVeque 1979 parallel algorithm).
 
     Use to fit a Normalizer over data that doesn't fit in memory:
         acc = _WelfordAccumulator(n_features)
@@ -211,13 +211,23 @@ class _WelfordAccumulator:
         self._M2 = np.zeros(n_features, dtype=np.float64)
 
     def update(self, X: np.ndarray) -> None:
+        # Computes the chunk's own local mean/M2 (two passes over X, no
+        # reference to the running mean) and merges it into the running
+        # totals with the O(F) Chan/Golub/LeVeque combination formula.
+        # Equivalent to the textbook single-pass streaming update (which
+        # instead re-derives two full (B, F) arrays from the running mean,
+        # before and after updating it) but ~40% cheaper here since it
+        # avoids one of those (B, F) passes and its temporary array.
         X = np.asarray(X, dtype=np.float64)
         B = X.shape[0]
+        mean_b = X.mean(0)
+        diff = X - mean_b
+        M2_b = np.einsum("ij,ij->j", diff, diff)
+
         new_n = self.n + B
-        delta = X - self._mean
-        self._mean += delta.sum(0) / new_n
-        delta2 = X - self._mean
-        self._M2 += (delta * delta2).sum(0)
+        delta = mean_b - self._mean
+        self._mean += delta * (B / new_n)
+        self._M2 += M2_b + delta * delta * (self.n * B / new_n)
         self.n = new_n
 
     def to_normalizer(self) -> "Normalizer":
@@ -276,6 +286,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 +386,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 +408,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 +636,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 +690,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 +717,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 +756,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 +784,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 +825,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_loader.py

@@ -1,7 +1,13 @@
+import numpy as np
 import pandas as pd
 import pytest
 
-from giant.data.loader import build_process_map_from_files, find_parquet_files
+from giant.data.loader import (
+    build_index_maps,
+    build_index_maps_from_files,
+    build_process_map_from_files,
+    find_parquet_files,
+)
 
 
 def _touch(path):
@@ -90,3 +96,186 @@ def test_build_process_map_from_files_spans_multiple_files(tmp_path):
     assert proc_map["phot"] == 0
     assert proc_map["eIoni"] == 1
     assert proc_map["compt"] == 2
+
+
+def test_build_process_map_from_files_tie_breaking_pins_first_seen_order(tmp_path):
+    """When two processes end up with equal total counts, ranking falls back
+    to whichever was accumulated first (`sorted(..., reverse=True)` is stable,
+    and `counts` is built in file/row-scan order) — this is implementation-
+    defined, not a documented contract, so pin it explicitly: a future
+    rewrite (e.g. a polars-based single-scan) that ties differently would
+    silently reshuffle which processes get their own expert slot across a
+    retrain, and this test is what should catch that."""
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"process": ["compt", "phot", "compt", "phot"]}).to_parquet(path)
+
+    proc_map = build_process_map_from_files([path], n_experts=3)
+
+    assert proc_map == {"compt": 0, "phot": 1}
+
+
+def test_build_process_map_from_files_tie_breaking_favors_first_scanned_file(
+    tmp_path,
+):
+    """Same total-count tie as above, but split across two files with equal
+    per-file counts — the file listed first wins the tie."""
+    path_a = tmp_path / "a.parquet"
+    path_b = tmp_path / "b.parquet"
+    pd.DataFrame({"process": ["zzz", "zzz"]}).to_parquet(path_a)
+    pd.DataFrame({"process": ["aaa", "aaa"]}).to_parquet(path_b)
+
+    forward = build_process_map_from_files([path_a, path_b], n_experts=3)
+    backward = build_process_map_from_files([path_b, path_a], n_experts=3)
+
+    assert forward == {"zzz": 0, "aaa": 1}
+    assert backward == {"aaa": 0, "zzz": 1}
+
+
+def test_build_process_map_from_files_fewer_processes_than_experts(tmp_path):
+    """When there are fewer distinct processes than expert slots, every
+    process gets its own index and the shared "other" bucket goes unused."""
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"process": ["eIoni", "phot"]}).to_parquet(path)
+
+    proc_map = build_process_map_from_files([path], n_experts=5)
+
+    assert proc_map == {"eIoni": 0, "phot": 1}
+    assert 4 not in proc_map.values()  # the "other" slot (n_experts - 1) is unused
+
+
+def test_build_process_map_from_files_n_experts_one_buckets_everything(tmp_path):
+    """n_experts=1 leaves no room for a "most frequent" slot — every process
+    (however frequent) is bucketed into the single shared index 0."""
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"process": ["eIoni"] * 10 + ["phot"] * 1}).to_parquet(path)
+
+    proc_map = build_process_map_from_files([path], n_experts=1)
+
+    assert proc_map == {"eIoni": 0, "phot": 0}
+
+
+def test_build_process_map_from_files_three_files_partial_overlap(tmp_path):
+    """Counts for a process appearing in only some of several files must sum
+    correctly, not just match the two-file case already covered above."""
+    path_a = tmp_path / "a.parquet"
+    path_b = tmp_path / "b.parquet"
+    path_c = tmp_path / "c.parquet"
+    pd.DataFrame({"process": ["eIoni"] * 2}).to_parquet(path_a)
+    pd.DataFrame({"process": ["phot"] * 3}).to_parquet(path_b)
+    pd.DataFrame({"process": ["eIoni"] * 2 + ["compt"] * 1}).to_parquet(path_c)
+
+    # eIoni: 2+2=4 > phot: 3 > compt: 1
+    proc_map = build_process_map_from_files([path_a, path_b, path_c], n_experts=3)
+
+    assert proc_map["eIoni"] == 0
+    assert proc_map["phot"] == 1
+    assert proc_map["compt"] == 2
+
+
+# ── build_index_maps (in-memory) ────────────────────────────────────────────
+
+
+def test_build_index_maps_sorts_numerically_not_lexicographically():
+    """10-digit nuclear/ion PDG codes must sort numerically — a lexicographic
+    sort would place "1000060120" before "22" since '1' < '2'."""
+    data = {
+        "pdg": np.array([22, 1000060120, 11], dtype=np.int64),
+        "material": np.array(["G4_AIR", "PbWO4", "G4_Fe"], dtype=object),
+    }
+    pdg_map, mat_map = build_index_maps(data)
+    assert list(pdg_map.keys()) == [11, 22, 1000060120]
+    assert mat_map == {"G4_AIR": 0, "G4_Fe": 1, "PbWO4": 2}
+
+
+def test_build_index_maps_dedups_repeated_values():
+    data = {
+        "pdg": np.array([11, 11, 22, 22, 22], dtype=np.int64),
+        "material": np.array(["PbWO4"] * 5, dtype=object),
+    }
+    pdg_map, mat_map = build_index_maps(data)
+    assert pdg_map == {11: 0, 22: 1}
+    assert mat_map == {"PbWO4": 0}
+
+
+def test_build_index_maps_handles_negative_pdg_codes():
+    """Antiparticle codes (negative) must sort numerically, not by magnitude."""
+    data = {
+        "pdg": np.array([-13, 11, -11, 13], dtype=np.int64),
+        "material": np.array(["X"] * 4, dtype=object),
+    }
+    pdg_map, _ = build_index_maps(data)
+    assert list(pdg_map.keys()) == [-13, -11, 11, 13]
+
+
+def test_build_index_maps_indices_are_dense_and_bijective():
+    data = {
+        "pdg": np.array([5, 1, 9, 1, 5], dtype=np.int64),
+        "material": np.array(["a", "b", "c", "a", "b"], dtype=object),
+    }
+    pdg_map, mat_map = build_index_maps(data)
+    assert sorted(pdg_map.values()) == list(range(len(pdg_map)))
+    assert sorted(mat_map.values()) == list(range(len(mat_map)))
+
+
+# ── build_index_maps_from_files ─────────────────────────────────────────────
+
+
+def test_build_index_maps_from_files_unions_and_dedups_across_files(tmp_path):
+    path_a = tmp_path / "a.parquet"
+    path_b = tmp_path / "b.parquet"
+    pd.DataFrame({"pdg": [11, 22], "material": ["G4_AIR", "PbWO4"]}).to_parquet(path_a)
+    pd.DataFrame({"pdg": [22, 2112], "material": ["PbWO4", "G4_Fe"]}).to_parquet(path_b)
+
+    pdg_map, mat_map = build_index_maps_from_files([path_a, path_b])
+
+    assert pdg_map == {11: 0, 22: 1, 2112: 2}
+    assert mat_map == {"G4_AIR": 0, "G4_Fe": 1, "PbWO4": 2}
+
+
+def test_build_index_maps_from_files_ordering_independent_of_file_order(tmp_path):
+    """Index assignment comes from the globally sorted union, not file-scan
+    order — swapping which file is scanned first must not change the map,
+    since the map is baked into a trained checkpoint's vocabulary."""
+    path_a = tmp_path / "a.parquet"
+    path_b = tmp_path / "b.parquet"
+    pd.DataFrame({"pdg": [22], "material": ["PbWO4"]}).to_parquet(path_a)
+    pd.DataFrame({"pdg": [11], "material": ["G4_AIR"]}).to_parquet(path_b)
+
+    forward = build_index_maps_from_files([path_a, path_b])
+    backward = build_index_maps_from_files([path_b, path_a])
+
+    assert forward == backward
+    assert forward == ({11: 0, 22: 1}, {"G4_AIR": 0, "PbWO4": 1})
+
+
+def test_build_index_maps_from_files_numeric_sort_for_nuclear_codes(tmp_path):
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"pdg": [22, 1000060120, 11], "material": ["X", "X", "X"]}).to_parquet(
+        path
+    )
+
+    pdg_map, _ = build_index_maps_from_files([path])
+    assert list(pdg_map.keys()) == [11, 22, 1000060120]
+
+
+def test_build_index_maps_from_files_single_file(tmp_path):
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"pdg": [11, 11, 22], "material": ["PbWO4"] * 3}).to_parquet(path)
+    pdg_map, mat_map = build_index_maps_from_files([path])
+    assert pdg_map == {11: 0, 22: 1}
+    assert mat_map == {"PbWO4": 0}
+
+
+def test_build_index_maps_from_files_matches_build_index_maps(tmp_path):
+    """Sanity-pin: the file-scanning and in-memory variants must agree on the
+    same data, since a future single-pass rewrite (pyarrow/polars) may
+    replace one but not the other."""
+    rng = np.random.default_rng(0)
+    pdg = rng.choice([11, -11, 22, 2112, 1000060120], size=200)
+    material = rng.choice(["G4_AIR", "PbWO4", "G4_Fe"], size=200)
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"pdg": pdg, "material": material}).to_parquet(path)
+
+    from_files = build_index_maps_from_files([path])
+    from_memory = build_index_maps({"pdg": pdg, "material": material})
+    assert from_files == from_memory

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,10 @@ from giant.data.transforms import (
     log_transform,
     Normalizer,
     reconstruct_post_pos,
+    sorted_membership,
     travel_direction,
+    _vectorized_map_lookup,
+    _WelfordAccumulator,
 )
 
 
@@ -440,3 +443,120 @@ 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)
+
+
+# ── _WelfordAccumulator ──────────────────────────────────────────────────────
+
+
+def test_welford_accumulator_matches_direct_mean_std_over_many_chunks():
+    rng = np.random.default_rng(5)
+    F = 4
+    chunks = [rng.standard_normal((rng.integers(1, 50), F)) * 10 + 3 for _ in range(20)]
+    full = np.concatenate(chunks, axis=0)
+
+    acc = _WelfordAccumulator(F)
+    for chunk in chunks:
+        acc.update(chunk)
+    norm = acc.to_normalizer()
+
+    assert norm.mean is not None and norm.std is not None
+    np.testing.assert_allclose(norm.mean, full.mean(axis=0), rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(norm.std, full.std(axis=0), rtol=1e-5, atol=1e-5)
+    assert acc.n == full.shape[0]
+
+
+def test_welford_accumulator_single_chunk():
+    rng = np.random.default_rng(6)
+    X = rng.standard_normal((100, 3)) * 5 - 2
+
+    acc = _WelfordAccumulator(3)
+    acc.update(X)
+    norm = acc.to_normalizer()
+
+    assert norm.mean is not None and norm.std is not None
+    np.testing.assert_allclose(norm.mean, X.mean(axis=0), rtol=1e-5)
+    np.testing.assert_allclose(norm.std, X.std(axis=0), rtol=1e-5)
+
+
+def test_welford_accumulator_matches_naive_running_mean_reference():
+    """The chunk-local-mean + Chan-merge formula must agree with the naive
+    textbook streaming update (subtract the *running* mean before and after
+    updating it) that it replaces, within float64 rounding tolerance."""
+    rng = np.random.default_rng(7)
+    F = 3
+    chunks = [rng.standard_normal((rng.integers(1, 40), F)) for _ in range(15)]
+
+    def naive_update(mean, M2, n, X):
+        X = np.asarray(X, dtype=np.float64)
+        B = X.shape[0]
+        new_n = n + B
+        delta = X - mean
+        mean = mean + delta.sum(0) / new_n
+        delta2 = X - mean
+        M2 = M2 + (delta * delta2).sum(0)
+        return mean, M2, new_n
+
+    naive_mean = np.zeros(F)
+    naive_M2 = np.zeros(F)
+    naive_n = 0
+    for chunk in chunks:
+        naive_mean, naive_M2, naive_n = naive_update(
+            naive_mean, naive_M2, naive_n, chunk
+        )
+
+    acc = _WelfordAccumulator(F)
+    for chunk in chunks:
+        acc.update(chunk)
+
+    assert acc.n == naive_n
+    np.testing.assert_allclose(acc._mean, naive_mean, rtol=1e-9, atol=1e-9)
+    np.testing.assert_allclose(acc._M2, naive_M2, rtol=1e-9, atol=1e-9)