giant/giant/model/routers.py

"""Mixture-of-experts routing: `Router` base + registry, the four concrete
router types, and composed/config-driven construction — self-contained, no
dependency on any other `giant.model` submodule (issues.md Issue 8)."""

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.cond_layout import CondLayout
from giant.constants import COND_DIM

# ---------------------------------------------------------------------------
# Routers — carried over unchanged from v0.2
# ---------------------------------------------------------------------------


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
        self.gumbel = False
        self.gumbel_tau = 1.0

    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 combine_weights(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        """(B, n_experts) train-time expert-combination weights.

        Default (`gumbel=False`): identical to `gate()`. Opt-in
        straight-through Gumbel-softmax (`gumbel=True`, train mode only):
        hardens the forward pass to a one-hot sample (matching eval-time
        top-1 dispatch) while keeping the soft sample's gradient on backward.

        Forced fp32 (`torch.autocast(..., enabled=False)`) regardless of the
        caller's ambient `train.precision` autocast region: `clamp_min(1e-8)`
        below sits under bf16's precision but *above* fp16's ~6e-8 subnormal
        floor, so `log_probs` degrading here is exactly the kind of quiet
        drift that cost a whole rollout benchmark before (see the MoE section
        of CLAUDE.md's Roadmap) — cheap to rule out (gitea #47).
        """
        with torch.autocast(cond_cont.device.type, enabled=False):
            probs = self.gate(cond_cont, cond_cat)
            if not (self.gumbel and self.training):
                return probs
            log_probs = torch.log(probs.clamp_min(1e-8))
            return F.gumbel_softmax(log_probs, tau=self.gumbel_tau, hard=True, dim=-1)

    def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
        """(B,) hard expert index, used for eval-time grouped dispatch."""
        with torch.autocast(cond_cont.device.type, enabled=False):
            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).

        Forced fp32 — `importance` sums `gate()` over the whole batch (a
        large-magnitude accumulation in reduced precision), then takes a
        `std/mean` ratio: a classic catastrophic-cancellation shape (gitea
        #47)."""
        with torch.autocast(cond_cont.device.type, enabled=False):
            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). Routers gating on an unobservable
        pre-step quantity (e.g. ProcessRouter) override this.
        """
        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."""
        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: `(norm_entropy, importance)` — see v0.2 docstring for
        the full explanation, unchanged in v0.3.0.

        Forced fp32, same rationale as `balance_loss`/`combine_weights`: the
        `+ 1e-8` epsilon here is `entropy_loss`'s training-loss path too, not
        just a diagnostic (gitea #47)."""
        with torch.autocast(cond_cont.device.type, enabled=False):
            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 `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 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 the initial effective width/temperature exactly equals `value`."""
    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). `gate(e) =
    softmax_i(-(e - c_i)^2 / tau)`; as tau -> 0 this hardens to
    nearest-center (Voronoi) selection, exactly what `top1` uses at eval.
    """

    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 ({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, 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:
        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 (own table, separate
    from the trunk's `ConditionEncoder`). No supervision needed — PDG code
    is already known at pre-step time."""

    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[:, CondLayout.PDG_COL])  # (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 — a post-step
    outcome, so a small classifier over pre-step conditioning predicts it
    (own pdg/material embeddings, separate from the trunk's ConditionEncoder).
    `n_experts` doubles as the number of process classes. Supervised via
    `classify_loss` against the true `process` label at train time only;
    `gate`/`top1` never see it."""

    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:
        pdg_e = self.pdg_emb(cond_cat[:, CondLayout.PDG_COL])
        mat_e = self.mat_emb(cond_cat[:, CondLayout.MAT_COL])
        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. Not registered in `ROUTER_REGISTRY`; use `build_composed_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:
        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 — see
    `_parse_composed_axes`."""
    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)


_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.

    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.
    """
    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))]


# Router types that read cond_cat's pdg index through their own
# nn.Embedding(pdg_vocab, ...), regardless of the trunk's particle
# conditioning mode — see _check_router_conditioning_compat.
_VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process")


def _check_router_conditioning_compat(router_types: list[str], particle_conditioning: str) -> None:
    """Reject a router axis that reintroduces a training-vocab PDG lookup
    under `conditioning.particle.type = "physical"`.

    `PdgRouter`/`ProcessRouter` always build their own dataset-scoped
    `nn.Embedding(pdg_vocab, ...)`, independent of `ConditionEncoder`'s
    particle mode. Pairing either with `"physical"` would silently
    reintroduce a training-menu-scoped lookup at the routing layer,
    defeating the point of physical-property conditioning. Raised loudly at
    model-build time.
    """
    bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES))
    if bad and particle_conditioning == "physical":
        raise ValueError(
            f"router type(s) {bad} always use a training-vocab PDG embedding, "
            "which is incompatible with conditioning.particle.type='physical' "
            "(whose whole point is generalizing beyond that vocab) — pick a "
            "different router type (e.g. 'energy') or use "
            "conditioning.particle.type='embedding'."
        )


def _build_router_from_cfg(
    router_cfg: dict,
    pdg_vocab: int,
    mat_vocab: int,
    particle_conditioning: str = "embedding",
) -> Router:
    """Resolve one stage's `router` config into a `Router`, single-axis or
    composed. `gumbel` is set as a post-construction attribute (shared by
    every router type, not a per-type constructor kwarg)."""
    shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
    if router_cfg["type"] == "composed":
        axes = _parse_composed_axes(router_cfg)
        _check_router_conditioning_compat([a["type"] for a in axes], particle_conditioning)
        router = build_composed_router(axes, **shared_vocab)
        router.gumbel = bool(router_cfg.get("gumbel", False))
        return router
    _check_router_conditioning_compat([router_cfg["type"]], particle_conditioning)
    router_kwargs = {k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")}
    router_kwargs.setdefault("pdg_vocab", pdg_vocab)
    router_kwargs.setdefault("mat_vocab", mat_vocab)
    router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
    router.gumbel = bool(router_cfg.get("gumbel", False))
    return router