MLXID/GridBased
0
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial commit: grid-based multi-photon reconstruction

Browse Source 
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>
This commit is contained in:
xie_x1andClaude Opus 5.5 committed 2026-10-20 07:36:49 +02:00
commit 0f9a3f496f
222 files changed
+53227
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

scripts/04_evaluate.py
+344
View File
@@ -0,0 +1,344 @@
#!/usr/bin/env python3
"""
WP4 -- evaluate a grid model on the fixed-occupancy MC scan.
python scripts/04_evaluate.py --checkpoint Results/<exp>/Models/best.pth
python scripts/04_evaluate.py --checkpoint ... --data <folder>/eval_scan
Produces, per occupancy point:
efficiency / purity at the working threshold, position RMS and sigma68,
energy bias/resolution, sub-pixel non-uniformity (risk R1),
plus the threshold ROC, the pair-resolution curve (efficiency vs distance to the
nearest other photon -- the figure of merit that replaces "efficiency vs
occupancy"), and residual vs conversion depth z.
"""
import argparse
import glob
import json
import os
import re
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
from torch.utils.data import DataLoader
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent / "src"))
from datasets import GridFrameDataset, collate_with_photons # noqa: E402
from decode import decode # noqa: E402
from metrics import (DetectionStats, greedy_match, # noqa: E402
nearest_neighbour_distance)
from models import build_model # noqa: E402
THRESHOLDS = np.array([.05, .1, .2, .3, .4, .5, .6, .7, .8, .9, .95, .98])
NEIGHBOUR_BIN_EDGES = np.array([0, .5, 1, 1.5, 2, 2.5, 3, 4, 6, 1e9])
ENERGY_BIN_WIDTH = 2.5 # keV; one bin only for a
# monochromatic set, which is fine
def energy_bin_edges(energies, width=ENERGY_BIN_WIDTH):
"""2.5 keV bins on a fixed grid, spanning the energies of the set."""
low = np.floor(energies.min() / width) * width
high = np.ceil(energies.max() / width) * width
return np.arange(low, high + width / 2, width)
def occupancy_of_filename(path):
match = re.search(r"occ([0-9p]+)pct", os.path.basename(path))
return float(match.group(1).replace("p", ".")) / 100 if match else float("nan")
@torch.no_grad()
def run_file(model, conf, path, halo, batch_size=512, workers=8, max_frames=None,
roc_frames=2000):
dataset = GridFrameDataset([path], augment=False, return_photons=True,
input_scale=float(conf.data.input_scale),
energy_scale=float(conf.data.energy_scale),
max_frames=max_frames, name=os.path.basename(path))
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False,
num_workers=workers, collate_fn=collate_with_photons,
pin_memory=True)
frame_size = dataset.frame_size
working_threshold = float(conf.eval.threshold)
lowest_threshold = float(THRESHOLDS.min())
stats = DetectionStats(pixel_size_um=float(conf.eval.pixel_size_um),
match_radius=float(conf.eval.match_radius),
halo=halo, frame_size=frame_size)
roc_predicted, roc_true = [], []
for frames, objectness, offset, energy, photons in loader:
frames = frames.cuda(non_blocking=True)
prediction = model(frames).float()
points = decode(prediction, threshold=lowest_threshold,
merge_distance=float(conf.eval.merge_distance), halo=halo,
predict_energy=bool(conf.model.predict_energy),
energy_scale=float(conf.data.energy_scale))
points = [frame_points.numpy() for frame_points in points]
truth = [frame_truth.numpy() for frame_truth in photons]
stats.update([p[p[:, 3] > working_threshold] for p in points], truth)
if len(roc_predicted) < roc_frames:
roc_predicted.extend(points)
roc_true.extend(truth)
summary = stats.summary()
differential = stats.differential()
# ---- threshold sweep on a subset -------------------------------------
roc = []
for threshold in THRESHOLDS:
subset_stats = DetectionStats(pixel_size_um=float(conf.eval.pixel_size_um),
match_radius=float(conf.eval.match_radius),
halo=halo, frame_size=frame_size)
subset_stats.update([p[p[:, 3] > threshold] for p in roc_predicted[:roc_frames]],
roc_true[:roc_frames])
subset_summary = subset_stats.summary()
roc.append((float(threshold), subset_summary["efficiency"],
subset_summary["purity"], subset_summary["rms_1d_um"]))
# ---- pair resolution: efficiency vs truth nearest-neighbour distance --
n_bins = len(NEIGHBOUR_BIN_EDGES) - 1
n_true_binned = np.zeros(n_bins)
n_matched_binned = np.zeros(n_bins)
for predicted, truth in zip(roc_predicted[:roc_frames], roc_true[:roc_frames]):
predicted = predicted[predicted[:, 3] > working_threshold]
if len(truth) == 0:
continue
inside = ((truth[:, 0] >= halo) & (truth[:, 0] < frame_size - halo) &
(truth[:, 1] >= halo) & (truth[:, 1] < frame_size - halo))
truth = truth[inside]
if len(truth) == 0:
continue
neighbour_distance = nearest_neighbour_distance(truth[:, :2])
bin_index = np.digitize(neighbour_distance, NEIGHBOUR_BIN_EDGES) - 1
for index in range(n_bins):
n_true_binned[index] += (bin_index == index).sum()
if len(predicted):
_, index_true = greedy_match(predicted[:, :2], truth[:, :2],
float(conf.eval.match_radius))
for index in range(n_bins):
n_matched_binned[index] += (bin_index[index_true] == index).sum()
pair_resolution = dict(
edges=NEIGHBOUR_BIN_EDGES[:-1].tolist(),
n_true=n_true_binned.tolist(),
n_match=n_matched_binned.tolist(),
efficiency=(n_matched_binned / np.maximum(n_true_binned, 1)).tolist())
return (summary, differential, roc, pair_resolution,
stats.versus_energy(energy_bin_edges(dataset.photons[:, 2])))
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--checkpoint", required=True)
parser.add_argument("--data", default=None,
help="folder with the eval_scan h5 files")
parser.add_argument("--max-frames", type=int, default=None)
parser.add_argument("--workers", type=int, default=8)
parser.add_argument("--outdir", default=None)
args = parser.parse_args()
checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
conf = OmegaConf.create(checkpoint["conf"])
model = build_model(conf).cuda().eval()
model.load_state_dict(checkpoint["model"])
halo = model.rf_radius
print(f"[eval] checkpoint epoch {checkpoint['epoch']}, R={halo}, "
f"{sum(p.numel() for p in model.parameters())} params, "
f"threshold {conf.eval.threshold}, merge {conf.eval.merge_distance} px")
data_folder = args.data or os.path.join(conf.data.folder, "eval_scan")
files = sorted(glob.glob(os.path.join(data_folder, "*.h5")),
key=occupancy_of_filename)
if not files:
sys.exit(f"no eval files in {data_folder}")
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 = occupancy_of_filename(path)
summary, differential, roc, pair_resolution, per_energy = run_file(
model, conf, path, halo, workers=args.workers, max_frames=args.max_frames)
results[f"{occupancy*100:g}%"] = dict(occupancy=occupancy, summary=summary,
roc=roc, pair=pair_resolution,
per_energy=per_energy,
_differential=differential)
print(f" occ {occupancy*100:5.2f}% eff {summary['efficiency']*100:6.2f}% "
f"pur {summary['purity']*100:6.2f}% RMS {summary['rms_1d_um']:.3f} um "
f"s68 {summary['sigma68_1d_um']:.3f} um "
f"dE {summary.get('energy_rms_keV', float('nan')):.3f} keV "
f"subpix-nonunif {summary['subpixel_nonuniformity']*100:.1f}%",
flush=True)
populated = [row for row in list(results.values())[len(results)//2]["per_energy"]
if row["n_true"] > 100]
if len(populated) > 1:
print(f"\n per energy bin at occupancy "
f"{list(results.values())[len(results)//2]['occupancy']*100:g}%:")
print(f" {'keV':>11} {'n':>9} {'eff%':>7} {'RMS um':>8} {'s68 um':>8} "
f"{'dE bias':>8} {'dE rms':>8}")
for row in populated:
print(f" {row['energy_low']:4.1f}-{row['energy_high']:4.1f} "
f"{row['n_true']:9d} {row['efficiency']*100:7.2f} "
f"{row['rms_1d_um']:8.3f} {row['sigma68_1d_um']:8.3f} "
f"{row['energy_bias_keV']:8.3f} {row['energy_rms_keV']:8.3f}")
plot_versus_energy(results, output_dir)
plot_all(results, output_dir, float(conf.eval.threshold))
serialisable = {label: {key: value for key, value in result.items()
if key != "_differential"}
for label, result in results.items()}
with open(output_dir / "eval_summary.json", "w") as handle:
json.dump(serialisable, handle, indent=1)
print(f"[eval] wrote {output_dir/'eval_summary.json'} and {output_dir/'Plots'}")
def plot_versus_energy(results, output_dir):
"""Efficiency, position resolution and energy resolution against photon energy."""
figure, axes = plt.subplots(1, 3, figsize=(15, 4.4))
for label, result in results.items():
rows = [row for row in result["per_energy"] if row["n_true"] > 100]
if len(rows) < 2:
continue
centre = np.array([0.5 * (r["energy_low"] + r["energy_high"]) for r in rows])
axes[0].plot(centre, [r["efficiency"] * 100 for r in rows], "o-", label=f"occ {label}")
axes[1].plot(centre, [r["sigma68_1d_um"] for r in rows], "o-", label=f"occ {label}")
axes[2].plot(centre, [r["energy_rms_keV"] for r in rows], "o-", label=f"occ {label}")
for axis, ylabel in zip(axes, ("efficiency [%]", r"position $\sigma_{68}$ [$\mu$m]",
r"energy RMS [keV]")):
axis.set_xlabel("photon energy [keV]")
axis.set_ylabel(ylabel)
axis.grid(alpha=.3)
axis.legend(fontsize=8)
figure.tight_layout()
figure.savefig(output_dir / "Plots" / "eval_vs_energy.png", dpi=130)
plt.close(figure)
def plot_all(results, output_dir, working_threshold):
occupancy = np.array([r["occupancy"] for r in results.values()]) * 100
efficiency = np.array([r["summary"]["efficiency"] for r in results.values()]) * 100
purity = np.array([r["summary"]["purity"] for r in results.values()]) * 100
rms = np.array([r["summary"]["rms_1d_um"] for r in results.values()])
sigma68 = np.array([r["summary"]["sigma68_1d_um"] for r in results.values()])
figure, axes = plt.subplots(2, 3, figsize=(16, 9))
axes[0, 0].plot(occupancy, efficiency, "o-", label="efficiency")
axes[0, 0].plot(occupancy, purity, "s-", label="purity")
axes[0, 0].set_xlabel("occupancy [%]")
axes[0, 0].set_ylabel("[%]")
axes[0, 0].set_title(f"working point p>{working_threshold}")
axes[0, 0].legend()
axes[0, 0].grid(alpha=.3)
axes[0, 1].plot(occupancy, rms, "o-", label="RMS")
axes[0, 1].plot(occupancy, sigma68, "s-", label=r"$\sigma_{68}$")
axes[0, 1].set_xlabel("occupancy [%]")
axes[0, 1].set_ylabel(r"residual [$\mu$m]")
axes[0, 1].legend()
axes[0, 1].grid(alpha=.3)
for label, result in results.items():
roc = np.array(result["roc"])
axes[0, 2].plot(roc[:, 2] * 100, roc[:, 1] * 100, ".-", label=f"occ {label}")
axes[0, 2].set_xlabel("purity [%]")
axes[0, 2].set_ylabel("efficiency [%]")
axes[0, 2].set_title("threshold sweep")
axes[0, 2].legend(fontsize=8)
axes[0, 2].grid(alpha=.3)
for label, result in results.items():
residual = result["_differential"]["residual_um"]
if len(residual):
axes[1, 0].hist(residual.ravel(), bins=200, range=(-6, 6),
histtype="step", density=True, label=f"occ {label}")
axes[1, 0].set_xlabel(r"residual [$\mu$m]")
axes[1, 0].set_yscale("log")
axes[1, 0].legend(fontsize=8)
axes[1, 0].grid(alpha=.3)
for label, result in results.items():
pair = result["pair"]
edges = np.array(pair["edges"])
pair_efficiency = np.array(pair["efficiency"]) * 100
populated = np.array(pair["n_true"]) > 50
axes[1, 1].plot(edges[populated], pair_efficiency[populated], "o-",
label=f"occ {label}")
axes[1, 1].set_xlabel("distance to nearest other photon [px]")
axes[1, 1].set_ylabel("efficiency [%]")
axes[1, 1].set_xlim(0, 6)
axes[1, 1].set_title("pair resolution")
axes[1, 1].legend(fontsize=8)
axes[1, 1].grid(alpha=.3)
middle = list(results.values())[len(results) // 2]
subpixel = middle["_differential"]["subpixel"]
if len(subpixel):
histogram, _, _ = np.histogram2d(subpixel[:, 0], subpixel[:, 1], bins=25,
range=[[0, 1], [0, 1]])
image = axes[1, 2].imshow(histogram.T / histogram.mean(), origin="lower",
extent=[0, 1, 0, 1], vmin=0.5, vmax=1.5,
cmap="viridis")
plt.colorbar(image, ax=axes[1, 2])
axes[1, 2].set_title(f"sub-pixel uniformity (occ {middle['occupancy']*100:g}%)\n"
"flat = no grid artefact (risk R1)")
axes[1, 2].set_xlabel("sub-pixel x")
axes[1, 2].set_ylabel("sub-pixel y")
figure.tight_layout()
figure.savefig(output_dir / "Plots" / "eval_scan.png", dpi=130)
plt.close(figure)
# ---- residual vs depth, vs isolation, and energy ---------------------
figure, axes = plt.subplots(1, 3, figsize=(16, 4.4))
depth_edges = np.linspace(0, 650, 14)
neighbour_edges = np.array([0, .5, 1, 1.5, 2, 2.5, 3, 4, 6, 10])
for label, result in results.items():
differential = result["_differential"]
if not len(differential["z_um"]):
continue
depth_bin = np.digitize(differential["z_um"], depth_edges) - 1
depth_rms = [differential["residual_um"][depth_bin == index].std()
if (depth_bin == index).sum() > 50 else np.nan
for index in range(len(depth_edges) - 1)]
axes[0].plot(0.5 * (depth_edges[:-1] + depth_edges[1:]), depth_rms, "o-",
label=f"occ {label}")
neighbour_bin = np.digitize(differential["nn_px"], neighbour_edges) - 1
neighbour_rms = [differential["residual_um"][neighbour_bin == index].std()
if (neighbour_bin == index).sum() > 50 else np.nan
for index in range(len(neighbour_edges) - 1)]
axes[1].plot(neighbour_edges[:-1], neighbour_rms, "o-", label=f"occ {label}")
if np.isfinite(differential["e_pred"]).any():
axes[2].hist(differential["e_pred"] - differential["e_true"], bins=150,
range=(-4, 4), histtype="step", density=True,
label=f"occ {label}")
axes[0].set_xlabel(r"conversion depth z [$\mu$m]")
axes[0].set_ylabel(r"RMS [$\mu$m]")
axes[0].set_title("resolution vs charge-cloud size")
axes[0].legend(fontsize=8)
axes[0].grid(alpha=.3)
axes[1].set_xlabel("distance to nearest other photon [px]")
axes[1].set_ylabel(r"RMS [$\mu$m]")
axes[1].legend(fontsize=8)
axes[1].grid(alpha=.3)
axes[2].set_xlabel(r"$E_{pred}-E_{true}$ [keV]")
axes[2].set_yscale("log")
axes[2].legend(fontsize=8)
axes[2].grid(alpha=.3)
figure.tight_layout()
figure.savefig(output_dir / "Plots" / "eval_differential.png", dpi=130)
plt.close(figure)
if __name__ == "__main__":
main()