Files
giant/scripts/geometry_oracle.py
lars 3faa272562 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>
2026-07-08 11:42:12 +02:00

75 lines
2.4 KiB
Python

"""Build a position -> (material, layer_id) geometry oracle from steps parquet.
Backs the `dwarf build-geometry-oracle` subcommand. The oracle is consumed by
`giant rollout` to supply the material/layer conditioning at each step, since the
surrogate does not predict them.
"""
from __future__ import annotations
from pathlib import Path
from giant.data.loader import find_parquet_files
from giant.geometry import build_geometry_oracle
def run_build_geometry_oracle(
data: Path,
out: Path,
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")
oracle = build_geometry_oracle(
files,
method=method,
k=k,
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("classes (material, layer_id):")
for material, layer_id in oracle.classes:
print(f" {material:<12} layer_id={layer_id}")
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) — "
"every rollout query would be flagged as escaped. Pass "
"`giant rollout --escape-threshold <mm>` to override, or use data "
"with distinct step positions."
)
out.parent.mkdir(parents=True, exist_ok=True)
oracle.save(out)
print(f"wrote oracle -> {out}")