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>

giant/validate.py

@@ -53,18 +53,21 @@ def validate_marginals(
     model.eval()
 
     all_real, all_gen = [], []
-    for i, (cond_cont, cond_cat, x1) in enumerate(val_loader):
+    for i, batch in enumerate(val_loader):
         if n_batches is not None and i >= n_batches:
             break
+        # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx);
+        # validate_marginals only needs the Stage-1 primary target.
+        cond_cont, cond_cat, x1 = batch[0], batch[1], batch[2]
         cond_cont = cond_cont.to(device)
         cond_cat = cond_cat.to(device)
 
         if mode == "flow":
-            gen = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
+            gen, _n_sec = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
         elif mode == "ddpm":
-            gen = sample_ddpm(model, cond_cont, cond_cat, schedule)
+            gen, _n_sec = sample_ddpm(model, cond_cont, cond_cat, schedule)
         else:
-            gen = sample_ddim(model, cond_cont, cond_cat, schedule, **_kw(steps))
+            gen, _n_sec = sample_ddim(model, cond_cont, cond_cat, schedule, **_kw(steps))
 
         all_real.append(x1.numpy())
         all_gen.append(gen.cpu().numpy())