Implement Phase 2: secondary particle prediction
Two-stage factorisation: Stage 1 predicts 9D primary kinematics + n_sec classification head (COND_DIM reduced to 8, dropping n_sec/e_sec inputs); Stage 2 (SecondaryDecoder) generates K_MAX=15 secondary slots via masked flow matching over (stick_logit, local_dir, type_emb) conditioned on Stage 1 output. Joint training with combined loss L_s1 + λ_nsec*L_nsec + λ_s2*L_s2. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
+24
-58
@@ -1,67 +1,33 @@
|
||||
import numpy as np
|
||||
from giant.data.dataset import StepsDataset, train_val_split
|
||||
from giant.data.dataset import make_event_split
|
||||
|
||||
|
||||
def _dummy(N=500, n_events=20):
|
||||
def test_make_event_split_sizes():
|
||||
rng = np.random.default_rng(42)
|
||||
data = {"event_id": rng.integers(0, n_events, size=N)}
|
||||
cond_cont = rng.standard_normal((N, 9)).astype(np.float32)
|
||||
cond_cat = rng.integers(0, 3, size=(N, 2)).astype(np.int64)
|
||||
target = rng.standard_normal((N, 6)).astype(np.float32)
|
||||
return data, cond_cont, cond_cat, target
|
||||
event_ids = rng.integers(0, 50, size=1000)
|
||||
train_set, val_set = make_event_split(event_ids, val_fraction=0.2)
|
||||
unique = np.unique(event_ids)
|
||||
assert len(train_set) + len(val_set) == len(unique)
|
||||
|
||||
|
||||
def test_dataset_length():
|
||||
data, cond_cont, cond_cat, target = _dummy()
|
||||
assert len(StepsDataset(cond_cont, cond_cat, target)) == len(target)
|
||||
|
||||
|
||||
def test_dataset_item_shapes():
|
||||
data, cond_cont, cond_cat, target = _dummy()
|
||||
c, k, t = StepsDataset(cond_cont, cond_cat, target)[0]
|
||||
assert c.shape == (9,)
|
||||
assert k.shape == (2,)
|
||||
assert t.shape == (6,)
|
||||
|
||||
|
||||
def test_split_sizes_sum_to_total():
|
||||
data, cond_cont, cond_cat, target = _dummy(N=500)
|
||||
train_ds, val_ds = train_val_split(
|
||||
data, cond_cont, cond_cat, target, val_fraction=0.2
|
||||
)
|
||||
assert len(train_ds) + len(val_ds) == 500
|
||||
|
||||
|
||||
def test_split_no_empty_sets():
|
||||
data, cond_cont, cond_cat, target = _dummy(N=500, n_events=20)
|
||||
train_ds, val_ds = train_val_split(
|
||||
data, cond_cont, cond_cat, target, val_fraction=0.2
|
||||
)
|
||||
assert len(val_ds) > 0
|
||||
assert len(train_ds) > 0
|
||||
|
||||
|
||||
def test_split_event_leakage():
|
||||
"""Train and val must not share any event_id."""
|
||||
N = 1000
|
||||
n_events = 50
|
||||
def test_make_event_split_no_overlap():
|
||||
rng = np.random.default_rng(7)
|
||||
event_ids = rng.integers(0, n_events, size=N)
|
||||
data = {"event_id": event_ids}
|
||||
cond_cont = rng.standard_normal((N, 9)).astype(np.float32)
|
||||
cond_cat = rng.integers(0, 3, size=(N, 2)).astype(np.int64)
|
||||
target = rng.standard_normal((N, 6)).astype(np.float32)
|
||||
event_ids = rng.integers(0, 50, size=1000)
|
||||
train_set, val_set = make_event_split(event_ids, val_fraction=0.2)
|
||||
assert train_set.isdisjoint(val_set)
|
||||
|
||||
train_ds, val_ds = train_val_split(
|
||||
data, cond_cont, cond_cat, target, val_fraction=0.2
|
||||
)
|
||||
|
||||
# Recover which event_ids ended up in each split via the indices
|
||||
# (The dataset doesn't store event_ids, so we check via the original mask logic)
|
||||
unique_events = np.unique(event_ids)
|
||||
rng2 = np.random.default_rng(42)
|
||||
rng2.shuffle(unique_events)
|
||||
n_val = max(1, int(len(unique_events) * 0.2))
|
||||
val_events = set(unique_events[:n_val].tolist())
|
||||
train_events = set(unique_events[n_val:].tolist())
|
||||
assert val_events.isdisjoint(train_events)
|
||||
def test_make_event_split_no_empty_sets():
|
||||
rng = np.random.default_rng(0)
|
||||
event_ids = rng.integers(0, 20, size=500)
|
||||
train_set, val_set = make_event_split(event_ids, val_fraction=0.2)
|
||||
assert len(train_set) > 0
|
||||
assert len(val_set) > 0
|
||||
|
||||
|
||||
def test_make_event_split_reproducible():
|
||||
event_ids = np.arange(100)
|
||||
a_tr, a_val = make_event_split(event_ids, val_fraction=0.1, seed=42)
|
||||
b_tr, b_val = make_event_split(event_ids, val_fraction=0.1, seed=42)
|
||||
assert a_tr == b_tr
|
||||
assert a_val == b_val
|
||||
|
||||
+6
-4
@@ -39,8 +39,9 @@ def test_sample_flow_shape():
|
||||
B = 6
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
out = sample_flow(_small_model(), cond_cont, cond_cat, steps=5)
|
||||
assert out.shape == (B, 9)
|
||||
sample, n_sec = sample_flow(_small_model(), cond_cont, cond_cat, steps=5)
|
||||
assert sample.shape == (B, 9)
|
||||
assert n_sec.shape == (B,)
|
||||
|
||||
|
||||
def test_ddpm_loss_nonneg():
|
||||
@@ -55,5 +56,6 @@ def test_sample_ddim_shape():
|
||||
schedule = CosineSchedule(T=50)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
out = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5)
|
||||
assert out.shape == (B, 9)
|
||||
sample, n_sec = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5)
|
||||
assert sample.shape == (B, 9)
|
||||
assert n_sec.shape == (B,)
|
||||
|
||||
@@ -39,7 +39,9 @@ def test_denoising_mlp_gradients_flow():
|
||||
t = torch.rand(B)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
loss = model(x_t, t, cond_cont, cond_cat).sum()
|
||||
loss.backward()
|
||||
# Both paths must be exercised to get gradients through all parameters.
|
||||
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}"
|
||||
|
||||
@@ -0,0 +1,222 @@
|
||||
"""Tests for Phase 2: secondary particle prediction."""
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_DIM, X_DIM
|
||||
from giant.model.network import DenoisingMLP, SecondaryDecoder
|
||||
from giant.model.schedule import flow_matching_loss_secondary
|
||||
from giant.sample import sample_secondaries, snap_type_to_pdg_idx
|
||||
|
||||
|
||||
# ── helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def _stage1(pdg=3, mat=2):
|
||||
return DenoisingMLP(pdg_vocab=pdg, mat_vocab=mat, hidden_dim=32, n_blocks=2)
|
||||
|
||||
|
||||
def _sec_decoder(pdg=3, mat=2):
|
||||
return SecondaryDecoder(pdg_vocab=pdg, mat_vocab=mat, hidden_dim=32, n_blocks=2)
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
# ── DenoisingMLP Phase-2 additions ───────────────────────────────────────────
|
||||
|
||||
def test_predict_n_sec_shape():
|
||||
B = 8
|
||||
model = _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_predict_n_sec_no_nan():
|
||||
B = 8
|
||||
model = _stage1()
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
logits = model.predict_n_sec(cond_cont, cond_cat)
|
||||
assert torch.isfinite(logits).all()
|
||||
|
||||
|
||||
def test_pdg_embedding_weight_shape():
|
||||
model = _stage1(pdg=5, mat=2)
|
||||
w = model.pdg_embedding_weight()
|
||||
assert w.shape == (5, EMB_DIM)
|
||||
|
||||
|
||||
# ── SecondaryDecoder ──────────────────────────────────────────────────────────
|
||||
|
||||
def test_sec_decoder_output_shape():
|
||||
B = 8
|
||||
decoder = _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)
|
||||
out = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
|
||||
assert out.shape == (B, SEC_DIM)
|
||||
|
||||
|
||||
def test_sec_decoder_no_nan():
|
||||
B = 4
|
||||
decoder = _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)
|
||||
out = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
|
||||
assert torch.isfinite(out).all()
|
||||
|
||||
|
||||
def test_sec_decoder_gradients():
|
||||
B = 4
|
||||
decoder = _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(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}"
|
||||
|
||||
|
||||
# ── masked flow matching loss ─────────────────────────────────────────────────
|
||||
|
||||
def test_flow_matching_loss_secondary_scalar():
|
||||
B, pdg, mat = 8, 3, 2
|
||||
decoder = _sec_decoder(pdg, mat)
|
||||
x1 = torch.randn(B, SEC_DIM)
|
||||
cond_cont, cond_cat = _cond(B, pdg, mat)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
sec_mask = torch.ones(B, K_MAX, dtype=torch.bool)
|
||||
loss = flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask)
|
||||
assert loss.shape == ()
|
||||
assert loss.item() >= 0.0
|
||||
|
||||
|
||||
def test_flow_matching_loss_secondary_mask_zeros_padding():
|
||||
"""Loss with all-zero mask (no valid secondaries) should be 0."""
|
||||
B, pdg, mat = 4, 3, 2
|
||||
decoder = _sec_decoder(pdg, mat)
|
||||
x1 = torch.randn(B, SEC_DIM)
|
||||
cond_cont, cond_cat = _cond(B, pdg, mat)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
sec_mask = torch.zeros(B, K_MAX, dtype=torch.bool)
|
||||
loss = flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask)
|
||||
assert loss.item() == pytest.approx(0.0, abs=1e-6)
|
||||
|
||||
|
||||
def test_flow_matching_loss_secondary_has_grad():
|
||||
B, pdg, mat = 4, 3, 2
|
||||
decoder = _sec_decoder(pdg, mat)
|
||||
x1 = torch.randn(B, SEC_DIM)
|
||||
cond_cont, cond_cat = _cond(B, pdg, mat)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
sec_mask = torch.ones(B, K_MAX, dtype=torch.bool)
|
||||
flow_matching_loss_secondary(
|
||||
decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
|
||||
).backward()
|
||||
assert any(p.grad is not None for p in decoder.parameters())
|
||||
|
||||
|
||||
# ── sampling ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def test_sample_secondaries_shapes():
|
||||
B, pdg, mat = 6, 3, 2
|
||||
decoder = _sec_decoder(pdg, mat)
|
||||
cond_cont, cond_cat = _cond(B, pdg, mat)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
|
||||
sec_cont, sec_type_emb, sec_valid = sample_secondaries(
|
||||
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=3
|
||||
)
|
||||
assert sec_cont.shape == (B, K_MAX, 4)
|
||||
assert sec_type_emb.shape == (B, K_MAX, EMB_DIM)
|
||||
assert sec_valid.shape == (B, K_MAX)
|
||||
assert sec_valid.dtype == torch.bool
|
||||
|
||||
|
||||
def test_sample_secondaries_valid_mask_matches_n_sec():
|
||||
B, pdg, mat = 4, 3, 2
|
||||
decoder = _sec_decoder(pdg, mat)
|
||||
cond_cont, cond_cat = _cond(B, pdg, mat)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
n_sec_pred = torch.tensor([0, 1, 3, K_MAX])
|
||||
_, _, sec_valid = sample_secondaries(
|
||||
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
|
||||
)
|
||||
for i, n in enumerate(n_sec_pred.tolist()):
|
||||
assert sec_valid[i, :n].all()
|
||||
assert not sec_valid[i, n:].any()
|
||||
|
||||
|
||||
def test_snap_type_to_pdg_idx_shape():
|
||||
B, pdg_vocab = 4, 5
|
||||
emb_weight = torch.randn(pdg_vocab, EMB_DIM)
|
||||
sec_type_emb = torch.randn(B, K_MAX, EMB_DIM)
|
||||
idx = snap_type_to_pdg_idx(sec_type_emb, emb_weight)
|
||||
assert idx.shape == (B, K_MAX)
|
||||
assert idx.dtype == torch.int64
|
||||
assert (idx >= 0).all() and (idx < pdg_vocab).all()
|
||||
|
||||
|
||||
# ── encode_secondaries round-trip ─────────────────────────────────────────────
|
||||
|
||||
def test_encode_secondaries_energy_conservation():
|
||||
"""Decoded stick-breaking fractions must sum to ≈ e_sec."""
|
||||
from giant.data.transforms import encode_secondaries
|
||||
|
||||
rng = np.random.default_rng(42)
|
||||
N = 50
|
||||
n_sec = rng.integers(1, 5, size=N)
|
||||
e_sec = rng.uniform(0.1, 10.0, size=N).astype(np.float32)
|
||||
|
||||
sec_E_list = np.zeros((N, K_MAX), dtype=np.float32)
|
||||
sec_dir_list = np.zeros((N, K_MAX, 3), dtype=np.float32)
|
||||
sec_dir_list[:, :, 2] = 1.0
|
||||
sec_valid = np.zeros((N, K_MAX), dtype=bool)
|
||||
for i in range(N):
|
||||
k = n_sec[i]
|
||||
energies = rng.dirichlet(np.ones(k)) * e_sec[i]
|
||||
energies = np.sort(energies)[::-1]
|
||||
sec_E_list[i, :k] = energies.astype(np.float32)
|
||||
sec_valid[i, :k] = True
|
||||
|
||||
pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
|
||||
pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
|
||||
|
||||
sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir)
|
||||
assert sec_cont.shape == (N, K_MAX, 4)
|
||||
assert np.isfinite(sec_cont).all()
|
||||
|
||||
|
||||
def test_encode_secondaries_direction_encoding():
|
||||
"""Local-frame secondary directions should be unit vectors for valid slots."""
|
||||
from giant.data.transforms import encode_secondaries
|
||||
|
||||
rng = np.random.default_rng(7)
|
||||
N = 20
|
||||
e_sec = np.ones(N, dtype=np.float32) * 5.0
|
||||
sec_E_list = np.zeros((N, K_MAX), dtype=np.float32)
|
||||
sec_E_list[:, 0] = 3.0
|
||||
sec_E_list[:, 1] = 2.0
|
||||
sec_dir_list = rng.standard_normal((N, K_MAX, 3)).astype(np.float32)
|
||||
norms = np.linalg.norm(sec_dir_list, axis=-1, keepdims=True)
|
||||
sec_dir_list /= np.where(norms > 0, norms, 1.0)
|
||||
sec_valid = np.zeros((N, K_MAX), dtype=bool)
|
||||
sec_valid[:, :2] = True
|
||||
|
||||
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
|
||||
sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir)
|
||||
|
||||
# dir columns are sec_cont[:, :, 1:4]
|
||||
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)
|
||||
Reference in New Issue
Block a user