lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
6744d05948bc0e3a58b03a3b366c2caa8419fbaf
giant/tests/test_analysis.py
T
lars 55c676fb9b Rewrite analysis module as a lean, fully-streaming pipeline 
Replace giant/analysis.py's dual numpy-SampleCollection + polars paths with a
single polars-streaming implementation that produces the validation notebook's
plots directly from a `giant predict --coord local` parquet, sized for files
larger than RAM.

- Drop the numpy SampleCollection path (load_predicted_local, marginal_table,
  correlation_matrices, direction_alignment, constraint_report, plot_kl_bars)
  and the rollout observables; the 5 remaining plotters now take a parquet
  path / LazyFrame and stream internally.
- Rewrite compute_event_observables_pl to aggregate in parallel streaming
  polars (post-pos reconstruction as expressions) instead of a serial
  pyarrow-batch + numpy loop, fixing a pre-existing OOM (holistic median +
  323M-row join in the bin-edge sizing). Medians are approximated from a
  streaming log-bin histogram with within-bin interpolation.
- Keep every full-file scan narrow (few columns): on a file larger than RAM,
  peak mmap memory, not scan count, is the binding constraint. Marginals run
  one dim at a time (~15GB peak) rather than a combined all-dims pass (OOM).
- Update analysis/validation.ipynb to the path-based API; delete the
  analysis/export_*.py and compare_ode_steps_*.py one-off scripts.
- Rewrite tests/test_analysis.py around parquet fixtures with an inline numpy
  oracle; add correlation/streaming-plotter and approx-median coverage.

Verified end-to-end on the 32GB predict file: full notebook completes at
~25GB peak (no OOM); event rollup runs at ~13 cores.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-10 14:47:46 +02:00

747 lines
27 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
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,
compute_event_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_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,
)
from giant.data.transforms import (
energy_simplex_decode,
inv_log_transform,
log_transform,
reconstruct_post_pos,
)
# ---------------------------------------------------------------------------
# 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)
table = obs.event_table.sort("event_id")
for i, eid in enumerate(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(
table["real_total_edep"][i], real_total_edep, rtol=1e-4
)
np.testing.assert_allclose(
table["gen_total_edep"][i], gen_total_edep, rtol=1e-4
)
np.testing.assert_allclose(
table["real_total_length"][i], real_total_length, rtol=1e-4
)
np.testing.assert_allclose(
table["gen_total_length"][i], gen_total_length, rtol=1e-4
)
np.testing.assert_allclose(
table["real_centroid_depth"][i], real_centroid, rtol=1e-3, atol=1e-4
)
np.testing.assert_allclose(
table["gen_centroid_depth"][i], gen_centroid, rtol=1e-3, atol=1e-4
)
np.testing.assert_allclose(
table["real_transverse_rms"][i], real_rms, rtol=1e-3, atol=1e-4
)
np.testing.assert_allclose(
table["gen_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.event_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.event_table.sort("event_id")["real_total_edep"].to_numpy(),
obs_from_path.event_table.sort("event_id")["real_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.event_table["real_median_edep"][0]
approx_length = obs.event_table["real_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
Reference in New Issue View Git Blame Copy Permalink