Merge branch '4-prototype-a-mixture-of-experts-routing-tree-architecture' into analysis-streaming-rewrite
# Conflicts: # giant/analysis.py # tests/test_analysis.py
This commit is contained in:
@@ -139,24 +139,29 @@ def test_plan_jobs_multiple_detectors_each_start_independently(tmp_path):
|
||||
|
||||
def test_job_seed_deterministic():
|
||||
job = SimJob(detector="pbwo4", config=None, shard_index=3)
|
||||
assert job_seed("steps", "gen1", job) == job_seed("steps", "gen1", job)
|
||||
assert job_seed("steps", "gen1", job, None) == job_seed("steps", "gen1", job, None)
|
||||
|
||||
|
||||
def test_job_seed_varies_by_shard_index():
|
||||
a = SimJob(detector="pbwo4", config=None, shard_index=0)
|
||||
b = SimJob(detector="pbwo4", config=None, shard_index=1)
|
||||
assert job_seed("steps", "gen1", a) != job_seed("steps", "gen1", b)
|
||||
assert job_seed("steps", "gen1", a, None) != job_seed("steps", "gen1", b, None)
|
||||
|
||||
|
||||
def test_job_seed_varies_by_detector():
|
||||
a = SimJob(detector="pbwo4", config=None, shard_index=0)
|
||||
b = SimJob(detector="sampling_pb_scint", config="pb_scint", shard_index=0)
|
||||
assert job_seed("steps", "gen1", a) != job_seed("steps", "gen1", b)
|
||||
assert job_seed("steps", "gen1", a, None) != job_seed("steps", "gen1", b, None)
|
||||
|
||||
|
||||
def test_job_seed_varies_by_gen():
|
||||
job = SimJob(detector="pbwo4", config=None, shard_index=0)
|
||||
assert job_seed("steps", "gen1", job) != job_seed("steps", "gen2", job)
|
||||
assert job_seed("steps", "gen1", job, None) != job_seed("steps", "gen2", job, None)
|
||||
|
||||
|
||||
def test_job_seed_varies_by_energy():
|
||||
job = SimJob(detector="pbwo4", config=None, shard_index=0)
|
||||
assert job_seed("steps", "gen1", job, 1.0) != job_seed("steps", "gen1", job, 10.0)
|
||||
|
||||
|
||||
def test_run_job_passes_deterministic_seed_env_var(tmp_path):
|
||||
@@ -166,11 +171,11 @@ def test_run_job_passes_deterministic_seed_env_var(tmp_path):
|
||||
tmp_root.mkdir()
|
||||
|
||||
job = SimJob(detector="pbwo4", config=None, shard_index=5)
|
||||
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
||||
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
|
||||
|
||||
assert result.dest is not None
|
||||
payload = json.loads(result.dest.read_text())
|
||||
assert payload["seed"] == str(job_seed("steps", "gen1", job))
|
||||
assert payload["seed"] == str(job_seed("steps", "gen1", job, None))
|
||||
|
||||
|
||||
def test_run_job_moves_output_to_correct_shard_path(tmp_path):
|
||||
@@ -181,7 +186,7 @@ def test_run_job_moves_output_to_correct_shard_path(tmp_path):
|
||||
tmp_root.mkdir()
|
||||
|
||||
job = SimJob(detector="pbwo4", config=None, shard_index=7)
|
||||
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
||||
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
|
||||
|
||||
assert result.ok
|
||||
assert result.dest == gen_dir / "pbwo4" / "shard-007.root"
|
||||
@@ -197,7 +202,7 @@ def test_run_job_passes_config_arg_and_isolates_cwd(tmp_path):
|
||||
tmp_root.mkdir()
|
||||
|
||||
job = SimJob(detector="sampling_pb_scint", config="pb_scint", shard_index=0)
|
||||
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
||||
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
|
||||
|
||||
assert result.ok
|
||||
assert result.dest is not None
|
||||
@@ -215,13 +220,27 @@ def test_run_job_omits_config_arg_when_none(tmp_path):
|
||||
tmp_root.mkdir()
|
||||
|
||||
job = SimJob(detector="pbwo4", config=None, shard_index=0)
|
||||
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
||||
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
|
||||
|
||||
assert result.dest is not None
|
||||
payload = json.loads(result.dest.read_text())
|
||||
assert payload["argv"] == ["10000"]
|
||||
|
||||
|
||||
def test_run_job_appends_energy_arg_when_given(tmp_path):
|
||||
fake = _write_fake_executable(tmp_path / "fake_exe.py")
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
tmp_root = tmp_path / ".sim-tmp"
|
||||
tmp_root.mkdir()
|
||||
|
||||
job = SimJob(detector="pbwo4_10gev", config=None, shard_index=0)
|
||||
result = run_job(job, fake, 10000, 10.0, tmp_path, "steps", "gen1", tmp_root)
|
||||
|
||||
assert result.dest is not None
|
||||
payload = json.loads(result.dest.read_text())
|
||||
assert payload["argv"] == ["10000", "10.0"]
|
||||
|
||||
|
||||
def test_run_job_fails_when_executable_errors(tmp_path):
|
||||
fake = _write_fake_executable(tmp_path / "fake_exe.py", exit_code=1)
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
@@ -229,7 +248,7 @@ def test_run_job_fails_when_executable_errors(tmp_path):
|
||||
tmp_root.mkdir()
|
||||
|
||||
job = SimJob(detector="pbwo4", config=None, shard_index=0)
|
||||
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
||||
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
|
||||
|
||||
assert not result.ok
|
||||
assert "exited 1" in result.message
|
||||
@@ -242,7 +261,7 @@ def test_run_job_fails_when_no_root_file_produced(tmp_path):
|
||||
tmp_root.mkdir()
|
||||
|
||||
job = SimJob(detector="pbwo4", config=None, shard_index=0)
|
||||
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
||||
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
|
||||
|
||||
assert not result.ok
|
||||
assert "found 0" in result.message
|
||||
@@ -255,7 +274,7 @@ def test_run_job_fails_when_multiple_root_files_produced(tmp_path):
|
||||
tmp_root.mkdir()
|
||||
|
||||
job = SimJob(detector="pbwo4", config=None, shard_index=0)
|
||||
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
||||
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
|
||||
|
||||
assert not result.ok
|
||||
assert "found 2" in result.message
|
||||
@@ -270,7 +289,7 @@ def test_run_job_refuses_to_overwrite_existing_shard(tmp_path):
|
||||
tmp_root.mkdir()
|
||||
|
||||
job = SimJob(detector="pbwo4", config=None, shard_index=0)
|
||||
result = run_job(job, fake, 10000, tmp_path, "steps", "gen1", tmp_root)
|
||||
result = run_job(job, fake, 10000, None, tmp_path, "steps", "gen1", tmp_root)
|
||||
|
||||
assert not result.ok
|
||||
assert "overwrite" in result.message
|
||||
@@ -285,7 +304,15 @@ def test_run_all_caps_concurrency(tmp_path):
|
||||
|
||||
jobs = [SimJob(detector="pbwo4", config=None, shard_index=i) for i in range(6)]
|
||||
results = run_all(
|
||||
jobs, fake, 10000, tmp_path, "steps", "gen1", max_workers=2, tmp_root=tmp_root
|
||||
jobs,
|
||||
fake,
|
||||
10000,
|
||||
None,
|
||||
tmp_path,
|
||||
"steps",
|
||||
"gen1",
|
||||
max_workers=2,
|
||||
tmp_root=tmp_root,
|
||||
)
|
||||
|
||||
assert all(r.ok and r.dest is not None for r in results)
|
||||
|
||||
+33
-1
@@ -1,6 +1,7 @@
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from giant.data.loader import find_parquet_files
|
||||
from giant.data.loader import build_process_map_from_files, find_parquet_files
|
||||
|
||||
|
||||
def _touch(path):
|
||||
@@ -58,3 +59,34 @@ def test_manifest_with_no_entries_raises(tmp_path):
|
||||
|
||||
with pytest.raises(FileNotFoundError):
|
||||
find_parquet_files(manifest)
|
||||
|
||||
|
||||
def test_build_process_map_from_files_keeps_most_frequent(tmp_path):
|
||||
"""process counts: eIoni=5, phot=3, compt=2, Rayl=1 — with n_experts=3, only
|
||||
the top 2 (eIoni, phot) get their own index; compt/Rayl share the "other"
|
||||
(last) index."""
|
||||
process = ["eIoni"] * 5 + ["phot"] * 3 + ["compt"] * 2 + ["Rayl"] * 1
|
||||
path = tmp_path / "shard-000.parquet"
|
||||
pd.DataFrame({"process": process}).to_parquet(path)
|
||||
|
||||
proc_map = build_process_map_from_files([path], n_experts=3)
|
||||
|
||||
assert proc_map["eIoni"] == 0
|
||||
assert proc_map["phot"] == 1
|
||||
assert proc_map["compt"] == 2
|
||||
assert proc_map["Rayl"] == 2
|
||||
assert set(proc_map.values()) <= {0, 1, 2}
|
||||
|
||||
|
||||
def test_build_process_map_from_files_spans_multiple_files(tmp_path):
|
||||
path_a = tmp_path / "a.parquet"
|
||||
path_b = tmp_path / "b.parquet"
|
||||
pd.DataFrame({"process": ["eIoni"] * 3 + ["phot"] * 1}).to_parquet(path_a)
|
||||
pd.DataFrame({"process": ["phot"] * 4 + ["compt"] * 1}).to_parquet(path_b)
|
||||
|
||||
# phot: 1+4=5 total > eIoni: 3 > compt: 1
|
||||
proc_map = build_process_map_from_files([path_a, path_b], n_experts=3)
|
||||
|
||||
assert proc_map["phot"] == 0
|
||||
assert proc_map["eIoni"] == 1
|
||||
assert proc_map["compt"] == 2
|
||||
|
||||
@@ -231,3 +231,130 @@ def test_encode_secondaries_direction_encoding():
|
||||
local_dirs = sec_cont[:, :2, 1:4] # (N, 2, 3) — valid slots only
|
||||
norms_out = np.linalg.norm(local_dirs, axis=-1)
|
||||
np.testing.assert_allclose(norms_out, 1.0, atol=1e-5)
|
||||
|
||||
|
||||
# ── decode_secondaries: exact energy conservation ────────────────────────────
|
||||
|
||||
|
||||
def _random_sec_cont(rng, N, stick_logit_scale=1.0):
|
||||
sec_cont = rng.standard_normal((N, K_MAX, 4)).astype(np.float32)
|
||||
sec_cont[:, :, 0] *= stick_logit_scale
|
||||
dirs = sec_cont[:, :, 1:]
|
||||
dirs /= np.linalg.norm(dirs, axis=-1, keepdims=True)
|
||||
return sec_cont
|
||||
|
||||
|
||||
def test_decode_secondaries_valid_slots_sum_to_e_sec():
|
||||
"""The valid slots' energies must sum to exactly e_sec, not just <= e_sec.
|
||||
|
||||
Rows with n_sec=0 are excluded: there's no slot to put the budget in, so
|
||||
valid_sum is correctly 0 regardless of e_sec there (see
|
||||
test_decode_secondaries_zero_n_sec_has_zero_energy) — the shortfall in
|
||||
that case is handled downstream (e.g. rollout.py dumps it into edep).
|
||||
"""
|
||||
from giant.data.transforms import decode_secondaries
|
||||
|
||||
rng = np.random.default_rng(0)
|
||||
N = 200
|
||||
sec_cont = _random_sec_cont(rng, N)
|
||||
sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)
|
||||
n_sec = rng.integers(0, K_MAX + 1, size=N)
|
||||
e_sec = rng.uniform(0.0, 50.0, size=N).astype(np.float32)
|
||||
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
|
||||
|
||||
sec_E, _sec_dir, _sec_pdg, sec_valid = decode_secondaries(
|
||||
sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22}
|
||||
)
|
||||
|
||||
valid_sum = (sec_E * sec_valid).sum(axis=1)
|
||||
has_secondaries = n_sec > 0
|
||||
np.testing.assert_allclose(
|
||||
valid_sum[has_secondaries],
|
||||
e_sec[has_secondaries],
|
||||
atol=1e-3,
|
||||
rtol=1e-5,
|
||||
)
|
||||
|
||||
|
||||
def test_decode_secondaries_zero_n_sec_has_zero_energy():
|
||||
"""n_sec=0 rows get no secondaries and no forced energy assignment."""
|
||||
from giant.data.transforms import decode_secondaries
|
||||
|
||||
rng = np.random.default_rng(1)
|
||||
N = 10
|
||||
sec_cont = _random_sec_cont(rng, N)
|
||||
sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)
|
||||
n_sec = np.zeros(N, dtype=np.int64)
|
||||
e_sec = rng.uniform(1.0, 10.0, size=N).astype(np.float32)
|
||||
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
|
||||
|
||||
sec_E, _sec_dir, _sec_pdg, sec_valid = decode_secondaries(
|
||||
sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22}
|
||||
)
|
||||
|
||||
assert not sec_valid.any()
|
||||
np.testing.assert_allclose(sec_E, 0.0)
|
||||
|
||||
|
||||
def test_decode_secondaries_degenerate_row_falls_back_to_even_split():
|
||||
"""All-zero stick fractions for the valid slots fall back to an even split."""
|
||||
from giant.data.transforms import decode_secondaries
|
||||
|
||||
rng = np.random.default_rng(2)
|
||||
N = 4
|
||||
sec_cont = _random_sec_cont(rng, N)
|
||||
# Drive every valid slot's stick-breaking fraction to ~0 (huge negative logit).
|
||||
n_sec = np.array([0, 1, 3, K_MAX])
|
||||
for i, k in enumerate(n_sec):
|
||||
sec_cont[i, :k, 0] = -80.0
|
||||
sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)
|
||||
e_sec = np.array([0.0, 4.0, 9.0, 30.0], dtype=np.float32)
|
||||
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
|
||||
|
||||
sec_E, _sec_dir, _sec_pdg, sec_valid = decode_secondaries(
|
||||
sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22}
|
||||
)
|
||||
|
||||
for i, k in enumerate(n_sec):
|
||||
if k == 0:
|
||||
continue
|
||||
np.testing.assert_allclose(sec_E[i, :k], e_sec[i] / k, atol=1e-4)
|
||||
np.testing.assert_allclose(sec_E[i, :k].sum(), e_sec[i], atol=1e-3)
|
||||
|
||||
|
||||
def test_decode_secondaries_rescale_preserves_relative_shares():
|
||||
"""Rescaling should keep each valid slot's *share* of the budget unchanged.
|
||||
|
||||
A shortfall shouldn't get dumped into whichever slot is last by energy
|
||||
rank — it should be spread proportionally, i.e. sec_E[i] / sec_E[j] for
|
||||
two valid slots must match before and after the e_sec rescale.
|
||||
"""
|
||||
from giant.data.transforms import decode_secondaries
|
||||
|
||||
rng = np.random.default_rng(3)
|
||||
N = 1
|
||||
sec_cont = _random_sec_cont(rng, N)
|
||||
n_sec = np.array([4])
|
||||
pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64)
|
||||
|
||||
sec_E_small, _, _, sec_valid = decode_secondaries(
|
||||
sec_cont,
|
||||
sec_pdg_pred,
|
||||
n_sec,
|
||||
np.array([5.0], dtype=np.float32),
|
||||
pre_dir,
|
||||
{0: 22},
|
||||
)
|
||||
sec_E_large, _, _, _ = decode_secondaries(
|
||||
sec_cont,
|
||||
sec_pdg_pred,
|
||||
n_sec,
|
||||
np.array([50.0], dtype=np.float32),
|
||||
pre_dir,
|
||||
{0: 22},
|
||||
)
|
||||
|
||||
ratio_small = sec_E_small[0, :4] / sec_E_small[0, 0]
|
||||
ratio_large = sec_E_large[0, :4] / sec_E_large[0, 0]
|
||||
np.testing.assert_allclose(ratio_small, ratio_large, rtol=1e-4)
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Tests for the autoregressive shower rollout driver."""
|
||||
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
@@ -161,3 +162,94 @@ def test_max_tracks_cap_conserves_energy():
|
||||
leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum()
|
||||
assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)
|
||||
assert len(np.unique(rec["track_id"][m])) <= 3
|
||||
|
||||
|
||||
# ── Streaming output (on_chunk) ──────────────────────────────────────────────
|
||||
|
||||
|
||||
def _run_streaming(on_chunk, **kwargs):
|
||||
torch.manual_seed(0)
|
||||
np.random.seed(0)
|
||||
s1, s2 = _models()
|
||||
cond, tgt = _norms()
|
||||
seeds = kwargs.pop("seeds", None) or _seeds()
|
||||
return rollout(
|
||||
s1,
|
||||
s2,
|
||||
_oracle(),
|
||||
seeds,
|
||||
cond,
|
||||
tgt,
|
||||
PDG_MAP,
|
||||
MAT_MAP,
|
||||
energy_cutoff=kwargs.pop("energy_cutoff", 1.0),
|
||||
max_steps=kwargs.pop("max_steps", 30),
|
||||
steps=4,
|
||||
batch_size=128,
|
||||
max_tracks_per_event=kwargs.pop("max_tracks_per_event", 300),
|
||||
escape_threshold=kwargs.pop("escape_threshold", 1e9),
|
||||
on_chunk=on_chunk,
|
||||
)
|
||||
|
||||
|
||||
def test_on_chunk_receives_every_row_exactly_once():
|
||||
"""Concatenating the streamed chunks must reproduce the buffered result."""
|
||||
from giant.rollout import _RECORD_KEYS
|
||||
|
||||
buffered = _run()
|
||||
|
||||
chunks: list[dict[str, np.ndarray]] = []
|
||||
summary = _run_streaming(chunks.append)
|
||||
|
||||
streamed = {k: np.concatenate([c[k] for c in chunks]) for k in _RECORD_KEYS}
|
||||
assert summary["n_rows"] == len(buffered["event_id"])
|
||||
assert len(streamed["event_id"]) == len(buffered["event_id"])
|
||||
for k in _RECORD_KEYS:
|
||||
np.testing.assert_array_equal(streamed[k], buffered[k])
|
||||
|
||||
|
||||
def test_on_chunk_summary_termination_reason_counts_match_buffered():
|
||||
buffered = _run()
|
||||
summary = _run_streaming(lambda row: None)
|
||||
|
||||
expected = Counter(r for r in buffered["termination_reason"].tolist() if r)
|
||||
assert summary["termination_reason_counts"] == dict(expected)
|
||||
|
||||
|
||||
def test_on_chunk_never_buffers_full_records():
|
||||
"""Streaming mode must not accumulate rows for later to_dict() retrieval."""
|
||||
from giant.rollout import _Recorder
|
||||
|
||||
rec = _Recorder(sink=lambda row: None)
|
||||
rec.add(
|
||||
event_id=np.array([0]),
|
||||
track_id=np.array([0]),
|
||||
parent_id=np.array([-1]),
|
||||
generation=np.array([0]),
|
||||
step_no=np.array([0]),
|
||||
pdg=np.array([11]),
|
||||
pre_x=np.array([0.0]),
|
||||
pre_y=np.array([0.0]),
|
||||
pre_z=np.array([0.0]),
|
||||
pre_E=np.array([1.0]),
|
||||
pre_dx=np.array([0.0]),
|
||||
pre_dy=np.array([0.0]),
|
||||
pre_dz=np.array([1.0]),
|
||||
post_x=np.array([0.0]),
|
||||
post_y=np.array([0.0]),
|
||||
post_z=np.array([1.0]),
|
||||
post_E=np.array([0.0]),
|
||||
post_dx=np.array([0.0]),
|
||||
post_dy=np.array([0.0]),
|
||||
post_dz=np.array([1.0]),
|
||||
edep=np.array([1.0]),
|
||||
step_length=np.array([1.0]),
|
||||
material=np.array(["G4_AIR"], dtype=object),
|
||||
layer_id=np.array([0]),
|
||||
n_sec_pred=np.array([0]),
|
||||
termination_reason=np.array(["natural_end"], dtype=object),
|
||||
)
|
||||
assert rec.n_rows == 1
|
||||
assert rec.termination_reason_counts == {"natural_end": 1}
|
||||
with pytest.raises(AssertionError):
|
||||
rec.to_dict()
|
||||
|
||||
@@ -0,0 +1,704 @@
|
||||
"""Tests for the mixture-of-experts routing prototype (giant/model/network.py)."""
|
||||
|
||||
import torch
|
||||
|
||||
from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
|
||||
from giant.model.network import (
|
||||
ComposedRouter,
|
||||
DenoisingMLP,
|
||||
EnergyRouter,
|
||||
PdgRouter,
|
||||
ProcessRouter,
|
||||
ROUTER_REGISTRY,
|
||||
RoutedDenoisingMLP,
|
||||
RoutedSecondaryDecoder,
|
||||
SecondaryDecoder,
|
||||
build_composed_router,
|
||||
build_models,
|
||||
build_router,
|
||||
)
|
||||
|
||||
|
||||
def _cond(B=8, pdg=3, mat=2):
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.stack(
|
||||
[torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
|
||||
)
|
||||
return cond_cont, cond_cat
|
||||
|
||||
|
||||
def _routed_stage1(n_experts=4, pdg=3, mat=2, **router_kwargs):
|
||||
router = build_router("energy", n_experts, **router_kwargs)
|
||||
return RoutedDenoisingMLP(
|
||||
pdg_vocab=pdg,
|
||||
mat_vocab=mat,
|
||||
router=router,
|
||||
expert_hidden_dim=16,
|
||||
expert_n_blocks=2,
|
||||
)
|
||||
|
||||
|
||||
def _routed_sec_decoder(n_experts=4, pdg=3, mat=2, **router_kwargs):
|
||||
router = build_router("energy", n_experts, **router_kwargs)
|
||||
return RoutedSecondaryDecoder(
|
||||
pdg_vocab=pdg,
|
||||
mat_vocab=mat,
|
||||
router=router,
|
||||
expert_hidden_dim=16,
|
||||
expert_n_blocks=2,
|
||||
)
|
||||
|
||||
|
||||
# ── Router / EnergyRouter contract ──────────────────────────────────────────
|
||||
|
||||
|
||||
def test_energy_router_registered():
|
||||
assert ROUTER_REGISTRY["energy"] is EnergyRouter
|
||||
|
||||
|
||||
def test_energy_router_gate_partition_of_unity():
|
||||
router = EnergyRouter(n_experts=4)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
g = router.gate(cond_cont, cond_cat)
|
||||
assert g.shape == (16, 4)
|
||||
torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
|
||||
|
||||
|
||||
def test_energy_router_top1_matches_gate_argmax():
|
||||
router = EnergyRouter(n_experts=4)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
assert torch.equal(
|
||||
router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
|
||||
)
|
||||
|
||||
|
||||
def test_energy_router_hardens_as_temperature_shrinks():
|
||||
"""As tau -> 0 the soft gate should converge to a one-hot at the argmax."""
|
||||
router = EnergyRouter(n_experts=4, temperature=1e-4)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
g = router.gate(cond_cont, cond_cat)
|
||||
top1 = router.top1(cond_cont, cond_cat)
|
||||
onehot = torch.nn.functional.one_hot(top1, num_classes=4).float()
|
||||
torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0)
|
||||
|
||||
|
||||
def test_energy_router_balance_loss_is_nonnegative_scalar():
|
||||
router = EnergyRouter(n_experts=4)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
loss = router.balance_loss(cond_cont, cond_cat)
|
||||
assert loss.shape == ()
|
||||
assert loss.item() >= 0.0
|
||||
|
||||
|
||||
def test_build_router_ignores_unrecognized_kwargs():
|
||||
# lambda_balance is a model_config.router key but not an EnergyRouter kwarg
|
||||
router = build_router("energy", 4, temperature=0.3, lambda_balance=0.5)
|
||||
assert isinstance(router, EnergyRouter)
|
||||
assert router.temperature == 0.3
|
||||
|
||||
|
||||
def test_build_router_unknown_type_raises():
|
||||
try:
|
||||
build_router("nonexistent", 4)
|
||||
except ValueError:
|
||||
return
|
||||
raise AssertionError("expected ValueError for unknown router type")
|
||||
|
||||
|
||||
# ── PdgRouter ────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_pdg_router_registered():
|
||||
assert ROUTER_REGISTRY["pdg"] is PdgRouter
|
||||
|
||||
|
||||
def test_pdg_router_gate_partition_of_unity():
|
||||
router = PdgRouter(n_experts=4, pdg_vocab=3)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
g = router.gate(cond_cont, cond_cat)
|
||||
assert g.shape == (16, 4)
|
||||
torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
|
||||
|
||||
|
||||
def test_pdg_router_top1_matches_gate_argmax():
|
||||
router = PdgRouter(n_experts=4, pdg_vocab=3)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
assert torch.equal(
|
||||
router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
|
||||
)
|
||||
|
||||
|
||||
def test_pdg_router_hardens_as_temperature_shrinks():
|
||||
"""As tau -> 0 the soft gate should converge to a one-hot at the argmax."""
|
||||
router = PdgRouter(n_experts=4, pdg_vocab=3, temperature=1e-4)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
g = router.gate(cond_cont, cond_cat)
|
||||
top1 = router.top1(cond_cont, cond_cat)
|
||||
onehot = torch.nn.functional.one_hot(top1, num_classes=4).float()
|
||||
torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0)
|
||||
|
||||
|
||||
def test_pdg_router_balance_loss_is_nonnegative_scalar():
|
||||
router = PdgRouter(n_experts=4, pdg_vocab=3)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
loss = router.balance_loss(cond_cont, cond_cat)
|
||||
assert loss.shape == ()
|
||||
assert loss.item() >= 0.0
|
||||
|
||||
|
||||
def test_pdg_router_classify_loss_defaults_to_zero():
|
||||
"""PDG is already known at gate time (unlike ProcessRouter's process
|
||||
label), so no supervision is needed — falls back to Router's default."""
|
||||
router = PdgRouter(n_experts=4, pdg_vocab=3)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
labels = torch.randint(0, 4, (16,))
|
||||
loss = router.classify_loss(cond_cont, cond_cat, labels)
|
||||
assert loss.shape == ()
|
||||
assert loss.item() == 0.0
|
||||
|
||||
|
||||
def test_pdg_router_only_reads_pdg_column():
|
||||
"""Gate must depend on cond_cat[:, 0] (pdg) only, not cond_cont or material."""
|
||||
router = PdgRouter(n_experts=4, pdg_vocab=3)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
g_before = router.gate(cond_cont, cond_cat)
|
||||
|
||||
cond_cont_perturbed = torch.randn_like(cond_cont)
|
||||
cond_cat_diff_mat = cond_cat.clone()
|
||||
cond_cat_diff_mat[:, 1] = (cond_cat_diff_mat[:, 1] + 1) % 2
|
||||
g_after = router.gate(cond_cont_perturbed, cond_cat_diff_mat)
|
||||
|
||||
torch.testing.assert_close(g_before, g_after, atol=1e-6, rtol=0)
|
||||
|
||||
|
||||
def test_build_router_pdg_type_uses_pdg_vocab():
|
||||
router = build_router("pdg", 4, pdg_vocab=5, mat_vocab=3, emb_dim=8)
|
||||
assert isinstance(router, PdgRouter)
|
||||
assert router.pdg_emb.num_embeddings == 5
|
||||
|
||||
|
||||
def test_build_models_routed_with_pdg_router():
|
||||
model_config = dict(
|
||||
pdg_vocab=4,
|
||||
mat_vocab=2,
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=16,
|
||||
expert_n_blocks=2,
|
||||
router={
|
||||
"enabled": True,
|
||||
"type": "pdg",
|
||||
"n_experts": 3,
|
||||
},
|
||||
)
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
assert isinstance(stage1, RoutedDenoisingMLP)
|
||||
assert isinstance(stage1.router, PdgRouter)
|
||||
assert len(stage1.experts) == 3
|
||||
assert stage1.router.pdg_emb.num_embeddings == 4
|
||||
|
||||
|
||||
# ── ProcessRouter ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_process_router_registered():
|
||||
assert ROUTER_REGISTRY["process"] is ProcessRouter
|
||||
|
||||
|
||||
def test_process_router_gate_partition_of_unity():
|
||||
router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
g = router.gate(cond_cont, cond_cat)
|
||||
assert g.shape == (16, 4)
|
||||
torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
|
||||
|
||||
|
||||
def test_process_router_top1_matches_gate_argmax():
|
||||
router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
assert torch.equal(
|
||||
router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
|
||||
)
|
||||
|
||||
|
||||
def test_process_router_balance_loss_is_nonnegative_scalar():
|
||||
router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
loss = router.balance_loss(cond_cont, cond_cat)
|
||||
assert loss.shape == ()
|
||||
assert loss.item() >= 0.0
|
||||
|
||||
|
||||
def test_process_router_classify_loss_decreases_with_training():
|
||||
"""The classifier should be able to fit an arbitrary label assignment —
|
||||
a sanity check that gradients actually flow to the process classifier."""
|
||||
torch.manual_seed(0)
|
||||
router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
|
||||
cond_cont, cond_cat = _cond(32)
|
||||
labels = torch.randint(0, 4, (32,))
|
||||
|
||||
opt = torch.optim.Adam(router.parameters(), lr=0.05)
|
||||
first = router.classify_loss(cond_cont, cond_cat, labels).item()
|
||||
for _ in range(50):
|
||||
opt.zero_grad()
|
||||
loss = router.classify_loss(cond_cont, cond_cat, labels)
|
||||
loss.backward()
|
||||
opt.step()
|
||||
last = loss.item()
|
||||
assert last < first
|
||||
|
||||
|
||||
def test_energy_router_classify_loss_defaults_to_zero():
|
||||
"""Routers with no supervised signal (EnergyRouter) fall back to the
|
||||
Router base class's zero-loss default."""
|
||||
router = EnergyRouter(n_experts=4)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
labels = torch.randint(0, 4, (16,))
|
||||
loss = router.classify_loss(cond_cont, cond_cat, labels)
|
||||
assert loss.shape == ()
|
||||
assert loss.item() == 0.0
|
||||
|
||||
|
||||
def test_build_router_process_type_uses_pdg_mat_vocab():
|
||||
router = build_router("process", 4, pdg_vocab=5, mat_vocab=3, emb_dim=8)
|
||||
assert isinstance(router, ProcessRouter)
|
||||
assert router.pdg_emb.num_embeddings == 5
|
||||
assert router.mat_emb.num_embeddings == 3
|
||||
|
||||
|
||||
def test_build_models_routed_with_process_router():
|
||||
model_config = dict(
|
||||
pdg_vocab=4,
|
||||
mat_vocab=2,
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=16,
|
||||
expert_n_blocks=2,
|
||||
router={
|
||||
"enabled": True,
|
||||
"type": "process",
|
||||
"n_experts": 3,
|
||||
"lambda_proc": 1.0,
|
||||
},
|
||||
)
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
assert isinstance(stage1, RoutedDenoisingMLP)
|
||||
assert isinstance(stage1.router, ProcessRouter)
|
||||
assert len(stage1.experts) == 3
|
||||
assert stage1.router.pdg_emb.num_embeddings == 4
|
||||
assert stage1.router.mat_emb.num_embeddings == 2
|
||||
|
||||
|
||||
# ── ComposedRouter ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_composed_router_n_experts_is_product():
|
||||
router = ComposedRouter(
|
||||
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
|
||||
)
|
||||
assert router.n_experts == 12
|
||||
|
||||
|
||||
def test_composed_router_gate_partition_of_unity():
|
||||
router = ComposedRouter(
|
||||
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
|
||||
)
|
||||
cond_cont, cond_cat = _cond(16, pdg=5)
|
||||
g = router.gate(cond_cont, cond_cat)
|
||||
assert g.shape == (16, 12)
|
||||
torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
|
||||
|
||||
|
||||
def test_composed_router_gate_is_outer_product_of_sub_gates():
|
||||
energy_router = EnergyRouter(n_experts=4)
|
||||
pdg_router = PdgRouter(n_experts=3, pdg_vocab=5)
|
||||
router = ComposedRouter([energy_router, pdg_router])
|
||||
cond_cont, cond_cat = _cond(16, pdg=5)
|
||||
|
||||
g_energy = energy_router.gate(cond_cont, cond_cat) # (16, 4)
|
||||
g_pdg = pdg_router.gate(cond_cont, cond_cat) # (16, 3)
|
||||
expected = (g_energy.unsqueeze(-1) * g_pdg.unsqueeze(1)).flatten(1) # (16, 12)
|
||||
|
||||
torch.testing.assert_close(router.gate(cond_cont, cond_cat), expected)
|
||||
|
||||
|
||||
def test_composed_router_top1_factors_into_per_axis_argmax():
|
||||
"""Joint argmax over the outer product must equal the pair of per-axis
|
||||
argmaxes, flattened with the same row-major index convention as gate()."""
|
||||
energy_router = EnergyRouter(n_experts=4)
|
||||
pdg_router = PdgRouter(n_experts=3, pdg_vocab=5)
|
||||
router = ComposedRouter([energy_router, pdg_router])
|
||||
cond_cont, cond_cat = _cond(16, pdg=5)
|
||||
|
||||
joint_idx = router.top1(cond_cont, cond_cat)
|
||||
energy_idx = energy_router.top1(cond_cont, cond_cat)
|
||||
pdg_idx = pdg_router.top1(cond_cont, cond_cat)
|
||||
expected = energy_idx * pdg_router.n_experts + pdg_idx
|
||||
|
||||
assert torch.equal(joint_idx, expected)
|
||||
|
||||
|
||||
def test_composed_router_supports_different_expert_counts_per_axis():
|
||||
router = ComposedRouter(
|
||||
[EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)]
|
||||
)
|
||||
assert router.n_experts == 10
|
||||
cond_cont, cond_cat = _cond(8, pdg=5)
|
||||
assert router.gate(cond_cont, cond_cat).shape == (8, 10)
|
||||
|
||||
|
||||
def test_composed_router_classify_loss_sums_sub_router_losses():
|
||||
"""energy/pdg both default to zero, so the composed loss should too."""
|
||||
router = ComposedRouter(
|
||||
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
|
||||
)
|
||||
cond_cont, cond_cat = _cond(16, pdg=5)
|
||||
labels = torch.randint(0, 4, (16,))
|
||||
loss = router.classify_loss(cond_cont, cond_cat, labels)
|
||||
assert loss.shape == ()
|
||||
assert loss.item() == 0.0
|
||||
|
||||
|
||||
def test_composed_router_rejects_empty_router_list():
|
||||
try:
|
||||
ComposedRouter([])
|
||||
except ValueError:
|
||||
return
|
||||
raise AssertionError("expected ValueError for empty router list")
|
||||
|
||||
|
||||
def test_composed_router_not_in_registry():
|
||||
assert "composed" not in ROUTER_REGISTRY
|
||||
|
||||
|
||||
# ── _parse_composed_axes (axis{i}_{field} flat-key config convention) ───────
|
||||
|
||||
|
||||
def test_parse_composed_axes_groups_indexed_keys():
|
||||
from giant.model.network import _parse_composed_axes
|
||||
|
||||
router_cfg = {
|
||||
"enabled": True,
|
||||
"type": "composed",
|
||||
"axis0_type": "energy",
|
||||
"axis0_n_experts": 4,
|
||||
"axis0_temperature": 0.3,
|
||||
"axis1_type": "pdg",
|
||||
"axis1_n_experts": 3,
|
||||
"axis1_emb_dim": 6,
|
||||
}
|
||||
axes = _parse_composed_axes(router_cfg)
|
||||
assert axes == [
|
||||
{"type": "energy", "n_experts": 4, "temperature": 0.3},
|
||||
{"type": "pdg", "n_experts": 3, "emb_dim": 6},
|
||||
]
|
||||
|
||||
|
||||
def test_parse_composed_axes_ignores_unrelated_keys():
|
||||
from giant.model.network import _parse_composed_axes
|
||||
|
||||
router_cfg = {
|
||||
"enabled": True,
|
||||
"type": "composed",
|
||||
"lambda_balance": 0.0,
|
||||
"axis0_type": "energy",
|
||||
"axis0_n_experts": 4,
|
||||
}
|
||||
axes = _parse_composed_axes(router_cfg)
|
||||
assert axes == [{"type": "energy", "n_experts": 4}]
|
||||
|
||||
|
||||
def test_parse_composed_axes_raises_on_index_gap():
|
||||
from giant.model.network import _parse_composed_axes
|
||||
|
||||
router_cfg = {
|
||||
"type": "composed",
|
||||
"axis0_type": "energy",
|
||||
"axis0_n_experts": 4,
|
||||
# axis1 missing entirely
|
||||
"axis2_type": "pdg",
|
||||
"axis2_n_experts": 3,
|
||||
}
|
||||
try:
|
||||
_parse_composed_axes(router_cfg)
|
||||
except ValueError:
|
||||
return
|
||||
raise AssertionError("expected ValueError for a gap in axis indices")
|
||||
|
||||
|
||||
def test_build_composed_router_resolves_per_axis_specs():
|
||||
router = build_composed_router(
|
||||
[
|
||||
{"type": "energy", "n_experts": 4, "temperature": 0.3},
|
||||
{"type": "pdg", "n_experts": 3, "emb_dim": 6},
|
||||
],
|
||||
pdg_vocab=5,
|
||||
mat_vocab=2,
|
||||
)
|
||||
assert isinstance(router, ComposedRouter)
|
||||
assert router.n_experts == 12
|
||||
energy_router, pdg_router = router.routers
|
||||
assert isinstance(energy_router, EnergyRouter)
|
||||
assert energy_router.temperature == 0.3
|
||||
assert isinstance(pdg_router, PdgRouter)
|
||||
assert pdg_router.pdg_emb.num_embeddings == 5
|
||||
assert pdg_router.pdg_emb.embedding_dim == 6
|
||||
|
||||
|
||||
def test_build_models_routed_with_composed_router():
|
||||
model_config = dict(
|
||||
pdg_vocab=5,
|
||||
mat_vocab=2,
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=16,
|
||||
expert_n_blocks=2,
|
||||
router={
|
||||
"enabled": True,
|
||||
"type": "composed",
|
||||
"axis0_type": "energy",
|
||||
"axis0_n_experts": 4,
|
||||
"axis1_type": "pdg",
|
||||
"axis1_n_experts": 3,
|
||||
},
|
||||
)
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
assert isinstance(stage1, RoutedDenoisingMLP)
|
||||
assert isinstance(stage1.router, ComposedRouter)
|
||||
assert len(stage1.experts) == 12
|
||||
assert len(sec_decoder.experts) == 12
|
||||
# stage1 and sec_decoder must not share router weights (same convention
|
||||
# as the single-axis routers built by build_models).
|
||||
assert stage1.router is not sec_decoder.router
|
||||
|
||||
|
||||
def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow():
|
||||
from giant.sample import sample_flow, sample_secondaries
|
||||
|
||||
model_config = dict(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=8,
|
||||
expert_n_blocks=1,
|
||||
router={
|
||||
"enabled": True,
|
||||
"type": "composed",
|
||||
"axis0_type": "energy",
|
||||
"axis0_n_experts": 2,
|
||||
"axis1_type": "pdg",
|
||||
"axis1_n_experts": 2,
|
||||
},
|
||||
)
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
B = 5
|
||||
cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
|
||||
stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
|
||||
assert stage1_norm.shape == (B, X_DIM)
|
||||
assert n_sec_pred.shape == (B,)
|
||||
|
||||
sec_cont, sec_type_emb, sec_valid = sample_secondaries(
|
||||
sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
|
||||
)
|
||||
assert sec_cont.shape == (B, K_MAX, 4)
|
||||
assert sec_valid.shape == (B, K_MAX)
|
||||
|
||||
|
||||
# ── RoutedDenoisingMLP ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_routed_denoising_mlp_output_shape_train_and_eval():
|
||||
B = 8
|
||||
model = _routed_stage1()
|
||||
x_t = torch.randn(B, X_DIM)
|
||||
t = torch.rand(B)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
|
||||
model.train()
|
||||
out_train = model(x_t, t, cond_cont, cond_cat)
|
||||
assert out_train.shape == (B, X_DIM)
|
||||
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
out_eval = model(x_t, t, cond_cont, cond_cat)
|
||||
assert out_eval.shape == (B, X_DIM)
|
||||
|
||||
|
||||
def test_routed_denoising_mlp_gradients_flow_in_train_mode():
|
||||
"""Soft mixture in train mode should touch every expert's parameters."""
|
||||
B = 8
|
||||
model = _routed_stage1(n_experts=3)
|
||||
x_t = torch.randn(B, X_DIM)
|
||||
t = torch.rand(B)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
model.train()
|
||||
flow_loss = model(x_t, t, cond_cont, cond_cat).sum()
|
||||
nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum()
|
||||
(flow_loss + nsec_loss).backward()
|
||||
for name, p in model.named_parameters():
|
||||
assert p.grad is not None, f"no grad for {name}"
|
||||
|
||||
|
||||
def test_routed_denoising_mlp_eval_dispatch_matches_manual_grouping():
|
||||
"""Eval-mode grouped top-1 dispatch must equal running each row through
|
||||
its assigned expert individually (batch order shouldn't matter)."""
|
||||
B = 12
|
||||
model = _routed_stage1(n_experts=4)
|
||||
model.eval()
|
||||
x_t = torch.randn(B, X_DIM)
|
||||
t = torch.rand(B)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
|
||||
with torch.no_grad():
|
||||
batched = model(x_t, t, cond_cont, cond_cat)
|
||||
|
||||
t_emb = model.time_emb(t)
|
||||
c_emb = model.cond_enc(cond_cont, cond_cat)
|
||||
cond = torch.cat([t_emb, c_emb], dim=-1)
|
||||
idx = model.router.top1(cond_cont, cond_cat)
|
||||
manual = torch.zeros_like(x_t)
|
||||
for i in range(B):
|
||||
manual[i] = model.experts[int(idx[i])](x_t[i : i + 1], cond[i : i + 1])[0]
|
||||
|
||||
torch.testing.assert_close(batched, manual, atol=1e-5, rtol=1e-4)
|
||||
|
||||
|
||||
def test_routed_denoising_mlp_predict_n_sec_shape():
|
||||
B = 6
|
||||
model = _routed_stage1()
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
logits = model.predict_n_sec(cond_cont, cond_cat)
|
||||
assert logits.shape == (B, K_MAX + 1)
|
||||
|
||||
|
||||
def test_routed_denoising_mlp_pdg_embedding_weight_shape():
|
||||
model = _routed_stage1(pdg=5, mat=2)
|
||||
from giant.constants import EMB_DIM
|
||||
|
||||
assert model.pdg_embedding_weight().shape == (5, EMB_DIM)
|
||||
|
||||
|
||||
# ── RoutedSecondaryDecoder ───────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_routed_secondary_decoder_output_shape_train_and_eval():
|
||||
B = 8
|
||||
decoder = _routed_sec_decoder()
|
||||
x_t = torch.randn(B, SEC_DIM)
|
||||
t = torch.rand(B)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
|
||||
decoder.train()
|
||||
out_train = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
|
||||
assert out_train.shape == (B, SEC_DIM)
|
||||
|
||||
decoder.eval()
|
||||
with torch.no_grad():
|
||||
out_eval = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
|
||||
assert out_eval.shape == (B, SEC_DIM)
|
||||
|
||||
|
||||
def test_routed_secondary_decoder_gradients_flow():
|
||||
B = 4
|
||||
decoder = _routed_sec_decoder(n_experts=3)
|
||||
x_t = torch.randn(B, SEC_DIM)
|
||||
t = torch.rand(B)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
decoder.train()
|
||||
decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().backward()
|
||||
for name, p in decoder.named_parameters():
|
||||
assert p.grad is not None, f"no grad for {name}"
|
||||
|
||||
|
||||
# ── build_models dispatch ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_build_models_monolith_when_router_absent():
|
||||
model_config = dict(
|
||||
pdg_vocab=4,
|
||||
mat_vocab=2,
|
||||
hidden_dim=32,
|
||||
n_blocks=2,
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
)
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
assert isinstance(stage1, DenoisingMLP)
|
||||
assert isinstance(sec_decoder, SecondaryDecoder)
|
||||
|
||||
|
||||
def test_build_models_monolith_when_router_disabled():
|
||||
model_config = dict(
|
||||
pdg_vocab=4,
|
||||
mat_vocab=2,
|
||||
hidden_dim=32,
|
||||
n_blocks=2,
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
router={"enabled": False, "type": "energy", "n_experts": 4},
|
||||
)
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
assert isinstance(stage1, DenoisingMLP)
|
||||
assert isinstance(sec_decoder, SecondaryDecoder)
|
||||
|
||||
|
||||
def test_build_models_routed_when_enabled():
|
||||
model_config = dict(
|
||||
pdg_vocab=4,
|
||||
mat_vocab=2,
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=16,
|
||||
expert_n_blocks=2,
|
||||
router={
|
||||
"enabled": True,
|
||||
"type": "energy",
|
||||
"n_experts": 4,
|
||||
"temperature": 0.5,
|
||||
"learn_centers": True,
|
||||
"lambda_balance": 0.0,
|
||||
},
|
||||
)
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
assert isinstance(stage1, RoutedDenoisingMLP)
|
||||
assert isinstance(sec_decoder, RoutedSecondaryDecoder)
|
||||
assert len(stage1.experts) == 4
|
||||
assert len(sec_decoder.experts) == 4
|
||||
|
||||
|
||||
def test_build_models_routed_pair_is_drop_in_for_sample_flow():
|
||||
"""Exercise the exact calling convention giant/sample.py uses."""
|
||||
from giant.sample import sample_flow, sample_secondaries
|
||||
|
||||
model_config = dict(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=8,
|
||||
expert_n_blocks=1,
|
||||
router={"enabled": True, "type": "energy", "n_experts": 2},
|
||||
)
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
B = 5
|
||||
cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
|
||||
stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
|
||||
assert stage1_norm.shape == (B, X_DIM)
|
||||
assert n_sec_pred.shape == (B,)
|
||||
|
||||
sec_cont, sec_type_emb, sec_valid = sample_secondaries(
|
||||
sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
|
||||
)
|
||||
assert sec_cont.shape == (B, K_MAX, 4)
|
||||
assert sec_valid.shape == (B, K_MAX)
|
||||
@@ -69,7 +69,9 @@ def test_orphaned_child_track_is_dropped_not_nulled():
|
||||
}
|
||||
)
|
||||
out, n_orphaned = steps_to_parquet._add_secondary_attributes(df)
|
||||
row = out.filter((pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0))
|
||||
row = out.filter(
|
||||
(pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0)
|
||||
)
|
||||
|
||||
assert n_orphaned == 1
|
||||
assert row["child_track_ids"].to_list() == [[2]]
|
||||
|
||||
@@ -235,25 +235,22 @@ def test_build_features_clamps_n_sec_label_to_k_max():
|
||||
pdg_map = {11: 0}
|
||||
mat_map = {"PbWO4": 0}
|
||||
|
||||
_, _, _, n_sec, _, _, _, _ = build_features(data, pdg_map, mat_map)
|
||||
_, _, _, n_sec, _, _, _, _, _ = build_features(data, pdg_map, mat_map)
|
||||
|
||||
assert n_sec.max() <= K_MAX
|
||||
np.testing.assert_array_equal(n_sec, [0, 5, K_MAX])
|
||||
|
||||
|
||||
def _step_data_no_sec_lists(n_sec: np.ndarray) -> dict:
|
||||
"""Minimal build_features input with n_sec but no per-secondary list columns
|
||||
(mimics a parquet that skipped the parent->child join)."""
|
||||
N = len(n_sec)
|
||||
def _minimal_step_data(N: int, process: np.ndarray | None = None) -> dict:
|
||||
rng = np.random.default_rng(0)
|
||||
return {
|
||||
data = {
|
||||
"pdg": np.full(N, 11, dtype=np.int32),
|
||||
"material": np.full(N, "PbWO4", dtype=object),
|
||||
"pre_pos": rng.standard_normal((N, 3)).astype(np.float32),
|
||||
"pre_E": np.full(N, 10.0, dtype=np.float32),
|
||||
"pre_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)),
|
||||
"layer_id": np.zeros(N, dtype=np.int32),
|
||||
"n_sec": np.asarray(n_sec, dtype=np.int32),
|
||||
"n_sec": np.zeros(N, dtype=np.int32),
|
||||
"e_sec": np.full(N, 1.0, dtype=np.float32),
|
||||
"step_length": np.full(N, 1.0, dtype=np.float32),
|
||||
"post_E": np.full(N, 9.0, dtype=np.float32),
|
||||
@@ -261,6 +258,40 @@ def _step_data_no_sec_lists(n_sec: np.ndarray) -> dict:
|
||||
"post_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)),
|
||||
"post_pos": rng.standard_normal((N, 3)).astype(np.float32),
|
||||
}
|
||||
if process is not None:
|
||||
data["process"] = process
|
||||
return data
|
||||
|
||||
|
||||
def _step_data_no_sec_lists(n_sec: np.ndarray) -> dict:
|
||||
"""Minimal build_features input with n_sec but no per-secondary list columns
|
||||
(mimics a parquet that skipped the parent->child join)."""
|
||||
data = _minimal_step_data(len(n_sec))
|
||||
data["n_sec"] = np.asarray(n_sec, dtype=np.int32)
|
||||
return data
|
||||
|
||||
|
||||
def test_build_features_proc_idx_zero_without_proc_map():
|
||||
data = _minimal_step_data(
|
||||
3, process=np.array(["compt", "phot", "eIoni"], dtype=object)
|
||||
)
|
||||
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
|
||||
|
||||
*_, proc_idx, _, _ = build_features(data, pdg_map, mat_map)
|
||||
|
||||
np.testing.assert_array_equal(proc_idx, [0, 0, 0])
|
||||
|
||||
|
||||
def test_build_features_proc_idx_looks_up_proc_map():
|
||||
data = _minimal_step_data(
|
||||
3, process=np.array(["compt", "phot", "eIoni"], dtype=object)
|
||||
)
|
||||
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
|
||||
proc_map = {"compt": 0, "phot": 1, "eIoni": 2}
|
||||
|
||||
*_, proc_idx, _, _ = build_features(data, pdg_map, mat_map, proc_map=proc_map)
|
||||
|
||||
np.testing.assert_array_equal(proc_idx, [0, 1, 2])
|
||||
|
||||
|
||||
def test_build_features_require_secondaries_raises_when_lists_missing():
|
||||
@@ -281,7 +312,7 @@ def test_build_features_require_secondaries_ok_when_no_secondaries():
|
||||
data = _step_data_no_sec_lists(np.zeros(3, dtype=np.int32))
|
||||
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
|
||||
|
||||
_, _, _, _, sec_cont, sec_pdg_idx, _, _ = build_features(
|
||||
_, _, _, _, sec_cont, sec_pdg_idx, *_ = build_features(
|
||||
data, pdg_map, mat_map, require_secondaries=True
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user