lars
26a176aeaa
Closes the loop from single-step prediction into full showers: - giant/geometry.py + `dwarf build-geometry-oracle`: learn position -> (material, layer_id) from data (KNN/SVM) to supply the conditioning the surrogate does not predict; flag detector escape by NN distance. - giant/rollout.py: breadth-first batched frontier that steps all active tracks, spawns secondaries as new tracks, and terminates on energy cutoff, per-track max steps, escape, or natural end. Energy is deposited locally on every stop except escape (leakage), so showers conserve energy exactly. - `giant rollout` CLI: seed from real events (argmax pre_E), load checkpoint, write a world-frame steps parquet + YAML sidecar. - giant/analysis.py: compute_rollout_observables + plot_rollout_* for single-sided longitudinal/transverse/total-energy shower profiles; analysis/export_rollout_observables.py driver. - scikit-learn added as an optional `geometry` extra (lazy-imported). - Tests: tests/test_geometry.py, tests/test_rollout.py. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
"""Tests for the geometry oracle (position -> material/layer_id + escape)."""
|
|
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from giant import geometry as g
|
|
|
|
# scikit-learn is an optional extra; skip the whole module if it's missing.
|
|
pytest.importorskip("sklearn")
|
|
|
|
|
|
def _box_batch(n, rng):
|
|
"""A labelled point cloud: inside a 100mm box -> PbWO4/0, else AIR/-1."""
|
|
pos = rng.uniform(-200, 200, (n, 3)).astype(np.float32)
|
|
inside = (np.abs(pos) < 100).all(axis=1)
|
|
mat = np.where(inside, "G4_PbWO4", "G4_AIR").astype(object)
|
|
lay = np.where(inside, 0, -1).astype(np.int64)
|
|
return pos, mat, lay
|
|
|
|
|
|
def _build(subsample=30000, method="knn", escape_factor=5.0):
|
|
rng = np.random.default_rng(0)
|
|
batches = [_box_batch(20000, rng) for _ in range(3)]
|
|
with patch.object(g, "_iter_point_batches", lambda p: iter(batches)):
|
|
return g.build_geometry_oracle(
|
|
[Path("x")], method=method, subsample=subsample,
|
|
escape_factor=escape_factor,
|
|
)
|
|
|
|
|
|
def test_classes_discovered():
|
|
orc = _build()
|
|
assert set(orc.classes) == {("G4_PbWO4", 0), ("G4_AIR", -1)}
|
|
|
|
|
|
def test_query_labels_inside_and_outside():
|
|
orc = _build()
|
|
pos = np.array([[0.0, 0.0, 0.0], [150.0, 150.0, 150.0]])
|
|
material, layer_id, _ = orc.query(pos)
|
|
assert material[0] == "G4_PbWO4" and layer_id[0] == 0
|
|
assert material[1] == "G4_AIR" and layer_id[1] == -1
|
|
|
|
|
|
def test_escape_flag_fires_far_from_data():
|
|
orc = _build()
|
|
pos = np.array([[0.0, 0.0, 0.0], [1e5, 0.0, 0.0]])
|
|
_, _, escaped = orc.query(pos)
|
|
assert not escaped[0]
|
|
assert escaped[1]
|
|
|
|
|
|
def test_query_empty():
|
|
orc = _build()
|
|
material, layer_id, escaped = orc.query(np.empty((0, 3)))
|
|
assert len(material) == len(layer_id) == len(escaped) == 0
|
|
|
|
|
|
def test_query_bad_shape_raises():
|
|
orc = _build()
|
|
with pytest.raises(ValueError):
|
|
orc.query(np.zeros((4, 2)))
|
|
|
|
|
|
def test_save_load_roundtrip(tmp_path):
|
|
orc = _build()
|
|
p = tmp_path / "oracle.pkl"
|
|
orc.save(p)
|
|
loaded = g.GeometryOracle.load(p)
|
|
|
|
pos = np.array([[0.0, 0.0, 0.0], [150.0, 150.0, 150.0], [1e5, 0.0, 0.0]])
|
|
m0, l0, e0 = orc.query(pos)
|
|
m1, l1, e1 = loaded.query(pos)
|
|
assert (m0 == m1).all() and (l0 == l1).all() and (e0 == e1).all()
|
|
assert loaded.escape_threshold == orc.escape_threshold
|
|
assert loaded.classes == orc.classes
|
|
|
|
|
|
def test_svm_method_has_escape_tree():
|
|
orc = _build(subsample=4000, method="svm")
|
|
# SVM cannot answer NN-distance, so a reference tree backs the escape test.
|
|
assert orc._ref_tree is not None
|
|
_, _, escaped = orc.query(np.array([[1e5, 0.0, 0.0]]))
|
|
assert escaped[0]
|