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>

tests/test_flow.py

@@ -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,)