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>
242 lines
11 KiB
Python
242 lines
11 KiB
Python
#!/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()
|