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,
)
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)["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 / "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