giant/giant/model/network.py

import inspect
import math
import re
from collections.abc import Sequence

import torch
import torch.nn as nn
import torch.nn.functional as F

from giant.constants import (
    COND_DIM,
    COND_DIM_BASE,
    EMB_DIM,
    K_MAX,
    MATERIAL_PHYS_DIM,
    PARTICLE_PHYS_DIM,
    SEC_DIM,
    X_DIM,
)


class SinusoidalEmbedding(nn.Module):
    def __init__(self, dim: int) -> None:
        super().__init__()
        assert dim % 2 == 0, "dim must be even"
        half = dim // 2
        freqs = torch.exp(
            -math.log(10000)
            * torch.arange(half, dtype=torch.float32)
            / max(half - 1, 1)
        )
        self.register_buffer("freqs", freqs)

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        t = t.reshape(-1, 1).float()
        args = t * self.freqs.unsqueeze(0)  # (B, half)
        return torch.cat([args.sin(), args.cos()], dim=-1)  # (B, dim)


class ConditionEncoder(nn.Module):
    """Fuses continuous conditioning with particle/material identity.

    Two mutually exclusive ways to turn (pdg, material) identity into the
    two `emb_dim`-wide vectors concatenated with the base continuous
    conditioning before the fusion MLP:
      - "embedding": a learned `nn.Embedding` lookup table per axis, indexed
        by `cond_cat`'s dense training-vocab index. Memorizes the training
        menu; the original Phase-2 design.
      - "physical": a small MLP per axis, mapping the axis's raw physical
        properties (already present in `cond_cont[:, COND_DIM_BASE:]` — see
        giant.data.transforms.build_features) to an `emb_dim`-wide vector —
        a drop-in replacement for the embedding lookup, computable for any
        PDG code / material name rather than only ones seen in training.
    Both modes produce the same `in_dim = COND_DIM_BASE + 2*emb_dim` for the
    fusion MLP, so only how the two vectors are produced differs.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        cont_dim: int = COND_DIM,
        emb_dim: int = 16,
        out_dim: int = 128,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        if conditioning not in ("embedding", "physical"):
            raise ValueError(f"unknown conditioning mode {conditioning!r}")
        self.conditioning = conditioning
        if conditioning == "embedding":
            self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
            self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
        else:
            self.particle_mlp = nn.Sequential(
                nn.Linear(PARTICLE_PHYS_DIM, emb_dim),
                nn.SiLU(),
                nn.Linear(emb_dim, emb_dim),
            )
            self.material_mlp = nn.Sequential(
                nn.Linear(MATERIAL_PHYS_DIM, emb_dim),
                nn.SiLU(),
                nn.Linear(emb_dim, emb_dim),
            )
        in_dim = COND_DIM_BASE + 2 * emb_dim
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, out_dim),
            nn.SiLU(),
            nn.Linear(out_dim, out_dim),
        )

    def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        if self.conditioning == "embedding":
            pdg_e = self.pdg_emb(cond_cat[:, 0])
            mat_e = self.mat_emb(cond_cat[:, 1])
        else:
            particle_phys = cond_cont[
                :, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM
            ]
            material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :]
            pdg_e = self.particle_mlp(particle_phys)
            mat_e = self.material_mlp(material_phys)
        x = torch.cat([cond_cont[:, :COND_DIM_BASE], pdg_e, mat_e], dim=-1)
        return self.mlp(x)


class ResBlock(nn.Module):
    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.1) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.linear1 = nn.Linear(dim, dim)
        self.cond_proj = nn.Linear(cond_dim, dim, bias=False)
        self.act = nn.SiLU()
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        h = self.norm(x)
        h = self.linear1(h) + self.cond_proj(cond)
        h = self.act(h)
        h = self.dropout(h)
        h = self.linear2(h)
        return x + h


class DenoisingMLP(nn.Module):
    """Stage-1 model: predicts the 9D primary post-step vector field + n_sec logits.

    The n_sec head runs on the condition encoding only (no diffusion noise),
    so it can be called at inference time independently via `predict_n_sec`.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        hidden_dim: int = 256,
        n_blocks: int = 6,
        emb_dim: int = 16,
        time_dim: int = 64,
        cond_out_dim: int = 128,
        x_dim: int = X_DIM,
        dropout: float = 0.1,
        k_max: int = K_MAX,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        self.time_emb = SinusoidalEmbedding(time_dim)
        self.cond_enc = ConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            out_dim=cond_out_dim,
            conditioning=conditioning,
        )
        merged_cond_dim = time_dim + cond_out_dim
        self.input_proj = nn.Linear(x_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_proj = nn.Linear(hidden_dim, x_dim)
        # Predicts n_sec as classification over {0, 1, ..., k_max}.
        # Applied to the condition encoding (not the diffused latent).
        self.n_sec_head = nn.Sequential(
            nn.Linear(cond_out_dim, hidden_dim // 2),
            nn.SiLU(),
            nn.Linear(hidden_dim // 2, k_max + 1),
        )

    def forward(
        self,
        x_t: torch.Tensor,
        t: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        t_emb = self.time_emb(t)  # (B, time_dim)
        c_emb = self.cond_enc(cond_cont, cond_cat)  # (B, cond_out_dim)
        cond = torch.cat([t_emb, c_emb], dim=-1)
        x = self.input_proj(x_t)
        for block in self.blocks:
            x = block(x, cond)
        return self.out_proj(x)

    def predict_n_sec(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        """Return n_sec logits (B, K_MAX+1) from conditioning alone."""
        c_emb = self.cond_enc(cond_cont, cond_cat)
        return self.n_sec_head(c_emb)


class SecondaryConditionEncoder(nn.Module):
    """Encodes pre-step conditioning + Stage-1 output for the secondary decoder."""

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        emb_dim: int = 16,
        cond_out_dim: int = 128,
        stage1_dim: int = X_DIM,
        stage1_proj_dim: int = 64,
        out_dim: int = 128,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        self.base = ConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            out_dim=cond_out_dim,
            conditioning=conditioning,
        )
        self.stage1_proj = nn.Linear(stage1_dim, stage1_proj_dim)
        fused_dim = cond_out_dim + stage1_proj_dim
        self.fuse = nn.Sequential(
            nn.Linear(fused_dim, out_dim),
            nn.SiLU(),
        )

    def forward(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
    ) -> torch.Tensor:
        base = self.base(cond_cont, cond_cat)  # (B, cond_out_dim)
        s1 = self.stage1_proj(stage1_out).tanh()  # (B, stage1_proj_dim)
        return self.fuse(torch.cat([base, s1], dim=-1))  # (B, out_dim)


class SecondaryDecoder(nn.Module):
    """Stage-2 model: predicts vector field over K_MAX secondary slots simultaneously.

    Each slot encodes (stick_break_logit, local_dir_3D, log_mass, charge) for
    one secondary ordered by descending energy — mass/charge are the
    secondary's predicted physical identity, regressed directly against real
    physics targets (see giant.data.transforms.encode_secondaries), used
    as-is with no snapping to a discrete PDG code. Padded slots are masked
    from loss.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        hidden_dim: int = 256,
        n_blocks: int = 6,
        emb_dim: int = 16,
        time_dim: int = 64,
        cond_out_dim: int = 128,
        stage1_proj_dim: int = 64,
        sec_dim: int = SEC_DIM,
        dropout: float = 0.1,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        self.time_emb = SinusoidalEmbedding(time_dim)
        self.cond_enc = SecondaryConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            cond_out_dim=cond_out_dim,
            stage1_proj_dim=stage1_proj_dim,
            out_dim=cond_out_dim,
            conditioning=conditioning,
        )
        merged_cond_dim = time_dim + cond_out_dim
        self.input_proj = nn.Linear(sec_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_proj = nn.Linear(hidden_dim, sec_dim)

    def forward(
        self,
        x_t: torch.Tensor,
        t: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
    ) -> torch.Tensor:
        t_emb = self.time_emb(t)
        c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out)
        cond = torch.cat([t_emb, c_emb], dim=-1)
        x = self.input_proj(x_t)
        for block in self.blocks:
            x = block(x, cond)
        return self.out_proj(x)


class WGANGenerator(nn.Module):
    """Stage-1 WGAN-GP generator: single forward pass, no diffusion/flow time.

    Same `ConditionEncoder` + `ResBlock` trunk as `DenoisingMLP`, but the
    input is a noise vector `z` (not a diffused/interpolated `x_t`) and the
    ResBlocks condition on the condition encoding alone (no time embedding to
    concatenate) — see `giant/model/wgan.py` for the adversarial losses, and
    `giant.sample.sample_wgan` for single-pass sampling.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        hidden_dim: int = 256,
        n_blocks: int = 6,
        emb_dim: int = 16,
        cond_out_dim: int = 128,
        x_dim: int = X_DIM,
        noise_dim: int = 64,
        dropout: float = 0.1,
        k_max: int = K_MAX,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        self.noise_dim = noise_dim
        self.cond_enc = ConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            out_dim=cond_out_dim,
            conditioning=conditioning,
        )
        self.input_proj = nn.Linear(noise_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_proj = nn.Linear(hidden_dim, x_dim)
        self.n_sec_head = nn.Sequential(
            nn.Linear(cond_out_dim, hidden_dim // 2),
            nn.SiLU(),
            nn.Linear(hidden_dim // 2, k_max + 1),
        )

    def forward(
        self,
        z: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        cond = self.cond_enc(cond_cont, cond_cat)
        x = self.input_proj(z)
        for block in self.blocks:
            x = block(x, cond)
        return self.out_proj(x)

    def predict_n_sec(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        """Return n_sec logits (B, K_MAX+1) from conditioning alone."""
        c_emb = self.cond_enc(cond_cont, cond_cat)
        return self.n_sec_head(c_emb)


class Critic(nn.Module):
    """Stage-1 WGAN-GP critic: scalar realism score, own `ConditionEncoder`.

    Kept structurally parallel to `WGANGenerator` (own condition encoder —
    separate weights from the generator's, standard GAN practice) but has no
    n_sec head: n_sec is never adversarial, it stays a plain classifier on
    the generator side.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        hidden_dim: int = 256,
        n_blocks: int = 6,
        emb_dim: int = 16,
        cond_out_dim: int = 128,
        x_dim: int = X_DIM,
        dropout: float = 0.1,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        self.cond_enc = ConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            out_dim=cond_out_dim,
            conditioning=conditioning,
        )
        self.input_proj = nn.Linear(x_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_norm = nn.LayerNorm(hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, 1)

    def forward(
        self,
        x: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        cond = self.cond_enc(cond_cont, cond_cat)
        h = self.input_proj(x)
        for block in self.blocks:
            h = block(h, cond)
        return self.out_proj(self.out_norm(h)).squeeze(-1)


class WGANSecondaryGenerator(nn.Module):
    """Stage-2 WGAN-GP generator: single forward pass over all K_MAX slots.

    Mirrors `SecondaryDecoder` minus the time embedding, the same way
    `WGANGenerator` mirrors `DenoisingMLP` — takes noise `z` instead of `x_t`,
    conditions on `SecondaryConditionEncoder`'s output alone.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        hidden_dim: int = 256,
        n_blocks: int = 6,
        emb_dim: int = 16,
        cond_out_dim: int = 128,
        stage1_proj_dim: int = 64,
        sec_dim: int = SEC_DIM,
        noise_dim: int = 64,
        dropout: float = 0.1,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        self.noise_dim = noise_dim
        self.cond_enc = SecondaryConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            cond_out_dim=cond_out_dim,
            stage1_proj_dim=stage1_proj_dim,
            out_dim=cond_out_dim,
            conditioning=conditioning,
        )
        self.input_proj = nn.Linear(noise_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_proj = nn.Linear(hidden_dim, sec_dim)

    def forward(
        self,
        z: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
    ) -> torch.Tensor:
        cond = self.cond_enc(cond_cont, cond_cat, stage1_out)
        x = self.input_proj(z)
        for block in self.blocks:
            x = block(x, cond)
        return self.out_proj(x)


class SecondaryCritic(nn.Module):
    """Stage-2 WGAN-GP critic: scalar realism score over the flattened 90D slots."""

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        hidden_dim: int = 256,
        n_blocks: int = 6,
        emb_dim: int = 16,
        cond_out_dim: int = 128,
        stage1_proj_dim: int = 64,
        sec_dim: int = SEC_DIM,
        dropout: float = 0.1,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        self.cond_enc = SecondaryConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            cond_out_dim=cond_out_dim,
            stage1_proj_dim=stage1_proj_dim,
            out_dim=cond_out_dim,
            conditioning=conditioning,
        )
        self.input_proj = nn.Linear(sec_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_norm = nn.LayerNorm(hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, 1)

    def forward(
        self,
        x: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
    ) -> torch.Tensor:
        cond = self.cond_enc(cond_cont, cond_cat, stage1_out)
        h = self.input_proj(x)
        for block in self.blocks:
            h = block(h, cond)
        return self.out_proj(self.out_norm(h)).squeeze(-1)


class Router(nn.Module):
    """Contract for a pluggable mixture-of-experts routing axis.

    Subclasses implement `gate` (soft partition-of-unity weights over
    experts, used in train mode for a fully differentiable mixture);
    `top1` and `balance_loss` have working defaults so a new routing axis
    is usually a one-method add. See `ROUTER_REGISTRY` / `build_router`.
    """

    def __init__(self, n_experts: int) -> None:
        super().__init__()
        self.n_experts = n_experts

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        """(B, n_experts) soft weights, rows summing to 1."""
        raise NotImplementedError

    def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        """(B,) hard expert index, used for eval-time grouped dispatch."""
        return self.gate(cond_cont, cond_cat).argmax(dim=-1)

    def balance_loss(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
    ) -> torch.Tensor:
        """Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017)."""
        importance = self.gate(cond_cont, cond_cat).sum(dim=0)  # (n_experts,)
        return (importance.std() / (importance.mean() + 1e-8)) ** 2

    def classify_loss(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
    ) -> torch.Tensor:
        """Optional supervised auxiliary loss shaping the router's own belief.

        Default: none (a scalar 0), for routers like EnergyRouter that read a
        quantity directly off cond_cont/cond_cat and need no label. Routers
        gating on an unobservable pre-step quantity (e.g. ProcessRouter,
        which predicts the physics process that will end the step) override
        this to supervise their internal classifier against the true label.
        """
        return torch.zeros((), device=cond_cont.device)

    def entropy_loss(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
    ) -> torch.Tensor:
        """Optional auxiliary loss rewarding sharper (lower-entropy) routing.

        Reuses `gate_stats`'s `norm_entropy` (already in [0, 1], 1.0 =
        uniform/collapsed) directly as the loss, so minimizing it pushes
        every router's gate toward decisiveness. A generic base-class
        default — works for any Router via gate_stats, no per-subclass
        override needed. Off by default (see `lambda_entropy` in
        giant.train): bounded width/temperature (EnergyRouter's
        `learn_width`/`learn_temperature`) is the primary defense against
        gate collapse; this is a secondary, use-with-caution lever, since
        indiscriminately penalizing entropy can also suppress legitimate
        soft ambiguity near a router's own decision boundary.
        """
        norm_entropy, _ = self.gate_stats(cond_cont, cond_cat)
        return norm_entropy

    def gate_stats(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Diagnostics for catching a router that fails to specialize.

        Returns `(norm_entropy, importance)`:
          - `norm_entropy`: scalar, the batch-mean of each row's gate entropy
            divided by `log(n_experts)`, in [0, 1] and comparable across
            routers with different `n_experts` (1.0 = uniform/collapsed
            gating, 0.0 = fully hard routing).
          - `importance`: (n_experts,) tensor, `gate(...).sum(dim=0)` — the
            *unnormalized* per-expert weight mass for this batch. Callers
            wanting a global utilization share across many batches must sum
            this across batches first and normalize once at the end;
            averaging per-batch shares instead would treat every batch as
            equally important regardless of size and understate a
            rarely-but-fully-used expert.
        """
        gate = self.gate(cond_cont, cond_cat)  # (B, n_experts)
        row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1)  # (B,)
        norm_entropy = row_entropy.mean() / math.log(self.n_experts)
        importance = gate.sum(dim=0)  # (n_experts,)
        return norm_entropy, importance


ROUTER_REGISTRY: dict[str, type[Router]] = {}


def register_router(name: str):
    def decorator(cls: type[Router]) -> type[Router]:
        ROUTER_REGISTRY[name] = cls
        return cls

    return decorator


def build_router(name: str, n_experts: int, **kwargs) -> Router:
    """Factory: look up a `Router` subclass by name from the registry.

    Every registered router type is fed the same `model.router` config
    dict; kwargs not declared by that type's constructor are silently
    dropped, so per-type hyperparameters (e.g. EnergyRouter's
    `temperature`) can coexist in one config without special-casing.
    """
    if name not in ROUTER_REGISTRY:
        raise ValueError(
            f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}"
        )
    cls = ROUTER_REGISTRY[name]
    accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"}
    filtered = {k: v for k, v in kwargs.items() if k in accepted}
    return cls(n_experts=n_experts, **filtered)


def _bounded_interp(raw: torch.Tensor, lo: float, hi: float) -> torch.Tensor:
    """Sigmoid interpolation into `[lo, hi]` — smooth, always-positive-gradient
    bound (unlike `clamp`, which zeroes gradient past the boundary) used for
    EnergyRouter's `learn_width`/`learn_temperature` modes."""
    return lo + (hi - lo) * torch.sigmoid(raw)


def _inverse_bounded_interp(value: float, lo: float, hi: float) -> float:
    """Inverse of `_bounded_interp`, used once at construction to warm-start
    `raw` so `_bounded_interp(raw, lo, hi) == value` — lets `learn_width`/
    `learn_temperature` start out exactly reproducing the fixed-`temperature`
    gate before any training moves them."""
    p = min(max((value - lo) / (hi - lo), 1e-6), 1 - 1e-6)
    return math.log(p / (1 - p))


@register_router("energy")
class EnergyRouter(Router):
    """Soft turn-on gate over normalized pre-step log-energy.

    Reads `cond_cont[:, energy_idx]` (ignores cond_cat). Learnable (or
    fixed) 1-D centers. By default initialized spread evenly across
    [-2, 2] — an assumed-uniform z-normalized energy range that may not
    match the true (often skewed) distribution and can leave experts
    overlapping instead of partitioning the range; pass `centers_init` to
    seed them from data (e.g. energy quantiles) instead.
    `gate(e) = softmax_i(-(e - c_i)^2 / tau)`, differentiable in e; as
    tau -> 0 this hardens to nearest-center (Voronoi) selection, which is
    exactly what `top1` uses at eval.

    `temperature` is normally a single fixed scalar shared by every expert.
    Two mutually exclusive optional modes generalize it:
      - `learn_width`: each expert gets its own learnable width, so
        `gate(e) = softmax_i(-(e - c_i)^2 / width_i)` — experts can learn
        independently how much of the energy axis they cover.
      - `learn_temperature`: the single shared `temperature` itself becomes
        learnable (still one scalar for every expert).
    Both parameterize their raw learnable value through a sigmoid bounded
    into `[width_min_ratio, width_max_ratio] * temperature` (see
    `_bounded_interp`), warm-started so the initial effective width/
    temperature exactly equals `temperature` — enabling either mode is a
    no-op at init. The bound is deliberately not raw `softplus`/`exp`
    (unbounded above): an unbounded width lets one expert's width run away
    to infinity, making its logit `-d2/width -> 0` almost everywhere so it
    wins nearly every row regardless of true distance to its center — the
    same "experts overlap instead of partitioning" failure this whole
    router design is trying to avoid, just via a new mechanism. See
    `Router.entropy_loss`/`giant.train`'s `lambda_entropy` for a secondary,
    optional guard against all experts' widths co-inflating together
    (which bounding caps but doesn't forbid, and which the load-balance
    loss alone can't see since usage shares stay even throughout).
    """

    def __init__(
        self,
        n_experts: int = 4,
        temperature: float = 0.5,
        learn_centers: bool = True,
        energy_idx: int = 3,
        centers_init: Sequence[float] | None = None,
        learn_width: bool = False,
        learn_temperature: bool = False,
        width_min_ratio: float = 0.1,
        width_max_ratio: float = 10.0,
    ) -> None:
        super().__init__(n_experts)
        if learn_width and learn_temperature:
            raise ValueError("learn_width and learn_temperature are mutually exclusive")
        self.temperature = temperature
        self.energy_idx = energy_idx
        self.learn_width = learn_width
        self.learn_temperature = learn_temperature
        if learn_width or learn_temperature:
            if not (width_min_ratio < 1.0 < width_max_ratio):
                raise ValueError(
                    f"width_min_ratio ({width_min_ratio}) and width_max_ratio "
                    f"({width_max_ratio}) must bracket 1.0"
                )
            self._width_lo = width_min_ratio * temperature
            self._width_hi = width_max_ratio * temperature
            raw0 = _inverse_bounded_interp(temperature, self._width_lo, self._width_hi)
            if learn_width:
                self.raw_width = nn.Parameter(torch.full((n_experts,), raw0))
            else:
                self.raw_temperature = nn.Parameter(torch.tensor(raw0))
        if centers_init is None:
            centers = torch.linspace(-2.0, 2.0, n_experts)
        else:
            if len(centers_init) != n_experts:
                raise ValueError(
                    f"centers_init has {len(centers_init)} values, "
                    f"expected n_experts={n_experts}"
                )
            centers = torch.tensor(list(centers_init), dtype=torch.float32)
        if learn_centers:
            self.centers = nn.Parameter(centers)
        else:
            self.register_buffer("centers", centers)

    def effective_width(self) -> torch.Tensor | float:
        """Softmax denominator used by `gate()`: a fixed scalar `temperature`
        (default), a per-expert `(n_experts,)` bounded width (`learn_width`),
        or a single bounded learnable scalar (`learn_temperature`)."""
        if self.learn_width:
            return _bounded_interp(self.raw_width, self._width_lo, self._width_hi)
        if self.learn_temperature:
            return _bounded_interp(self.raw_temperature, self._width_lo, self._width_hi)
        return self.temperature

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        e = cond_cont[:, self.energy_idx].unsqueeze(-1)  # (B, 1)
        d2 = (e - self.centers.unsqueeze(0)) ** 2  # (B, n_experts)
        return torch.softmax(-d2 / self.effective_width(), 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.

    Unlike EnergyRouter (which reads a quantity that's already known at
    pre-step time), the process — Compton, photoelectric, brems, ... — is a
    *post-step outcome*: it can't be read off cond_cont/cond_cat directly.
    Instead this router runs a small classifier over pre-step conditioning
    (its own pdg/material embeddings, kept separate from the trunk's
    ConditionEncoder) that predicts it, one class per expert slot
    (`n_experts` doubles as the number of process classes — see
    `build_process_map_from_files`, which caps the process vocabulary to
    exactly this many classes, bucketing rare processes into a shared
    "other" slot).

    The classifier is supervised by `classify_loss` against the true
    `process` label (see `giant/train.py`) — a *training-time* signal only;
    `gate`/`top1` never see it, so eval-time dispatch (rollout, predict)
    needs no ground truth, same as every other Router. This sidesteps the
    gradient/differentiability problem that sank the earlier
    process-conditioned-flow proposal (see the archived decision doc): the
    hard categorical choice only ever feeds a non-differentiable expert
    *dispatch*, never the flow's own conditioning path.
    """

    def __init__(
        self,
        n_experts: int,
        pdg_vocab: int,
        mat_vocab: int,
        emb_dim: int = 8,
        hidden_dim: int = 64,
    ) -> None:
        super().__init__(n_experts)
        self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
        self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
        self.classifier = nn.Sequential(
            nn.Linear(COND_DIM + 2 * emb_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, n_experts),
        )

    def logits(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        """(B, n_experts) raw process-classifier logits, one class per expert."""
        pdg_e = self.pdg_emb(cond_cat[:, 0])
        mat_e = self.mat_emb(cond_cat[:, 1])
        h = torch.cat([cond_cont, pdg_e, mat_e], dim=-1)
        return self.classifier(h)

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1)

    def classify_loss(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
    ) -> torch.Tensor:
        return F.cross_entropy(self.logits(cond_cont, cond_cat), labels)


class ComposedRouter(Router):
    """Joint router over independent axes (e.g. energy x pdg), outer-product gated.

    Wraps N already-built sub-routers, each free to have its own
    `n_experts` and hyperparameters (an `EnergyRouter(n_experts=4, ...)`
    composed with a `PdgRouter(n_experts=3, ...)` needs no axis to match
    the other's expert count). The joint gate is the outer product of the
    per-axis softmax gates, flattened to `(B, prod(n_experts_i))` — still a
    partition of unity, since each factor is one. Because the axes are
    routed independently, the joint argmax factors into the per-axis
    argmaxes, so `top1` (inherited from `Router`) costs no more than
    routing each axis alone despite the multiplicative expert count; the
    same is true of `balance_loss` (inherited, computed on the flattened
    joint gate — now one importance term per *joint* expert cell).

    Not registered in `ROUTER_REGISTRY` / buildable via `build_router`,
    since those assume one `n_experts` int shared by a single router type;
    use `build_composed_router` instead, which resolves a list of per-axis
    specs (each independently typed and sized) through `build_router`.
    """

    def __init__(self, routers: list[Router]) -> None:
        if not routers:
            raise ValueError("ComposedRouter needs at least one sub-router")
        n_experts = 1
        for r in routers:
            n_experts *= r.n_experts
        super().__init__(n_experts)
        self.routers = nn.ModuleList(routers)

    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        joint = self.routers[0].gate(cond_cont, cond_cat)  # (B, n_0)
        for router in self.routers[1:]:
            g = router.gate(cond_cont, cond_cat)  # (B, n_i)
            joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(
                1
            )  # (B, prod so far)
        return joint

    def classify_loss(
        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
    ) -> torch.Tensor:
        """Sum of each sub-router's own classify_loss (0 for unsupervised axes)."""
        total = torch.zeros((), device=cond_cont.device)
        for router in self.routers:
            total = total + router.classify_loss(cond_cont, cond_cat, labels)
        return total


def build_composed_router(specs: list[dict], **shared_kwargs) -> ComposedRouter:
    """Build a `ComposedRouter` from a list of per-axis router specs.

    Each spec is a `{"type": ..., "n_experts": ..., ...per-axis kwargs}`
    dict resolved through `build_router` exactly like a single-axis router
    config, so axes can differ in both expert count and hyperparameters
    (e.g. an energy axis's `temperature` vs a pdg axis's `emb_dim`).
    `shared_kwargs` (`pdg_vocab`, `mat_vocab`, ...) are merged under each
    spec, with the spec's own keys taking precedence.
    """
    routers = [
        build_router(
            spec["type"],
            spec["n_experts"],
            **{
                **shared_kwargs,
                **{k: v for k, v in spec.items() if k not in ("type", "n_experts")},
            },
        )
        for spec in specs
    ]
    return ComposedRouter(routers)


class ExpertTrunk(nn.Module):
    """One small expert: `input_proj -> ResBlock stack -> out_proj`.

    Same shape as the monolithic DenoisingMLP/SecondaryDecoder trunk, but
    intended to be narrower/shallower (per-call cost is the whole point).
    """

    def __init__(
        self,
        in_dim: int,
        hidden_dim: int,
        n_blocks: int,
        merged_cond_dim: int,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()
        self.input_proj = nn.Linear(in_dim, hidden_dim)
        self.blocks = nn.ModuleList(
            [
                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
                for _ in range(n_blocks)
            ]
        )
        self.out_proj = nn.Linear(hidden_dim, in_dim)

    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        x = self.input_proj(x)
        for block in self.blocks:
            x = block(x, cond)
        return self.out_proj(x)


def _route_forward(
    experts: nn.ModuleList,
    router: Router,
    x: torch.Tensor,
    cond: torch.Tensor,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    training: bool,
) -> torch.Tensor:
    """Shared dispatch for both Routed* trunks.

    Train mode: full soft mixture `sum_i gate_i * expert_i(x)` — fully
    differentiable, N-expert compute. Eval mode: grouped top-1 dispatch —
    each row runs exactly one (small) expert, which is the actual source
    of the per-call speedup this architecture is for.
    """
    if training:
        weights = router.gate(cond_cont, cond_cat)  # (B, n_experts)
        out = torch.zeros_like(x)
        for i, expert in enumerate(experts):
            out = out + weights[:, i : i + 1] * expert(x, cond)
        return out

    idx = router.top1(cond_cont, cond_cat)  # (B,)
    out = torch.zeros_like(x)
    for i, expert in enumerate(experts):
        mask = idx == i
        if mask.any():
            out[mask] = expert(x[mask], cond[mask])
    return out


class RoutedDenoisingMLP(nn.Module):
    """Routed drop-in for `DenoisingMLP`.

    Shares the time embedding, `ConditionEncoder`, and `n_sec_head` (all
    tiny) across experts and routes only the trunk (where the FLOPs are).
    Same `forward`/`predict_n_sec` signatures as `DenoisingMLP`, so
    sample.py/rollout.py/validate.py need no changes.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        router: Router,
        expert_hidden_dim: int = 128,
        expert_n_blocks: int = 3,
        emb_dim: int = EMB_DIM,
        time_dim: int = 64,
        cond_out_dim: int = 128,
        x_dim: int = X_DIM,
        dropout: float = 0.1,
        k_max: int = K_MAX,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        self.router = router
        self.time_emb = SinusoidalEmbedding(time_dim)
        self.cond_enc = ConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            out_dim=cond_out_dim,
            conditioning=conditioning,
        )
        merged_cond_dim = time_dim + cond_out_dim
        self.experts = nn.ModuleList(
            [
                ExpertTrunk(
                    x_dim,
                    expert_hidden_dim,
                    expert_n_blocks,
                    merged_cond_dim,
                    dropout=dropout,
                )
                for _ in range(router.n_experts)
            ]
        )
        self.n_sec_head = nn.Sequential(
            nn.Linear(cond_out_dim, cond_out_dim),
            nn.SiLU(),
            nn.Linear(cond_out_dim, k_max + 1),
        )

    def forward(
        self,
        x_t: torch.Tensor,
        t: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        t_emb = self.time_emb(t)
        c_emb = self.cond_enc(cond_cont, cond_cat)
        cond = torch.cat([t_emb, c_emb], dim=-1)
        return _route_forward(
            self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
        )

    def predict_n_sec(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        """Return n_sec logits (B, K_MAX+1) from conditioning alone."""
        c_emb = self.cond_enc(cond_cont, cond_cat)
        return self.n_sec_head(c_emb)


class RoutedSecondaryDecoder(nn.Module):
    """Routed drop-in for `SecondaryDecoder`.

    Shares the time embedding and `SecondaryConditionEncoder` across
    experts and routes only the trunk. Same `forward` signature as
    `SecondaryDecoder`.
    """

    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        router: Router,
        expert_hidden_dim: int = 128,
        expert_n_blocks: int = 3,
        emb_dim: int = EMB_DIM,
        time_dim: int = 64,
        cond_out_dim: int = 128,
        stage1_proj_dim: int = 64,
        sec_dim: int = SEC_DIM,
        dropout: float = 0.1,
        conditioning: str = "embedding",
    ) -> None:
        super().__init__()
        self.router = router
        self.time_emb = SinusoidalEmbedding(time_dim)
        self.cond_enc = SecondaryConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            cond_out_dim=cond_out_dim,
            stage1_proj_dim=stage1_proj_dim,
            out_dim=cond_out_dim,
            conditioning=conditioning,
        )
        merged_cond_dim = time_dim + cond_out_dim
        self.experts = nn.ModuleList(
            [
                ExpertTrunk(
                    sec_dim,
                    expert_hidden_dim,
                    expert_n_blocks,
                    merged_cond_dim,
                    dropout=dropout,
                )
                for _ in range(router.n_experts)
            ]
        )

    def forward(
        self,
        x_t: torch.Tensor,
        t: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_out: torch.Tensor,
    ) -> torch.Tensor:
        t_emb = self.time_emb(t)
        c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out)
        cond = torch.cat([t_emb, c_emb], dim=-1)
        return _route_forward(
            self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
        )


_STAGE1_MODEL_KEYS = {
    "pdg_vocab",
    "mat_vocab",
    "hidden_dim",
    "n_blocks",
    "emb_dim",
    "dropout",
    "k_max",
    "conditioning",
}
_SEC_DECODER_MODEL_KEYS = {
    "pdg_vocab",
    "mat_vocab",
    "hidden_dim",
    "n_blocks",
    "emb_dim",
    "dropout",
    "conditioning",
}
_WGAN_GENERATOR_MODEL_KEYS = _STAGE1_MODEL_KEYS | {"noise_dim"}
_WGAN_SEC_GENERATOR_MODEL_KEYS = _SEC_DECODER_MODEL_KEYS | {"noise_dim"}
# Critic has no n_sec head (n_sec is never adversarial), so it doesn't accept
# k_max the way DenoisingMLP/WGANGenerator do.
_CRITIC_MODEL_KEYS = _STAGE1_MODEL_KEYS - {"k_max"}


_AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$")


def _parse_composed_axes(router_cfg: dict) -> list[dict]:
    """Regroup `axis{i}_{field}` flat keys into a list of per-axis spec dicts.

    Flat keys (rather than a nested list-of-dicts) keep composed-router
    config expressible in the same one-level-of-nesting TOML/CLI shape as
    every other router option (`model.router` stays a flat table of
    scalars) — e.g. `axis0_type = "energy"`, `axis0_n_experts = 4`,
    `axis1_type = "pdg"`, `axis1_n_experts = 3`, `axis1_emb_dim = 8`.
    Axis indices must be contiguous from 0; order follows the index, not
    dict insertion order (TOML/CLI merging doesn't preserve it reliably).
    """
    axes: dict[int, dict] = {}
    for key, value in router_cfg.items():
        m = _AXIS_KEY_RE.match(key)
        if m is None:
            continue
        idx, field = int(m.group(1)), m.group(2)
        axes.setdefault(idx, {})[field] = value
    missing = set(range(len(axes))) - axes.keys()
    if missing:
        raise ValueError(f"composed router config has gaps at axis indices {missing}")
    return [axes[i] for i in range(len(axes))]


def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) -> Router:
    """Resolve one `model.router` config into a `Router`, single-axis or composed.

    `router_cfg["type"] == "composed"` reads `axis{i}_{field}` flat keys
    (see `_parse_composed_axes`) instead of a single `type`/`n_experts` pair.
    """
    shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
    if router_cfg["type"] == "composed":
        return build_composed_router(_parse_composed_axes(router_cfg), **shared_vocab)
    router_kwargs = {
        k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")
    }
    # Not every router needs these (EnergyRouter doesn't declare them, so
    # build_router's kwarg filtering drops them silently) but ProcessRouter
    # needs its own pdg/material embeddings sized to match the checkpoint's
    # vocab, same as the trunk's ConditionEncoder.
    router_kwargs.setdefault("pdg_vocab", pdg_vocab)
    router_kwargs.setdefault("mat_vocab", mat_vocab)
    return build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)


def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
    """Construct (stage1, sec_decoder) from a persisted/CLI model_config dict.

    Dispatches to the routed pair when `model_config["router"]["enabled"]`
    is truthy; a missing/absent "router" key (pre-routing checkpoints)
    falls back to the monolithic pair unchanged, so this is a drop-in
    replacement for the ad-hoc constructions it replaces.

    `model_config.get("conditioning", "embedding")` — old checkpoints have no
    "conditioning" key and must keep loading with their original embedding
    tables, so the default here is "embedding", not the training-time
    default (which is "physical" — see giant.config.DEFAULT_CONFIG). Read
    once and passed to both stage1/sec_decoder, so they structurally always
    share one mode.
    """
    if model_config.get("mode") == "wgan":
        stage1 = WGANGenerator(
            **{k: v for k, v in model_config.items() if k in _WGAN_GENERATOR_MODEL_KEYS}
        )
        sec_decoder = WGANSecondaryGenerator(
            **{
                k: v
                for k, v in model_config.items()
                if k in _WGAN_SEC_GENERATOR_MODEL_KEYS
            }
        )
        return stage1, sec_decoder

    router_cfg = model_config.get("router")
    if router_cfg and router_cfg.get("enabled"):
        pdg_vocab = model_config["pdg_vocab"]
        mat_vocab = model_config["mat_vocab"]
        shared = dict(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            expert_hidden_dim=model_config.get("expert_hidden_dim")
            or model_config.get("hidden_dim", 128),
            expert_n_blocks=model_config.get("expert_n_blocks")
            or model_config.get("n_blocks", 3),
            emb_dim=model_config.get("emb_dim", EMB_DIM),
            dropout=model_config.get("dropout", 0.1),
            conditioning=model_config.get("conditioning", "embedding"),
        )
        stage1 = RoutedDenoisingMLP(
            router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab),
            k_max=model_config.get("k_max", K_MAX),
            **shared,
        )
        sec_decoder = RoutedSecondaryDecoder(
            router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab),
            **shared,
        )
        return stage1, sec_decoder

    stage1 = DenoisingMLP(
        **{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS}
    )
    sec_decoder = SecondaryDecoder(
        **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
    )
    return stage1, sec_decoder


def build_critics(model_config: dict) -> tuple[nn.Module, nn.Module]:
    """Construct (critic, sec_critic) for `--mode wgan` training.

    Training-only — never persisted for inference the way `build_models`'s
    pair is, since `predict`/`rollout` only ever run the generators.
    """
    critic = Critic(
        **{k: v for k, v in model_config.items() if k in _CRITIC_MODEL_KEYS}
    )
    sec_critic = SecondaryCritic(
        **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
    )
    return critic, sec_critic