giant/giant/model/network.py

import inspect
import math

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

from giant.constants import COND_DIM, EMB_DIM, K_MAX, 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):
    def __init__(
        self,
        pdg_vocab: int,
        mat_vocab: int,
        cont_dim: int = COND_DIM,
        emb_dim: int = 16,
        out_dim: int = 128,
    ) -> None:
        super().__init__()
        self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
        self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
        in_dim = cont_dim + 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:
        pdg_e = self.pdg_emb(cond_cat[:, 0])
        mat_e = self.mat_emb(cond_cat[:, 1])
        x = torch.cat([cond_cont, 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,
    ) -> 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,
        )
        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)

    def pdg_embedding_weight(self) -> torch.Tensor:
        """Return the PDG embedding table weights for secondary type targets."""
        return self.cond_enc.pdg_emb.weight


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,
    ) -> None:
        super().__init__()
        self.base = ConditionEncoder(
            pdg_vocab=pdg_vocab,
            mat_vocab=mat_vocab,
            emb_dim=emb_dim,
            out_dim=cond_out_dim,
        )
        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, type_emb) for one
    secondary ordered by descending energy. 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,
    ) -> 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,
        )
        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 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)


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)


@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, initialized spread across [-2, 2] — roughly the
    z-normalized energy range. `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.
    """

    def __init__(
        self,
        n_experts: int = 4,
        temperature: float = 0.5,
        learn_centers: bool = True,
        energy_idx: int = 3,
    ) -> None:
        super().__init__(n_experts)
        self.temperature = temperature
        self.energy_idx = energy_idx
        centers = torch.linspace(-2.0, 2.0, n_experts)
        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 = cond_cont[:, self.energy_idx].unsqueeze(-1)  # (B, 1)
        d2 = (e - self.centers.unsqueeze(0)) ** 2  # (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 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`/`pdg_embedding_weight` 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,
    ) -> 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,
        )
        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)

    def pdg_embedding_weight(self) -> torch.Tensor:
        """Return the PDG embedding table weights for secondary type targets."""
        return self.cond_enc.pdg_emb.weight


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,
    ) -> 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,
        )
        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",
}
_SEC_DECODER_MODEL_KEYS = {
    "pdg_vocab",
    "mat_vocab",
    "hidden_dim",
    "n_blocks",
    "emb_dim",
    "dropout",
}


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.
    """
    router_cfg = model_config.get("router")
    if router_cfg and router_cfg.get("enabled"):
        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", model_config["pdg_vocab"])
        router_kwargs.setdefault("mat_vocab", model_config["mat_vocab"])
        shared = dict(
            pdg_vocab=model_config["pdg_vocab"],
            mat_vocab=model_config["mat_vocab"],
            expert_hidden_dim=model_config.get("expert_hidden_dim", 128),
            expert_n_blocks=model_config.get("expert_n_blocks", 3),
            emb_dim=model_config.get("emb_dim", EMB_DIM),
            dropout=model_config.get("dropout", 0.1),
        )
        stage1 = RoutedDenoisingMLP(
            router=build_router(
                router_cfg["type"], router_cfg["n_experts"], **router_kwargs
            ),
            k_max=model_config.get("k_max", K_MAX),
            **shared,
        )
        sec_decoder = RoutedSecondaryDecoder(
            router=build_router(
                router_cfg["type"], router_cfg["n_experts"], **router_kwargs
            ),
            **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