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>
This commit is contained in:
+8
-1
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.")
|
||||
|
||||
Reference in New Issue
Block a user