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>

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)