lars
55332db67a
CI / Lint (ruff check) (push) Successful in 31s
CI / Format (ruff format) (push) Successful in 32s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 26s
CI / Type check (ty) (push) Successful in 29s
CI / Format (ruff format) (pull_request) Successful in 33s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 35s
CI / Tests (pull_request) Successful in 3m47s
CI / Tests (push) Successful in 3m55s
Rejoins lines that only wrapped because they exceeded the old 88-char limit; ruff check and the full test suite (725 passed) are unaffected.
75 lines
3.0 KiB
Python
75 lines
3.0 KiB
Python
import polars as pl
|
|
|
|
from scripts import steps_to_parquet
|
|
|
|
|
|
def _frame() -> pl.DataFrame:
|
|
# event 0: step (1,0) spawns track 2 (first-step pre_E=15) → e_sec=15.
|
|
# event 1: step (1,0) spawns tracks 2 & 3 (20 + 30) → e_sec=50.
|
|
return pl.DataFrame(
|
|
{
|
|
"event_id": [0, 0, 0, 1, 1, 1],
|
|
"track_id": [1, 1, 2, 1, 2, 3],
|
|
"step_no": [0, 1, 0, 0, 0, 0],
|
|
"pre_E": [100.0, 80.0, 15.0, 200.0, 20.0, 30.0],
|
|
"pdg": [11, 11, 22, 11, 22, 22],
|
|
"pre_dx": [0.0, 0.0, 1.0, 0.0, 1.0, 0.0],
|
|
"pre_dy": [0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
|
|
"pre_dz": [1.0, 1.0, 0.0, 1.0, 0.0, 0.0],
|
|
"child_track_ids": [[2], [], [], [2, 3], [], []],
|
|
}
|
|
)
|
|
|
|
|
|
def test_e_sec_sums_child_first_step_energy():
|
|
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"]))
|
|
assert e_sec[(1, 0, 0)] == 15.0 # one child, first-step pre_E 15
|
|
assert e_sec[(1, 0, 1)] == 50.0 # two children, 20 + 30
|
|
|
|
|
|
def test_e_sec_zero_when_no_children():
|
|
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))
|
|
assert childless["e_sec"].item() == 0.0
|
|
|
|
|
|
def test_e_sec_preserves_row_count_and_order():
|
|
df = _frame()
|
|
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()
|