Add regression coverage for vocab/process index-map builders

build_index_maps, build_index_maps_from_files, and (mostly)
build_process_map_from_files had no test pinning their sort order,
tie-breaking, or cross-file union behavior — all load-bearing for a
trained checkpoint's vocabulary, and all at risk of silently changing
under a future single-pass (pyarrow/polars) rewrite of the setup-stage
scan. Add tests for numeric-vs-lexicographic PDG sort (nuclear/ion
codes), negative PDG codes, dedup/bijective indices, file-order
independence, and process-map tie-breaking/boundary conditions
(n_experts=1, fewer processes than experts, 3-file partial overlap).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

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