Drop orphaned child tracks instead of nulling secondary targets

A listed child_track_id can fail to match any first-step row (e.g. a
secondary absorbed below the tracking threshold at birth). The
parent->child left join in _add_secondary_attributes left these as
nulls, which silently became NaN once the parquet round-tripped
through the loader's float32 padding — poisoning every later secondary
slot in that step via the cumulative "remaining budget" in
encode_secondaries, while e_sec quietly undercounted and n_sec (from
len(child_track_ids)) overcounted relative to the actual lists.

Drop orphans from both the per-secondary lists and child_track_ids
itself so downstream counts stay consistent, and thread the per-file
orphaned count back through convert_steps_to_parquet so both the
sequential and --jobs>1 batch paths in `dwarf convert` can report an
aggregate total instead of relying on grepping printed output.

Also floors encode_secondaries' slot-0 budget to _EPS (matching the
i>0 branch), fixing a harmless but noisy 0/0 divide warning on
zero-secondary steps.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

giant/data/transforms.py

@@ -310,7 +310,7 @@ def encode_secondaries(
     stick_logits = np.zeros((N, K), dtype=np.float32)
     for i in range(K):
         if i == 0:
-            remaining = e_sec
+            remaining = np.maximum(e_sec, _EPS)
         else:
             remaining = np.maximum(e_sec - sec_E_list[:, :i].sum(axis=1), _EPS)
         f = np.clip(sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS)

scripts/dwarf.py

@@ -116,14 +116,21 @@ def convert(
                 "error: --output can only be used with a single input file", err=True
             )
             raise typer.Exit(1)
+        total_orphaned = 0
         for root_file in root_files:
-            convert_steps_to_parquet(
+            _, n_orphaned = convert_steps_to_parquet(
                 root_file,
                 output_path=output,
                 batch_size=batch_size,
                 tree_name=tree,
                 compression=compression_value,
             )
+            total_orphaned += n_orphaned
+        if total_orphaned:
+            typer.echo(
+                f"\n{total_orphaned} orphaned child track(s) dropped across "
+                f"{len(root_files)} file(s)."
+            )
         return
 
     if output is not None:

scripts/steps_to_parquet.py

@@ -4,7 +4,7 @@ See `uv run dwarf convert --help` for the CLI.
 """
 
 from pathlib import Path
-from typing import Literal
+from typing import Literal, cast
 
 import awkward as ak
 import polars as pl
@@ -13,7 +13,7 @@ import uproot
 ParquetCompression = Literal["lz4", "uncompressed", "snappy", "gzip", "brotli", "zstd"]
 
 
-def _add_secondary_attributes(df: pl.DataFrame) -> pl.DataFrame:
+def _add_secondary_attributes(df: pl.DataFrame) -> tuple[pl.DataFrame, int]:
     """Add per-step secondary attributes via the parent→child track join.
 
     For each step that spawns secondaries, collects each child track's birth
@@ -27,6 +27,18 @@ def _add_secondary_attributes(df: pl.DataFrame) -> pl.DataFrame:
 
     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).
+
+    A listed child_track_id can fail to match any row in `first_step` — the
+    child track never took a recorded step (e.g. absorbed below the tracking
+    threshold at birth). Such orphans carry no physical secondary data, so
+    they're dropped from child_track_ids/sec_*_list rather than left as nulls:
+    a null in a float32 list silently becomes NaN once the parquet round-trips
+    through the loader (`giant/data/loader.py:_pad_list_col`), and that NaN
+    poisons every later secondary slot in the same step via the cumulative-sum
+    "remaining budget" in `encode_secondaries`.
+
+    Returns (df, n_orphaned) — the caller uses the count to report/aggregate
+    across files rather than relying solely on the printed message here.
     """
     first_step = (
         df.sort("step_no")
@@ -41,6 +53,8 @@ def _add_secondary_attributes(df: pl.DataFrame) -> pl.DataFrame:
         .rename({"track_id": "child_track_id"})
     )
 
+    child_track_id_dtype = cast(pl.List, df.schema["child_track_ids"]).inner
+
     exploded = (
         df.select(["event_id", "child_track_ids"])
         .with_row_index("_step_row")
@@ -51,11 +65,20 @@ def _add_secondary_attributes(df: pl.DataFrame) -> pl.DataFrame:
 
     joined = exploded.join(first_step, on=["event_id", "child_track_id"], how="left")
 
+    n_orphaned = joined["child_E"].null_count()
+    if n_orphaned:
+        print(
+            f"  dropping {n_orphaned} orphaned child_track_id(s) with no "
+            "recorded first step (absorbed below tracking threshold?)"
+        )
+    joined = joined.drop_nulls("child_E")
+
     # 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_track_id").alias("child_track_ids"),
             pl.col("child_E").sum().alias("e_sec"),
             pl.col("child_E").alias("sec_E_list"),
             pl.col("child_pdg").alias("sec_pdg_list"),
@@ -67,11 +90,14 @@ def _add_secondary_attributes(df: pl.DataFrame) -> pl.DataFrame:
 
     empty_list_f64 = pl.Series("x", [[]], dtype=pl.List(pl.Float64))
     empty_list_i32 = pl.Series("x", [[]], dtype=pl.List(pl.Int32))
+    empty_list_child_id = pl.Series("x", [[]], dtype=pl.List(child_track_id_dtype))
 
-    return (
-        df.with_row_index("_step_row")
+    out = (
+        df.drop("child_track_ids")
+        .with_row_index("_step_row")
         .join(per_step, on="_step_row", how="left")
         .with_columns(
+            pl.col("child_track_ids").fill_null(empty_list_child_id),
             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),
@@ -81,6 +107,7 @@ def _add_secondary_attributes(df: pl.DataFrame) -> pl.DataFrame:
         )
         .drop("_step_row")
     )
+    return out, n_orphaned
 
 
 def _batch_to_polars(batch: ak.Array) -> pl.DataFrame:
@@ -106,7 +133,7 @@ def convert_steps_to_parquet(
     batch_size: str = "100 MB",
     tree_name: str = "Steps",
     compression: ParquetCompression = "snappy",
-) -> Path:
+) -> tuple[Path, int]:
     """Read *tree_name* from *root_path* and write it to a Parquet file.
 
     Reads in batches of *batch_size* so that peak ROOT-deserialization memory
@@ -122,6 +149,11 @@ def convert_steps_to_parquet(
                  integer row count (500_000).
     tree_name:   Name of the TTree inside the ROOT file.
     compression: Parquet compression codec (snappy | lz4 | zstd | gzip | none).
+
+    Returns (output_path, n_orphaned) — n_orphaned is the count of dropped
+    orphaned child_track_ids (see `_add_secondary_attributes`), 0 if the tree
+    has no child_track_ids column at all. Callers converting many files use
+    it to aggregate a total instead of grepping the printed per-file message.
     """
     root_path = Path(root_path)
     if output_path is None:
@@ -144,11 +176,12 @@ def convert_steps_to_parquet(
     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.
+    n_orphaned = 0
     if "child_track_ids" in df.columns:
         print("\nComputing per-step secondary attributes …", end=" ", flush=True)
-        df = _add_secondary_attributes(df)
+        df, n_orphaned = _add_secondary_attributes(df)
 
     print(f"\nWriting {output_path} …", end=" ", flush=True)
     df.write_parquet(output_path, compression=compression)
     print(f"done  ({output_path.stat().st_size / 1e6:.1f} MB)")
-    return output_path
+    return output_path, n_orphaned

scripts/steps_to_parquet_parallel.py

@@ -26,6 +26,12 @@ from pathlib import Path
 GEN_RE = re.compile(r"^gen\d+$")
 SCHEMA_RE = re.compile(r"^schema(\d+)$")
 
+# Matches the per-file orphan-drop message printed by
+# steps_to_parquet._add_secondary_attributes — each subprocess's count is
+# parsed back out of its captured stdout since there's no in-process return
+# value across the subprocess boundary.
+_ORPHAN_RE = re.compile(r"dropping (\d+) orphaned child_track_id")
+
 
 class DestinationError(ValueError):
     pass
@@ -210,4 +216,9 @@ def run_parallel_job(
             print(f"  {root_file}", file=sys.stderr)
         raise SystemExit(1)
 
+    total_orphaned = sum(
+        int(m.group(1)) for _, _, stdout, _ in results for m in _ORPHAN_RE.finditer(stdout)
+    )
+    if total_orphaned:
+        print(f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s).")
     print(f"\nAll {len(results)} conversion(s) completed.")

tests/test_steps_to_parquet.py

@@ -22,7 +22,8 @@ def _frame() -> pl.DataFrame:
 
 
 def test_e_sec_sums_child_first_step_energy():
-    out = steps_to_parquet._add_secondary_attributes(_frame())
+    out, n_orphaned = steps_to_parquet._add_secondary_attributes(_frame())
+    assert n_orphaned == 0
     e_sec = dict(
         zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"])
     )
@@ -31,7 +32,7 @@ def test_e_sec_sums_child_first_step_energy():
 
 
 def test_e_sec_zero_when_no_children():
-    out = steps_to_parquet._add_secondary_attributes(_frame())
+    out, _ = steps_to_parquet._add_secondary_attributes(_frame())
     childless = out.filter(
         (pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 1)
     )
@@ -40,6 +41,38 @@ def test_e_sec_zero_when_no_children():
 
 def test_e_sec_preserves_row_count_and_order():
     df = _frame()
-    out = steps_to_parquet._add_secondary_attributes(df)
+    out, _ = steps_to_parquet._add_secondary_attributes(df)
     assert out.height == df.height
     assert out["pre_E"].to_list() == df["pre_E"].to_list()
+
+
+def test_orphaned_child_track_is_dropped_not_nulled():
+    """A listed child_track_id with no first step of its own (e.g. absorbed
+    below the tracking threshold at birth) must not leave a null in
+    sec_E_list/sec_pdg_list/etc: that null turns into NaN once the parquet
+    round-trips through the loader, poisoning every later secondary slot in
+    the step via encode_secondaries' cumulative "remaining budget". It must
+    also be dropped from child_track_ids itself, so n_sec (len(child_track_ids)
+    downstream) matches the actual, orphan-free secondary lists."""
+    df = pl.DataFrame(
+        {
+            "event_id": [0, 0, 0],
+            "track_id": [1, 1, 2],
+            "step_no": [0, 1, 0],
+            "pre_E": [100.0, 80.0, 15.0],
+            "pdg": [11, 11, 22],
+            "pre_dx": [0.0, 0.0, 1.0],
+            "pre_dy": [0.0, 0.0, 0.0],
+            "pre_dz": [1.0, 1.0, 0.0],
+            # track 3 is listed as a child but never appears with its own step.
+            "child_track_ids": [[2, 3], [], []],
+        }
+    )
+    out, n_orphaned = steps_to_parquet._add_secondary_attributes(df)
+    row = out.filter((pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0))
+
+    assert n_orphaned == 1
+    assert row["child_track_ids"].to_list() == [[2]]
+    assert row["e_sec"].item() == 15.0
+    assert row["sec_E_list"].to_list() == [[15.0]]
+    assert None not in row["sec_E_list"].item()