lars
60c2ca1985
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>
127 lines
3.8 KiB
Python
127 lines
3.8 KiB
Python
"""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
|