From 47a6c9db1f8844cb024de790236493fa7def2706 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 29 Jul 2026 13:45:33 +0200 Subject: [PATCH] Add regression coverage for vocab/process index-map builders MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- tests/test_loader.py | 191 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 190 insertions(+), 1 deletion(-) diff --git a/tests/test_loader.py b/tests/test_loader.py index af609dc..b6af7bb 100644 --- a/tests/test_loader.py +++ b/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