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>

giant/config.py

@@ -34,10 +34,10 @@ DEFAULT_CONFIG: dict = {
             "n_experts": 4,
             "expert_hidden_dim": 128,
             "expert_n_blocks": 3,
-            "temperature": 0.5,  # energy-router kwarg
-            "learn_centers": True,  # energy-router kwarg
+            "temperature": 0.5,  # energy/pdg-router kwarg
+            "learn_centers": True,  # energy/pdg-router kwarg
             "lambda_balance": 0.0,  # optional load-balance aux loss weight
-            "emb_dim": 8,  # process-router kwarg: its own pdg/mat embedding width
+            "emb_dim": 8,  # process/pdg-router kwarg: own pdg(/mat) embedding width
             "hidden_dim": 64,  # process-router kwarg: its classifier's hidden width
             "lambda_proc": 0.0,  # process-router kwarg: supervised process-CE weight
             # (0.0 still trains a working router — the gate gets gradient

giant/model/network.py

@@ -344,6 +344,48 @@ class EnergyRouter(Router):
         return torch.softmax(-d2 / self.temperature, dim=-1)
 
 
+@register_router("pdg")
+class PdgRouter(Router):
+    """Soft turn-on gate over a learned PDG embedding.
+
+    Unlike ProcessRouter's process label, PDG code is already known at
+    pre-step time (it's a conditioning input, `cond_cat[:, 0]`), so no
+    supervision is needed — `classify_loss` falls back to the Router base
+    class's zero-loss default, same as EnergyRouter. Because PDG is
+    categorical rather than a scalar, this generalizes EnergyRouter's
+    soft-turn-on-then-Voronoi trick from a 1-D distance to a distance in a
+    small embedding space: its own embedding table (kept separate from the
+    trunk's ConditionEncoder, same reasoning as ProcessRouter's own
+    pdg/mat embeddings) maps each PDG code to a point, and `n_experts`
+    learnable (or fixed) centers partition that space.
+    `gate(pdg) = softmax_i(-||emb(pdg) - c_i||^2 / tau)`.
+    """
+
+    def __init__(
+        self,
+        n_experts: int,
+        pdg_vocab: int,
+        emb_dim: int = 8,
+        temperature: float = 0.5,
+        learn_centers: bool = True,
+    ) -> None:
+        super().__init__(n_experts)
+        self.temperature = temperature
+        self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
+        centers = torch.randn(n_experts, emb_dim) * 0.1
+        if learn_centers:
+            self.centers = nn.Parameter(centers)
+        else:
+            self.register_buffer("centers", centers)
+
+    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        e = self.pdg_emb(cond_cat[:, 0])  # (B, emb_dim)
+        d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(
+            -1
+        )  # (B, n_experts)
+        return torch.softmax(-d2 / self.temperature, dim=-1)
+
+
 @register_router("process")
 class ProcessRouter(Router):
     """Routes on the physics process expected to end the step.

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 ────────────────────────────────────────────────────────────