giant/tests/test_router_gating.py

"""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,
    compute_router_specialization,
)
from giant.analysis.sources import RolloutSide
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, name: str = "ckpt.pt") -> str:
    cfg = _model_cfg()
    stage1 = build_models(cfg)["stage1"]
    assert stage1 is not None
    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 / name
    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 _side(checkpoint: str | None, lf: pl.LazyFrame) -> RolloutSide:
    return RolloutSide(all=lf, phys=lf, checkpoint=checkpoint)


def test_compute_router_gating_shapes(tmp_path):
    checkpoint = _write_checkpoint(tmp_path)
    lf = _steps_frame()
    r = compute_router_gating({"rollout": _side(checkpoint, lf)}, lf)
    assert r.kind == "router_gating"
    assert list(r.payload["series"]) == ["rollout"]
    entry = r.payload["series"]["rollout"]
    assert entry["n_experts"] == 2
    for side in ("rollout", "reference"):
        means = entry[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({"rollout": _side(None, lf)}, lf)
    assert r.kind == "unavailable"
    assert "note" in r.payload
    assert r.title


def test_compute_router_gating_two_rollouts_only_moe_ones_included(tmp_path):
    lf = _steps_frame()
    ckpt = _write_checkpoint(tmp_path)
    rollouts = {"flow": _side(None, lf), "moe": _side(ckpt, lf)}
    r = compute_router_gating(rollouts, lf)
    assert list(r.payload["series"]) == ["moe"]


def test_compute_router_specialization_two_rollouts(tmp_path):
    lf = _steps_frame()
    ckpt_a = _write_checkpoint(tmp_path, "a.pt")
    ckpt_b = _write_checkpoint(tmp_path, "b.pt")
    rollouts = {"a": _side(ckpt_a, lf), "b": _side(ckpt_b, lf)}
    r = compute_router_specialization(rollouts, lf)
    assert r.kind == "router_specialization"
    assert list(r.payload["series"]) == ["a", "b"]
    for entry in r.payload["series"].values():
        assert entry["chance_level"] == 0.5
        assert len(entry["rollout"]["centers"]) == len(entry["rollout"]["score"])


def test_compute_router_share_by_pdg(tmp_path):
    checkpoint = _write_checkpoint(tmp_path)
    lf = _steps_frame()
    r = compute_router_share_by_pdg({"rollout": _side(checkpoint, lf)}, lf, top_pdgs=[11, 22])
    assert r.kind == "router_share"
    entry = r.payload["series"]["rollout"]
    for side in ("rollout", "reference"):
        assert set(entry[side]) == {"e-", "gamma"}
        for shares in entry[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({"rollout": _side(checkpoint, lf)}, lf)
    assert r.kind == "router_share"
    entry = r.payload["series"]["rollout"]
    assert set(entry["categories"]) <= {"eIoni", "compt"}
    for shares in entry["reference"].values():
        assert abs(sum(shares) - 1.0) < 1e-5


def test_no_moe_rollouts_are_unavailable(tmp_path):
    lf = _steps_frame()
    rollouts = {"flow": _side(None, lf), "wgan": _side(None, lf)}
    assert compute_router_gating(rollouts, lf).kind == "unavailable"
    assert compute_router_share_by_pdg(rollouts, lf, top_pdgs=[11, 22]).kind == "unavailable"
    assert compute_router_share_by_process(rollouts, lf).kind == "unavailable"
    assert compute_router_specialization(rollouts, lf).kind == "unavailable"