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>

scripts/dwarf.py

@@ -387,6 +387,7 @@ def make_root(
 
 
 class OracleMethod(str, Enum):
+    slab = "slab"
     knn = "knn"
     svm = "svm"
 
@@ -396,13 +397,18 @@ def build_geometry_oracle(
     data: Annotated[
         Path, typer.Argument(help="Steps parquet file or directory of steps files")
     ],
-    out: Annotated[
-        Path, typer.Option("--out", "-o", help="Output oracle .pkl path")
-    ],
+    out: Annotated[Path, typer.Option("--out", "-o", help="Output oracle .pkl path")],
     method: Annotated[
         OracleMethod,
-        typer.Option("--method", help="Classifier: knn (default) or svm"),
-    ] = OracleMethod.knn,
+        typer.Option(
+            "--method",
+            help=(
+                "Lookup strategy: slab (default; exact O(log #segments) fast "
+                "path for the layered-slab detector geometry), knn, or svm "
+                "(generic fallbacks for non-slab geometries)"
+            ),
+        ),
+    ] = OracleMethod.slab,
     k: Annotated[
         int, typer.Option("--k", help="Neighbours for the knn classifier")
     ] = 1,
@@ -414,10 +420,24 @@ def build_geometry_oracle(
         float,
         typer.Option(
             "--escape-factor",
-            help="escape_threshold = this x median NN spacing of reference points",
+            help="escape_threshold = this x median spacing of reference points",
         ),
     ] = 5.0,
     seed: Annotated[int, typer.Option("--seed", help="Sampling seed")] = 0,
+    depth_axis: Annotated[
+        int,
+        typer.Option(
+            "--depth-axis",
+            help="0/1/2 -> x/y/z axis the layers stack along (method=slab only)",
+        ),
+    ] = 2,
+    n_bins: Annotated[
+        int,
+        typer.Option(
+            "--n-bins",
+            help="Depth-axis resolution, finer than the thinnest layer (method=slab only)",
+        ),
+    ] = 2000,
 ) -> None:
     """Fit a position -> (material, layer_id) oracle for `giant rollout`."""
     run_build_geometry_oracle(
@@ -428,6 +448,8 @@ def build_geometry_oracle(
         subsample=subsample,
         escape_factor=escape_factor,
         seed=seed,
+        depth_axis=depth_axis,
+        n_bins=n_bins,
     )
 
 

scripts/geometry_oracle.py

@@ -16,11 +16,13 @@ from giant.geometry import build_geometry_oracle
 def run_build_geometry_oracle(
     data: Path,
     out: Path,
-    method: str = "knn",
+    method: str = "slab",
     k: int = 1,
     subsample: int = 500_000,
     escape_factor: float = 5.0,
     seed: int = 0,
+    depth_axis: int = 2,
+    n_bins: int = 2000,
 ) -> None:
     files = find_parquet_files(data)
     print(f"found {len(files)} parquet file(s); sampling up to {subsample:,} points")
@@ -32,17 +34,33 @@ def run_build_geometry_oracle(
         subsample=subsample,
         escape_factor=escape_factor,
         seed=seed,
+        depth_axis=depth_axis,
+        n_bins=n_bins,
     )
 
-    print(f"method: {method}  reference points: {oracle.metadata['n_reference_points']:,}")
+    print(
+        f"method: {method}  reference points: {oracle.metadata['n_reference_points']:,}"
+    )
     print("classes (material, layer_id):")
     for material, layer_id in oracle.classes:
         print(f"  {material:<12} layer_id={layer_id}")
-    print(
-        f"median NN spacing: {oracle.metadata['median_nn_dist']:.3f}  "
-        f"escape_threshold: {oracle.escape_threshold:.3f} "
-        f"(= {escape_factor}x spacing)"
-    )
+
+    if method == "slab":
+        z_lo, z_hi = oracle.metadata["z_range"]
+        print(
+            f"depth axis: {'xyz'[depth_axis]}  segments: {oracle.metadata['n_segments']}  "
+            f"z range: [{z_lo:.3f}, {z_hi:.3f}]  radius_max: {oracle.metadata['radius_max']:.3f}"
+        )
+        print(
+            f"median depth spacing: {oracle.metadata['median_z_spacing']:.3f}  "
+            f"escape_threshold: {oracle.escape_threshold:.3f} (= {escape_factor}x spacing)"
+        )
+    else:
+        print(
+            f"median NN spacing: {oracle.metadata['median_nn_dist']:.3f}  "
+            f"escape_threshold: {oracle.escape_threshold:.3f} "
+            f"(= {escape_factor}x spacing)"
+        )
     if oracle.escape_threshold <= 0.0:
         print(
             "warning: escape_threshold is 0 (reference points are coincident) — "