analyze: add MoE router gating/share diagnostic plots

New "model" family in the gallery: router_gating (mean soft gate weight
vs. pre-step energy, showing the router's soft decision boundaries) and
router_share_by_pdg/router_share_by_process (stacked top-1 dispatch share
by species / true physics process). Needs a live checkpoint's Router, so
it's a documented exception to the rest of the package's polars/numpy-only
contract; gracefully degrades to a placeholder for non-MoE checkpoints.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

giant/analysis/catalog.py

@@ -37,6 +37,11 @@ from giant.analysis.reduce import (
     weighted_profile,
 )
 from giant.analysis.reduced import Reduced
+from giant.analysis.router_gating import (
+    compute_router_gating,
+    compute_router_share_by_pdg,
+    compute_router_share_by_process,
+)
 from giant.analysis.sources import Side, open_side, physical_steps, secondaries
 from giant.analysis.variables import RANGED_VARS, cos_scatter_expr
 
@@ -50,9 +55,10 @@ class Bundle:
     t_all: pl.LazyFrame  # reference, all rows
     r_phys: pl.LazyFrame  # rollout, physical steps only
     t_phys: pl.LazyFrame  # reference, physical steps only
+    checkpoint: str | None = None  # from the rollout YAML; router_gating only
 
     @classmethod
-    def open(cls, rollout, reference, ctx: Context) -> "Bundle":
+    def open(cls, rollout, reference, ctx: Context, checkpoint=None) -> "Bundle":
         r_all = open_side(rollout, Side.rollout)
         t_all = open_side(reference, Side.reference)
         return cls(
@@ -61,6 +67,7 @@ class Bundle:
             t_all=t_all,
             r_phys=physical_steps(r_all, Side.rollout),
             t_phys=physical_steps(t_all, Side.reference),
+            checkpoint=checkpoint,
         )
 
 
@@ -559,6 +566,23 @@ def build_catalog() -> list[PlotSpec]:
         PlotSpec("sec_count_per_species", "secondaries", _sec_count_per_species),
         PlotSpec("sec_energy", "secondaries", _sec_energy),
         PlotSpec("sec_cos_angle", "secondaries", _sec_cos_angle),
+        PlotSpec(
+            "router_gating",
+            "model",
+            lambda b: compute_router_gating(b.checkpoint, b.r_phys, b.t_phys),
+        ),
+        PlotSpec(
+            "router_share_by_pdg",
+            "model",
+            lambda b: compute_router_share_by_pdg(
+                b.checkpoint, b.r_phys, b.t_phys, b.ctx.top_pdgs
+            ),
+        ),
+        PlotSpec(
+            "router_share_by_process",
+            "model",
+            lambda b: compute_router_share_by_process(b.checkpoint, b.t_phys),
+        ),
     ]
     return specs
 

giant/analysis/condor.py

@@ -152,10 +152,11 @@ def compute_reduced(
     reference: str | Path,
     shared: str | Path,
     out: str | Path,
+    checkpoint: str | None = None,
 ) -> Path:
     """Core: run one plot's reduction against explicit paths → ``Reduced`` JSON."""
     ctx = Context.load(shared)
-    bundle = Bundle.open(rollout, reference, ctx)
+    bundle = Bundle.open(rollout, reference, ctx, checkpoint=checkpoint)
     reduced = get_spec(spec_id).compute(bundle)
     out = Path(out)
     reduced.save(out)
@@ -172,6 +173,7 @@ def compute_one(spec_id: str, run_dir: str | Path) -> Path:
         meta.reference,
         run_path / "shared.json",
         run_path / "reduced" / f"{spec_id}.json",
+        checkpoint=meta.plot_meta.get("checkpoint"),
     )
 
 

giant/analysis/reduced.py

@@ -12,11 +12,14 @@ from dataclasses import asdict, dataclass, field
 from pathlib import Path
 
 # Reduced.kind values:
-#   "overlay_hist"  rollout vs reference density histogram over shared edges
-#   "grouped_hist"  one panel per group (energy/pdg/material), each an overlay
-#   "profile"       edep-weighted mean +/- event-RMS vs depth/radius, two series
-#   "bar"           per-category rollout vs reference bars (share / counts)
-#   "single_hist"   one series only (e.g. rollout leakage; reference has none)
+#   "overlay_hist"   rollout vs reference density histogram over shared edges
+#   "grouped_hist"   one panel per group (energy/pdg/material), each an overlay
+#   "profile"        edep-weighted mean +/- event-RMS vs depth/radius, two series
+#   "bar"            per-category rollout vs reference bars (share / counts)
+#   "single_hist"    one series only (e.g. rollout leakage; reference has none)
+#   "router_gating"  stacked mean MoE gate weight vs energy, rollout + reference
+#   "router_share"   stacked bar of MoE top-1 dispatch share by category
+#   "unavailable"    plot not applicable to this run (e.g. non-MoE checkpoint)
 
 
 @dataclass

giant/analysis/render.py

@@ -137,12 +137,89 @@ def _render_bar(r: Reduced, params: dict):
     return fig
 
 
+def _render_router_gating(r: Reduced, params: dict):
+    n_experts = r.payload["n_experts"]
+    log_x = r.payload.get("log_x", False)
+    fig, axes = ps.new_figure(
+        "slide-16x9", title=r.title, params=params, nrows=1, ncols=2, squeeze=False
+    )
+    flat = axes.ravel()
+    for ax, key in zip(flat, ("rollout", "reference")):
+        side = r.payload.get(key, {})
+        centers = np.asarray(side.get("centers", []))
+        means = np.asarray(side.get("means", []))
+        if len(centers) and means.size:
+            cum = np.zeros(len(centers))
+            for i in range(n_experts):
+                ax.fill_between(
+                    centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}"
+                )
+                cum = cum + means[:, i]
+        if log_x:
+            ax.set_xscale("log")
+        ax.set_ylim(0, 1)
+        ax.set_title(_SERIES_LABELS[key], fontsize=8)
+        ax.set_xlabel(r.xlabel)
+    flat[0].set_ylabel("mean gate weight")
+    ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
+    return fig
+
+
+def _render_router_share(r: Reduced, params: dict):
+    categories = r.payload["categories"]
+    n_experts = r.payload["n_experts"]
+    x = np.arange(len(categories))
+    present = [k for k in ("rollout", "reference") if k in r.payload]
+    fig, axes = ps.new_figure(
+        "slide-16x9",
+        title=r.title,
+        params=params,
+        nrows=1,
+        ncols=len(present),
+        squeeze=False,
+    )
+    flat = axes.ravel()
+    for ax, key in zip(flat, present):
+        side = r.payload[key]
+        shares = np.array([side[c] for c in categories])  # (n_cat, n_experts)
+        bottom = np.zeros(len(categories))
+        for i in range(n_experts):
+            ax.bar(x, shares[:, i], bottom=bottom, label=f"expert {i}")
+            bottom += shares[:, i]
+        ax.set_xticks(x)
+        ax.set_xticklabels(categories, rotation=45, ha="right")
+        ax.set_ylim(0, 1)
+        ax.set_title(_SERIES_LABELS[key], fontsize=8)
+    flat[0].set_ylabel("share of rows dispatched to expert")
+    ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
+    return fig
+
+
+def _render_unavailable(r: Reduced, params: dict):
+    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
+    ax.axis("off")
+    ax.text(
+        0.5,
+        0.5,
+        r.payload.get("note", "not available"),
+        ha="center",
+        va="center",
+        wrap=True,
+        fontsize=10,
+        transform=ax.transAxes,
+    )
+    return fig
+
+
 _RENDERERS = {
     "overlay_hist": _render_overlay,
     "single_hist": _render_single,
     "grouped_hist": _render_grouped,
     "profile": _render_profile,
     "bar": _render_bar,
+    "router_gating": _render_router_gating,
+    "router_share": _render_router_share,
+    "unavailable": _render_unavailable,
 }
 
 

giant/analysis/router_gating.py

@@ -0,0 +1,339 @@
+"""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"
+    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.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 {}
+    router_cfg = model_cfg.get("router")
+    if not router_cfg or not router_cfg.get("enabled"):
+        return None
+
+    stage1, _ = build_models(model_cfg)
+    stage1.load_state_dict(ckpt["model"])
+    stage1.eval()
+
+    return _RouterHandle(
+        router=stage1.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"]),
+        conditioning=model_cfg.get("conditioning", "embedding"),
+        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,
+        conditioning=handle.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},
+        },
+    )

tests/test_catalog.py

@@ -44,6 +44,9 @@ def test_every_spec_computes_valid_reduced(bundle: Bundle):
             "profile",
             "bar",
             "single_hist",
+            "router_gating",
+            "router_share",
+            "unavailable",
         }
         assert r.title and r.xlabel
         _validate_payload(r)
@@ -67,3 +70,14 @@ def _validate_payload(r) -> None:
             assert len(p[k]) == n
     elif r.kind == "bar":
         assert len(p["labels"]) == len(p["rollout"]) == len(p["reference"])
+    elif r.kind == "unavailable":
+        assert p["note"]
+    elif r.kind == "router_gating":
+        for side in ("rollout", "reference"):
+            if side in p:
+                assert len(p[side]["centers"]) == len(p[side]["means"])
+    elif r.kind == "router_share":
+        for cat in p["categories"]:
+            for side in ("rollout", "reference"):
+                if side in p:
+                    assert cat in p[side]

tests/test_router_gating.py

@@ -0,0 +1,126 @@
+"""Tests for the MoE router-gating diagnostic (giant.analysis.router_gating)."""
+
+from __future__ import annotations
+
+import numpy as np
+import polars as pl
+import torch
+
+from giant.analysis.router_gating import (
+    compute_router_gating,
+    compute_router_share_by_pdg,
+    compute_router_share_by_process,
+)
+from giant.data.transforms import Normalizer
+from giant.model.network import build_models
+
+_PDG_MAP = {11: 0, 22: 1}
+_MAT_MAP = {"G4_PbWO4": 0, "G4_Pb": 1}
+
+
+def _model_cfg() -> dict:
+    return {
+        "router": {
+            "enabled": True,
+            "type": "energy",
+            "n_experts": 2,
+            "temperature": 0.5,
+            "learn_centers": True,
+            "energy_idx": 3,
+        },
+        "pdg_vocab": len(_PDG_MAP),
+        "mat_vocab": len(_MAT_MAP),
+        "conditioning": "embedding",
+    }
+
+
+def _write_checkpoint(tmp_path) -> str:
+    cfg = _model_cfg()
+    stage1, _ = build_models(cfg)
+    norm = Normalizer()
+    norm.mean = np.zeros(15, dtype=np.float32)
+    norm.std = np.ones(15, dtype=np.float32)
+    ckpt = {
+        "model_config": cfg,
+        "model": stage1.state_dict(),
+        "pdg_map": _PDG_MAP,
+        "mat_map": _MAT_MAP,
+        "normalizer": {"cond": norm.to_dict()},
+    }
+    path = tmp_path / "ckpt.pt"
+    torch.save(ckpt, path)
+    return str(path)
+
+
+def _steps_frame(process: bool = False) -> pl.LazyFrame:
+    n = 40
+    rng = np.random.default_rng(0)
+    pre_e = np.concatenate([rng.uniform(1, 10, n // 2), rng.uniform(100, 1000, n // 2)])
+    pdg = np.where(np.arange(n) % 2 == 0, 11, 22)
+    material = np.where(np.arange(n) % 3 == 0, "G4_Pb", "G4_PbWO4")
+    data = {
+        "event_id": np.arange(n),
+        "pdg": pdg,
+        "pre_x": np.zeros(n),
+        "pre_y": np.zeros(n),
+        "pre_z": np.zeros(n),
+        "pre_E": pre_e,
+        "pre_dx": np.zeros(n),
+        "pre_dy": np.zeros(n),
+        "pre_dz": np.ones(n),
+        "post_x": np.zeros(n),
+        "post_y": np.zeros(n),
+        "post_z": np.ones(n),
+        "post_E": pre_e * 0.5,
+        "post_dx": np.zeros(n),
+        "post_dy": np.zeros(n),
+        "post_dz": np.ones(n),
+        "edep": pre_e * 0.5,
+        "step_length": np.ones(n),
+        "material": material,
+        "layer_id": np.zeros(n, dtype=np.int64),
+    }
+    if process:
+        data["process"] = np.where(pdg == 11, "eIoni", "compt")
+    return pl.DataFrame(data).lazy()
+
+
+def test_compute_router_gating_shapes(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    lf = _steps_frame()
+    r = compute_router_gating(checkpoint, lf, lf)
+    assert r.kind == "router_gating"
+    assert r.payload["n_experts"] == 2
+    for side in ("rollout", "reference"):
+        means = r.payload[side]["means"]
+        assert means, f"{side} produced no bins"
+        assert all(abs(sum(row) - 1.0) < 1e-5 for row in means)
+
+
+def test_compute_router_gating_missing_checkpoint_is_unavailable():
+    lf = _steps_frame()
+    r = compute_router_gating(None, lf, lf)
+    assert r.kind == "unavailable"
+    assert "note" in r.payload
+    assert r.title
+
+
+def test_compute_router_share_by_pdg(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    lf = _steps_frame()
+    r = compute_router_share_by_pdg(checkpoint, lf, lf, top_pdgs=[11, 22])
+    assert r.kind == "router_share"
+    for side in ("rollout", "reference"):
+        assert set(r.payload[side]) == {"e-", "gamma"}
+        for shares in r.payload[side].values():
+            assert abs(sum(shares) - 1.0) < 1e-5
+
+
+def test_compute_router_share_by_process(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    lf = _steps_frame(process=True)
+    r = compute_router_share_by_process(checkpoint, lf)
+    assert r.kind == "router_share"
+    assert set(r.payload["categories"]) <= {"eIoni", "compt"}
+    for shares in r.payload["reference"].values():
+        assert abs(sum(shares) - 1.0) < 1e-5