lars
3faa272562
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>
487 lines
18 KiB
Python
487 lines
18 KiB
Python
"""Geometry oracle: learn a position -> (material, layer_id) map from step data.
|
|
|
|
The surrogate conditions on `material` and `layer_id`, but does not predict
|
|
them — during a shower rollout they must be looked up from the new position.
|
|
There is no in-repo detector geometry (it lives in external miniCaloSim), so we
|
|
learn it from positions sampled from a real steps dataset.
|
|
|
|
miniCaloSim's detector is a stack of planar layer slabs along one axis (see
|
|
`physics/detector-design/minicalosim-geometry.md`), so `material`/`layer_id`
|
|
are a pure function of depth. The default ("slab") method exploits this: fit a
|
|
1D lookup table of depth-axis segment boundaries and do an exact O(log
|
|
#segments) binary search per query, with escape decided by depth/transverse
|
|
bounds — far cheaper per call than a nearest-neighbour search over hundreds of
|
|
thousands of reference points, which matters because this oracle is queried on
|
|
every autoregressive step of a shower rollout. "knn"/"svm" remain as generic
|
|
fallbacks (a classifier over 3D positions, escape decided by distance to the
|
|
nearest reference point) for geometries that aren't simple slab stacks.
|
|
|
|
scikit-learn / joblib are an optional dependency (the `geometry` extra) needed
|
|
by "knn"/"svm" and by `save`/`load`; they are imported lazily so the core
|
|
install stays lean.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any, Iterable
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow.parquet as pq
|
|
|
|
_INSTALL_HINT = (
|
|
"the geometry oracle needs scikit-learn — install it with "
|
|
"`uv sync --extra cpu --extra geometry`"
|
|
)
|
|
|
|
|
|
def _require_sklearn():
|
|
try:
|
|
import joblib # noqa: F401
|
|
from sklearn.neighbors import KNeighborsClassifier # noqa: F401
|
|
from sklearn.svm import SVC # noqa: F401
|
|
except ImportError as exc: # pragma: no cover - exercised only without extra
|
|
raise ImportError(_INSTALL_HINT) from exc
|
|
|
|
|
|
@dataclass
|
|
class _SlabLookup:
|
|
"""Fast path for a detector that is a stack of planar layer slabs along one
|
|
axis (miniCaloSim's actual geometry — see `giant/geometry.py` module docstring
|
|
and `physics/detector-design/minicalosim-geometry.md`). `material`/`layer_id`
|
|
are then a pure function of depth, found by binary search over `z_edges`
|
|
instead of a nearest-neighbour search over the whole reference point cloud —
|
|
O(log(#segments)) instead of O(log(#reference points)), with a far smaller
|
|
constant factor, and exact rather than approximate.
|
|
"""
|
|
|
|
axis: int # which of the 3 position components is the depth axis
|
|
z_edges: np.ndarray # (n_segments + 1,) sorted boundaries between segments
|
|
materials: np.ndarray # (n_segments,) object, material of each segment
|
|
layer_ids: np.ndarray # (n_segments,) int64, layer_id of each segment
|
|
radius_max: float # largest transverse radius seen in training data
|
|
|
|
def query(
|
|
self, pos: np.ndarray, margin: float
|
|
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
other = [i for i in range(3) if i != self.axis]
|
|
z = pos[:, self.axis]
|
|
radius = np.sqrt(pos[:, other[0]] ** 2 + pos[:, other[1]] ** 2)
|
|
|
|
idx = np.searchsorted(self.z_edges, z, side="right") - 1
|
|
idx = np.clip(idx, 0, len(self.materials) - 1)
|
|
material = self.materials[idx]
|
|
layer_id = self.layer_ids[idx]
|
|
|
|
escaped = (
|
|
(z < self.z_edges[0] - margin)
|
|
| (z > self.z_edges[-1] + margin)
|
|
| (radius > self.radius_max + margin)
|
|
)
|
|
return material, layer_id, escaped
|
|
|
|
|
|
@dataclass
|
|
class GeometryOracle:
|
|
"""Maps world-frame position -> (material, layer_id, escaped).
|
|
|
|
Two lookup strategies are supported (`metadata["method"]`):
|
|
- `"slab"`: exact O(log #segments) binary search exploiting the known
|
|
layered-slab detector geometry (see `_SlabLookup`). Fast and preferred.
|
|
- `"knn"` / `"svm"`: a generic sklearn classifier over 3D positions,
|
|
predicting a class index into `classes` (a list of (material, layer_id)
|
|
pairs). Kept as a fallback for geometries that aren't simple slab stacks.
|
|
|
|
`escape_threshold` is a distance in position units (mm). For knn/svm it's
|
|
compared against the nearest training reference point. For slab it's the
|
|
slack allowed beyond the observed depth range / transverse radius before a
|
|
point is flagged `escaped`.
|
|
"""
|
|
|
|
estimator: Any
|
|
classes: list[tuple[str, int]]
|
|
escape_threshold: float
|
|
metadata: dict
|
|
# Only populated for non-neighbour estimators (SVM) to answer the escape
|
|
# distance query; KNeighborsClassifier answers it directly.
|
|
_ref_tree: Any = field(default=None)
|
|
# Only populated when metadata["method"] == "slab".
|
|
_slab: _SlabLookup | None = field(default=None)
|
|
|
|
def query(self, pos: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
"""Return (material (N,) str, layer_id (N,) int, escaped (N,) bool)."""
|
|
pos = np.ascontiguousarray(np.asarray(pos, dtype=np.float64))
|
|
if pos.ndim != 2 or pos.shape[1] != 3:
|
|
raise ValueError(f"pos must be (N, 3), got {pos.shape}")
|
|
if len(pos) == 0:
|
|
return (
|
|
np.empty(0, dtype=object),
|
|
np.empty(0, dtype=np.int64),
|
|
np.empty(0, dtype=bool),
|
|
)
|
|
|
|
if self._slab is not None:
|
|
return self._slab.query(pos, self.escape_threshold)
|
|
|
|
# kneighbors gives the distance to the nearest reference point, which is
|
|
# what the escape test needs; it exists on both KNeighborsClassifier and
|
|
# (via a stored reference tree) our SVM wrapper below.
|
|
dist = self._nearest_distance(pos)
|
|
escaped = dist > self.escape_threshold
|
|
|
|
cls_idx = self.estimator.predict(pos).astype(np.int64)
|
|
material = np.array([self.classes[i][0] for i in cls_idx], dtype=object)
|
|
layer_id = np.array([self.classes[i][1] for i in cls_idx], dtype=np.int64)
|
|
return material, layer_id, escaped
|
|
|
|
def _nearest_distance(self, pos: np.ndarray) -> np.ndarray:
|
|
from sklearn.neighbors import KNeighborsClassifier
|
|
|
|
if isinstance(self.estimator, KNeighborsClassifier):
|
|
dist, _ = self.estimator.kneighbors(pos, n_neighbors=1)
|
|
return dist[:, 0]
|
|
# SVM (or any non-neighbour estimator): use the separately stored
|
|
# NearestNeighbors index purely for the escape distance.
|
|
dist, _ = self._ref_tree.kneighbors(pos, n_neighbors=1)
|
|
return dist[:, 0]
|
|
|
|
def save(self, path: str | Path) -> None:
|
|
_require_sklearn()
|
|
import joblib
|
|
|
|
joblib.dump(
|
|
{
|
|
"estimator": self.estimator,
|
|
"classes": self.classes,
|
|
"escape_threshold": self.escape_threshold,
|
|
"metadata": self.metadata,
|
|
"ref_tree": getattr(self, "_ref_tree", None),
|
|
"slab": getattr(self, "_slab", None),
|
|
},
|
|
path,
|
|
)
|
|
|
|
@classmethod
|
|
def load(cls, path: str | Path) -> "GeometryOracle":
|
|
_require_sklearn()
|
|
import joblib
|
|
|
|
d = joblib.load(path)
|
|
obj = cls(
|
|
estimator=d["estimator"],
|
|
classes=[tuple(c) for c in d["classes"]],
|
|
escape_threshold=float(d["escape_threshold"]),
|
|
metadata=d.get("metadata", {}),
|
|
)
|
|
obj._ref_tree = d.get("ref_tree")
|
|
obj._slab = d.get("slab")
|
|
return obj
|
|
|
|
|
|
_PRE_COLS = ["pre_x", "pre_y", "pre_z"]
|
|
_POST_COLS = ["post_x", "post_y", "post_z"]
|
|
_LABEL_COLS = ["material", "layer_id"]
|
|
|
|
|
|
def _iter_point_batches(path: Path, batch_size: int = 1_000_000):
|
|
"""Yield (pos (M,3), material (M,), layer_id (M,)) from any parquet with a
|
|
position + material + layer_id schema (raw steps *or* predict output).
|
|
|
|
Only the needed columns are read. post_pos points are included when present
|
|
(they share their step's label) so boundary regions are densely sampled.
|
|
"""
|
|
pf = pq.ParquetFile(path)
|
|
have = set(pf.schema_arrow.names)
|
|
missing = [c for c in (*_PRE_COLS, *_LABEL_COLS) if c not in have]
|
|
if missing:
|
|
raise ValueError(
|
|
f"{path} is missing columns {missing} needed to build a geometry "
|
|
"oracle (expected pre_x/y/z, material, layer_id)"
|
|
)
|
|
has_post = all(c in have for c in _POST_COLS)
|
|
cols = [*_PRE_COLS, *_LABEL_COLS] + (_POST_COLS if has_post else [])
|
|
|
|
for batch in pf.iter_batches(batch_size=batch_size, columns=cols):
|
|
d = batch.to_pydict()
|
|
pre = np.array([d[c] for c in _PRE_COLS], dtype=np.float32).T
|
|
mat = np.array([str(m) for m in d["material"]], dtype=object)
|
|
lay = np.asarray(d["layer_id"], dtype=np.int64)
|
|
if has_post:
|
|
post = np.array([d[c] for c in _POST_COLS], dtype=np.float32).T
|
|
yield (
|
|
np.concatenate([pre, post], axis=0),
|
|
np.concatenate([mat, mat], axis=0),
|
|
np.concatenate([lay, lay], axis=0),
|
|
)
|
|
else:
|
|
yield pre, mat, lay
|
|
|
|
|
|
def _collect_points(
|
|
files: Iterable[Path],
|
|
subsample: int,
|
|
seed: int,
|
|
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
"""Stream files, reservoir-sample (pos, material, layer_id) points.
|
|
|
|
Reservoir sampling keeps memory bounded regardless of total file size.
|
|
"""
|
|
rng = np.random.default_rng(seed)
|
|
res_pos = np.empty((subsample, 3), dtype=np.float32)
|
|
res_mat = np.empty(subsample, dtype=object)
|
|
res_lay = np.empty(subsample, dtype=np.int64)
|
|
seen = 0
|
|
|
|
for path in files:
|
|
for pos, mat, lay in _iter_point_batches(path):
|
|
m = len(pos)
|
|
|
|
if seen < subsample:
|
|
take = min(subsample - seen, m)
|
|
res_pos[seen : seen + take] = pos[:take]
|
|
res_mat[seen : seen + take] = mat[:take]
|
|
res_lay[seen : seen + take] = lay[:take]
|
|
seen += take
|
|
if take < m:
|
|
# Reservoir is now full; run the standard replacement rule
|
|
# on the remainder of this chunk.
|
|
_reservoir_replace(
|
|
res_pos,
|
|
res_mat,
|
|
res_lay,
|
|
pos[take:],
|
|
mat[take:],
|
|
lay[take:],
|
|
seen,
|
|
rng,
|
|
)
|
|
seen += m - take
|
|
else:
|
|
_reservoir_replace(res_pos, res_mat, res_lay, pos, mat, lay, seen, rng)
|
|
seen += m
|
|
|
|
n = min(seen, subsample)
|
|
return res_pos[:n], res_mat[:n], res_lay[:n]
|
|
|
|
|
|
def _reservoir_replace(res_pos, res_mat, res_lay, pos, mat, lay, seen, rng) -> None:
|
|
"""Vectorized reservoir replacement for a batch of incoming points."""
|
|
m = len(pos)
|
|
k = res_pos.shape[0]
|
|
# For incoming global index j (seen..seen+m-1), keep with prob k/(j+1),
|
|
# replacing a uniformly-chosen reservoir slot.
|
|
idx = seen + np.arange(m)
|
|
j = rng.integers(0, idx + 1) # j in [0, global_index]
|
|
keep = j < k
|
|
slots = j[keep]
|
|
res_pos[slots] = pos[keep]
|
|
res_mat[slots] = mat[keep]
|
|
res_lay[slots] = lay[keep]
|
|
|
|
|
|
def _fit_slab_lookup(
|
|
pos: np.ndarray,
|
|
mat: np.ndarray,
|
|
lay: np.ndarray,
|
|
axis: int,
|
|
n_bins: int,
|
|
) -> tuple[_SlabLookup, dict]:
|
|
"""Fit a `_SlabLookup` assuming material/layer_id are a function of depth
|
|
(`pos[:, axis]`) alone — true for a stack of planar layer slabs.
|
|
|
|
Bins the depth axis into `n_bins` equal-width bins, takes the majority
|
|
(material, layer_id) label per bin (robust to the handful of points near a
|
|
boundary whose true label is ambiguous at bin resolution), fills any empty
|
|
bins from the nearest populated bin, then run-length-encodes consecutive
|
|
bins sharing a label into segments. Binary search over the segment
|
|
boundaries then answers a query in O(log #segments).
|
|
"""
|
|
other = [i for i in range(3) if i != axis]
|
|
z = pos[:, axis].astype(np.float64)
|
|
radius = np.sqrt(
|
|
pos[:, other[0]].astype(np.float64) ** 2
|
|
+ pos[:, other[1]].astype(np.float64) ** 2
|
|
)
|
|
|
|
z_min, z_max = float(z.min()), float(z.max())
|
|
if z_min == z_max:
|
|
raise ValueError(
|
|
"all points share the same depth-axis coordinate — pick a "
|
|
"different `depth_axis` or use method='knn'/'svm'"
|
|
)
|
|
edges = np.linspace(z_min, z_max, n_bins + 1)
|
|
bin_idx = np.clip(np.searchsorted(edges, z, side="right") - 1, 0, n_bins - 1)
|
|
|
|
counts = (
|
|
pd.DataFrame({"bin": bin_idx, "material": mat, "layer_id": lay})
|
|
.groupby(["bin", "material", "layer_id"])
|
|
.size()
|
|
.to_frame("n")
|
|
.reset_index()
|
|
.sort_values("n", ascending=False)
|
|
.drop_duplicates("bin")
|
|
)
|
|
|
|
bin_material = np.full(n_bins, "", dtype=object)
|
|
bin_layer = np.full(n_bins, -1, dtype=np.int64)
|
|
has_data = np.zeros(n_bins, dtype=bool)
|
|
idx = counts["bin"].to_numpy()
|
|
bin_material[idx] = counts["material"].to_numpy()
|
|
bin_layer[idx] = counts["layer_id"].to_numpy()
|
|
has_data[idx] = True
|
|
|
|
# Forward/backward-fill bins with no samples from the nearest populated one.
|
|
fill_from = np.where(has_data, np.arange(n_bins), -1)
|
|
for b in range(1, n_bins):
|
|
if fill_from[b] == -1:
|
|
fill_from[b] = fill_from[b - 1]
|
|
for b in range(n_bins - 2, -1, -1):
|
|
if fill_from[b] == -1:
|
|
fill_from[b] = fill_from[b + 1]
|
|
bin_material = bin_material[fill_from]
|
|
bin_layer = bin_layer[fill_from]
|
|
|
|
# Run-length-encode consecutive bins sharing a label into segments.
|
|
changed = (
|
|
np.flatnonzero(
|
|
(bin_material[1:] != bin_material[:-1]) | (bin_layer[1:] != bin_layer[:-1])
|
|
)
|
|
+ 1
|
|
)
|
|
seg_starts = np.concatenate([[0], changed])
|
|
z_edges = np.concatenate([edges[seg_starts], edges[-1:]])
|
|
materials = bin_material[seg_starts]
|
|
layer_ids = bin_layer[seg_starts]
|
|
|
|
unique_z = np.unique(z)
|
|
median_spacing = float(np.median(np.diff(unique_z))) if len(unique_z) > 1 else 1.0
|
|
radius_max = float(radius.max())
|
|
|
|
slab = _SlabLookup(
|
|
axis=axis,
|
|
z_edges=z_edges,
|
|
materials=materials,
|
|
layer_ids=layer_ids,
|
|
radius_max=radius_max,
|
|
)
|
|
info = {
|
|
"n_segments": int(len(materials)),
|
|
"z_range": (z_min, z_max),
|
|
"median_z_spacing": median_spacing,
|
|
"radius_max": radius_max,
|
|
}
|
|
return slab, info
|
|
|
|
|
|
def build_geometry_oracle(
|
|
files: list[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,
|
|
) -> GeometryOracle:
|
|
"""Fit a position -> (material, layer_id) classifier from steps files.
|
|
|
|
method: "slab" (default-recommended fast path exploiting the known
|
|
layered-slab detector geometry — see `_SlabLookup`), "knn"
|
|
(KNeighborsClassifier), or "svm" (SVC). "slab" is O(log #segments) per
|
|
query and exact; "knn"/"svm" are generic fallbacks for geometries that
|
|
aren't simple slab stacks, at the cost of a much slower query (a
|
|
nearest-neighbour or kernel evaluation against up to `subsample`
|
|
reference points) and, for "svm", occasional misclassification.
|
|
k: neighbours for the knn classifier (ignored otherwise).
|
|
subsample: max reference points held in memory / used for the fit.
|
|
escape_factor: escape_threshold = escape_factor * median spacing of the
|
|
reference points along the relevant axis/axes, so it scales with the
|
|
sampling density of the data.
|
|
depth_axis: index (0/1/2 -> x/y/z) of the position component that layers
|
|
stack along. Only used by method="slab"; default 2 (z) matches
|
|
miniCaloSim's beam-axis-aligned layer stack.
|
|
n_bins: depth-axis resolution for method="slab" — should be finer than the
|
|
thinnest layer.
|
|
"""
|
|
pos, mat, lay = _collect_points(files, subsample, seed)
|
|
if len(pos) == 0:
|
|
raise ValueError("no points collected — are these steps parquet files?")
|
|
|
|
# Combined (material, layer_id) class label, used for `classes` regardless
|
|
# of method (informational for slab; the actual classifier index for
|
|
# knn/svm).
|
|
pairs = list(zip((str(m) for m in mat), (int(v) for v in lay)))
|
|
classes = sorted(set(pairs))
|
|
|
|
if method == "slab":
|
|
slab, info = _fit_slab_lookup(pos, mat, lay, axis=depth_axis, n_bins=n_bins)
|
|
escape_threshold = escape_factor * info["median_z_spacing"]
|
|
oracle = GeometryOracle(
|
|
estimator=None,
|
|
classes=classes,
|
|
escape_threshold=escape_threshold,
|
|
metadata={
|
|
"method": "slab",
|
|
"depth_axis": depth_axis,
|
|
"n_bins": n_bins,
|
|
"n_reference_points": int(len(pos)),
|
|
"escape_factor": escape_factor,
|
|
"n_files": len(files),
|
|
**info,
|
|
},
|
|
)
|
|
oracle._slab = slab
|
|
return oracle
|
|
|
|
_require_sklearn()
|
|
from sklearn.neighbors import KNeighborsClassifier, NearestNeighbors
|
|
from sklearn.svm import SVC
|
|
|
|
class_to_idx = {c: i for i, c in enumerate(classes)}
|
|
y = np.array([class_to_idx[p] for p in pairs], dtype=np.int64)
|
|
|
|
X = pos.astype(np.float64)
|
|
|
|
if method == "knn":
|
|
estimator = KNeighborsClassifier(n_neighbors=k)
|
|
estimator.fit(X, y)
|
|
ref_tree = None
|
|
elif method == "svm":
|
|
estimator = SVC(kernel="rbf")
|
|
estimator.fit(X, y)
|
|
# SVM cannot answer nearest-neighbour distance queries, so keep a light
|
|
# reference tree alongside it purely for the escape test.
|
|
ref_tree = NearestNeighbors(n_neighbors=1).fit(X)
|
|
else:
|
|
raise ValueError(f"unknown method {method!r}; use 'slab', 'knn', or 'svm'")
|
|
|
|
# Escape threshold from the reference point spacing. Sample a subset for the
|
|
# median 2-NN distance (the 1st neighbour of a training point is itself).
|
|
nn = NearestNeighbors(n_neighbors=2).fit(X)
|
|
probe = (
|
|
X
|
|
if len(X) <= 20_000
|
|
else X[np.random.default_rng(seed).choice(len(X), 20_000, replace=False)]
|
|
)
|
|
d2, _ = nn.kneighbors(probe, n_neighbors=2)
|
|
median_nn = float(np.median(d2[:, 1]))
|
|
escape_threshold = escape_factor * median_nn
|
|
|
|
oracle = GeometryOracle(
|
|
estimator=estimator,
|
|
classes=classes,
|
|
escape_threshold=escape_threshold,
|
|
metadata={
|
|
"method": method,
|
|
"k": k,
|
|
"n_reference_points": int(len(X)),
|
|
"median_nn_dist": median_nn,
|
|
"escape_factor": escape_factor,
|
|
"n_files": len(files),
|
|
},
|
|
)
|
|
oracle._ref_tree = ref_tree
|
|
return oracle
|