Add PdgRouter for particle-type-based expert gating

Routes on the pre-step PDG code, which — unlike ProcessRouter's process
label — is already known at gate time (a conditioning input), so no
supervision is needed and classify_loss falls back to the zero default.
Generalizes EnergyRouter's soft-turn-on-then-Voronoi trick from a 1-D
distance to a small learned PDG embedding space: its own embedding table
maps each PDG code to a point, and n_experts learnable centers partition
that space.

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

tests/test_router.py

@@ -6,6 +6,7 @@ from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
 from giant.model.network import (
     DenoisingMLP,
     EnergyRouter,
+    PdgRouter,
     ProcessRouter,
     ROUTER_REGISTRY,
     RoutedDenoisingMLP,
@@ -102,6 +103,100 @@ def test_build_router_unknown_type_raises():
     raise AssertionError("expected ValueError for unknown router type")
 
 
+# ── PdgRouter ────────────────────────────────────────────────────────────────
+
+
+def test_pdg_router_registered():
+    assert ROUTER_REGISTRY["pdg"] is PdgRouter
+
+
+def test_pdg_router_gate_partition_of_unity():
+    router = PdgRouter(n_experts=4, pdg_vocab=3)
+    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_pdg_router_top1_matches_gate_argmax():
+    router = PdgRouter(n_experts=4, pdg_vocab=3)
+    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_pdg_router_hardens_as_temperature_shrinks():
+    """As tau -> 0 the soft gate should converge to a one-hot at the argmax."""
+    router = PdgRouter(n_experts=4, pdg_vocab=3, temperature=1e-4)
+    cond_cont, cond_cat = _cond(16)
+    g = router.gate(cond_cont, cond_cat)
+    top1 = router.top1(cond_cont, cond_cat)
+    onehot = torch.nn.functional.one_hot(top1, num_classes=4).float()
+    torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0)
+
+
+def test_pdg_router_balance_loss_is_nonnegative_scalar():
+    router = PdgRouter(n_experts=4, pdg_vocab=3)
+    cond_cont, cond_cat = _cond(16)
+    loss = router.balance_loss(cond_cont, cond_cat)
+    assert loss.shape == ()
+    assert loss.item() >= 0.0
+
+
+def test_pdg_router_classify_loss_defaults_to_zero():
+    """PDG is already known at gate time (unlike ProcessRouter's process
+    label), so no supervision is needed — falls back to Router's default."""
+    router = PdgRouter(n_experts=4, pdg_vocab=3)
+    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_pdg_router_only_reads_pdg_column():
+    """Gate must depend on cond_cat[:, 0] (pdg) only, not cond_cont or material."""
+    router = PdgRouter(n_experts=4, pdg_vocab=3)
+    cond_cont, cond_cat = _cond(16)
+    g_before = router.gate(cond_cont, cond_cat)
+
+    cond_cont_perturbed = torch.randn_like(cond_cont)
+    cond_cat_diff_mat = cond_cat.clone()
+    cond_cat_diff_mat[:, 1] = (cond_cat_diff_mat[:, 1] + 1) % 2
+    g_after = router.gate(cond_cont_perturbed, cond_cat_diff_mat)
+
+    torch.testing.assert_close(g_before, g_after, atol=1e-6, rtol=0)
+
+
+def test_build_router_pdg_type_uses_pdg_vocab():
+    router = build_router("pdg", 4, pdg_vocab=5, mat_vocab=3, emb_dim=8)
+    assert isinstance(router, PdgRouter)
+    assert router.pdg_emb.num_embeddings == 5
+
+
+def test_build_models_routed_with_pdg_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": "pdg",
+            "n_experts": 3,
+        },
+    )
+    stage1, sec_decoder = build_models(model_config)
+    assert isinstance(stage1, RoutedDenoisingMLP)
+    assert isinstance(stage1.router, PdgRouter)
+    assert len(stage1.experts) == 3
+    assert stage1.router.pdg_emb.num_embeddings == 4
+
+
 # ── ProcessRouter ────────────────────────────────────────────────────────────