Add fast slab lookup for the GeometryOracle, replacing knn as the default

miniCaloSim's detector is a stack of planar layer slabs along one axis, so
material/layer_id are a pure function of depth. The new "slab" method
exploits this with an exact O(log #segments) binary search over
depth-axis segment boundaries, instead of a nearest-neighbour search over
hundreds of thousands of reference points — much cheaper per call, which
matters since the oracle is queried on every autoregressive rollout step.
"knn"/"svm" remain as fallbacks for non-slab geometries.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

tests/test_geometry.py

@@ -26,7 +26,9 @@ def _build(subsample=30000, method="knn", escape_factor=5.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,
+            [Path("x")],
+            method=method,
+            subsample=subsample,
             escape_factor=escape_factor,
         )
 
@@ -84,3 +86,93 @@ def test_svm_method_has_escape_tree():
     assert orc._ref_tree is not None
     _, _, escaped = orc.query(np.array([[1e5, 0.0, 0.0]]))
     assert escaped[0]
+
+
+def _layer_batch(n, rng):
+    """Two 100mm slabs along z (with an air gap between/around them), bounded
+    to a 100x100mm transverse footprint — miniCaloSim's actual layer-stack
+    shape."""
+    z = rng.uniform(-20.0, 220.0, n).astype(np.float32)
+    x = rng.uniform(-50.0, 50.0, n).astype(np.float32)
+    y = rng.uniform(-50.0, 50.0, n).astype(np.float32)
+    material = np.full(n, "G4_AIR", dtype=object)
+    layer_id = np.full(n, -1, dtype=np.int64)
+    in_l0 = (z >= 0.0) & (z < 100.0)
+    in_l1 = (z >= 110.0) & (z < 210.0)
+    material[in_l0] = "G4_PbWO4"
+    layer_id[in_l0] = 0
+    material[in_l1] = "G4_W"
+    layer_id[in_l1] = 1
+    pos = np.stack([x, y, z], axis=1)
+    return pos, material, layer_id
+
+
+def _build_slab(subsample=60000, n_bins=500, escape_factor=5.0):
+    rng = np.random.default_rng(0)
+    batches = [_layer_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="slab",
+            subsample=subsample,
+            escape_factor=escape_factor,
+            depth_axis=2,
+            n_bins=n_bins,
+        )
+
+
+def test_slab_is_default_method():
+    rng = np.random.default_rng(0)
+    batches = [_layer_batch(20000, rng)]
+    with patch.object(g, "_iter_point_batches", lambda p: iter(batches)):
+        orc = g.build_geometry_oracle([Path("x")], subsample=20000)
+    assert orc.metadata["method"] == "slab"
+    assert orc._slab is not None
+
+
+def test_slab_classes_discovered():
+    orc = _build_slab()
+    assert set(orc.classes) == {("G4_PbWO4", 0), ("G4_W", 1), ("G4_AIR", -1)}
+
+
+def test_slab_query_labels_by_depth():
+    orc = _build_slab()
+    pos = np.array(
+        [[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0]]
+    )  # layer 0, gap, layer 1
+    material, layer_id, escaped = orc.query(pos)
+    assert list(material) == ["G4_PbWO4", "G4_AIR", "G4_W"]
+    assert list(layer_id) == [0, -1, 1]
+    assert not escaped.any()
+
+
+def test_slab_escape_beyond_depth_range():
+    orc = _build_slab()
+    pos = np.array([[0.0, 0.0, 50.0], [0.0, 0.0, 1e5]])
+    _, _, escaped = orc.query(pos)
+    assert not escaped[0]
+    assert escaped[1]
+
+
+def test_slab_escape_beyond_transverse_radius():
+    orc = _build_slab()
+    pos = np.array([[0.0, 0.0, 50.0], [1e5, 1e5, 50.0]])
+    _, _, escaped = orc.query(pos)
+    assert not escaped[0]
+    assert escaped[1]
+
+
+def test_slab_save_load_roundtrip(tmp_path):
+    orc = _build_slab()
+    p = tmp_path / "slab_oracle.pkl"
+    orc.save(p)
+    loaded = g.GeometryOracle.load(p)
+
+    pos = np.array(
+        [[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0], [0.0, 0.0, 1e5]]
+    )
+    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._slab is not None

tests/test_rollout.py

@@ -39,7 +39,11 @@ def _oracle():
     mat = np.where(inside, "G4_PbWO4", "G4_AIR").astype(object)
     lay = np.where(inside, 0, -1).astype(np.int64)
     with patch.object(g, "_iter_point_batches", lambda p: iter([(pos, mat, lay)])):
-        return g.build_geometry_oracle([Path("x")], subsample=20000)
+        # Pinned to "knn" explicitly: this test's escape-threshold semantics
+        # (tiny threshold -> escape even at a valid interior point, because no
+        # training point is that close) are KNN-specific, and the fixture's
+        # box geometry isn't a layer stack the "slab" method could fit anyway.
+        return g.build_geometry_oracle([Path("x")], method="knn", subsample=20000)
 
 
 def _seeds(n=6):