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/data/loader.py

@@ -40,8 +40,52 @@ def find_parquet_files(path: str | Path) -> list[Path]:
     return [p]
 
 
+def _pad_list_col(series: pd.Series, K: int, fill: float = 0.0) -> np.ndarray:
+    """Pad / truncate a list-valued Series to fixed width K → (N, K) float32."""
+    out = np.full((len(series), K), fill, dtype=np.float32)
+    for i, lst in enumerate(series):
+        if lst is not None and len(lst) > 0:
+            n = min(len(lst), K)
+            out[i, :n] = lst[:n]
+    return out
+
+
+def _pad_list_col_int(series: pd.Series, K: int, fill: int = 0) -> np.ndarray:
+    """Pad / truncate a list-valued integer Series to fixed width K → (N, K) int64."""
+    out = np.full((len(series), K), fill, dtype=np.int64)
+    for i, lst in enumerate(series):
+        if lst is not None and len(lst) > 0:
+            n = min(len(lst), K)
+            out[i, :n] = lst[:n]
+    return out
+
+
+def _pad_dir_col(
+    dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int
+) -> np.ndarray:
+    """Pad three list-valued direction columns → (N, K, 3) float32.
+
+    Padding direction defaults to (0,0,1) (forward) so it is a valid unit vector.
+    """
+    N = len(dx)
+    out = np.zeros((N, K, 3), dtype=np.float32)
+    out[:, :, 2] = 1.0
+    for i in range(N):
+        lx, ly, lz = dx.iloc[i], dy.iloc[i], dz.iloc[i]
+        if lx is not None and len(lx) > 0:
+            n = min(len(lx), K)
+            out[i, :n, 0] = lx[:n]
+            out[i, :n, 1] = ly[:n]
+            out[i, :n, 2] = lz[:n]
+    return out
+
+
 def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
-    return {
+    from giant.constants import K_MAX
+
+    has_sec_lists = "sec_E_list" in df.columns
+
+    d: dict[str, np.ndarray] = {
         "event_id": df["event_id"].to_numpy(),
         "pdg": df["pdg"].to_numpy(dtype=np.int32),
         "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
@@ -59,6 +103,15 @@ def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
         "post_pos": df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32),
     }
 
+    if has_sec_lists:
+        d["sec_E_list"] = _pad_list_col(df["sec_E_list"], K_MAX)
+        d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], K_MAX)
+        d["sec_dir_list"] = _pad_dir_col(
+            df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], K_MAX
+        )
+
+    return d
+
 
 def load_steps(path: str | Path) -> dict[str, np.ndarray]:
     return _df_to_dict(pd.read_parquet(path))