xie_x1andClaude Opus 5.5
0f9a3f496f
Code, configs, charge LUTs, best.pth and evaluation records for each training run, the knife-edge / Siemens-star / XRF notebooks with outputs, and the overnight run reports. Epoch checkpoints, inference results and HTML exports stay local (see .gitignore). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
322 lines
15 KiB
Python
322 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
WP4 -- baseline comparison on identical MC frames (gate G2).
|
|
|
|
conda activate mlxid_demo # this env has both torch and aare
|
|
python scripts/06_baseline_comparison.py \
|
|
--checkpoint Results/<exp>/Models/E100.pth --n-frames 4000
|
|
|
|
Three reconstruction chains are run over the *same* simulated frames and scored
|
|
against the same truth with the same matcher:
|
|
|
|
grid GridNet + position-space decode
|
|
cluster+dl aare VarClusterFinder -> cluster energy picks the multiplicity
|
|
-> the pretrained 1/2/3-photon CNNs on 3x3 / 6x6 / 9x9 crops
|
|
cluster+eta the same cluster finder, 1-photon clusters only, classical eta
|
|
interpolation with a Rosenblatt LUT fitted on the same events
|
|
|
|
Cluster cropping, energy windows and coordinate conventions follow
|
|
``DataProcess/UsefulFuncs.py::_processFrames``; the CNN input conventions follow
|
|
``DeepLearning/src/datasets.py``. Every chain is scored on the same valid
|
|
region (frame minus the grid model's halo), so the numbers are comparable.
|
|
"""
|
|
import argparse
|
|
import importlib.util
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import matplotlib
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import torch
|
|
from omegaconf import OmegaConf
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
# GridBased/src and MLXIDDemo/src both define a `models` module. MLXIDDemo's
|
|
# model_zoo does `from models import ...`, so that name is left to the demo and
|
|
# the GridBased one is loaded by path under its own name.
|
|
sys.path.insert(0, str(HERE.parent / "src"))
|
|
sys.path.insert(0, str(HERE.parent.parent / "MLXIDDemo" / "src"))
|
|
sys.path.append("/home/xie_x1/software/aareBuild")
|
|
|
|
import aare # noqa: E402
|
|
from decode import decode # noqa: E402
|
|
from eta_interpolation_functions import (bilinear_xy_lookup, # noqa: E402
|
|
build_xy_lut_Rosenblatt)
|
|
from metrics import DetectionStats # noqa: E402
|
|
from model_zoo import load_pretrained # noqa: E402
|
|
|
|
|
|
def _load_by_path(name, path):
|
|
spec = importlib.util.spec_from_file_location(name, path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
build_model = _load_by_path("grid_models", HERE.parent / "src" / "models.py").build_model
|
|
|
|
CLUSTER_SIZE = {1: 3, 2: 6, 3: 9}
|
|
ETA_BINS = 201
|
|
ETA_POSITION_WEIGHTS = np.array([-1.0, 0.0, 1.0]) + 0.5
|
|
|
|
|
|
def coordinate_channels(size):
|
|
"""The two coordinate planes singlePhotonDataset appends to its input."""
|
|
axis = np.linspace(-0.5, 0.5, size)
|
|
x_grid, y_grid = np.meshgrid(axis, axis, indexing="ij")
|
|
return torch.from_numpy(np.stack([x_grid, y_grid])).float()
|
|
|
|
|
|
def find_clusters(frames, noise_keV, threshold_sigma=5, peripheral_factor=3,
|
|
n_neighbours=4):
|
|
"""Run the variable cluster finder frame by frame.
|
|
|
|
Returns a list per frame of dicts with the seed pixel, the charge-weighted
|
|
centre and the cluster energy.
|
|
"""
|
|
frame_size = frames.shape[-1]
|
|
noise_map = np.ones((frame_size, frame_size)) * noise_keV
|
|
|
|
def new_finder():
|
|
# A finder is built per frame. Reusing one across frames is what
|
|
# DataProcess/UsefulFuncs._processFrames does and is fine at its
|
|
# 101x101 ROI, but on these 32x32 frames the internal hit buffer
|
|
# overflows after a few frames and hits() starts returning stale
|
|
# entries (and can abort the process).
|
|
finder = aare.VarClusterFinder((frame_size, frame_size), threshold_sigma)
|
|
finder.set_peripheralThresholdFactor(peripheral_factor)
|
|
finder.set_noiseMap(noise_map)
|
|
finder.set_numberOfNeighbours(n_neighbours)
|
|
finder.set_empty_surroundingPixels(False)
|
|
return finder
|
|
|
|
per_frame = []
|
|
for frame in frames:
|
|
finder = new_finder()
|
|
working_copy = np.array(frame, dtype=np.float64) # the finder consumes it
|
|
finder.find_clusters_X(working_copy)
|
|
hits = finder.hits()
|
|
clusters = []
|
|
for index in range(len(hits)):
|
|
size = hits["size"][index]
|
|
energies = hits["enes"][index][:size]
|
|
total = energies.sum()
|
|
if total <= 0:
|
|
continue
|
|
clusters.append(dict(
|
|
energy=float(hits["energy"][index]),
|
|
seed_col=int(hits["col"][index]), seed_row=int(hits["row"][index]),
|
|
centre_col=float((hits["cols"][index][:size] * energies).sum() / total),
|
|
centre_row=float((hits["rows"][index][:size] * energies).sum() / total),
|
|
))
|
|
per_frame.append(clusters)
|
|
return per_frame
|
|
|
|
|
|
def collect_crops(frames, clusters_per_frame, beam_energy, selection_range):
|
|
"""Sort clusters into 1/2/3-photon crops, following UsefulFuncs._processFrames."""
|
|
frame_size = frames.shape[-1]
|
|
crops = {n: [] for n in CLUSTER_SIZE}
|
|
origins = {n: [] for n in CLUSTER_SIZE} # (frame_index, ref_col, ref_row)
|
|
for frame_index, clusters in enumerate(clusters_per_frame):
|
|
for cluster in clusters:
|
|
energy = cluster["energy"]
|
|
for n_photons, size in CLUSTER_SIZE.items():
|
|
window = n_photons * selection_range
|
|
if not (n_photons * beam_energy - window < energy
|
|
< n_photons * beam_energy + window):
|
|
continue
|
|
if n_photons == 1:
|
|
ref_col = cluster["seed_col"] - size // 2
|
|
ref_row = cluster["seed_row"] - size // 2
|
|
else:
|
|
ref_col = int(cluster["centre_col"] - size / 2) + 1
|
|
ref_row = int(cluster["centre_row"] - size / 2) + 1
|
|
if (ref_col < 0 or ref_col + size > frame_size
|
|
or ref_row < 0 or ref_row + size > frame_size):
|
|
continue
|
|
crops[n_photons].append(
|
|
frames[frame_index, ref_row:ref_row + size, ref_col:ref_col + size])
|
|
origins[n_photons].append((frame_index, ref_col, ref_row))
|
|
return ({n: (np.asarray(v, np.float32) if v else np.zeros((0, CLUSTER_SIZE[n],
|
|
CLUSTER_SIZE[n]), np.float32))
|
|
for n, v in crops.items()},
|
|
{n: (np.asarray(v, np.int64) if v else np.zeros((0, 3), np.int64))
|
|
for n, v in origins.items()})
|
|
|
|
|
|
@torch.no_grad()
|
|
def cluster_dl_points(crops, origins, n_frames, device="cuda", batch_size=8192):
|
|
"""Run the pretrained 1/2/3-photon CNNs -> per-frame photon lists."""
|
|
points = [[] for _ in range(n_frames)]
|
|
for n_photons, size in CLUSTER_SIZE.items():
|
|
sample = crops[n_photons]
|
|
if len(sample) == 0:
|
|
continue
|
|
model = load_pretrained(n_photons, device)
|
|
centre = size / 2.0
|
|
extra = coordinate_channels(size).to(device) if n_photons == 1 else None
|
|
for start in range(0, len(sample), batch_size):
|
|
batch = torch.from_numpy(sample[start:start + batch_size]).to(device)
|
|
batch = batch.unsqueeze(1)
|
|
if extra is not None:
|
|
batch = torch.cat([batch, extra.expand(batch.shape[0], -1, -1, -1)], dim=1)
|
|
output = model(batch).cpu().numpy().reshape(len(batch), n_photons, 2)
|
|
for row, (frame_index, ref_col, ref_row) in enumerate(
|
|
origins[n_photons][start:start + batch_size]):
|
|
for photon in range(n_photons):
|
|
points[frame_index].append((
|
|
ref_col + centre + output[row, photon, 0],
|
|
ref_row + centre + output[row, photon, 1],
|
|
float(n_photons and 12.0), 1.0))
|
|
del model
|
|
torch.cuda.empty_cache()
|
|
return [np.asarray(p, np.float64).reshape(-1, 4) for p in points]
|
|
|
|
|
|
def cluster_eta_points(crops, origins, n_frames):
|
|
"""Classical eta interpolation on the 1-photon clusters."""
|
|
sample = crops[1]
|
|
points = [[] for _ in range(n_frames)]
|
|
if len(sample) == 0:
|
|
return [np.zeros((0, 4)) for _ in range(n_frames)]
|
|
totals = sample.sum(axis=(1, 2))
|
|
totals[totals == 0] = 1.0
|
|
eta_x = np.clip((sample.sum(axis=1) * ETA_POSITION_WEIGHTS).sum(axis=-1) / totals, 0, 1)
|
|
eta_y = np.clip((sample.sum(axis=2) * ETA_POSITION_WEIGHTS).sum(axis=-1) / totals, 0, 1)
|
|
histogram, _, _ = np.histogram2d(eta_x, eta_y, bins=ETA_BINS, range=[[0, 1], [0, 1]])
|
|
x_table, y_table = build_xy_lut_Rosenblatt(histogram)
|
|
x_fraction, y_fraction = bilinear_xy_lookup(eta_x, eta_y, x_table, y_table)
|
|
for row, (frame_index, ref_col, ref_row) in enumerate(origins[1]):
|
|
points[frame_index].append((ref_col + 1.0 + x_fraction[row],
|
|
ref_row + 1.0 + y_fraction[row], 12.0, 1.0))
|
|
return [np.asarray(p, np.float64).reshape(-1, 4) for p in points]
|
|
|
|
|
|
@torch.no_grad()
|
|
def grid_points(model, frames, train_conf, halo, threshold, merge_distance,
|
|
batch_size=512):
|
|
input_scale = float(train_conf.data.input_scale)
|
|
energy_scale = float(train_conf.data.energy_scale)
|
|
points = []
|
|
for start in range(0, len(frames), batch_size):
|
|
batch = torch.from_numpy(frames[start:start + batch_size]).cuda().unsqueeze(1)
|
|
prediction = model(batch / input_scale).float()
|
|
points.extend(p.numpy() for p in decode(
|
|
prediction, threshold=threshold, merge_distance=merge_distance,
|
|
halo=halo, predict_energy=bool(train_conf.model.predict_energy),
|
|
energy_scale=energy_scale))
|
|
return points
|
|
|
|
|
|
def score(points_per_frame, truth_per_frame, halo, frame_size, match_radius,
|
|
pixel_size_um):
|
|
stats = DetectionStats(pixel_size_um=pixel_size_um, match_radius=match_radius,
|
|
halo=halo, frame_size=frame_size)
|
|
stats.update(points_per_frame, truth_per_frame)
|
|
return stats.summary()
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--checkpoint", required=True)
|
|
parser.add_argument("--data", default=None, help="eval_scan folder")
|
|
parser.add_argument("--n-frames", type=int, default=4000)
|
|
parser.add_argument("--beam-energy", type=float, default=12.0)
|
|
parser.add_argument("--selection-range", type=float, default=3.0)
|
|
parser.add_argument("--outdir", default=None)
|
|
args = parser.parse_args()
|
|
|
|
import h5py
|
|
checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
|
|
train_conf = OmegaConf.create(checkpoint["conf"])
|
|
model = build_model(train_conf).cuda().eval()
|
|
model.load_state_dict(checkpoint["model"])
|
|
halo = model.rf_radius
|
|
threshold = float(train_conf.eval.threshold)
|
|
merge_distance = float(train_conf.eval.merge_distance)
|
|
match_radius = float(train_conf.eval.match_radius)
|
|
pixel_size_um = float(train_conf.eval.pixel_size_um)
|
|
noise_keV = 0.13
|
|
print(f"[baseline] grid model {args.checkpoint} (epoch {checkpoint['epoch']}), "
|
|
f"R={halo}, threshold {threshold}, merge {merge_distance} px")
|
|
|
|
data_folder = Path(args.data or Path(train_conf.data.folder) / "eval_scan")
|
|
files = sorted(data_folder.glob("*.h5"),
|
|
key=lambda p: float(p.stem.replace("occ", "").replace("pct", "")
|
|
.replace("p", ".")))
|
|
output_dir = Path(args.outdir or Path(args.checkpoint).parent.parent)
|
|
(output_dir / "Plots").mkdir(parents=True, exist_ok=True)
|
|
|
|
results = {}
|
|
for path in files:
|
|
occupancy = float(path.stem.replace("occ", "").replace("pct", "")
|
|
.replace("p", ".")) / 100
|
|
with h5py.File(path, "r") as handle:
|
|
frames = handle["frames"][:args.n_frames].astype(np.float32)
|
|
pointer = handle["frame_ptr"][:args.n_frames + 1]
|
|
photons = handle["photons"][:pointer[-1]]
|
|
truth = [photons[pointer[i]:pointer[i + 1]] for i in range(len(frames))]
|
|
frame_size = frames.shape[-1]
|
|
|
|
clusters_per_frame = find_clusters(frames.astype(np.float64), noise_keV)
|
|
crops, origins = collect_crops(frames, clusters_per_frame,
|
|
args.beam_energy, args.selection_range)
|
|
chains = {
|
|
"grid": grid_points(model, frames, train_conf, halo, threshold,
|
|
merge_distance),
|
|
"cluster+dl": cluster_dl_points(crops, origins, len(frames)),
|
|
"cluster+eta": cluster_eta_points(crops, origins, len(frames)),
|
|
}
|
|
results[f"{occupancy*100:g}%"] = dict(
|
|
occupancy=occupancy,
|
|
n_clusters={str(n): int(len(crops[n])) for n in CLUSTER_SIZE},
|
|
**{name: score(pts, truth, halo, frame_size, match_radius, pixel_size_um)
|
|
for name, pts in chains.items()})
|
|
|
|
print(f"\n occupancy {occupancy*100:g}% ({len(frames)} frames, clusters "
|
|
f"1ph={len(crops[1])} 2ph={len(crops[2])} 3ph={len(crops[3])})")
|
|
print(f" {'chain':12s} {'eff%':>7} {'pur%':>7} {'RMS um':>8} {'s68 um':>8}")
|
|
for name in chains:
|
|
s = results[f"{occupancy*100:g}%"][name]
|
|
print(f" {name:12s} {s['efficiency']*100:7.2f} {s['purity']*100:7.2f} "
|
|
f"{s['rms_1d_um']:8.3f} {s['sigma68_1d_um']:8.3f}", flush=True)
|
|
|
|
with open(output_dir / "baseline_comparison.json", "w") as handle:
|
|
json.dump(results, handle, indent=1)
|
|
plot(results, output_dir, checkpoint["epoch"])
|
|
print(f"\n[baseline] wrote {output_dir/'baseline_comparison.json'}")
|
|
|
|
|
|
def plot(results, output_dir, epoch):
|
|
occupancy = np.array([r["occupancy"] for r in results.values()]) * 100
|
|
figure, axes = plt.subplots(1, 3, figsize=(15, 4.4))
|
|
styles = {"grid": ("o-", "#10695f"), "cluster+dl": ("s--", "#4b2a6b"),
|
|
"cluster+eta": ("^:", "#a0451c")}
|
|
for name, (style, colour) in styles.items():
|
|
efficiency = np.array([r[name]["efficiency"] for r in results.values()]) * 100
|
|
purity = np.array([r[name]["purity"] for r in results.values()]) * 100
|
|
rms = np.array([r[name]["rms_1d_um"] for r in results.values()])
|
|
axes[0].plot(occupancy, efficiency, style, color=colour, label=name)
|
|
axes[1].plot(occupancy, purity, style, color=colour, label=name)
|
|
axes[2].plot(occupancy, rms, style, color=colour, label=name)
|
|
for axis, label in zip(axes, ("efficiency [%]", "purity [%]",
|
|
r"position RMS [$\mu$m]")):
|
|
axis.set_xlabel("occupancy [%]")
|
|
axis.set_ylabel(label)
|
|
axis.set_xscale("log")
|
|
axis.grid(alpha=.3)
|
|
axis.legend()
|
|
figure.suptitle(f"same MC frames, same matcher (grid model at epoch {epoch})")
|
|
figure.tight_layout()
|
|
figure.savefig(output_dir / "Plots" / "baseline_comparison.png", dpi=130)
|
|
plt.close(figure)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|