Initial commit: grid-based multi-photon reconstruction
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
commit
0f9a3f496f
222 files changed
+53227
No files matched your search
@@ -0,0 +1,321 @@
|
||||
#!/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()
|
||||
Reference in new issue
Block a user