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.
2026-07-08 16:15:09 +02:00
parent bac541240f
commit 0b3ece52ed
12 changed files with 385 additions and 31 deletions
@@ -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
@@ -6,6 +6,7 @@ from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
 from giant.model.network import (
    DenoisingMLP,
    EnergyRouter,
+    ProcessRouter,
    ROUTER_REGISTRY,
    RoutedDenoisingMLP,
    RoutedSecondaryDecoder,
@@ -101,6 +102,98 @@ def test_build_router_unknown_type_raises():
    raise AssertionError("expected ValueError for unknown router type")


+# ── ProcessRouter ────────────────────────────────────────────────────────────
+
+
+def test_process_router_registered():
+    assert ROUTER_REGISTRY["process"] is ProcessRouter
+
+
+def test_process_router_gate_partition_of_unity():
+    router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
+    cond_cont, cond_cat = _cond(16)
+    g = router.gate(cond_cont, cond_cat)
+    assert g.shape == (16, 4)
+    torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
+
+
+def test_process_router_top1_matches_gate_argmax():
+    router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
+    cond_cont, cond_cat = _cond(16)
+    assert torch.equal(
+        router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
+    )
+
+
+def test_process_router_balance_loss_is_nonnegative_scalar():
+    router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
+    cond_cont, cond_cat = _cond(16)
+    loss = router.balance_loss(cond_cont, cond_cat)
+    assert loss.shape == ()
+    assert loss.item() >= 0.0
+
+
+def test_process_router_classify_loss_decreases_with_training():
+    """The classifier should be able to fit an arbitrary label assignment —
+    a sanity check that gradients actually flow to the process classifier."""
+    torch.manual_seed(0)
+    router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
+    cond_cont, cond_cat = _cond(32)
+    labels = torch.randint(0, 4, (32,))
+
+    opt = torch.optim.Adam(router.parameters(), lr=0.05)
+    first = router.classify_loss(cond_cont, cond_cat, labels).item()
+    for _ in range(50):
+        opt.zero_grad()
+        loss = router.classify_loss(cond_cont, cond_cat, labels)
+        loss.backward()
+        opt.step()
+    last = loss.item()
+    assert last < first
+
+
+def test_energy_router_classify_loss_defaults_to_zero():
+    """Routers with no supervised signal (EnergyRouter) fall back to the
+    Router base class's zero-loss default."""
+    router = EnergyRouter(n_experts=4)
+    cond_cont, cond_cat = _cond(16)
+    labels = torch.randint(0, 4, (16,))
+    loss = router.classify_loss(cond_cont, cond_cat, labels)
+    assert loss.shape == ()
+    assert loss.item() == 0.0
+
+
+def test_build_router_process_type_uses_pdg_mat_vocab():
+    router = build_router("process", 4, pdg_vocab=5, mat_vocab=3, emb_dim=8)
+    assert isinstance(router, ProcessRouter)
+    assert router.pdg_emb.num_embeddings == 5
+    assert router.mat_emb.num_embeddings == 3
+
+
+def test_build_models_routed_with_process_router():
+    model_config = dict(
+        pdg_vocab=4,
+        mat_vocab=2,
+        emb_dim=16,
+        dropout=0.1,
+        k_max=K_MAX,
+        expert_hidden_dim=16,
+        expert_n_blocks=2,
+        router={
+            "enabled": True,
+            "type": "process",
+            "n_experts": 3,
+            "lambda_proc": 1.0,
+        },
+    )
+    stage1, sec_decoder = build_models(model_config)
+    assert isinstance(stage1, RoutedDenoisingMLP)
+    assert isinstance(stage1.router, ProcessRouter)
+    assert len(stage1.experts) == 3
+    assert stage1.router.pdg_emb.num_embeddings == 4
+    assert stage1.router.mat_emb.num_embeddings == 2
+
+
 # ── RoutedDenoisingMLP ───────────────────────────────────────────────────────


@@ -235,7 +235,48 @@ def test_build_features_clamps_n_sec_label_to_k_max():
    pdg_map = {11: 0}
    mat_map = {"PbWO4": 0}

-    _, _, _, n_sec, _, _, _, _ = build_features(data, pdg_map, mat_map)
+    _, _, _, n_sec, _, _, _, _, _ = build_features(data, pdg_map, mat_map)

    assert n_sec.max() <= K_MAX
    np.testing.assert_array_equal(n_sec, [0, 5, K_MAX])
+
+
+def _minimal_step_data(N: int, process: np.ndarray | None = None) -> dict:
+    rng = np.random.default_rng(0)
+    data = {
+        "pdg": np.full(N, 11, dtype=np.int32),
+        "material": np.full(N, "PbWO4", dtype=object),
+        "pre_pos": rng.standard_normal((N, 3)).astype(np.float32),
+        "pre_E": np.full(N, 10.0, dtype=np.float32),
+        "pre_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)),
+        "layer_id": np.zeros(N, dtype=np.int32),
+        "n_sec": np.zeros(N, dtype=np.int32),
+        "e_sec": np.full(N, 1.0, dtype=np.float32),
+        "step_length": np.full(N, 1.0, dtype=np.float32),
+        "post_E": np.full(N, 9.0, dtype=np.float32),
+        "edep": np.full(N, 1.0, dtype=np.float32),
+        "post_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)),
+        "post_pos": rng.standard_normal((N, 3)).astype(np.float32),
+    }
+    if process is not None:
+        data["process"] = process
+    return data
+
+
+def test_build_features_proc_idx_zero_without_proc_map():
+    data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object))
+    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
+
+    *_, proc_idx, _, _ = build_features(data, pdg_map, mat_map)
+
+    np.testing.assert_array_equal(proc_idx, [0, 0, 0])
+
+
+def test_build_features_proc_idx_looks_up_proc_map():
+    data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object))
+    pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
+    proc_map = {"compt": 0, "phot": 1, "eIoni": 2}
+
+    *_, proc_idx, _, _ = build_features(data, pdg_map, mat_map, proc_map=proc_map)
+
+    np.testing.assert_array_equal(proc_idx, [0, 1, 2])