From 740ebdf6b6d09cf104a41a6aeacdd2c6729fc129 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 9 Jul 2026 14:36:06 +0200 Subject: [PATCH] Drop orphaned child tracks instead of nulling secondary targets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- giant/data/transforms.py | 2 +- scripts/dwarf.py | 9 +++++- scripts/steps_to_parquet.py | 47 +++++++++++++++++++++++----- scripts/steps_to_parquet_parallel.py | 11 +++++++ tests/test_steps_to_parquet.py | 39 +++++++++++++++++++++-- 5 files changed, 96 insertions(+), 12 deletions(-) diff --git a/giant/data/transforms.py b/giant/data/transforms.py index fe3a81d..2994393 100644 --- a/giant/data/transforms.py +++ b/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) diff --git a/scripts/dwarf.py b/scripts/dwarf.py index f728bcf..e096834 100644 --- a/scripts/dwarf.py +++ b/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: diff --git a/scripts/steps_to_parquet.py b/scripts/steps_to_parquet.py index cb74fe5..880c965 100644 --- a/scripts/steps_to_parquet.py +++ b/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 diff --git a/scripts/steps_to_parquet_parallel.py b/scripts/steps_to_parquet_parallel.py index d89a607..dc1c5fe 100644 --- a/scripts/steps_to_parquet_parallel.py +++ b/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.") diff --git a/tests/test_steps_to_parquet.py b/tests/test_steps_to_parquet.py index 6ec6918..9810d6e 100644 --- a/tests/test_steps_to_parquet.py +++ b/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()