giant/giant/analysis/router_gating.py

"""Router gating diagnostic: where a MoE checkpoint's decision boundaries sit.

Unlike everything else in this package, this reduction needs a live PyTorch
model — soft expert gate weights aren't columns in a rollout/predict parquet,
they only exist by calling `Router.gate(cond_cont, cond_cat)` (see
`giant.model.network.Router`) against the checkpoint that produced the
rollout. That's a deliberate, narrow exception to the rest of the catalog's
"polars/numpy only" contract; it still runs fine as a `compute-one` HTCondor
job since torch is already installed there (the same env trains checkpoints).

The routing axis is fixed to pre-step energy: every router type at least
indirectly depends on it (`EnergyRouter` reads it directly; `PdgRouter` and
`ProcessRouter` correlate with it through the physics), and it's the one axis
a reader can interpret without knowing the checkpoint's specific router
config. `x` is binned into equal-population (quantile) bins rather than
equal-width ones, since energy is heavy-tailed and equal-width bins would
leave the upper end almost empty. Mean gate weight per bin is stacked as
filled areas per expert — since `gate` rows are a partition of unity, the
stack always fills exactly to 1, and the crossover bands are the router's
soft decision boundaries (where two experts' means cross ~0.5).
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np
import polars as pl

from giant.analysis.grouping import pdg_label
from giant.analysis.reduced import Reduced

if TYPE_CHECKING:
    import torch

    from giant.data.transforms import Normalizer

_SAMPLE_ROWS = 200_000
_N_BINS = 40
_TOP_K_PROCESS = 8

_COLS = (
    "pre_x",
    "pre_y",
    "pre_z",
    "pre_E",
    "pre_dx",
    "pre_dy",
    "pre_dz",
    "layer_id",
    "pdg",
    "material",
)


@dataclass
class _RouterHandle:
    router: "torch.nn.Module"
    pdg_map: dict[int, int]
    mat_map: dict[str, int]
    cond_normalizer: "Normalizer"
    particle_conditioning: str
    material_conditioning: str
    router_type: str


def load_router(checkpoint: str | Path) -> _RouterHandle | None:
    """Load a checkpoint's Stage-1 router, or None if it isn't a MoE checkpoint."""
    import torch

    from giant.checkpoint_io import conditioning_axes
    from giant.data.transforms import Normalizer
    from giant.model.network import build_models

    ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
    model_cfg = ckpt.get("model_config") or {}
    # New nested shape (has a "stage1_model" key) vs. a v0.2 checkpoint's
    # flat model_config.
    router_cfg = (
        (model_cfg.get("stage1_model") or {}).get("router") if "stage1_model" in model_cfg else model_cfg.get("router")
    )
    if not router_cfg or not router_cfg.get("enabled"):
        return None

    built = build_models(model_cfg)
    stage1 = built["stage1"]
    if stage1 is None:
        return None
    stage1.load_state_dict(ckpt["model"])
    stage1.eval()
    router = stage1.trunk.router
    if router is None:
        return None

    particle_conditioning, material_conditioning = conditioning_axes(model_cfg)
    return _RouterHandle(
        router=router,
        pdg_map={int(k): v for k, v in ckpt["pdg_map"].items()},
        mat_map={str(k): v for k, v in ckpt["mat_map"].items()},
        cond_normalizer=Normalizer.from_dict(ckpt["normalizer"]["cond"]),
        particle_conditioning=particle_conditioning,
        material_conditioning=material_conditioning,
        router_type=router_cfg["type"],
    )


def _subsample(lf: pl.LazyFrame, n: int, seed: int, extra_cols: tuple = ()) -> pl.DataFrame:
    total = lf.select(pl.len()).collect(engine="streaming").item()
    if total > n:
        threshold = int(n / total * 2**32)
        lf = lf.filter((pl.col("pre_E").hash(seed=seed) % 2**32) < threshold)
    return lf.select(*_COLS, *extra_cols).collect(engine="streaming")


def _gate_for_df(handle: _RouterHandle, df: pl.DataFrame) -> tuple[pl.DataFrame, np.ndarray]:
    """(filtered df, gate_weights) for rows in ``df`` with a known pdg/material.

    Rows whose species or material never appeared in the checkpoint's
    training vocab can't be embedded — dropped here the same way
    `giant.rollout`'s own known-pdg gate drops them at inference. The
    returned df keeps every original column (filtered to the same rows), so
    callers can key gate weights by any of them (energy, pdg, process, ...).
    """
    import torch

    from giant.data.transforms import build_cond_features

    known = np.array(
        [
            int(p) in handle.pdg_map and str(m) in handle.mat_map
            for p, m in zip(df["pdg"].to_list(), df["material"].to_list())
        ]
    )
    if not known.any():
        return df.clear(), np.zeros((0, handle.router.n_experts))
    df = df.filter(pl.Series(known, dtype=pl.Boolean))

    data = {
        "pre_pos": np.column_stack([df["pre_x"].to_numpy(), df["pre_y"].to_numpy(), df["pre_z"].to_numpy()]),
        "pre_E": df["pre_E"].to_numpy(),
        "pre_dir": np.column_stack([df["pre_dx"].to_numpy(), df["pre_dy"].to_numpy(), df["pre_dz"].to_numpy()]),
        "layer_id": df["layer_id"].to_numpy(),
        "pdg": df["pdg"].to_numpy(),
        "material": df["material"].to_numpy(),
    }
    cond_cont, cond_cat = build_cond_features(
        data,
        handle.pdg_map,
        handle.mat_map,
        cond_normalizer=handle.cond_normalizer,
        particle_conditioning=handle.particle_conditioning,
        material_conditioning=handle.material_conditioning,
    )
    with torch.no_grad():
        gate = handle.router.gate(torch.from_numpy(cond_cont).float(), torch.from_numpy(cond_cat).long()).numpy()
    return df, gate


def _quantile_bins(x: np.ndarray, gate: np.ndarray, n_bins: int) -> dict:
    order = np.argsort(x)
    x_sorted, g_sorted = x[order], gate[order]
    edges = np.quantile(x_sorted, np.linspace(0, 1, n_bins + 1))
    edges[-1] = np.nextafter(edges[-1], np.inf)  # include the max value
    bin_idx = np.clip(np.digitize(x_sorted, edges[1:-1]), 0, n_bins - 1)

    n_experts = gate.shape[1]
    centers = np.full(n_bins, np.nan)
    means = np.full((n_bins, n_experts), np.nan)
    for b in range(n_bins):
        mask = bin_idx == b
        if mask.any():
            centers[b] = x_sorted[mask].mean()
            means[b] = g_sorted[mask].mean(axis=0)
    valid = ~np.isnan(centers)
    return {"centers": centers[valid].tolist(), "means": means[valid].tolist()}


def _top1_shares(categories: np.ndarray, idx: np.ndarray, order: list, n_experts: int) -> dict[str, list[float]]:
    """Fraction of each category's rows hard-dispatched to each expert.

    Uses `Router.top1` (argmax), not the soft `gate` mean — grouped top-1
    dispatch is what `_route_forward` actually runs in eval mode (rollout,
    predict), so this answers "which expert does a photon/Compton step
    actually go through", not just its average soft weight.
    """
    shares: dict[str, list[float]] = {}
    for key in order:
        mask = categories == key
        total = int(mask.sum())
        if total == 0:
            shares[str(key)] = [0.0] * n_experts
            continue
        counts = np.bincount(idx[mask], minlength=n_experts)
        shares[str(key)] = (counts / total).tolist()
    return shares


_NOTE_NOT_MOE = "checkpoint has no enabled MoE router (model.router.enabled is false/absent) — nothing to show"

_TITLES = {
    "router_gating": "Router gating (mixture-of-experts decision boundaries)",
    "router_share_by_pdg": "Router expert share by particle species",
    "router_share_by_process": "Router expert share by physics process",
}


def _unavailable(spec_id: str) -> Reduced:
    return Reduced(
        id=spec_id,
        family="model",
        kind="unavailable",
        title=_TITLES[spec_id],
        xlabel="n/a",
        payload={"note": _NOTE_NOT_MOE},
    )


def compute_router_gating(
    checkpoint: str | Path | None,
    r_phys: pl.LazyFrame,
    t_phys: pl.LazyFrame,
    seed: int = 0,
) -> Reduced:
    """`Reduced` for the router-gating figure, or an explanatory note if n/a."""
    handle = load_router(checkpoint) if checkpoint else None
    if handle is None:
        return _unavailable("router_gating")

    sides: dict[str, dict] = {}
    for name, lf in (("rollout", r_phys), ("reference", t_phys)):
        df = _subsample(lf, _SAMPLE_ROWS, seed)
        df, gate = _gate_for_df(handle, df)
        x = df["pre_E"].to_numpy()
        sides[name] = _quantile_bins(x, gate, _N_BINS) if len(x) else {"centers": [], "means": []}

    return Reduced(
        id="router_gating",
        family="model",
        kind="router_gating",
        title=_TITLES["router_gating"],
        xlabel="pre-step energy [MeV]",
        payload={
            "router_type": handle.router_type,
            "n_experts": handle.router.n_experts,
            "log_x": True,
            **sides,
        },
    )


def compute_router_share_by_pdg(
    checkpoint: str | Path | None,
    r_phys: pl.LazyFrame,
    t_phys: pl.LazyFrame,
    top_pdgs: list[int],
    seed: int = 0,
) -> Reduced:
    """Stacked-bar share of each particle species dispatched to each expert."""
    handle = load_router(checkpoint) if checkpoint else None
    if handle is None:
        return _unavailable("router_share_by_pdg")

    labels = [pdg_label(p) for p in top_pdgs]
    sides: dict[str, dict] = {}
    for name, lf in (("rollout", r_phys), ("reference", t_phys)):
        df = _subsample(lf, _SAMPLE_ROWS, seed)
        df, gate = _gate_for_df(handle, df)
        if len(df):
            idx = gate.argmax(axis=1)
            shares = _top1_shares(df["pdg"].to_numpy(), idx, top_pdgs, handle.router.n_experts)
        else:
            shares = {str(p): [0.0] * handle.router.n_experts for p in top_pdgs}
        sides[name] = {labels[i]: shares[str(p)] for i, p in enumerate(top_pdgs)}

    return Reduced(
        id="router_share_by_pdg",
        family="model",
        kind="router_share",
        title=_TITLES["router_share_by_pdg"],
        xlabel="particle species",
        payload={
            "router_type": handle.router_type,
            "n_experts": handle.router.n_experts,
            "categories": labels,
            **sides,
        },
    )


def compute_router_share_by_process(
    checkpoint: str | Path | None,
    t_phys: pl.LazyFrame,
    seed: int = 0,
    top_k: int = _TOP_K_PROCESS,
) -> Reduced:
    """Stacked-bar share of each physics process dispatched to each expert.

    Reference-only: ``process`` is the true post-step physics process — a
    label the rollout side has no equivalent of (see
    `giant.model.network.ProcessRouter`, which predicts it from pre-step
    conditioning alone, never observes it at eval time). This plot instead
    checks *after the fact*, on real data, how well the router's conditioning
    -based dispatch lines up with the true process.
    """
    handle = load_router(checkpoint) if checkpoint else None
    if handle is None:
        return _unavailable("router_share_by_process")

    df = _subsample(t_phys, _SAMPLE_ROWS, seed, extra_cols=("process",))
    df, gate = _gate_for_df(handle, df)
    if len(df):
        counts = df["process"].value_counts().sort("count", descending=True)
        order = counts["process"].to_list()[:top_k]
        idx = gate.argmax(axis=1)
        shares = _top1_shares(df["process"].to_numpy(), idx, order, handle.router.n_experts)
    else:
        order, shares = [], {}

    return Reduced(
        id="router_share_by_process",
        family="model",
        kind="router_share",
        title=_TITLES["router_share_by_process"],
        xlabel="physics process",
        payload={
            "router_type": handle.router_type,
            "n_experts": handle.router.n_experts,
            "categories": order,
            "reference": {p: shares[p] for p in order},
        },
    )