Encode edep/secondary/post energy as a conservation-constrained simplex

Replaces the independent log_delta_e/log_edep targets with 2 additive-log-ratio
coordinates over the deposit/secondary/post-energy simplex (fractions of pre_E
summing to 1), so edep + e_sec + post_E == pre_E holds by construction after
decoding (softmax) rather than being learned approximately. Requires e_sec
(secondary energy) as a new conditioning input and a steps_to_parquet.py pass
to derive it from child track first-step energies.

scripts/steps_to_parquet.py

@@ -19,6 +19,43 @@ 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.
+
+    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.
+    """
+    first_E = (
+        df.sort("step_no")
+        .group_by(["event_id", "track_id"])
+        .agg(pl.col("pre_E").first().alias("child_E"))
+        .rename({"track_id": "child_track_id"})
+    )
+
+    exploded = (
+        df.select(["event_id", "child_track_ids"])
+        .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
+    )
+    summed = (
+        exploded.join(first_E, on=["event_id", "child_track_id"], how="left")
+        .group_by("_step_row")
+        .agg(pl.col("child_E").sum().alias("e_sec"))
+    )
+    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))
+        .drop("_step_row")
+    )
+
+
 def _batch_to_polars(batch: ak.Array) -> pl.DataFrame:
     """Convert one awkward-array batch to a Polars DataFrame.
 
@@ -77,8 +114,15 @@ def convert_steps_to_parquet(
             rows_done += len(batch)
             print(f"  {rows_done:,} / {n_entries:,} rows read", end="\r", flush=True)
 
+    df = pl.concat(batches)
+    # 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(f"\nWriting {output_path} …", end=" ", flush=True)
-    pl.concat(batches).write_parquet(output_path, compression=compression)
+    df.write_parquet(output_path, compression=compression)
     print(f"done  ({output_path.stat().st_size / 1e6:.1f} MB)")
     return output_path