MLXID/GridBased
0
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
7de2d7d170aa2abe41e46640e77c7d143a28dfd8
GridBased/scripts/05_infer_measurement.py
T
xie_x1andClaude Opus 5.5 0f9a3f496f 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>
2026-10-20 07:36:49 +02:00

493 lines
24 KiB
Python
Raw Blame History

#!/usr/bin/env python3
"""
WP5.2 -- run the grid model over a measured run, straight from the raw files.
python scripts/05_infer_measurement.py --config Configs/infer_measurement_12keV.yaml \
--run Edge_3filters_pos0_12keV
python scripts/05_infer_measurement.py --config ... --run ... --n-frames 200000
python scripts/05_infer_measurement.py --config ... --run ... --save-photons
raw ADU -> pedestal -> common mode -> 3D-LUT -> GridNet -> position-space merge
-> sub-pixel binned image + energy spectrum + per-frame photon count.
Only the valid region (the frame minus the halo R) is kept: near the ROI border
the network sees zero padding instead of the neighbouring pixels that exist on
the chip, so those cells are not trustworthy.
With `inference.precision: fp16` the network runs under FP16 autocast in
channels-last layout, 2.3x faster for the poly model. The offsets and energy of
the cells above threshold are then recomputed in FP32 from the same FP16
features (`predict_fp16`), so energies are not rounded to the FP16 grid. The models are trained
with AMP, and on MC and measured frames FP16 reproduces the FP32 photon list
(efficiency, purity, position and energy resolution unchanged; 99.997 % of the
photons identical, positions within ~20 nm).
Images written per run, each as .npy (analysis) and .png (looking at it):
binned_image reconstructed photons at pixel_size / binning_factor
native_image the same photons rebinned back to the 25 um pitch
analog_image mean calibrated energy per frame -- the classical "native
readout" picture
plus comparison.png, the two side by side.
With `inference.self_flattening` the sub-pixel phase of the reconstructed
positions is flattened with the Rosenblatt transform of
`DeepLearning/src/etaInterpolationFromPoints.py`: the (frac_x, frac_y) histogram
of the whole run becomes a look-up table that maps every position onto the one a
uniform illumination would have put it at. A per-bin flat-field division cannot
do this -- it corrects intensities, while a grid artefact is a displacement --
so the flattened image is produced here alongside the raw one and both are kept.
"""
import argparse
import json
import math
import os
import queue
import sys
import threading
import time
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
sys.path.insert(0, str(HERE.parent / "src"))
from decode import decode # noqa: E402
from measurement import (GpuCalibrator, ImageAccumulator, # noqa: E402
RawFrameReader, estimate_pedestal)
from models import build_model # noqa: E402
from etaInterpolationFromPoints import (NEtaBins, # noqa: E402
interpolate_eta_from_points)
from flattening import StreamingFlattener # noqa: E402
@torch.no_grad()
def predict_fp16(model, inputs, threshold):
"""FP16 backbone and head; offsets and energy of candidate cells in FP32.
An FP16 head rounds every energy to the FP16 grid (7.3 meV below 15 keV,
14.6 meV above), which shows up as a comb in a 50 eV spectrum. Running the
whole head in FP32 means converting the full feature map and costs ~15 %,
so only the cells above the objectness threshold -- the only ones decode can
report -- get channels 1-3 recomputed in FP32 from the same FP16 features.
"""
with torch.autocast("cuda", dtype=torch.float16):
features = model.backbone(inputs.contiguous(memory_format=torch.channels_last))
prediction = model.head(features)
prediction = prediction.float()
logit_threshold = math.log(threshold / (1.0 - threshold))
batch_of, row_of, column_of = torch.nonzero(prediction[:, 0] > logit_threshold,
as_tuple=True)
if len(batch_of):
weight = model.head.weight[1:, :, 0, 0].float() # (3, channels)
bias = model.head.bias[1:].float()
cells = features[batch_of, :, row_of, column_of].float() # (n, channels)
values = cells @ weight.T + bias # (n, 3)
for channel in range(values.shape[1]):
prediction[batch_of, channel + 1, row_of, column_of] = values[:, channel]
return prediction
def load_model(checkpoint_path):
checkpoint = torch.load(checkpoint_path, 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"])
return model, train_conf, checkpoint
def prefetch(iterator, depth=4):
"""Overlap file reading with GPU work."""
output = queue.Queue(maxsize=depth)
def worker():
try:
for item in iterator:
output.put(item)
finally:
output.put(None)
threading.Thread(target=worker, daemon=True).start()
while True:
item = output.get()
if item is None:
return
yield item
def exposed_occupancy(density, fraction=0.5):
"""Occupancy inside the illuminated area only.
A knife-edge or slit run leaves most of the ROI in shadow, so an average
over the whole frame badly understates the flux the model actually sees:
on Edge_3filters_pos0_12keV it gives 0.16 % against the 0.67 % the
illuminated pixels really see.
"""
if density.size == 0 or density.max() <= 0:
return float("nan"), 0
reference = np.percentile(density, 99)
exposed = density > fraction * reference
if not exposed.any():
return float("nan"), 0
return float(density[exposed].mean()), int(exposed.sum())
def subpixel_histogram(x, y, bins=NEtaBins, chunk=50_000_000):
"""2D histogram of the sub-pixel phase of the reconstructed positions.
In slices: histogram2d stacks its input as float64 and adds int64 bin
indices, ~40 bytes per position, which ran out of memory on the 1e9-photon
flat fields. The histogram is additive, so the result is the same.
"""
histogram = np.zeros((bins, bins))
for start in range(0, len(x), chunk):
xs, ys = x[start:start + chunk], y[start:start + chunk]
histogram += np.histogram2d(xs - np.floor(xs), ys - np.floor(ys),
bins=bins, range=[[0, 1], [0, 1]])[0]
return histogram
def subpixel_nonuniformity(histogram):
"""Excess rms/mean over the Poisson floor; 0 means a flat sub-pixel phase."""
mean = histogram.mean()
if mean <= 0:
return float("nan")
relative = histogram.std() / mean
return float(np.sqrt(max(relative ** 2 - 1.0 / mean, 0.0)))
def save_subpixel_figure(output_dir, before, after, run_name):
panels = [("before flattening", before)]
if after is not None:
panels.append(("after flattening", after))
figure, axes = plt.subplots(1, len(panels) + 1,
figsize=(5.2 * (len(panels) + 1), 4.2))
for axis, (title, histogram) in zip(axes, panels):
image = axis.imshow(histogram.T / histogram.mean(), origin="lower",
extent=[0, 1, 0, 1], vmin=0.9, vmax=1.1, cmap="viridis")
axis.set_title(f"{title}\nexcess rms "
f"{subpixel_nonuniformity(histogram)*100:.2f} %", fontsize=10)
axis.set_xlabel("sub-pixel x")
axis.set_ylabel("sub-pixel y")
plt.colorbar(image, ax=axis, fraction=0.046)
axis = axes[-1]
for title, histogram in panels:
profile = histogram.sum(axis=1)
axis.plot(np.linspace(0, 1, len(profile)), profile / profile.mean(),
lw=1, label=title)
axis.axhline(1.0, color="k", lw=.6)
axis.set_xlabel("sub-pixel x")
axis.set_ylabel("relative density")
axis.legend(fontsize=8)
axis.grid(alpha=.3)
figure.suptitle(run_name)
figure.tight_layout()
figure.savefig(output_dir / "subpixel_phase.png", dpi=150)
plt.close(figure)
def save_images(output_dir, binned, native, analog, run_name, binning_factor,
pixel_size_um, halo):
"""Write the reconstruction and the analog picture as .npy and .png."""
np.save(output_dir / "binned_image.npy", binned)
np.save(output_dir / "native_image.npy", native)
np.save(output_dir / "analog_image.npy", analog)
bin_um = pixel_size_um / binning_factor
extent_binned = [0, binned.shape[1] * bin_um, 0, binned.shape[0] * bin_um]
extent_native = [0, native.shape[1] * pixel_size_um,
0, native.shape[0] * pixel_size_um]
def draw(axis, image, extent, title, unit):
finite = image[np.isfinite(image)]
high = np.percentile(finite, 99.5) if finite.size else 1.0
handle = axis.imshow(image, origin="lower", extent=extent, cmap="gray",
vmin=0, vmax=high if high > 0 else 1.0,
interpolation="nearest")
axis.set_title(title, fontsize=10)
axis.set_xlabel(r"x [$\mu$m]")
axis.set_ylabel(r"y [$\mu$m]")
plt.colorbar(handle, ax=axis, label=unit, fraction=0.046)
panels = (
(binned, extent_binned, "binned_image",
f"reconstruction, {bin_um:g} " + r"$\mu$m bins", "photons"),
(native, extent_native, "native_image",
f"reconstructed photons, {pixel_size_um:g} " + r"$\mu$m pitch", "photons"),
(analog, extent_native, "analog_image",
"native readout (mean energy per frame)", "keV/frame"),
)
for image, extent, name, title, unit in panels:
figure, axis = plt.subplots(figsize=(6.4, 5.4))
draw(axis, image, extent, f"{run_name}\n{title}", unit)
figure.tight_layout()
figure.savefig(output_dir / f"{name}.png", dpi=150)
plt.close(figure)
figure, axes = plt.subplots(1, 2, figsize=(12.4, 5.4))
draw(axes[0], analog, extent_native,
"native readout (mean energy per frame)", "keV/frame")
draw(axes[1], binned, extent_binned,
f"GridBased reconstruction, {bin_um:g} " + r"$\mu$m bins", "photons")
figure.suptitle(f"{run_name} (valid region = frame minus halo R={halo})")
figure.tight_layout()
figure.savefig(output_dir / "comparison.png", dpi=150)
plt.close(figure)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--config", required=True)
parser.add_argument("--run", required=True)
parser.add_argument("--n-frames", type=int, default=None)
parser.add_argument("--checkpoint", default=None)
parser.add_argument("--tag", default=None)
parser.add_argument("--save-photons", action="store_true",
help="write photons.npy (x, y, E, p) even if the config says no")
args = parser.parse_args()
conf = OmegaConf.load(args.config)
if args.run not in conf.runs:
sys.exit(f"unknown run {args.run!r}; have {list(conf.runs.keys())}")
run_conf = conf.runs[args.run]
# the stored frame is n_x x n_y; detector.crop = [x0, x1, y0, y1] reads and
# reconstructs only that window of it (the halo R is lost at its border, so
# the crop should reach R pixels past the region of interest)
frame_n_x, frame_n_y = int(conf.detector.n_x), int(conf.detector.n_y)
crop = conf.detector.get("crop", None)
crop = None if crop is None else [int(v) for v in crop]
if crop is None:
n_x, n_y, crop_x0, crop_y0 = frame_n_x, frame_n_y, 0, 0
else:
n_x, n_y, crop_x0, crop_y0 = crop[1] - crop[0], crop[3] - crop[2], crop[0], crop[2]
pixel_size_um = float(conf.detector.pixel_size_um)
checkpoint_path = args.checkpoint or conf.model.checkpoint
model, train_conf, checkpoint = load_model(checkpoint_path)
precision = str(conf.inference.get("precision", "fp32"))
if precision not in ("fp32", "fp16"):
sys.exit(f"unknown inference.precision {precision!r}")
if precision == "fp16":
model = model.to(memory_format=torch.channels_last)
torch.backends.cudnn.benchmark = True
halo = model.rf_radius
input_scale = float(train_conf.data.input_scale)
energy_scale = float(train_conf.data.energy_scale)
print(f"[infer] run {args.run}")
print(f"[infer] model {checkpoint_path} (epoch {checkpoint['epoch']}), "
f"R={halo} -> valid region {n_x-2*halo}x{n_y-2*halo} of {n_x}x{n_y}")
first_file, last_file = run_conf.signal_files
file_step = int(run_conf.get("signal_file_step", 1))
signal_files = [run_conf.signal_pattern.format(i)
for i in range(int(first_file), int(last_file) + 1, file_step)]
reader = RawFrameReader(signal_files, frame_n_x, frame_n_y,
header_words=int(conf.detector.header_words), crop=crop)
n_frames = args.n_frames or conf.inference.n_frames or reader.n_frames_available
n_frames = min(int(n_frames), reader.n_frames_available)
corner = (int(conf.detector.corner[0]) + crop_x0, int(conf.detector.corner[1]) + crop_y0)
calibrator = GpuCalibrator(conf.calibration.file, n_x, n_y, corner=corner,
common_mode=conf.calibration.common_mode)
if conf.calibration.pedestal_source == "pedestal_run":
pedestal_reader = RawFrameReader([run_conf.pedestal], frame_n_x, frame_n_y,
header_words=int(conf.detector.header_words),
crop=crop)
pedestal = estimate_pedestal(pedestal_reader,
n_frames=int(conf.calibration.pedestal_frames))
elif conf.calibration.pedestal_source == "per_file":
pedestal = estimate_pedestal(reader,
n_frames=int(conf.calibration.pedestal_frames))
else:
sys.exit(f"unknown pedestal_source {conf.calibration.pedestal_source!r}")
calibrator.set_pedestal(pedestal)
print(f"[infer] pedestal from {conf.calibration.pedestal_source}, "
f"common mode '{conf.calibration.common_mode}', "
f"{n_frames} frames of {reader.n_frames_available} "
f"({len(signal_files)} files, step {file_step})")
if crop is not None:
print(f"[infer] crop x {crop[0]}-{crop[1]}, y {crop[2]}-{crop[3]} of the "
f"{frame_n_x}x{frame_n_y} frame -> {n_x}x{n_y}, valid region "
f"x {crop[0]+halo}-{crop[1]-halo}, y {crop[2]+halo}-{crop[3]-halo}")
binning_factor = int(conf.inference.binning_factor)
accumulator = ImageAccumulator(n_x, n_y, binning_factor=binning_factor, halo=halo)
energy_sum = torch.zeros((n_y, n_x), dtype=torch.float64, device="cuda")
energy_low, energy_high, energy_n = conf.inference.energy_bins
energy_edges = np.linspace(float(energy_low), float(energy_high), int(energy_n) + 1)
energy_histogram = np.zeros(int(energy_n), np.int64)
photons_per_frame = []
save_photons = args.save_photons or bool(conf.inference.save_photons)
photon_dump = [] if save_photons else None
# self_flattening: true keeps every position and remaps it (the published
# runs); "streaming" histograms them on the GPU instead (src/flattening.py),
# for runs whose photon list would not fit in memory
flattening_mode = conf.inference.get("self_flattening", False)
streaming = str(flattening_mode).lower() == "streaming"
self_flattening = bool(flattening_mode) and not streaming
position_dump = [] if self_flattening else None
occupancy_edges = conf.inference.get("flattening_occupancy_edges", None)
flattener = (StreamingFlattener(
n_x, n_y, binning_factor, halo=halo,
lut_photons=int(conf.inference.get("flattening_lut_photons", 100_000_000)),
occupancy_edges=(None if occupancy_edges is None
else [float(e) for e in occupancy_edges]))
if streaming else None)
output_dir = Path(conf.inference.output_folder) / (
args.run + (f"_{args.tag}" if args.tag else ""))
output_dir.mkdir(parents=True, exist_ok=True)
batch_size = int(conf.inference.batch_size)
threshold = float(conf.model.threshold)
merge_distance = float(conf.model.merge_distance)
started = time.time()
n_done = 0
batches = reader.batches(batch_size, n_frames=n_frames)
with torch.no_grad():
for file_index, frame_index, adu_frames in prefetch(batches):
if len(adu_frames) == 0:
continue
energy_frames = calibrator.to_energy(torch.from_numpy(adu_frames).cuda())
energy_sum += energy_frames.sum(dim=0, dtype=torch.float64)
inputs = (energy_frames / input_scale).unsqueeze(1)
if precision == "fp16":
prediction = predict_fp16(model, inputs, threshold)
else:
prediction = model(inputs).float()
points_per_frame = decode(prediction, threshold=threshold,
merge_distance=merge_distance, halo=halo,
predict_energy=bool(train_conf.model.predict_energy),
energy_scale=energy_scale)
points_per_frame = [p.numpy() for p in points_per_frame]
accumulator.add(points_per_frame)
photons_per_frame.extend(len(p) for p in points_per_frame)
batch_energies = [p[:, 2] for p in points_per_frame if len(p)]
if batch_energies:
energy_histogram += np.histogram(np.concatenate(batch_energies),
bins=energy_edges)[0]
if photon_dump is not None:
photon_dump.extend(points_per_frame)
if position_dump is not None or flattener is not None:
# one array per batch, not per frame: 14 M small arrays would
# cost more in list overhead than the positions themselves
batch_points = [p[:, :2] for p in points_per_frame if len(p)]
if batch_points:
batch_points = np.concatenate(batch_points).astype(np.float32)
if flattener is not None:
flattener.add(batch_points, n_frames=len(adu_frames))
else:
position_dump.append(batch_points)
n_done += len(adu_frames)
if n_done % (batch_size * 200) < batch_size:
elapsed = time.time() - started
rate = n_done / elapsed
print(f" {n_done}/{n_frames} frames {rate:.0f} frame/s "
f"eta {(n_frames-n_done)/rate/60:.1f} min "
f"<photons/frame> {np.mean(photons_per_frame):.2f}", flush=True)
elapsed = time.time() - started
photons_per_frame = np.asarray(photons_per_frame)
binned = accumulator.image_2d
native = binned.reshape(n_y, binning_factor, n_x, binning_factor).sum(axis=(1, 3))
analog = (energy_sum / max(n_done, 1)).cpu().numpy()
phase_before = phase_after = None
flattening_classes = None
if self_flattening and position_dump:
positions = np.concatenate(position_dump)
position_dump.clear()
phase_before = subpixel_histogram(positions[:, 0], positions[:, 1])
print(f"[infer] self-flattening {len(positions)} positions, sub-pixel "
f"excess rms {subpixel_nonuniformity(phase_before)*100:.2f} %", flush=True)
corrected = interpolate_eta_from_points(positions)
phase_after = subpixel_histogram(corrected[:, 0], corrected[:, 1])
print(f"[infer] after flattening, sub-pixel excess rms "
f"{subpixel_nonuniformity(phase_after)*100:.2f} %", flush=True)
flattened_accumulator = ImageAccumulator(n_x, n_y,
binning_factor=binning_factor,
halo=halo)
# in slices: one call on 1e9 positions builds three int64 index arrays
# of the full length (~28 GB at the 0-filter flat field)
for start in range(0, len(corrected), 50_000_000):
flattened_accumulator.add([corrected[start:start + 50_000_000]])
np.save(output_dir / "binned_image_flattened.npy",
flattened_accumulator.image_2d)
np.save(output_dir / "subpixel_phase_before.npy", phase_before)
np.save(output_dir / "subpixel_phase_after.npy", phase_after)
save_subpixel_figure(output_dir, phase_before, phase_after, args.run)
del positions, corrected, flattened_accumulator
if flattener is not None and flattener.n_photons:
flattened, phase_before, phase_after = flattener.finalize()
print(f"[infer] streaming self-flattening of {flattener.n_photons} positions "
f"(table from the first {flattener.table_photons}): sub-pixel excess rms "
f"{subpixel_nonuniformity(phase_before)*100:.2f} % -> "
f"{subpixel_nonuniformity(phase_after)*100:.2f} %", flush=True)
for c in flattener.classes:
print(f"[infer] occupancy {c['occupancy_from']*100:5.2f}-{c['occupancy_to']*100:6.2f} %: "
f"{c['pixels']:6d} px, table from {c['table_photons']} photons", flush=True)
flattening_classes = flattener.classes
np.save(output_dir / "binned_image_flattened.npy", flattened)
np.save(output_dir / "subpixel_phase_before.npy", phase_before)
np.save(output_dir / "subpixel_phase_after.npy", phase_after)
save_subpixel_figure(output_dir, phase_before, phase_after, args.run)
del flattened, flattener
valid = (slice(halo, n_y - halo), slice(halo, n_x - halo))
density = native[valid] / max(n_done, 1) # photons per pixel per frame
occupancy_exposed, n_exposed = exposed_occupancy(density)
n_valid_pixels = (n_x - 2 * halo) * (n_y - 2 * halo)
save_images(output_dir, binned, native, analog, args.run, binning_factor,
pixel_size_um, halo)
np.save(output_dir / "energy_spectrum.npy",
np.stack([0.5 * (energy_edges[:-1] + energy_edges[1:]),
energy_histogram.astype(float)]))
np.save(output_dir / "photons_per_frame.npy", photons_per_frame)
if photon_dump is not None:
np.save(output_dir / "photons.npy", np.concatenate(photon_dump))
summary = dict(
run=args.run, checkpoint=str(checkpoint_path), epoch=int(checkpoint["epoch"]),
n_frames=int(n_done), n_photons=int(accumulator.n_photons),
photons_per_frame=float(photons_per_frame.mean()),
occupancy_valid_region=float(photons_per_frame.mean() / n_valid_pixels),
occupancy_exposed=occupancy_exposed,
exposed_pixels=n_exposed, valid_pixels=n_valid_pixels,
threshold=threshold, merge_distance=merge_distance, halo=halo,
common_mode=str(conf.calibration.common_mode),
pedestal_source=str(conf.calibration.pedestal_source),
binning_factor=binning_factor,
crop=crop, frame_size=[frame_n_x, frame_n_y],
signal_files=[int(first_file), int(last_file), file_step],
self_flattening="streaming" if streaming else self_flattening,
flattening_classes=flattening_classes,
precision=precision,
subpixel_nonuniformity_before=(subpixel_nonuniformity(phase_before)
if phase_before is not None else None),
subpixel_nonuniformity_after=(subpixel_nonuniformity(phase_after)
if phase_after is not None else None),
seconds=elapsed, frames_per_second=float(n_done / elapsed),
)
with open(output_dir / "summary.json", "w") as handle:
json.dump(summary, handle, indent=1)
print(f"[infer] {n_done} frames in {elapsed:.0f}s "
f"({n_done/elapsed:.0f} frame/s), {accumulator.n_photons} photons, "
f"{summary['photons_per_frame']:.2f}/frame")
print(f"[infer] occupancy: {summary['occupancy_valid_region']*100:.3f}% averaged "
f"over the valid region, {occupancy_exposed*100:.3f}% inside the "
f"illuminated area ({n_exposed} of {n_valid_pixels} px)")
print(f"[infer] wrote {output_dir}")
if __name__ == "__main__":
main()
Reference in new issue View Git Blame Copy Permalink