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,241 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
WP4/5.3 -- MC energy response of a polychromatic grid model.
|
||||
|
||||
python scripts/07_energy_response.py \
|
||||
--checkpoint Results/260901_poly5to30keV_Moench014_00/Models/best.pth
|
||||
|
||||
The energy head is a regression trained on a flat 5-30 keV spectrum, so it
|
||||
returns something close to the posterior mean: every prediction is pulled
|
||||
towards the centre of the training range, most strongly within ~1 keV of the
|
||||
edges where the prior is one-sided. A spectral line at E therefore shows up at
|
||||
g(E) != E. This script measures g on MC in fine true-energy bins -- the
|
||||
position of the Gaussian core of E_pred, i.e. what a line fit on measured data
|
||||
returns -- together with the core width and the detection efficiency, and
|
||||
writes
|
||||
|
||||
Results/<exp>/energy_response.npz per-bin table + all matched pairs
|
||||
Results/<exp>/Plots/energy_response.png
|
||||
|
||||
The correction applied to measured photons is E_corr = g^-1(E_pred) (see
|
||||
src/xrf_analysis.py: EnergyCorrection). Inverting the forward curve, rather
|
||||
than regressing E_true on E_pred, is what puts a line back at its energy: the
|
||||
reverse regression is the posterior mean again and would shrink a second time.
|
||||
|
||||
Only frames with occupancy <= --max-occupancy are used, so that the response
|
||||
matches the low-flux XRF runs it is applied to (0.7-1.2 %).
|
||||
"""
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
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
|
||||
sys.path.insert(0, str(HERE.parent / "src"))
|
||||
from decode import decode # noqa: E402
|
||||
from metrics import greedy_match # noqa: E402
|
||||
from models import build_model # noqa: E402
|
||||
from xrf_analysis import fit_gaussian_core # noqa: E402
|
||||
|
||||
DEFAULT_DATA = [
|
||||
"{folder}/eval_scan/occ0p5pct.h5",
|
||||
"{folder}/eval_scan/occ1pct.h5",
|
||||
"{folder}/eval_scan/occ2pct.h5",
|
||||
"{folder}/test/shard_00.h5",
|
||||
"{folder}/test/shard_01.h5",
|
||||
]
|
||||
|
||||
|
||||
def load_frames(path, max_occupancy):
|
||||
import h5py
|
||||
with h5py.File(path, "r") as handle:
|
||||
occupancy = handle["occupancy"][:]
|
||||
keep = np.nonzero(occupancy <= max_occupancy)[0]
|
||||
frames = handle["frames"][:][keep]
|
||||
photons = handle["photons"][:]
|
||||
pointer = handle["frame_ptr"][:]
|
||||
truth = [photons[pointer[i]:pointer[i + 1]] for i in keep]
|
||||
return frames, truth, occupancy[keep]
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def collect(model, conf, files, halo, max_occupancy, batch_size=1024):
|
||||
"""Matched (E_true, E_pred) pairs, and every valid truth photon with a flag."""
|
||||
input_scale = float(conf.data.input_scale)
|
||||
energy_scale = float(conf.data.energy_scale)
|
||||
threshold = float(conf.eval.threshold)
|
||||
merge_distance = float(conf.eval.merge_distance)
|
||||
match_radius = float(conf.eval.match_radius)
|
||||
e_true, e_pred, truth_energy, truth_found, fake_energy = [], [], [], [], []
|
||||
n_frames = 0
|
||||
for path in files:
|
||||
frames, truth, occupancy = load_frames(path, max_occupancy)
|
||||
size = frames.shape[-1]
|
||||
print(f"[response] {os.path.basename(path)}: {len(frames)} frames "
|
||||
f"with occupancy <= {max_occupancy*100:.1f} %", flush=True)
|
||||
n_frames += len(frames)
|
||||
for start in range(0, len(frames), batch_size):
|
||||
batch = torch.from_numpy(frames[start:start + batch_size] / input_scale)
|
||||
prediction = model(batch.float().unsqueeze(1).cuda()).float()
|
||||
points = decode(prediction, threshold=threshold,
|
||||
merge_distance=merge_distance, halo=halo,
|
||||
predict_energy=True, energy_scale=energy_scale)
|
||||
for predicted, true in zip(points, truth[start:start + batch_size]):
|
||||
predicted = predicted.numpy().astype(np.float64)
|
||||
inside = ((true[:, 0] >= halo) & (true[:, 0] < size - halo) &
|
||||
(true[:, 1] >= halo) & (true[:, 1] < size - halo))
|
||||
true = true[inside].astype(np.float64)
|
||||
found = np.zeros(len(true), bool)
|
||||
matched_prediction = np.zeros(len(predicted), bool)
|
||||
if len(true) and len(predicted):
|
||||
index_pred, index_true = greedy_match(predicted[:, :2], true[:, :2],
|
||||
match_radius)
|
||||
found[index_true] = True
|
||||
matched_prediction[index_pred] = True
|
||||
e_true.append(true[index_true, 2])
|
||||
e_pred.append(predicted[index_pred, 2])
|
||||
truth_energy.append(true[:, 2])
|
||||
truth_found.append(found)
|
||||
fake_energy.append(predicted[~matched_prediction, 2])
|
||||
return (np.concatenate(e_true), np.concatenate(e_pred),
|
||||
np.concatenate(truth_energy), np.concatenate(truth_found),
|
||||
np.concatenate(fake_energy), n_frames)
|
||||
|
||||
|
||||
def response_table(e_true, e_pred, truth_energy, truth_found, edges):
|
||||
rows = {key: [] for key in ("centre", "n", "efficiency", "median", "mean",
|
||||
"core_mu", "core_mu_err", "core_sigma",
|
||||
"core_sigma_err", "sigma68", "rms")}
|
||||
for low, high in zip(edges[:-1], edges[1:]):
|
||||
selected = (e_true >= low) & (e_true < high)
|
||||
in_bin = (truth_energy >= low) & (truth_energy < high)
|
||||
residual_reference = e_true[selected]
|
||||
predicted = e_pred[selected]
|
||||
core = fit_gaussian_core(predicted)
|
||||
rows["centre"].append(0.5 * (low + high))
|
||||
rows["n"].append(int(selected.sum()))
|
||||
rows["efficiency"].append(float(truth_found[in_bin].mean()) if in_bin.any()
|
||||
else np.nan)
|
||||
rows["median"].append(float(np.median(predicted)) if len(predicted) else np.nan)
|
||||
rows["mean"].append(float(predicted.mean()) if len(predicted) else np.nan)
|
||||
rows["core_mu"].append(core["mu"])
|
||||
rows["core_mu_err"].append(core["mu_err"])
|
||||
rows["core_sigma"].append(core["sigma"])
|
||||
rows["core_sigma_err"].append(core["sigma_err"])
|
||||
residual = predicted - residual_reference
|
||||
rows["sigma68"].append(float(np.percentile(np.abs(residual - np.median(residual)),
|
||||
68.27)) if len(residual) else np.nan)
|
||||
rows["rms"].append(float(residual.std()) if len(residual) else np.nan)
|
||||
return {key: np.asarray(value) for key, value in rows.items()}
|
||||
|
||||
|
||||
def plot(table, e_true, e_pred, output_path, title):
|
||||
centre = table["centre"]
|
||||
low, high = float(np.floor(e_true.min())), float(np.ceil(e_true.max()))
|
||||
figure, axes = plt.subplots(2, 2, figsize=(12, 8.5))
|
||||
axis = axes[0, 0]
|
||||
histogram, x_edges, y_edges = np.histogram2d(
|
||||
e_true, e_pred, bins=[int(5 * (high - low)), int(5 * (high - low + 5))],
|
||||
range=[[low, high], [low - 2.5, high + 2.5]])
|
||||
axis.imshow(np.log10(histogram.T + 1), origin="lower", aspect="auto",
|
||||
extent=[x_edges[0], x_edges[-1], y_edges[0], y_edges[-1]],
|
||||
cmap="viridis")
|
||||
axis.plot([low, high], [low, high], "w--", lw=.8)
|
||||
axis.set_xlabel("E_true [keV]")
|
||||
axis.set_ylabel("E_pred [keV]")
|
||||
axis.set_title("matched photons, log10 counts")
|
||||
|
||||
axis = axes[0, 1]
|
||||
axis.errorbar(centre, table["core_mu"] - centre, table["core_mu_err"], fmt=".",
|
||||
ms=3, label="Gaussian core position")
|
||||
axis.plot(centre, table["median"] - centre, lw=1, label="median")
|
||||
axis.plot(centre, table["mean"] - centre, lw=1, label="mean")
|
||||
axis.axhline(0, color="k", lw=.6)
|
||||
axis.set_xlabel("E_true [keV]")
|
||||
axis.set_ylabel("E_pred - E_true [keV]")
|
||||
axis.set_title("g(E) - E: pull towards the centre")
|
||||
axis.legend(fontsize=8)
|
||||
axis.grid(alpha=.3)
|
||||
|
||||
axis = axes[1, 0]
|
||||
axis.plot(centre, table["core_sigma"], ".", ms=4, label=r"core $\sigma$")
|
||||
axis.plot(centre, table["sigma68"], lw=1, label=r"$\sigma_{68}$ of residual")
|
||||
axis.plot(centre, table["rms"], lw=1, label="rms of residual")
|
||||
axis.set_xlabel("E_true [keV]")
|
||||
axis.set_ylabel("keV")
|
||||
axis.set_ylim(0, 1.6)
|
||||
axis.set_title("energy resolution")
|
||||
axis.legend(fontsize=8)
|
||||
axis.grid(alpha=.3)
|
||||
|
||||
axis = axes[1, 1]
|
||||
axis.plot(centre, table["efficiency"] * 100, ".", ms=4)
|
||||
axis.set_xlabel("E_true [keV]")
|
||||
axis.set_ylabel("efficiency [%]")
|
||||
axis.set_title("detection efficiency")
|
||||
axis.grid(alpha=.3)
|
||||
figure.suptitle(title)
|
||||
figure.tight_layout()
|
||||
figure.savefig(output_path, dpi=130)
|
||||
plt.close(figure)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--checkpoint", required=True)
|
||||
parser.add_argument("--data", nargs="*", default=None)
|
||||
parser.add_argument("--max-occupancy", type=float, default=0.025)
|
||||
parser.add_argument("--bin-width", type=float, default=0.25)
|
||||
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
|
||||
files = args.data or [pattern.format(folder=conf.data.folder)
|
||||
for pattern in DEFAULT_DATA]
|
||||
files = [f for pattern in files for f in sorted(glob.glob(pattern))]
|
||||
|
||||
e_true, e_pred, truth_energy, truth_found, fake_energy, n_frames = collect(
|
||||
model, conf, files, halo, args.max_occupancy)
|
||||
# the true-energy range of the MC set: 5-30 keV or 3-30 keV
|
||||
edges = np.arange(np.floor(truth_energy.min()), np.ceil(truth_energy.max()) + 1e-9,
|
||||
args.bin_width)
|
||||
table = response_table(e_true, e_pred, truth_energy, truth_found, edges)
|
||||
|
||||
experiment_dir = Path(args.checkpoint).resolve().parent.parent
|
||||
output = experiment_dir / "energy_response.npz"
|
||||
np.savez_compressed(
|
||||
output, edges=edges, **{f"table_{k}": v for k, v in table.items()},
|
||||
e_true=e_true.astype(np.float32), e_pred=e_pred.astype(np.float32),
|
||||
truth_energy=truth_energy.astype(np.float32), truth_found=truth_found,
|
||||
fake_energy=fake_energy.astype(np.float32),
|
||||
checkpoint=str(Path(args.checkpoint).resolve()), epoch=int(checkpoint["epoch"]),
|
||||
max_occupancy=args.max_occupancy, n_frames=n_frames,
|
||||
files=np.array([str(f) for f in files]))
|
||||
(experiment_dir / "Plots").mkdir(exist_ok=True)
|
||||
plot(table, e_true, e_pred, experiment_dir / "Plots" / "energy_response.png",
|
||||
f"{experiment_dir.name} (epoch {checkpoint['epoch']}), "
|
||||
f"occupancy <= {args.max_occupancy*100:.1f} %, {len(e_true)} matched photons")
|
||||
print(f"[response] {n_frames} frames, {len(truth_energy)} truth photons, "
|
||||
f"{len(e_true)} matched, {len(fake_energy)} unmatched predictions")
|
||||
print(f"[response] efficiency {truth_found.mean()*100:.2f} %")
|
||||
for energy in [e for e in (3.5, 4.5, 5.5, 7.0, 8.0, 12.0, 17.5, 22.0, 26.0, 29.5)
|
||||
if edges[0] < e < edges[-1]]:
|
||||
i = np.searchsorted(edges, energy) - 1
|
||||
print(f" E={table['centre'][i]:6.3f} g-E={table['core_mu'][i]-table['centre'][i]:+.3f}"
|
||||
f" core sigma={table['core_sigma'][i]:.3f} eff={table['efficiency'][i]*100:.2f} %")
|
||||
print(f"[response] wrote {output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in new issue
Block a user