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/sample.py

@@ -1,6 +1,6 @@
 import torch
 
-from giant.constants import X_DIM
+from giant.constants import K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM
 
 
 @torch.no_grad()
@@ -9,8 +9,13 @@ def sample_flow(
     cond_cont: torch.Tensor,
     cond_cat: torch.Tensor,
     steps: int = 10,
-) -> torch.Tensor:
-    """Euler integration of the learned vector field from t=0 to t=1."""
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Euler integration of the Stage-1 vector field from t=0 to t=1.
+
+    Returns (primary_sample, n_sec_pred):
+        primary_sample: (B, X_DIM) — normalised 9D primary post-step output
+        n_sec_pred: (B,) int64 — predicted secondary count
+    """
     model.eval()
     B = cond_cont.size(0)
     device = cond_cont.device
@@ -20,7 +25,63 @@ def sample_flow(
         t = torch.full((B,), i * dt, device=device)
         v = model(x, t, cond_cont, cond_cat)
         x = x + v * dt
-    return x
+    n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
+    n_sec_pred = n_sec_logits.argmax(dim=-1)
+    return x, n_sec_pred
+
+
+@torch.no_grad()
+def sample_secondaries(
+    sec_decoder: torch.nn.Module,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    stage1_out: torch.Tensor,
+    n_sec_pred: torch.Tensor,
+    steps: int = 10,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Euler integration of the Stage-2 vector field; return raw slot outputs.
+
+    n_sec_pred: (B,) int64 — number of valid secondaries per step
+
+    Returns (sec_cont, sec_type_emb, sec_valid):
+        sec_cont:     (B, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z]
+        sec_type_emb: (B, K_MAX, emb_dim) — predicted type embedding per slot
+        sec_valid:    (B, K_MAX) bool — True for slots i < n_sec_pred
+    """
+    sec_decoder.eval()
+    B = cond_cont.size(0)
+    device = cond_cont.device
+
+    x = torch.randn(B, SEC_DIM, device=device)
+    dt = 1.0 / steps
+    for i in range(steps):
+        t = torch.full((B,), i * dt, device=device)
+        v = sec_decoder(x, t, cond_cont, cond_cat, stage1_out)
+        x = x + v * dt
+
+    x_slots = x.view(B, K_MAX, SEC_SLOT_DIM)
+    sec_cont = x_slots[:, :, :4]
+    sec_type_emb = x_slots[:, :, 4:]
+
+    sec_valid = (
+        torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1)
+    )
+    return sec_cont, sec_type_emb, sec_valid
+
+
+def snap_type_to_pdg_idx(
+    sec_type_emb: torch.Tensor,
+    pdg_emb_weight: torch.Tensor,
+) -> torch.Tensor:
+    """Nearest-neighbour snap: predicted type embedding → PDG model-index.
+
+    sec_type_emb: (B, K_MAX, emb_dim)
+    Returns (B, K_MAX) int64 with model-indices.
+    """
+    B, K, D = sec_type_emb.shape
+    flat = sec_type_emb.reshape(-1, D)
+    dists = torch.cdist(flat.float(), pdg_emb_weight.float())
+    return dists.argmin(dim=-1).reshape(B, K)
 
 
 @torch.no_grad()
@@ -29,8 +90,8 @@ def sample_ddpm(
     cond_cont: torch.Tensor,
     cond_cat: torch.Tensor,
     schedule,
-) -> torch.Tensor:
-    """Full DDPM ancestral sampling (T reverse steps)."""
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Full DDPM ancestral sampling (T reverse steps). Returns (sample, n_sec_pred)."""
     model.eval()
     B = cond_cont.size(0)
     device = cond_cont.device
@@ -46,7 +107,9 @@ def sample_ddpm(
         x = (1.0 / alpha.sqrt()) * (
             x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred
         ) + beta.sqrt() * z
-    return x
+    n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
+    n_sec_pred = n_sec_logits.argmax(dim=-1)
+    return x, n_sec_pred
 
 
 @torch.no_grad()
@@ -56,8 +119,8 @@ def sample_ddim(
     cond_cat: torch.Tensor,
     schedule,
     steps: int = 50,
-) -> torch.Tensor:
-    """DDIM deterministic sampling (Song et al. 2020) with `steps` substeps."""
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """DDIM deterministic sampling (Song et al. 2020). Returns (sample, n_sec_pred)."""
     model.eval()
     B = cond_cont.size(0)
     device = cond_cont.device
@@ -75,4 +138,6 @@ def sample_ddim(
             ab_prev = torch.ones(1, device=device)
         x0_pred = (x - (1.0 - ab_t).sqrt() * eps_pred) / ab_t.sqrt()
         x = ab_prev.sqrt() * x0_pred + (1.0 - ab_prev).sqrt() * eps_pred
-    return x
+    n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
+    n_sec_pred = n_sec_logits.argmax(dim=-1)
+    return x, n_sec_pred