lars
c1e6ffd8c6
CI / Lint (ruff check) (push) Successful in 34s
CI / Format (ruff format) (push) Successful in 34s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 28s
CI / Type check (ty) (push) Successful in 30s
CI / Format (ruff format) (pull_request) Successful in 38s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 39s
CI / Tests (push) Successful in 4m54s
CI / Tests (pull_request) Successful in 4m55s
giant/analysis/sources.py's open_side scanned a reference directory of parquet shards with a bare glob and never offset event_id across them. Each shard is a separate Geant4 job whose own event_id numbering restarts from 0, so events from different shards collided on the same event_id, corrupting every downstream per-event grouping and the event_id % n_chunks condor chunking — the same root cause already fixed on the training/rollout side via giant/data/loader.py's per-file event_id_offset. open_side's reference branch now uses find_parquet_files (the same deterministically ordered file lister giant rollout's _seed_from_data uses) and offsets each shard's event_id via a join on polars' include_file_paths, so both sides of a comparison agree on what an event_id means. Two incidental behaviour changes come along for free: .manifest references now work (they crashed before), and the directory glob narrows from recursive **/*.parquet to top-level *.parquet, matching the file list rollout itself used to assign offsets — a deliberate choice, since a differing file list would make the two sides' offsets disagree again in a subtler way. No overflow guard on the per-shard offset stride (unlike loader's _offset_event_id): checking it here would cost an eager event_id-column read per shard in every condor compute job, and giant rollout already runs that check over the same file list when producing the seed. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
231 lines
8.9 KiB
Python
231 lines
8.9 KiB
Python
"""Tests for the streaming compute primitives (giant.analysis.reduce/sources/grouping)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import polars as pl
|
|
|
|
from giant.analysis import grouping as G
|
|
from giant.analysis import reduce as R
|
|
from giant.analysis.sources import (
|
|
SYNTHETIC_TERMINATION_REASONS,
|
|
Side,
|
|
open_side,
|
|
physical_steps,
|
|
secondaries,
|
|
)
|
|
from giant.data.loader import EVENT_ID_FILE_STRIDE
|
|
|
|
|
|
def _rollout_frame() -> pl.LazyFrame:
|
|
# event 1: primary (2 steps) + 1 secondary track + 1 escaped bookkeeping row
|
|
# event 2: primary (1 step)
|
|
return pl.DataFrame(
|
|
{
|
|
"event_id": [1, 1, 1, 1, 2],
|
|
"track_id": [0, 0, 1, 0, 0],
|
|
"parent_id": [-1, -1, 0, -1, -1],
|
|
"generation": [0, 0, 1, 0, 0],
|
|
"step_no": [0, 1, 0, 99, 0],
|
|
"pdg": [11, 11, 22, 11, 11],
|
|
"pre_x": [0.0, 0.0, 0.0, 0.0, 0.0],
|
|
"pre_y": [0.0, 0.0, 0.0, 0.0, 0.0],
|
|
"pre_z": [0.0, 1.0, 1.0, 2.0, 0.0],
|
|
"pre_E": [100.0, 60.0, 20.0, 30.0, 50.0],
|
|
"pre_dx": [0.0, 0.0, 1.0, 0.0, 0.0],
|
|
"pre_dy": [0.0, 0.0, 0.0, 0.0, 0.0],
|
|
"pre_dz": [1.0, 1.0, 0.0, 1.0, 1.0],
|
|
"post_x": [0.0, 0.0, 1.0, 0.0, 0.0],
|
|
"post_y": [0.0, 0.0, 0.0, 0.0, 0.0],
|
|
"post_z": [1.0, 2.0, 1.0, 2.0, 1.0],
|
|
"post_E": [60.0, 30.0, 0.0, 0.0, 20.0],
|
|
"post_dx": [0.0, 0.0, 1.0, 0.0, 0.0],
|
|
"post_dy": [0.0, 0.0, 0.0, 0.0, 0.0],
|
|
"post_dz": [1.0, 1.0, 0.0, 1.0, 1.0],
|
|
"edep": [40.0, 30.0, 20.0, 0.0, 30.0],
|
|
"step_length": [1.0, 1.0, 1.0, 0.0, 1.0],
|
|
"material": ["G4_PbWO4"] * 5,
|
|
"layer_id": [0, 1, 1, -1, 0],
|
|
"n_sec_pred": [1, 0, 0, 0, 0],
|
|
"termination_reason": [
|
|
"",
|
|
"natural_end",
|
|
"natural_end",
|
|
"escaped",
|
|
"natural_end",
|
|
],
|
|
}
|
|
).lazy()
|
|
|
|
|
|
def _reference_frame() -> pl.LazyFrame:
|
|
return pl.DataFrame(
|
|
{
|
|
"event_id": [1, 1, 2],
|
|
"track_id": [0, 0, 0],
|
|
"step_no": [0, 1, 0],
|
|
"pdg": [11, 11, 11],
|
|
"pre_x": [0.0, 0.0, 0.0],
|
|
"pre_y": [0.0, 0.0, 0.0],
|
|
"pre_z": [0.0, 1.0, 0.0],
|
|
"pre_E": [100.0, 60.0, 50.0],
|
|
"pre_dx": [0.0, 0.0, 0.0],
|
|
"pre_dy": [0.0, 0.0, 0.0],
|
|
"pre_dz": [1.0, 1.0, 1.0],
|
|
"post_x": [0.0, 0.0, 0.0],
|
|
"post_y": [0.0, 0.0, 0.0],
|
|
"post_z": [1.0, 2.0, 1.0],
|
|
"post_E": [60.0, 30.0, 20.0],
|
|
"post_dx": [0.0, 0.0, 0.0],
|
|
"post_dy": [0.0, 0.0, 0.0],
|
|
"post_dz": [1.0, 1.0, 1.0],
|
|
"edep": [40.0, 30.0, 30.0],
|
|
"step_length": [1.0, 1.0, 1.0],
|
|
"material": ["G4_PbWO4", "G4_PbWO4", "G4_Pb"],
|
|
"layer_id": [0, 1, 0],
|
|
"sec_E_list": [[20.0], [], [10.0]],
|
|
"sec_pdg_list": [[22], [], [22]],
|
|
"sec_dx_list": [[1.0], [], [0.0]],
|
|
"sec_dy_list": [[0.0], [], [0.0]],
|
|
"sec_dz_list": [[0.0], [], [1.0]],
|
|
}
|
|
).lazy()
|
|
|
|
|
|
def test_hist1d_overall_and_grouped():
|
|
lf = _rollout_frame()
|
|
edges = np.linspace(0.0, 50.0, 6) # width 10
|
|
h = R.hist1d(lf, pl.col("edep"), edges)
|
|
# edep values: 40,30,20,0,30 -> bins [0),[10),[20),[30),[40)
|
|
assert h[0].tolist() == [1, 0, 1, 2, 1]
|
|
# grouped by pdg: pdg 22 has a single edep=20
|
|
hg = R.hist1d(lf, pl.col("edep"), edges, group=pl.col("pdg"))
|
|
assert hg[22].tolist() == [0, 0, 1, 0, 0]
|
|
assert hg[11].sum() == 4
|
|
|
|
|
|
def test_hist1d_clamps_extreme_values_and_drops_nan():
|
|
# A rollout can emit a wildly out-of-range step_length (or an inf/NaN); the
|
|
# fixed-edge binning must clamp rather than overflow the i32 bin cast.
|
|
lf = pl.DataFrame({"x": [5.0, 1.0725e10, float("inf"), -float("inf"), float("nan"), None]}).lazy()
|
|
edges = np.linspace(0.0, 50.0, 6) # width 10
|
|
h = R.hist1d(lf, pl.col("x"), edges)
|
|
# 5 -> bin 0; 1e10 and +inf -> top bin; -inf -> bin 0; NaN/null dropped
|
|
assert h[0].tolist() == [2, 0, 0, 0, 2]
|
|
|
|
|
|
def test_profile_partial_clamps_extreme_values_and_drops_nan():
|
|
lf = pl.DataFrame(
|
|
{
|
|
"event_id": [1, 1, 1, 1],
|
|
"z": [5.0, 1.0725e10, float("nan"), 45.0],
|
|
"w": [1.0, 2.0, 4.0, 8.0],
|
|
}
|
|
).lazy()
|
|
edges = np.linspace(0.0, 50.0, 6)
|
|
ev, mat = R.profile_partial(lf, pl.col("z"), edges, pl.col("w"))
|
|
assert ev.tolist() == [1]
|
|
# 1e10 clamps into the top bin alongside 45; the NaN row's weight is dropped
|
|
assert mat[0].tolist() == [1.0, 0.0, 0.0, 0.0, 10.0]
|
|
|
|
|
|
def test_physical_steps_drops_synthetic_rollout_rows_only():
|
|
lf = _rollout_frame()
|
|
phys = physical_steps(lf, Side.rollout).collect()
|
|
assert phys.height == 4 # dropped the escaped bookkeeping row
|
|
assert "escaped" not in phys["termination_reason"].to_list()
|
|
assert SYNTHETIC_TERMINATION_REASONS # non-empty guard
|
|
# reference passes through unchanged
|
|
ref = _reference_frame()
|
|
assert physical_steps(ref, Side.reference).collect().height == ref.collect().height
|
|
|
|
|
|
def test_event_scalars_totals_include_all_rows():
|
|
lf = _rollout_frame()
|
|
es = R.event_scalars(lf).sort("event_id")
|
|
row1 = es.filter(pl.col("event_id") == 1).to_dicts()[0]
|
|
assert row1["total_edep"] == 90.0 # 40+30+20+0
|
|
assert row1["incident_E"] == 100.0
|
|
assert row1["n_steps"] == 4
|
|
|
|
|
|
def test_secondaries_rollout_vs_reference_align():
|
|
r = secondaries(_rollout_frame(), Side.rollout).collect().sort("event_id")
|
|
assert r["energy"].to_list() == [20.0] # only the generation>0, step_no==0 row
|
|
assert r["pdg"].to_list() == [22]
|
|
t = secondaries(_reference_frame(), Side.reference).collect().sort("event_id")
|
|
# two secondaries (event 1 and event 2); empty list dropped
|
|
assert sorted(t["energy"].to_list()) == [10.0, 20.0]
|
|
assert t["pdg"].to_list() == [22, 22]
|
|
|
|
|
|
def test_leakage_fraction():
|
|
frac = R.leakage_fraction(_rollout_frame())
|
|
# event 1: escaped pre_E=30, deposited=90 -> 30/120 = 0.25; event 2: 0
|
|
assert sorted(round(f, 6) for f in frac) == [0.0, 0.25]
|
|
|
|
|
|
def test_weighted_profile_matches_manual_bincount():
|
|
lf = _rollout_frame()
|
|
ea = R.entry_axis(lf)
|
|
lf2 = R.attach_entry_axis(lf, ea)
|
|
edges = np.linspace(0.0, 3.0, 4) # depth bins along +z
|
|
mean, std = R.weighted_profile(lf2, R.depth_expr(), edges, pl.col("edep"))
|
|
assert mean.shape == (3,)
|
|
# totals conserved: sum over bins == mean total edep per event
|
|
assert np.isclose(mean.sum() * 1, (90.0 + 30.0) / 2) # 2 events
|
|
|
|
|
|
def test_energy_bins_edges_and_event_map():
|
|
incident = np.array([100.0, 100.0, 1000.0, 1000.0])
|
|
edges = G.energy_bin_edges(incident, n_bins=2)
|
|
assert len(edges) == 3 and edges[0] <= 100.0 < edges[-1]
|
|
ids, bins = G.event_energy_bins(_rollout_frame(), edges)
|
|
assert set(bins.tolist()) <= {0, 1}
|
|
assert len(ids) == 2
|
|
|
|
|
|
def test_digitize_expr_matches_numpy():
|
|
edges = np.array([0.0, 10.0, 100.0, 1000.0])
|
|
df = pl.DataFrame({"v": [5.0, 50.0, 500.0, 2000.0]})
|
|
got = df.select(G.digitize_expr(pl.col("v"), edges).alias("b"))["b"].to_list()
|
|
assert got == np.digitize([5.0, 50.0, 500.0, 2000.0], edges[1:-1]).tolist()
|
|
|
|
|
|
def test_pdg_and_material_labels():
|
|
assert G.pdg_label(22) == "gamma"
|
|
assert G.pdg_label(999999) == "999999"
|
|
assert G.material_label("G4_PbWO4") == "PbWO4"
|
|
|
|
|
|
def _write_shard(path, event_ids, edeps):
|
|
pl.DataFrame({"event_id": event_ids, "pdg": [11] * len(event_ids), "edep": edeps}).write_parquet(path)
|
|
|
|
|
|
def test_open_side_reference_offsets_event_ids_across_shards(tmp_path):
|
|
# Each shard is a separate Geant4 job whose own event_id numbering restarts
|
|
# from 0 — a naive multi-shard scan collides on event_id across shards.
|
|
_write_shard(tmp_path / "a.parquet", [0, 1], [1.0, 2.0])
|
|
_write_shard(tmp_path / "b.parquet", [0, 1], [3.0, 4.0])
|
|
df = open_side(tmp_path, Side.reference).sort("event_id").collect()
|
|
assert df["event_id"].to_list() == [0, 1, EVENT_ID_FILE_STRIDE, EVENT_ID_FILE_STRIDE + 1]
|
|
assert df["edep"].to_list() == [1.0, 2.0, 3.0, 4.0]
|
|
assert "__source_path" not in df.columns
|
|
|
|
|
|
def test_open_side_reference_single_file_unchanged(tmp_path):
|
|
_write_shard(tmp_path / "only.parquet", [0, 1], [1.0, 2.0])
|
|
df = open_side(tmp_path / "only.parquet", Side.reference).sort("event_id").collect()
|
|
assert df["event_id"].to_list() == [0, 1]
|
|
assert "__source_path" not in df.columns
|
|
|
|
|
|
def test_open_side_reference_manifest(tmp_path):
|
|
_write_shard(tmp_path / "a.parquet", [0, 1], [1.0, 2.0])
|
|
_write_shard(tmp_path / "b.parquet", [0, 1], [3.0, 4.0])
|
|
manifest = tmp_path / "shards.manifest"
|
|
manifest.write_text("a.parquet\nb.parquet\n")
|
|
df = open_side(manifest, Side.reference).sort("event_id").collect()
|
|
assert df["event_id"].to_list() == [0, 1, EVENT_ID_FILE_STRIDE, EVENT_ID_FILE_STRIDE + 1]
|
|
assert df["edep"].to_list() == [1.0, 2.0, 3.0, 4.0]
|