lars
81ec225b86
Merged in every non-analysis change from the MoE-prototype branch (routing, training, data pipeline, streaming rollout output), keeping this branch's lean streaming giant/analysis.py and rebuilding the rollout-vs-truth feature natively on it instead of resurrecting the old numpy SampleCollection path. - Add RolloutVsTruth, accepted anywhere Tier 1-3 functions take a predict-parquet source: decodes a giant rollout file and a held-out truth file into RAW_TARGET_NAMES space via a polars port of the forward local-frame rotation, fully streaming (no SampleCollection, no eager materialization). - Add compute_rollout_vs_truth_observables_pl for Tier 4, reusing EventObservables (now backed by independent real_table/gen_table to support unequal rollout/truth event counts) so every existing shower-observable plot function works unchanged for both one-step and full-rollout comparisons. - Update analysis/rollout_validation.ipynb to the new API and CLAUDE.md's architecture description; add test coverage for the new source type. - Fix a pre-existing return-type mismatch in giant.rollout.rollout() (found by `ty check`): the on_chunk summary-dict branch didn't match the declared dict[str, np.ndarray] return type, now expressed as a RolloutSummary TypedDict. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
1176 lines
42 KiB
Python
1176 lines
42 KiB
Python
import matplotlib
|
||
|
||
matplotlib.use("Agg") # no display needed for plot smoke tests
|
||
|
||
import numpy as np
|
||
import pandas as pd
|
||
import polars as pl
|
||
import pyarrow as pa
|
||
import pyarrow.parquet as pq
|
||
import pytest
|
||
|
||
from giant.analysis import (
|
||
RAW_TARGET_NAMES,
|
||
RolloutVsTruth,
|
||
compute_event_observables_pl,
|
||
compute_rollout_vs_truth_observables_pl,
|
||
constraint_report_pl,
|
||
correlation_matrices_pl,
|
||
marginal_table_pl,
|
||
pdg_contribution_table_pl,
|
||
plot_constraint_violations,
|
||
plot_correlation_matrices,
|
||
plot_direction_alignment,
|
||
plot_kl_bars_pl,
|
||
plot_longitudinal_profile,
|
||
plot_marginals,
|
||
plot_mean_energy_per_step,
|
||
plot_mean_length_per_step,
|
||
plot_pairwise,
|
||
plot_pdg_energy_share,
|
||
plot_pdg_length_share,
|
||
plot_shower_max_depth,
|
||
plot_total_energy,
|
||
plot_total_length,
|
||
plot_transverse_profile,
|
||
)
|
||
from giant.constants import (
|
||
LOCAL_TARGET_NAMES,
|
||
PREDICT_COORD_METADATA_KEY,
|
||
PREDICT_SCHEMA_VERSION,
|
||
PREDICT_SCHEMA_VERSION_KEY,
|
||
ROLLOUT_COORD_VALUE,
|
||
TERM_ENERGY_CUTOFF,
|
||
TERM_NATURAL_END,
|
||
)
|
||
from giant.data.transforms import (
|
||
energy_simplex_decode,
|
||
inv_log_transform,
|
||
local_frame_rotation,
|
||
log_transform,
|
||
reconstruct_post_pos,
|
||
travel_direction,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Fixtures: predict `--coord local` parquet writers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _unit_vectors(rng, n):
|
||
v = rng.standard_normal((n, 3)).astype(np.float32)
|
||
return v / np.linalg.norm(v, axis=1, keepdims=True)
|
||
|
||
|
||
def _write_predicted_local_parquet(path, n=50, metadata=None, rng=None):
|
||
"""Mimic `giant predict --coord local`'s output schema for the loader tests.
|
||
|
||
Column 0 is a log-scaled step_length; columns 1–2 are the deposit/secondary
|
||
ALR energy logits (unconstrained reals, decoded against pre_E); columns 3–8
|
||
are direction components.
|
||
"""
|
||
rng = rng or np.random.default_rng(0)
|
||
true_log_local = rng.standard_normal((n, 9)).astype(np.float32)
|
||
true_log_local[:, 0] = log_transform(rng.uniform(0.1, 5.0, n).astype(np.float32))
|
||
pred_log_local = true_log_local + rng.normal(0, 0.01, (n, 9)).astype(np.float32)
|
||
pre_E = rng.uniform(1.0, 100.0, n).astype(np.float32)
|
||
|
||
table = pa.table(
|
||
{
|
||
"event_id": rng.integers(0, 10, n),
|
||
"pdg": rng.choice([11, -11, 22], n),
|
||
"pre_x": rng.standard_normal(n).astype(np.float32),
|
||
"pre_y": rng.standard_normal(n).astype(np.float32),
|
||
"pre_z": rng.standard_normal(n).astype(np.float32),
|
||
"pre_E": pre_E,
|
||
"pre_dx": rng.standard_normal(n).astype(np.float32),
|
||
"pre_dy": rng.standard_normal(n).astype(np.float32),
|
||
"pre_dz": rng.standard_normal(n).astype(np.float32),
|
||
"material": rng.choice(["W", "Pb"], n),
|
||
"layer_id": rng.integers(0, 10, n).astype(np.int32),
|
||
"n_sec": rng.integers(0, 3, n).astype(np.int32),
|
||
**{
|
||
f"pred_{name}": pred_log_local[:, j]
|
||
for j, name in enumerate(LOCAL_TARGET_NAMES)
|
||
},
|
||
**{
|
||
f"true_{name}": true_log_local[:, j]
|
||
for j, name in enumerate(LOCAL_TARGET_NAMES)
|
||
},
|
||
}
|
||
)
|
||
if metadata is not None:
|
||
table = table.replace_schema_metadata(metadata)
|
||
pq.write_table(table, path)
|
||
return true_log_local, pred_log_local, pre_E
|
||
|
||
|
||
def _predicted_local_path(tmp_path, n=200, seed=0):
|
||
path = tmp_path / "predicted_local.parquet"
|
||
_write_predicted_local_parquet(
|
||
path,
|
||
n=n,
|
||
rng=np.random.default_rng(seed),
|
||
metadata={
|
||
PREDICT_COORD_METADATA_KEY: "local",
|
||
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
|
||
},
|
||
)
|
||
return path
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Inline numpy oracle: the ground truth the streaming functions are checked
|
||
# against. This is the raw-space decoding the old `SampleCollection` path did,
|
||
# recomputed directly from the small fixture rather than in the module.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _raw_from_parquet(path):
|
||
"""(real_raw, gen_raw, pdg, material, pre_E) in physical units, via numpy."""
|
||
df = pl.read_parquet(path)
|
||
pre_E = df["pre_E"].to_numpy().astype(np.float32)
|
||
|
||
def decode(prefix):
|
||
log_local = np.column_stack(
|
||
[df[f"{prefix}_{name}"].to_numpy() for name in LOCAL_TARGET_NAMES]
|
||
).astype(np.float32)
|
||
raw = log_local.copy()
|
||
raw[:, 0] = inv_log_transform(log_local[:, 0])
|
||
edep, _e_sec, _post_E, delta_e = energy_simplex_decode(log_local[:, 1:3], pre_E)
|
||
raw[:, 1] = delta_e
|
||
raw[:, 2] = edep
|
||
return raw
|
||
|
||
return (
|
||
decode("true"),
|
||
decode("pred"),
|
||
df["pdg"].to_numpy(),
|
||
df["material"].to_numpy(),
|
||
pre_E,
|
||
)
|
||
|
||
|
||
def _histogram_kl(p_samples, q_samples, bins=50, eps=1e-8):
|
||
"""KL(P || Q) between two 1D samples via a shared histogram (numpy oracle)."""
|
||
lo = min(p_samples.min(), q_samples.min())
|
||
hi = max(p_samples.max(), q_samples.max())
|
||
if hi <= lo:
|
||
return 0.0
|
||
edges = np.linspace(lo, hi, bins + 1)
|
||
p_hist, _ = np.histogram(p_samples, bins=edges)
|
||
q_hist, _ = np.histogram(q_samples, bins=edges)
|
||
p = p_hist.astype(np.float64) + eps
|
||
q = q_hist.astype(np.float64) + eps
|
||
p /= p.sum()
|
||
q /= q.sum()
|
||
return float(np.sum(p * np.log(p / q)))
|
||
|
||
|
||
def _group_masks(pdg, material, pre_E, group_by, n_energy_bins=4):
|
||
"""Replicate the module's `_group` labelling so oracle labels line up."""
|
||
n = len(pdg)
|
||
if group_by is None:
|
||
return [("all", np.ones(n, dtype=bool))]
|
||
if group_by == "pdg":
|
||
return [(f"pdg={int(v)}", pdg == v) for v in np.unique(pdg)]
|
||
if group_by == "material":
|
||
return [(f"material={v}", material == v) for v in np.unique(material)]
|
||
if group_by == "energy":
|
||
edges = np.quantile(pre_E, np.linspace(0, 1, n_energy_bins + 1))
|
||
edges[-1] += 1e-6
|
||
bin_idx = np.digitize(pre_E, edges[1:-1])
|
||
return [
|
||
(f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})", bin_idx == i)
|
||
for i in range(n_energy_bins)
|
||
]
|
||
raise ValueError(group_by)
|
||
|
||
|
||
def _numpy_marginal_table(real, gen, pdg, material, pre_E, group_by, bins=50):
|
||
rows = []
|
||
for label, mask in _group_masks(pdg, material, pre_E, group_by):
|
||
if mask.sum() < 2:
|
||
continue
|
||
r, g = real[mask], gen[mask]
|
||
for j, name in enumerate(RAW_TARGET_NAMES):
|
||
rows.append(
|
||
{
|
||
"group": label,
|
||
"dim": name,
|
||
"n": int(mask.sum()),
|
||
"real_mean": r[:, j].mean(),
|
||
"gen_mean": g[:, j].mean(),
|
||
"real_std": r[:, j].std(),
|
||
"gen_std": g[:, j].std(),
|
||
"kl_real_gen": _histogram_kl(r[:, j], g[:, j], bins=bins),
|
||
}
|
||
)
|
||
return pd.DataFrame(rows).sort_values(["group", "dim"]).reset_index(drop=True)
|
||
|
||
|
||
def _numpy_constraint_report(gen, norm_tol=0.05):
|
||
post_norm = np.linalg.norm(gen[:, 3:6], axis=1)
|
||
travel_norm = np.linalg.norm(gen[:, 6:9], axis=1)
|
||
rows = [
|
||
{
|
||
"check": "post_dir unit norm",
|
||
"violation_rate": float(np.mean(np.abs(post_norm - 1) > norm_tol)),
|
||
"mean_abs_error": float(np.mean(np.abs(post_norm - 1))),
|
||
},
|
||
{
|
||
"check": "travel_dir unit norm",
|
||
"violation_rate": float(np.mean(np.abs(travel_norm - 1) > norm_tol)),
|
||
"mean_abs_error": float(np.mean(np.abs(travel_norm - 1))),
|
||
},
|
||
]
|
||
for j, name in enumerate(RAW_TARGET_NAMES[:3]):
|
||
rows.append(
|
||
{
|
||
"check": f"{name} >= 0",
|
||
"violation_rate": float(np.mean(gen[:, j] < 0)),
|
||
"mean_abs_error": float(np.mean(np.clip(-gen[:, j], 0, None))),
|
||
}
|
||
)
|
||
return pd.DataFrame(rows)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tier 1: stratified marginals
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
|
||
def test_marginal_table_pl_matches_numpy_oracle(tmp_path, group_by):
|
||
path = _predicted_local_path(tmp_path)
|
||
real, gen, pdg, material, pre_E = _raw_from_parquet(path)
|
||
|
||
expected = _numpy_marginal_table(real, gen, pdg, material, pre_E, group_by)
|
||
actual = (
|
||
marginal_table_pl(path, group_by=group_by).sort(["group", "dim"]).to_pandas()
|
||
)
|
||
|
||
assert list(expected["group"]) == list(actual["group"])
|
||
assert list(expected["n"]) == list(actual["n"])
|
||
for col in ["real_mean", "gen_mean", "real_std", "gen_std"]:
|
||
np.testing.assert_allclose(
|
||
expected[col].to_numpy(), actual[col].to_numpy(), atol=1e-4, rtol=1e-4
|
||
)
|
||
# KL uses np.histogram (oracle) vs polars binning (lazy path); the two bin the
|
||
# boundary (min/max) sample differently, so allow a small absolute discrepancy
|
||
# rather than requiring bit-identical estimates.
|
||
np.testing.assert_allclose(
|
||
expected["kl_real_gen"].to_numpy(), actual["kl_real_gen"].to_numpy(), atol=2e-2
|
||
)
|
||
|
||
|
||
def test_marginal_table_pl_aggregate_has_all_dims(tmp_path):
|
||
table = marginal_table_pl(_predicted_local_path(tmp_path))
|
||
assert set(table["dim"].to_list()) == set(RAW_TARGET_NAMES)
|
||
assert (table["group"] == "all").all()
|
||
|
||
|
||
def test_marginal_table_pl_accepts_lazyframe(tmp_path):
|
||
path = _predicted_local_path(tmp_path)
|
||
from_path = marginal_table_pl(path).sort(["group", "dim"])
|
||
from_lf = marginal_table_pl(pl.scan_parquet(path)).sort(["group", "dim"])
|
||
np.testing.assert_allclose(
|
||
from_lf["kl_real_gen"].to_numpy(), from_path["kl_real_gen"].to_numpy()
|
||
)
|
||
|
||
|
||
def test_marginal_table_pl_rejects_missing_metadata(tmp_path):
|
||
path = tmp_path / "no_metadata.parquet"
|
||
_write_predicted_local_parquet(path, metadata=None)
|
||
with pytest.raises(ValueError, match="no '.*' parquet metadata"):
|
||
marginal_table_pl(path)
|
||
|
||
|
||
def test_marginal_table_pl_rejects_global_coord(tmp_path):
|
||
path = tmp_path / "global.parquet"
|
||
_write_predicted_local_parquet(
|
||
path,
|
||
metadata={
|
||
PREDICT_COORD_METADATA_KEY: "global",
|
||
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
|
||
},
|
||
)
|
||
with pytest.raises(ValueError, match="coord=local"):
|
||
marginal_table_pl(path)
|
||
|
||
|
||
def test_marginal_table_pl_rejects_mismatched_schema_version(tmp_path):
|
||
path = tmp_path / "old_version.parquet"
|
||
_write_predicted_local_parquet(
|
||
path,
|
||
metadata={
|
||
PREDICT_COORD_METADATA_KEY: "local",
|
||
PREDICT_SCHEMA_VERSION_KEY: "999",
|
||
},
|
||
)
|
||
with pytest.raises(ValueError, match="schema version"):
|
||
marginal_table_pl(path)
|
||
|
||
|
||
@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
|
||
def test_plot_kl_bars_pl_runs_without_error(tmp_path, group_by):
|
||
fig = plot_kl_bars_pl(_predicted_local_path(tmp_path), group_by=group_by)
|
||
assert fig is not None
|
||
|
||
|
||
def test_plot_marginals_runs_without_error(tmp_path):
|
||
fig = plot_marginals(_predicted_local_path(tmp_path))
|
||
assert len(fig.axes) == len(RAW_TARGET_NAMES) # single "all" row
|
||
|
||
|
||
def test_plot_marginals_grouped_runs_without_error(tmp_path):
|
||
fig = plot_marginals(_predicted_local_path(tmp_path), group_by="material")
|
||
assert fig is not None
|
||
|
||
|
||
def test_plot_marginals_caps_groups(tmp_path):
|
||
path = _predicted_local_path(tmp_path, n=400)
|
||
df = pl.read_parquet(path).with_columns(
|
||
pl.Series("pdg", np.arange(400) % 8, dtype=pl.Int64)
|
||
)
|
||
fig = plot_marginals(df.lazy(), group_by="pdg", max_groups=3)
|
||
n_rows = len(fig.axes) // len(RAW_TARGET_NAMES)
|
||
assert n_rows <= 3
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tier 2: joint structure
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_correlation_matrices_pl_matches_numpy(tmp_path):
|
||
path = _predicted_local_path(tmp_path)
|
||
real, gen, *_ = _raw_from_parquet(path)
|
||
|
||
real_corr, gen_corr = correlation_matrices_pl(path)
|
||
|
||
np.testing.assert_allclose(
|
||
real_corr, np.corrcoef(real, rowvar=False), atol=1e-4, rtol=1e-4
|
||
)
|
||
np.testing.assert_allclose(
|
||
gen_corr, np.corrcoef(gen, rowvar=False), atol=1e-4, rtol=1e-4
|
||
)
|
||
for corr in (real_corr, gen_corr):
|
||
np.testing.assert_allclose(np.diag(corr), 1.0, atol=1e-6)
|
||
np.testing.assert_allclose(corr, corr.T, atol=1e-6)
|
||
|
||
|
||
def test_plot_correlation_matrices_runs_without_error(tmp_path):
|
||
fig = plot_correlation_matrices(_predicted_local_path(tmp_path))
|
||
assert len(fig.axes) >= 3 # real, generated, difference (+ colorbars)
|
||
|
||
|
||
def test_plot_pairwise_runs_without_error(tmp_path):
|
||
fig = plot_pairwise(_predicted_local_path(tmp_path), n_sample=100)
|
||
assert len(fig.axes) == 6 # 2 rows (real/gen) × 3 default pairs
|
||
|
||
|
||
def test_plot_direction_alignment_runs_without_error(tmp_path):
|
||
fig = plot_direction_alignment(_predicted_local_path(tmp_path))
|
||
assert fig is not None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tier 3: physical constraints
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_constraint_report_pl_matches_numpy_oracle(tmp_path):
|
||
path = _predicted_local_path(tmp_path)
|
||
_real, gen, *_ = _raw_from_parquet(path)
|
||
|
||
expected = _numpy_constraint_report(gen)
|
||
actual = constraint_report_pl(path).to_pandas()
|
||
|
||
assert list(expected["check"]) == list(actual["check"])
|
||
np.testing.assert_allclose(
|
||
expected["violation_rate"].to_numpy(),
|
||
actual["violation_rate"].to_numpy(),
|
||
atol=1e-6,
|
||
)
|
||
np.testing.assert_allclose(
|
||
expected["mean_abs_error"].to_numpy(),
|
||
actual["mean_abs_error"].to_numpy(),
|
||
atol=1e-4,
|
||
)
|
||
|
||
|
||
def test_constraint_report_pl_flags_bad_direction_norms(tmp_path):
|
||
path = _predicted_local_path(tmp_path)
|
||
# Force the generated post_dir off the unit sphere for every row: the fixed
|
||
# (1, 1, 1) vector has norm √3 ≈ 1.73, well outside the tolerance.
|
||
df = pl.read_parquet(path).with_columns(
|
||
pl.lit(1.0).alias("pred_post_dx"),
|
||
pl.lit(1.0).alias("pred_post_dy"),
|
||
pl.lit(1.0).alias("pred_post_dz"),
|
||
)
|
||
report = constraint_report_pl(df.lazy()).to_pandas()
|
||
rate = dict(zip(report["check"], report["violation_rate"]))
|
||
assert rate["post_dir unit norm"] == 1.0
|
||
|
||
|
||
def test_plot_constraint_violations_runs_without_error(tmp_path):
|
||
fig = plot_constraint_violations(_predicted_local_path(tmp_path))
|
||
assert len(fig.axes) == 2 + 3 # 2 direction norms + 3 scalar dims
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tier 4: event-level (shower) observables
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _make_event_level_arrays(rng):
|
||
"""3 events (3/2/4 steps), each with an unambiguous highest-pre_E row.
|
||
|
||
The forced max-pre_E rows (indices 1, 3, 7) fix a known shower axis/entry
|
||
point per event, so the expected event_table can be re-derived independently
|
||
in the test without depending on compute_event_observables_pl.
|
||
"""
|
||
event_id = np.array([0, 0, 0, 1, 1, 2, 2, 2, 2], dtype=np.int64)
|
||
n = len(event_id)
|
||
pre_pos = rng.uniform(-5.0, 5.0, (n, 3)).astype(np.float32)
|
||
pre_dir = _unit_vectors(rng, n)
|
||
pre_E = rng.uniform(1.0, 50.0, n).astype(np.float32)
|
||
pre_E[1] = 100.0 # event 0's entry step
|
||
pre_E[3] = 100.0 # event 1's entry step
|
||
pre_E[7] = 100.0 # event 2's entry step
|
||
|
||
def _local_block():
|
||
block = rng.standard_normal((n, 9)).astype(np.float32)
|
||
block[:, 0] = log_transform(rng.uniform(0.1, 5.0, n).astype(np.float32))
|
||
# cols 1–2 stay as random ALR energy logits (decoded against pre_E)
|
||
block[:, 3:6] = _unit_vectors(rng, n)
|
||
block[:, 6:9] = _unit_vectors(rng, n)
|
||
return block
|
||
|
||
true_log_local = _local_block()
|
||
pred_log_local = _local_block()
|
||
return event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local
|
||
|
||
|
||
def _write_event_level_parquet(path, rng=None):
|
||
rng = rng or np.random.default_rng(7)
|
||
event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local = (
|
||
_make_event_level_arrays(rng)
|
||
)
|
||
n = len(event_id)
|
||
pdg = rng.choice([11, -11, 22], n)
|
||
table = pa.table(
|
||
{
|
||
"event_id": event_id,
|
||
"pdg": pdg,
|
||
"pre_x": pre_pos[:, 0],
|
||
"pre_y": pre_pos[:, 1],
|
||
"pre_z": pre_pos[:, 2],
|
||
"pre_E": pre_E,
|
||
"pre_dx": pre_dir[:, 0],
|
||
"pre_dy": pre_dir[:, 1],
|
||
"pre_dz": pre_dir[:, 2],
|
||
"material": rng.choice(["W", "Pb"], n),
|
||
"layer_id": rng.integers(0, 10, n).astype(np.int32),
|
||
"n_sec": rng.integers(0, 3, n).astype(np.int32),
|
||
**{
|
||
f"pred_{name}": pred_log_local[:, j]
|
||
for j, name in enumerate(LOCAL_TARGET_NAMES)
|
||
},
|
||
**{
|
||
f"true_{name}": true_log_local[:, j]
|
||
for j, name in enumerate(LOCAL_TARGET_NAMES)
|
||
},
|
||
}
|
||
)
|
||
table = table.replace_schema_metadata(
|
||
{
|
||
PREDICT_COORD_METADATA_KEY: "local",
|
||
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
|
||
}
|
||
)
|
||
pq.write_table(table, path)
|
||
return event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local, pdg
|
||
|
||
|
||
def _expected_event_table(
|
||
event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local
|
||
):
|
||
"""Independent re-derivation of total/centroid/RMS per event, for comparison."""
|
||
expected = {}
|
||
for e in sorted(np.unique(event_id).tolist()):
|
||
mask = event_id == e
|
||
entry_idx = np.where(mask)[0][np.argmax(pre_E[mask])]
|
||
entry_pos = pre_pos[entry_idx]
|
||
axis_dir = pre_dir[entry_idx]
|
||
|
||
def agg(log_local, mask=mask, entry_pos=entry_pos, axis_dir=axis_dir):
|
||
step_length = inv_log_transform(log_local[mask, 0])
|
||
edep, _e_sec, _post_E, _delta_e = energy_simplex_decode(
|
||
log_local[mask, 1:3], pre_E[mask]
|
||
)
|
||
travel_dir_local = log_local[mask, 6:9]
|
||
post_pos = reconstruct_post_pos(
|
||
pre_pos[mask], pre_dir[mask], step_length, travel_dir_local
|
||
)
|
||
disp = post_pos - entry_pos
|
||
depth = disp @ axis_dir
|
||
transverse = np.linalg.norm(disp - depth[:, None] * axis_dir, axis=1)
|
||
total_edep = float(edep.sum())
|
||
total_length = float(step_length.sum())
|
||
centroid = float((edep * depth).sum() / total_edep)
|
||
rms = float(np.sqrt((edep * transverse**2).sum() / total_edep))
|
||
return total_edep, total_length, centroid, rms
|
||
|
||
real_total_edep, real_total_length, real_centroid, real_rms = agg(
|
||
true_log_local
|
||
)
|
||
gen_total_edep, gen_total_length, gen_centroid, gen_rms = agg(pred_log_local)
|
||
expected[e] = (
|
||
real_total_edep,
|
||
gen_total_edep,
|
||
real_total_length,
|
||
gen_total_length,
|
||
real_centroid,
|
||
gen_centroid,
|
||
real_rms,
|
||
gen_rms,
|
||
)
|
||
return expected
|
||
|
||
|
||
def test_compute_event_observables_pl_matches_manual_reconstruction(tmp_path):
|
||
path = tmp_path / "event_level.parquet"
|
||
event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local, _pdg = (
|
||
_write_event_level_parquet(path)
|
||
)
|
||
expected = _expected_event_table(
|
||
event_id, pre_pos, pre_dir, pre_E, true_log_local, pred_log_local
|
||
)
|
||
|
||
obs = compute_event_observables_pl(path, depth_bins=5, transverse_bins=5)
|
||
real_table = obs.real_table.sort("event_id")
|
||
gen_table = obs.gen_table.sort("event_id")
|
||
|
||
for i, eid in enumerate(real_table["event_id"].to_list()):
|
||
(
|
||
real_total_edep,
|
||
gen_total_edep,
|
||
real_total_length,
|
||
gen_total_length,
|
||
real_centroid,
|
||
gen_centroid,
|
||
real_rms,
|
||
gen_rms,
|
||
) = expected[eid]
|
||
np.testing.assert_allclose(
|
||
real_table["total_edep"][i], real_total_edep, rtol=1e-4
|
||
)
|
||
np.testing.assert_allclose(
|
||
gen_table["total_edep"][i], gen_total_edep, rtol=1e-4
|
||
)
|
||
np.testing.assert_allclose(
|
||
real_table["total_length"][i], real_total_length, rtol=1e-4
|
||
)
|
||
np.testing.assert_allclose(
|
||
gen_table["total_length"][i], gen_total_length, rtol=1e-4
|
||
)
|
||
np.testing.assert_allclose(
|
||
real_table["centroid_depth"][i], real_centroid, rtol=1e-3, atol=1e-4
|
||
)
|
||
np.testing.assert_allclose(
|
||
gen_table["centroid_depth"][i], gen_centroid, rtol=1e-3, atol=1e-4
|
||
)
|
||
np.testing.assert_allclose(
|
||
real_table["transverse_rms"][i], real_rms, rtol=1e-3, atol=1e-4
|
||
)
|
||
np.testing.assert_allclose(
|
||
gen_table["transverse_rms"][i], gen_rms, rtol=1e-3, atol=1e-4
|
||
)
|
||
|
||
|
||
def test_compute_event_observables_pl_profile_shapes(tmp_path):
|
||
path = tmp_path / "event_level.parquet"
|
||
_write_event_level_parquet(path)
|
||
obs = compute_event_observables_pl(path, depth_bins=7, transverse_bins=4)
|
||
|
||
assert obs.depth_edges.shape == (8,)
|
||
assert obs.transverse_edges.shape == (5,)
|
||
assert obs.real_depth_profile.shape == (7,)
|
||
assert obs.gen_depth_profile.shape == (7,)
|
||
assert obs.real_transverse_profile.shape == (4,)
|
||
assert obs.gen_transverse_profile.shape == (4,)
|
||
assert len(obs.real_table) == 3
|
||
assert len(obs.gen_table) == 3
|
||
|
||
|
||
def test_compute_event_observables_pl_accepts_lazyframe(tmp_path):
|
||
path = tmp_path / "event_level.parquet"
|
||
_write_event_level_parquet(path)
|
||
|
||
obs_from_path = compute_event_observables_pl(path, depth_bins=5, transverse_bins=5)
|
||
obs_from_lf = compute_event_observables_pl(
|
||
pl.scan_parquet(path), depth_bins=5, transverse_bins=5
|
||
)
|
||
|
||
np.testing.assert_allclose(
|
||
obs_from_lf.real_table.sort("event_id")["total_edep"].to_numpy(),
|
||
obs_from_path.real_table.sort("event_id")["total_edep"].to_numpy(),
|
||
)
|
||
|
||
|
||
def test_compute_event_observables_pl_approx_median(tmp_path):
|
||
"""The streaming log-bin median approximates the true per-event median."""
|
||
rng = np.random.default_rng(11)
|
||
n = 4000 # one event, many steps → a well-defined median
|
||
true_log = rng.standard_normal((n, 9)).astype(np.float32)
|
||
true_log[:, 0] = log_transform(rng.uniform(0.1, 5.0, n).astype(np.float32))
|
||
for s in (3, 6):
|
||
v = rng.standard_normal((n, 3)).astype(np.float32)
|
||
true_log[:, s : s + 3] = v / np.linalg.norm(v, axis=1, keepdims=True)
|
||
pre_E = rng.uniform(1.0, 100.0, n).astype(np.float32)
|
||
pre = rng.standard_normal((n, 3)).astype(np.float32)
|
||
pdir = _unit_vectors(rng, n)
|
||
pre_E[0] = 1000.0 # entry step
|
||
|
||
path = tmp_path / "one_event.parquet"
|
||
table = pa.table(
|
||
{
|
||
"event_id": np.zeros(n, dtype=np.int64),
|
||
"pdg": rng.choice([11, 22], n),
|
||
"pre_x": pre[:, 0],
|
||
"pre_y": pre[:, 1],
|
||
"pre_z": pre[:, 2],
|
||
"pre_E": pre_E,
|
||
"pre_dx": pdir[:, 0],
|
||
"pre_dy": pdir[:, 1],
|
||
"pre_dz": pdir[:, 2],
|
||
"material": rng.choice(["W", "Pb"], n),
|
||
"layer_id": rng.integers(0, 10, n).astype(np.int32),
|
||
"n_sec": rng.integers(0, 3, n).astype(np.int32),
|
||
**{f"pred_{nm}": true_log[:, j] for j, nm in enumerate(LOCAL_TARGET_NAMES)},
|
||
**{f"true_{nm}": true_log[:, j] for j, nm in enumerate(LOCAL_TARGET_NAMES)},
|
||
}
|
||
).replace_schema_metadata(
|
||
{
|
||
PREDICT_COORD_METADATA_KEY: "local",
|
||
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
|
||
}
|
||
)
|
||
pq.write_table(table, path)
|
||
|
||
true_edep = energy_simplex_decode(true_log[:, 1:3], pre_E)[0]
|
||
true_length = inv_log_transform(true_log[:, 0])
|
||
exact_edep_median = np.median(true_edep)
|
||
exact_length_median = np.median(true_length)
|
||
|
||
obs = compute_event_observables_pl(path)
|
||
approx_edep = obs.real_table["median_edep"][0]
|
||
approx_length = obs.real_table["median_length"][0]
|
||
|
||
# log-bin interpolation: expect within a few percent of the true median.
|
||
np.testing.assert_allclose(approx_edep, exact_edep_median, rtol=0.05)
|
||
np.testing.assert_allclose(approx_length, exact_length_median, rtol=0.05)
|
||
|
||
|
||
def test_event_level_plots_run_without_error(tmp_path):
|
||
path = tmp_path / "event_level.parquet"
|
||
_write_event_level_parquet(path)
|
||
obs = compute_event_observables_pl(path, depth_bins=5, transverse_bins=5)
|
||
|
||
assert plot_total_energy(obs) is not None
|
||
assert plot_total_length(obs) is not None
|
||
assert plot_longitudinal_profile(obs) is not None
|
||
assert plot_transverse_profile(obs) is not None
|
||
assert plot_shower_max_depth(obs) is not None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Particle-species (pdg) contribution shares
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_pdg_contribution_table_pl_matches_manual_sums(tmp_path):
|
||
path = tmp_path / "event_level.parquet"
|
||
_, _, _, pre_E, true_log_local, pred_log_local, pdg = _write_event_level_parquet(
|
||
path
|
||
)
|
||
|
||
real_edep = energy_simplex_decode(true_log_local[:, 1:3], pre_E)[0]
|
||
gen_edep = energy_simplex_decode(pred_log_local[:, 1:3], pre_E)[0]
|
||
real_length = inv_log_transform(true_log_local[:, 0])
|
||
gen_length = inv_log_transform(pred_log_local[:, 0])
|
||
|
||
expected = {}
|
||
for p in np.unique(pdg):
|
||
mask = pdg == p
|
||
expected[int(p)] = (
|
||
float(real_edep[mask].sum()),
|
||
float(gen_edep[mask].sum()),
|
||
float(real_length[mask].sum()),
|
||
float(gen_length[mask].sum()),
|
||
)
|
||
|
||
table = pdg_contribution_table_pl(path).sort("pdg")
|
||
for i, p in enumerate(table["pdg"].to_list()):
|
||
real_e, gen_e, real_l, gen_l = expected[int(p)]
|
||
np.testing.assert_allclose(table["real_total_edep"][i], real_e, rtol=1e-4)
|
||
np.testing.assert_allclose(table["gen_total_edep"][i], gen_e, rtol=1e-4)
|
||
np.testing.assert_allclose(table["real_total_length"][i], real_l, rtol=1e-4)
|
||
np.testing.assert_allclose(table["gen_total_length"][i], gen_l, rtol=1e-4)
|
||
|
||
|
||
def test_pdg_contribution_table_pl_accepts_lazyframe(tmp_path):
|
||
path = tmp_path / "event_level.parquet"
|
||
_write_event_level_parquet(path)
|
||
|
||
from_path = pdg_contribution_table_pl(path).sort("pdg")
|
||
from_lf = pdg_contribution_table_pl(pl.scan_parquet(path)).sort("pdg")
|
||
|
||
np.testing.assert_allclose(
|
||
from_lf["real_total_edep"].to_numpy(), from_path["real_total_edep"].to_numpy()
|
||
)
|
||
|
||
|
||
def test_pdg_pie_plots_run_without_error(tmp_path):
|
||
path = tmp_path / "event_level.parquet"
|
||
_write_event_level_parquet(path)
|
||
table = pdg_contribution_table_pl(path)
|
||
|
||
assert plot_pdg_energy_share(table) is not None
|
||
assert plot_pdg_length_share(table) is not None
|
||
|
||
|
||
def test_plot_pdg_energy_share_caps_slices():
|
||
table = pl.DataFrame(
|
||
{
|
||
"pdg": list(range(10)),
|
||
"real_total_edep": [float(10 - i) for i in range(10)],
|
||
"gen_total_edep": [float(10 - i) for i in range(10)],
|
||
"real_total_length": [float(10 - i) for i in range(10)],
|
||
"gen_total_length": [float(10 - i) for i in range(10)],
|
||
}
|
||
)
|
||
fig = plot_pdg_energy_share(table, max_slices=4)
|
||
for ax in fig.axes:
|
||
assert len(ax.patches) == 4
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Rollout vs. held-out truth: unpaired `RolloutVsTruth` source (Tier 1-3) and
|
||
# `compute_rollout_vs_truth_observables_pl` (Tier 4)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_WORLD_FRAME_COLS = [
|
||
"event_id",
|
||
"pdg",
|
||
"material",
|
||
"layer_id",
|
||
"pre_x",
|
||
"pre_y",
|
||
"pre_z",
|
||
"pre_E",
|
||
"pre_dx",
|
||
"pre_dy",
|
||
"pre_dz",
|
||
"post_x",
|
||
"post_y",
|
||
"post_z",
|
||
"post_E",
|
||
"post_dx",
|
||
"post_dy",
|
||
"post_dz",
|
||
"edep",
|
||
"step_length",
|
||
]
|
||
|
||
|
||
def _write_world_frame_parquet(
|
||
path, n=200, rng=None, rollout=True, termination_reasons=None
|
||
):
|
||
"""A `giant rollout` output or truth-schema steps file, in raw world-frame units.
|
||
|
||
`rollout=True` adds `track_id`/`termination_reason`/`n_sec_pred` (rollout-
|
||
only columns `RolloutVsTruth` reads) and tags `ROLLOUT_COORD_VALUE`
|
||
metadata; `rollout=False` is a plain truth-schema file (no metadata tag,
|
||
matching a real held-out/val parquet). `termination_reasons`, if given,
|
||
overrides the (rollout-only) `termination_reason` column — pass a fixed
|
||
array to test synthetic-row dropping deterministically.
|
||
"""
|
||
rng = rng or np.random.default_rng(0)
|
||
event_id = rng.integers(0, 5, n)
|
||
pre_pos = rng.standard_normal((n, 3)).astype(np.float32) * 5
|
||
pre_dir = _unit_vectors(rng, n)
|
||
pre_E = rng.uniform(1.0, 100.0, n).astype(np.float32)
|
||
post_dir = _unit_vectors(rng, n)
|
||
step_length = rng.uniform(0.1, 5.0, n).astype(np.float32)
|
||
travel_dir_world = _unit_vectors(rng, n)
|
||
post_pos = pre_pos + step_length[:, None] * travel_dir_world
|
||
post_E = (pre_E * rng.uniform(0.3, 0.99, n)).astype(np.float32)
|
||
edep = rng.uniform(0, 1, n).astype(np.float32)
|
||
|
||
columns = {
|
||
"event_id": event_id,
|
||
"pdg": rng.choice([11, -11, 22], n),
|
||
"material": rng.choice(["W", "Pb"], n),
|
||
"layer_id": rng.integers(0, 10, n).astype(np.int32),
|
||
"pre_x": pre_pos[:, 0],
|
||
"pre_y": pre_pos[:, 1],
|
||
"pre_z": pre_pos[:, 2],
|
||
"pre_E": pre_E,
|
||
"pre_dx": pre_dir[:, 0],
|
||
"pre_dy": pre_dir[:, 1],
|
||
"pre_dz": pre_dir[:, 2],
|
||
"post_x": post_pos[:, 0],
|
||
"post_y": post_pos[:, 1],
|
||
"post_z": post_pos[:, 2],
|
||
"post_E": post_E,
|
||
"post_dx": post_dir[:, 0],
|
||
"post_dy": post_dir[:, 1],
|
||
"post_dz": post_dir[:, 2],
|
||
"edep": edep,
|
||
"step_length": step_length,
|
||
}
|
||
metadata = None
|
||
if rollout:
|
||
columns["track_id"] = rng.integers(0, 3, n).astype(np.int64)
|
||
columns["termination_reason"] = (
|
||
termination_reasons
|
||
if termination_reasons is not None
|
||
else rng.choice(["", TERM_NATURAL_END], n)
|
||
)
|
||
columns["n_sec_pred"] = rng.integers(0, 3, n).astype(np.int32)
|
||
metadata = {PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE}
|
||
|
||
table = pa.table(columns)
|
||
if metadata is not None:
|
||
table = table.replace_schema_metadata(metadata)
|
||
pq.write_table(table, path)
|
||
return columns
|
||
|
||
|
||
def _raw_from_world_frame(columns):
|
||
"""(N, 9) `RAW_TARGET_NAMES` oracle for a world-frame steps dict, via numpy.
|
||
|
||
Independent re-derivation of `_world_frame_local_exprs`'s forward Rodrigues
|
||
rotation, using `giant.data.transforms.local_frame_rotation`/
|
||
`travel_direction` directly rather than the module's polars expressions.
|
||
"""
|
||
pre_pos = np.column_stack(
|
||
[columns["pre_x"], columns["pre_y"], columns["pre_z"]]
|
||
).astype(np.float32)
|
||
pre_dir = np.column_stack(
|
||
[columns["pre_dx"], columns["pre_dy"], columns["pre_dz"]]
|
||
).astype(np.float32)
|
||
post_pos = np.column_stack(
|
||
[columns["post_x"], columns["post_y"], columns["post_z"]]
|
||
).astype(np.float32)
|
||
post_dir = np.column_stack(
|
||
[columns["post_dx"], columns["post_dy"], columns["post_dz"]]
|
||
).astype(np.float32)
|
||
pre_E = np.asarray(columns["pre_E"], dtype=np.float32)
|
||
post_E = np.asarray(columns["post_E"], dtype=np.float32)
|
||
|
||
post_dir_local = local_frame_rotation(pre_dir, post_dir)
|
||
travel_dir_local = local_frame_rotation(
|
||
pre_dir, travel_direction(pre_pos, post_pos)
|
||
)
|
||
return np.column_stack(
|
||
[
|
||
np.asarray(columns["step_length"], dtype=np.float32),
|
||
pre_E - post_E,
|
||
np.asarray(columns["edep"], dtype=np.float32),
|
||
post_dir_local,
|
||
travel_dir_local,
|
||
]
|
||
).astype(np.float32)
|
||
|
||
|
||
def test_rollout_vs_truth_decodes_raw_targets_correctly(tmp_path):
|
||
rng = np.random.default_rng(1)
|
||
rollout_path = tmp_path / "rollout.parquet"
|
||
truth_path = tmp_path / "truth.parquet"
|
||
rollout_cols = _write_world_frame_parquet(
|
||
rollout_path, n=150, rng=rng, rollout=True
|
||
)
|
||
truth_cols = _write_world_frame_parquet(truth_path, n=200, rng=rng, rollout=False)
|
||
|
||
expected_gen = _raw_from_world_frame(rollout_cols)
|
||
expected_real = _raw_from_world_frame(truth_cols)
|
||
|
||
source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
|
||
table = marginal_table_pl(source).to_pandas().set_index("dim")
|
||
|
||
assert (table["n"] == len(truth_cols["pdg"])).all()
|
||
assert (table["n_gen"] == len(rollout_cols["pdg"])).all()
|
||
for j, name in enumerate(RAW_TARGET_NAMES):
|
||
np.testing.assert_allclose(
|
||
table.loc[name, "real_mean"], expected_real[:, j].mean(), atol=1e-3
|
||
)
|
||
np.testing.assert_allclose(
|
||
table.loc[name, "gen_mean"], expected_gen[:, j].mean(), atol=1e-3
|
||
)
|
||
|
||
|
||
def test_rollout_vs_truth_allows_unpaired_lengths(tmp_path):
|
||
rng = np.random.default_rng(2)
|
||
rollout_path = tmp_path / "rollout.parquet"
|
||
truth_path = tmp_path / "truth.parquet"
|
||
_write_world_frame_parquet(rollout_path, n=30, rng=rng, rollout=True)
|
||
_write_world_frame_parquet(truth_path, n=500, rng=rng, rollout=False)
|
||
|
||
source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
|
||
table = marginal_table_pl(source)
|
||
assert (table["n"] == 500).all()
|
||
assert (table["n_gen"] == 30).all()
|
||
|
||
|
||
def test_rollout_vs_truth_drops_synthetic_termination_rows(tmp_path):
|
||
rng = np.random.default_rng(3)
|
||
rollout_path = tmp_path / "rollout.parquet"
|
||
truth_path = tmp_path / "truth.parquet"
|
||
n = 100
|
||
# Half the rollout rows are synthetic bookkeeping (dropped), half are real.
|
||
reasons = np.where(np.arange(n) % 2 == 0, TERM_ENERGY_CUTOFF, "")
|
||
rollout_cols = _write_world_frame_parquet(
|
||
rollout_path, n=n, rng=rng, rollout=True, termination_reasons=reasons
|
||
)
|
||
_write_world_frame_parquet(truth_path, n=50, rng=rng, rollout=False)
|
||
|
||
kept = reasons != TERM_ENERGY_CUTOFF
|
||
expected_gen = _raw_from_world_frame(
|
||
{k: np.asarray(v)[kept] for k, v in rollout_cols.items()}
|
||
)
|
||
|
||
source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
|
||
table = marginal_table_pl(source)
|
||
assert (table["n_gen"] == kept.sum()).all()
|
||
step_length_row = table.filter(pl.col("dim") == "step_length")
|
||
np.testing.assert_allclose(
|
||
step_length_row["gen_mean"][0], expected_gen[:, 0].mean(), atol=1e-3
|
||
)
|
||
|
||
|
||
def test_rollout_vs_truth_rejects_wrong_coord_metadata(tmp_path):
|
||
rollout_path = tmp_path / "rollout.parquet"
|
||
truth_path = tmp_path / "truth.parquet"
|
||
_write_world_frame_parquet(rollout_path, rollout=True)
|
||
_write_world_frame_parquet(truth_path, rollout=False)
|
||
# Overwrite with a mismatched coord tag (predict-schema "local").
|
||
table = pq.read_table(rollout_path).replace_schema_metadata(
|
||
{PREDICT_COORD_METADATA_KEY: "local"}
|
||
)
|
||
pq.write_table(table, rollout_path)
|
||
|
||
source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
|
||
with pytest.raises(ValueError, match="not a rollout file"):
|
||
marginal_table_pl(source)
|
||
|
||
|
||
@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
|
||
def test_rollout_vs_truth_downstream_plots_run_without_error(tmp_path, group_by):
|
||
rollout_path = tmp_path / "rollout.parquet"
|
||
truth_path = tmp_path / "truth.parquet"
|
||
_write_world_frame_parquet(rollout_path, n=150, rollout=True)
|
||
_write_world_frame_parquet(truth_path, n=150, rollout=False)
|
||
source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
|
||
|
||
assert plot_marginals(source, group_by=group_by) is not None
|
||
assert plot_kl_bars_pl(source, group_by=group_by) is not None
|
||
|
||
|
||
def test_rollout_vs_truth_joint_and_constraint_checks_run(tmp_path):
|
||
rollout_path = tmp_path / "rollout.parquet"
|
||
truth_path = tmp_path / "truth.parquet"
|
||
_write_world_frame_parquet(rollout_path, n=150, rollout=True)
|
||
_write_world_frame_parquet(truth_path, n=150, rollout=False)
|
||
source = RolloutVsTruth(rollout=rollout_path, truth=truth_path)
|
||
|
||
assert plot_correlation_matrices(source) is not None
|
||
assert plot_pairwise(source, n_sample=50) is not None
|
||
assert plot_direction_alignment(source) is not None
|
||
assert plot_constraint_violations(source) is not None
|
||
report = constraint_report_pl(source).to_pandas()
|
||
assert set(report["check"]) == {
|
||
"post_dir unit norm",
|
||
"travel_dir unit norm",
|
||
"step_length >= 0",
|
||
"delta_e >= 0",
|
||
"edep >= 0",
|
||
}
|
||
|
||
|
||
def _write_truth_shower(path, event_id, n_steps_per_event, rng):
|
||
"""One-track-per-event world-frame shower: highest-`pre_E` row fixes the axis."""
|
||
rows_per_event = n_steps_per_event
|
||
rows = []
|
||
for i, eid in enumerate(event_id):
|
||
n = rows_per_event
|
||
pre_pos = rng.normal(size=(n, 3)).astype(np.float32) * 5
|
||
pre_E = rng.uniform(1.0, 50.0, n).astype(np.float32)
|
||
pre_E[0] = 1000.0 # entry step
|
||
axis = _unit_vectors(rng, 1)[0]
|
||
pre_dir = np.tile(axis, (n, 1)).astype(np.float32)
|
||
step_length = rng.uniform(0.1, 5.0, n).astype(np.float32)
|
||
post_pos = pre_pos + step_length[:, None] * axis
|
||
edep = rng.uniform(0, 1, n).astype(np.float32)
|
||
for j in range(n):
|
||
rows.append(
|
||
{
|
||
"event_id": eid,
|
||
"pdg": 11,
|
||
"material": "Pb",
|
||
"layer_id": 0,
|
||
"pre_x": pre_pos[j, 0],
|
||
"pre_y": pre_pos[j, 1],
|
||
"pre_z": pre_pos[j, 2],
|
||
"pre_E": pre_E[j],
|
||
"pre_dx": pre_dir[j, 0],
|
||
"pre_dy": pre_dir[j, 1],
|
||
"pre_dz": pre_dir[j, 2],
|
||
"post_x": post_pos[j, 0],
|
||
"post_y": post_pos[j, 1],
|
||
"post_z": post_pos[j, 2],
|
||
"edep": edep[j],
|
||
"step_length": step_length[j],
|
||
"track_id": i,
|
||
"termination_reason": TERM_NATURAL_END if j == n - 1 else "",
|
||
"n_sec_pred": 0,
|
||
}
|
||
)
|
||
return rows
|
||
|
||
|
||
def test_compute_rollout_vs_truth_observables_pl_matches_manual_reconstruction(
|
||
tmp_path,
|
||
):
|
||
rng = np.random.default_rng(4)
|
||
rollout_rows = _write_truth_shower(
|
||
tmp_path / "_r", event_id=[0, 1, 2], n_steps_per_event=15, rng=rng
|
||
)
|
||
truth_rows = _write_truth_shower(
|
||
tmp_path / "_t", event_id=[0, 1, 2], n_steps_per_event=20, rng=rng
|
||
)
|
||
|
||
rollout_path = tmp_path / "rollout.parquet"
|
||
truth_path = tmp_path / "truth.parquet"
|
||
pq.write_table(
|
||
pa.table(
|
||
{k: [r[k] for r in rollout_rows] for k in rollout_rows[0]}
|
||
).replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE}),
|
||
rollout_path,
|
||
)
|
||
truth_cols = [
|
||
"event_id",
|
||
"pdg",
|
||
"material",
|
||
"layer_id",
|
||
"pre_x",
|
||
"pre_y",
|
||
"pre_z",
|
||
"pre_E",
|
||
"pre_dx",
|
||
"pre_dy",
|
||
"pre_dz",
|
||
"post_x",
|
||
"post_y",
|
||
"post_z",
|
||
"edep",
|
||
"step_length",
|
||
]
|
||
pq.write_table(
|
||
pa.table({k: [r[k] for r in truth_rows] for k in truth_cols}), truth_path
|
||
)
|
||
|
||
obs = compute_rollout_vs_truth_observables_pl(
|
||
rollout_path, truth_path, depth_bins=5, transverse_bins=5
|
||
)
|
||
|
||
rollout_df = pd.DataFrame(rollout_rows)
|
||
truth_df = pd.DataFrame(truth_rows)
|
||
expected_gen_total = rollout_df.groupby("event_id")["edep"].sum().sort_index()
|
||
expected_real_total = truth_df.groupby("event_id")["edep"].sum().sort_index()
|
||
|
||
gen_table = obs.gen_table.sort("event_id")
|
||
real_table = obs.real_table.sort("event_id")
|
||
np.testing.assert_allclose(
|
||
gen_table["total_edep"].to_numpy(), expected_gen_total.to_numpy(), rtol=1e-4
|
||
)
|
||
np.testing.assert_allclose(
|
||
real_table["total_edep"].to_numpy(), expected_real_total.to_numpy(), rtol=1e-4
|
||
)
|
||
assert (gen_table["n_tracks"].to_numpy() == 1).all()
|
||
assert (gen_table["leaked_E"].to_numpy() == 0).all()
|
||
|
||
|
||
def test_compute_rollout_vs_truth_observables_pl_profile_shapes(tmp_path):
|
||
rng = np.random.default_rng(5)
|
||
rollout_rows = _write_truth_shower(
|
||
tmp_path / "_r", event_id=[0, 1], n_steps_per_event=10, rng=rng
|
||
)
|
||
truth_rows = _write_truth_shower(
|
||
tmp_path / "_t", event_id=[0, 1, 2], n_steps_per_event=10, rng=rng
|
||
)
|
||
rollout_path = tmp_path / "rollout.parquet"
|
||
truth_path = tmp_path / "truth.parquet"
|
||
pq.write_table(
|
||
pa.table(
|
||
{k: [r[k] for r in rollout_rows] for k in rollout_rows[0]}
|
||
).replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE}),
|
||
rollout_path,
|
||
)
|
||
truth_cols = [
|
||
"event_id",
|
||
"pdg",
|
||
"material",
|
||
"layer_id",
|
||
"pre_x",
|
||
"pre_y",
|
||
"pre_z",
|
||
"pre_E",
|
||
"pre_dx",
|
||
"pre_dy",
|
||
"pre_dz",
|
||
"post_x",
|
||
"post_y",
|
||
"post_z",
|
||
"edep",
|
||
"step_length",
|
||
]
|
||
pq.write_table(
|
||
pa.table({k: [r[k] for r in truth_rows] for k in truth_cols}), truth_path
|
||
)
|
||
|
||
obs = compute_rollout_vs_truth_observables_pl(
|
||
rollout_path, truth_path, depth_bins=6, transverse_bins=4
|
||
)
|
||
assert obs.depth_edges.shape == (7,)
|
||
assert obs.transverse_edges.shape == (5,)
|
||
assert len(obs.gen_table) == 2
|
||
assert len(obs.real_table) == 3
|
||
|
||
assert plot_total_energy(obs) is not None
|
||
assert plot_total_length(obs) is not None
|
||
assert plot_mean_energy_per_step(obs) is not None
|
||
assert plot_mean_length_per_step(obs) is not None
|
||
assert plot_longitudinal_profile(obs) is not None
|
||
assert plot_transverse_profile(obs) is not None
|
||
assert plot_shower_max_depth(obs) is not None
|
||
|
||
|
||
def test_compute_rollout_vs_truth_observables_pl_rejects_wrong_coord_metadata(
|
||
tmp_path,
|
||
):
|
||
rollout_path = tmp_path / "rollout.parquet"
|
||
truth_path = tmp_path / "truth.parquet"
|
||
_write_world_frame_parquet(rollout_path, rollout=True)
|
||
_write_world_frame_parquet(truth_path, rollout=False)
|
||
table = pq.read_table(rollout_path).replace_schema_metadata(
|
||
{PREDICT_COORD_METADATA_KEY: "local"}
|
||
)
|
||
pq.write_table(table, rollout_path)
|
||
|
||
with pytest.raises(ValueError, match="not a rollout file"):
|
||
compute_rollout_vs_truth_observables_pl(rollout_path, truth_path)
|