Add ProcessRouter for physics-process-based expert gating
Routes on the physics process (Compton, phot, brems, ...) that ends a step, supervised by a small classifier since process is a post-step outcome unobservable at gate time. Threads a process label end-to-end through the data pipeline (loader, build_features, dataset batches, training loss/checkpointing) alongside the existing EnergyRouter.
This commit is contained in:
+35
-1
@@ -1,6 +1,7 @@
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from giant.data.loader import find_parquet_files
|
||||
from giant.data.loader import build_process_map_from_files, find_parquet_files
|
||||
|
||||
|
||||
def _touch(path):
|
||||
@@ -58,3 +59,36 @@ def test_manifest_with_no_entries_raises(tmp_path):
|
||||
|
||||
with pytest.raises(FileNotFoundError):
|
||||
find_parquet_files(manifest)
|
||||
|
||||
|
||||
def test_build_process_map_from_files_keeps_most_frequent(tmp_path):
|
||||
"""process counts: eIoni=5, phot=3, compt=2, Rayl=1 — with n_experts=3, only
|
||||
the top 2 (eIoni, phot) get their own index; compt/Rayl share the "other"
|
||||
(last) index."""
|
||||
process = (
|
||||
["eIoni"] * 5 + ["phot"] * 3 + ["compt"] * 2 + ["Rayl"] * 1
|
||||
)
|
||||
path = tmp_path / "shard-000.parquet"
|
||||
pd.DataFrame({"process": process}).to_parquet(path)
|
||||
|
||||
proc_map = build_process_map_from_files([path], n_experts=3)
|
||||
|
||||
assert proc_map["eIoni"] == 0
|
||||
assert proc_map["phot"] == 1
|
||||
assert proc_map["compt"] == 2
|
||||
assert proc_map["Rayl"] == 2
|
||||
assert set(proc_map.values()) <= {0, 1, 2}
|
||||
|
||||
|
||||
def test_build_process_map_from_files_spans_multiple_files(tmp_path):
|
||||
path_a = tmp_path / "a.parquet"
|
||||
path_b = tmp_path / "b.parquet"
|
||||
pd.DataFrame({"process": ["eIoni"] * 3 + ["phot"] * 1}).to_parquet(path_a)
|
||||
pd.DataFrame({"process": ["phot"] * 4 + ["compt"] * 1}).to_parquet(path_b)
|
||||
|
||||
# phot: 1+4=5 total > eIoni: 3 > compt: 1
|
||||
proc_map = build_process_map_from_files([path_a, path_b], n_experts=3)
|
||||
|
||||
assert proc_map["phot"] == 0
|
||||
assert proc_map["eIoni"] == 1
|
||||
assert proc_map["compt"] == 2
|
||||
|
||||
Reference in New Issue
Block a user