lars
e6e0eb22bf
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>
34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
import numpy as np
|
|
from giant.data.dataset import make_event_split
|
|
|
|
|
|
def test_make_event_split_sizes():
|
|
rng = np.random.default_rng(42)
|
|
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_make_event_split_no_overlap():
|
|
rng = np.random.default_rng(7)
|
|
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)
|
|
|
|
|
|
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
|