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>

scripts/steps_to_parquet.py

@@ -19,20 +19,31 @@ import uproot
 ParquetCompression = Literal["lz4", "uncompressed", "snappy", "gzip", "brotli", "zstd"]
 
 
-def _add_secondary_energy(df: pl.DataFrame) -> pl.DataFrame:
-    """Add per-step `e_sec`: total initial kinetic energy of the secondaries born in it.
+def _add_secondary_attributes(df: pl.DataFrame) -> pl.DataFrame:
+    """Add per-step secondary attributes via the parent→child track join.
 
-    Each secondary's creation energy is the `pre_E` of that child track's first step
-    (min `step_no`) in the same event, so for a parent step
-    `e_sec = Σ over child_track_ids of the child track's first-step pre_E`. Steps that
-    spawn nothing get 0.0. The full event must be present in `df` (it is — the writer
-    concatenates every batch before this runs), since a child track's first step can
-    live in a different read batch than its parent step.
+    For each step that spawns secondaries, collects each child track's birth
+    state (from the child track's first step in the same event) and emits:
+      e_sec        float64  — total secondary energy (sum of child first-step pre_E)
+      sec_E_list   list[f64] — per-secondary energy, sorted descending
+      sec_pdg_list list[i32] — per-secondary PDG code, same order
+      sec_dx_list  list[f64] — per-secondary birth direction x, same order
+      sec_dy_list  list[f64] — per-secondary birth direction y, same order
+      sec_dz_list  list[f64] — per-secondary birth direction z, same order
+
+    Steps with no children get 0.0 / empty lists. The full event must be
+    present in `df` (it is — the writer concatenates before calling this).
     """
-    first_E = (
+    first_step = (
         df.sort("step_no")
         .group_by(["event_id", "track_id"])
-        .agg(pl.col("pre_E").first().alias("child_E"))
+        .agg(
+            pl.col("pre_E").first().alias("child_E"),
+            pl.col("pdg").first().alias("child_pdg"),
+            pl.col("pre_dx").first().alias("child_dx"),
+            pl.col("pre_dy").first().alias("child_dy"),
+            pl.col("pre_dz").first().alias("child_dz"),
+        )
         .rename({"track_id": "child_track_id"})
     )
 
@@ -41,17 +52,39 @@ def _add_secondary_energy(df: pl.DataFrame) -> pl.DataFrame:
         .with_row_index("_step_row")
         .explode("child_track_ids")
         .rename({"child_track_ids": "child_track_id"})
-        .drop_nulls("child_track_id")  # steps with no children explode to a null row
+        .drop_nulls("child_track_id")
     )
-    summed = (
-        exploded.join(first_E, on=["event_id", "child_track_id"], how="left")
+
+    joined = exploded.join(first_step, on=["event_id", "child_track_id"], how="left")
+
+    # Sort each step's secondaries by descending energy, then aggregate into lists
+    per_step = (
+        joined.sort("child_E", descending=True)
         .group_by("_step_row")
-        .agg(pl.col("child_E").sum().alias("e_sec"))
+        .agg(
+            pl.col("child_E").sum().alias("e_sec"),
+            pl.col("child_E").alias("sec_E_list"),
+            pl.col("child_pdg").alias("sec_pdg_list"),
+            pl.col("child_dx").alias("sec_dx_list"),
+            pl.col("child_dy").alias("sec_dy_list"),
+            pl.col("child_dz").alias("sec_dz_list"),
+        )
     )
+
+    empty_list_f64 = pl.Series("x", [[]], dtype=pl.List(pl.Float64))
+    empty_list_i32 = pl.Series("x", [[]], dtype=pl.List(pl.Int32))
+
     return (
         df.with_row_index("_step_row")
-        .join(summed, on="_step_row", how="left")
-        .with_columns(pl.col("e_sec").fill_null(0.0).cast(pl.Float64))
+        .join(per_step, on="_step_row", how="left")
+        .with_columns(
+            pl.col("e_sec").fill_null(0.0).cast(pl.Float64),
+            pl.col("sec_E_list").fill_null(empty_list_f64),
+            pl.col("sec_pdg_list").fill_null(empty_list_i32),
+            pl.col("sec_dx_list").fill_null(empty_list_f64),
+            pl.col("sec_dy_list").fill_null(empty_list_f64),
+            pl.col("sec_dz_list").fill_null(empty_list_f64),
+        )
         .drop("_step_row")
     )
 
@@ -118,8 +151,8 @@ def convert_steps_to_parquet(
     # Steps tree carries the parent→child links needed to derive secondary energy;
     # other trees (e.g. Hits) don't, so only augment when the column is present.
     if "child_track_ids" in df.columns:
-        print("\nComputing per-step secondary energy (e_sec) …", end=" ", flush=True)
-        df = _add_secondary_energy(df)
+        print("\nComputing per-step secondary attributes …", end=" ", flush=True)
+        df = _add_secondary_attributes(df)
 
     print(f"\nWriting {output_path} …", end=" ", flush=True)
     df.write_parquet(output_path, compression=compression)