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>
364 lines
16 KiB
Python
364 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Compare poly GridNet models: MC metrics, energy response, XRF lines, efficiency
|
||
against the classical 3x3 sum, fake rate and speed, in one report.
|
||
|
||
python scripts/08_compare_models.py --out Results/<dir>/report \\
|
||
--model "5-30 keV [32,64,128,128]" Results/260901_poly5to30keV_Moench014_00 - \\
|
||
--model "3-30 keV [32,64,96,96]" Results/<exp> c96
|
||
|
||
Each model is (label, experiment dir, XRF tag). The tag is the --tag its XRF
|
||
runs were reconstructed with (`-` for none), so the reconstructions are read
|
||
from <inference.output_folder>/<run>[_<tag>]. Expected inputs per model:
|
||
eval_summary.json and energy_response.npz in the experiment dir (04, 07), and
|
||
AgMo_1_factor_10, CuCrCo_1_factor_10 and Pedestal_CuCrCo_1 reconstructions (05).
|
||
Whatever is missing shows up as n/a, with the reason listed at the end.
|
||
|
||
Writes report.md, report.json and report_response.png.
|
||
"""
|
||
import argparse
|
||
import json
|
||
import sys
|
||
import time
|
||
import traceback
|
||
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 models import build_model # noqa: E402
|
||
from xrf_analysis import classical_and_model, fit_lines, line_centroid # noqa: E402
|
||
|
||
import importlib # noqa: E402
|
||
sys.path.insert(0, str(HERE))
|
||
predict_fp16 = importlib.import_module("05_infer_measurement").predict_fp16
|
||
|
||
RESPONSE_ENERGIES = (3.5, 4.5, 5.4, 6.9, 8.0, 12.0, 17.4, 22.1, 27.5, 29.4)
|
||
XRF_FITS = {
|
||
"AgMo_1_factor_10": [
|
||
dict(alpha=["Mo_Ka"], window=(16.0, 18.9)),
|
||
dict(alpha=["Mo_Kb"], window=(18.75, 20.7)),
|
||
dict(alpha=["Ag_Ka"], window=(20.6, 23.6)),
|
||
],
|
||
"CuCrCo_1_factor_10": [
|
||
dict(alpha=["Cr_Ka", "Co_Ka", "Cu_Ka"],
|
||
beta={"Cr_Kb": "Cr_Ka", "Co_Kb": "Co_Ka", "Cu_Kb": "Cu_Ka"},
|
||
ratio_fixed={"Co_Kb": 0.135}, fit_stretch=False, window=(4.85, 9.6)),
|
||
],
|
||
}
|
||
XRF_LINES = ["Cr_Ka", "Co_Ka", "Cu_Ka", "Mo_Ka", "Mo_Kb", "Ag_Ka"]
|
||
DARK_RUN = "Pedestal_CuCrCo_1"
|
||
|
||
|
||
def fmt(value, spec, missing="n/a"):
|
||
if value is None or (isinstance(value, float) and not np.isfinite(value)):
|
||
return missing
|
||
return format(value, spec)
|
||
|
||
|
||
def model_meta(experiment, errors):
|
||
out = {}
|
||
try:
|
||
checkpoint = torch.load(experiment / "Models" / "best.pth", map_location="cpu",
|
||
weights_only=False)
|
||
conf = OmegaConf.create(checkpoint["conf"])
|
||
model = build_model(conf)
|
||
out.update(channels=list(conf.model.channels), R=model.rf_radius,
|
||
params=sum(p.numel() for p in model.parameters()),
|
||
best_epoch=int(checkpoint["epoch"]))
|
||
history = json.load(open(experiment / "history.json"))
|
||
out.update(epochs=len(history["val_loss"]),
|
||
best_val_loss=float(min(history["val_loss"])))
|
||
out["speed_fp16"] = speed(model)
|
||
except Exception as error: # noqa: BLE001
|
||
errors.append(f"meta: {error!r}")
|
||
return out
|
||
|
||
|
||
@torch.no_grad()
|
||
def speed(model, size=400, batch=32, n_iter=20):
|
||
torch.backends.cudnn.benchmark = True
|
||
model = model.cuda().eval().to(memory_format=torch.channels_last)
|
||
x = (torch.randn(batch, 1, size, size, device="cuda") * 0.05).to(
|
||
memory_format=torch.channels_last)
|
||
def run(): # exactly the FP16 path of 05_infer_measurement
|
||
return predict_fp16(model, x, 0.3)
|
||
|
||
for _ in range(5):
|
||
run()
|
||
torch.cuda.synchronize()
|
||
started = time.perf_counter()
|
||
for _ in range(n_iter):
|
||
run()
|
||
torch.cuda.synchronize()
|
||
return batch * n_iter / (time.perf_counter() - started)
|
||
|
||
|
||
def mc_metrics(experiment, errors):
|
||
try:
|
||
data = json.load(open(experiment / "eval_summary.json"))
|
||
return {key: value["summary"] for key, value in data.items()}
|
||
except Exception as error: # noqa: BLE001
|
||
errors.append(f"eval_summary: {error!r}")
|
||
return {}
|
||
|
||
|
||
def response(experiment, errors):
|
||
out = {}
|
||
try:
|
||
data = np.load(experiment / "energy_response.npz")
|
||
centre = data["table_centre"]
|
||
e_true, e_pred = data["e_true"], data["e_pred"]
|
||
for energy in RESPONSE_ENERGIES:
|
||
if not centre[0] - 0.13 <= energy <= centre[-1] + 0.13:
|
||
continue
|
||
i = int(np.argmin(np.abs(centre - energy)))
|
||
near = np.abs(e_true - centre[i]) < 0.125
|
||
out[energy] = dict(
|
||
shift_meV=float((data["table_core_mu"][i] - centre[i]) * 1000),
|
||
core_sigma=float(data["table_core_sigma"][i]),
|
||
tail_percent=float((np.abs(e_pred[near] - e_true[near]) > 1).mean() * 100),
|
||
efficiency=float(data["table_efficiency"][i]))
|
||
out["_curve"] = dict(centre=centre.tolist(),
|
||
shift=(data["table_core_mu"] - centre).tolist(),
|
||
sigma=data["table_core_sigma"].tolist(),
|
||
efficiency=data["table_efficiency"].tolist())
|
||
except Exception as error: # noqa: BLE001
|
||
errors.append(f"energy_response: {error!r}")
|
||
return out
|
||
|
||
|
||
def xrf_lines(folder_of, errors):
|
||
out = {}
|
||
for run, groups in XRF_FITS.items():
|
||
try:
|
||
centres, counts = np.load(folder_of(run) / "energy_spectrum.npy")
|
||
for group in groups:
|
||
fit = fit_lines(centres, counts, group["alpha"], group.get("beta"),
|
||
window=group["window"],
|
||
beta_ratio_fixed=group.get("ratio_fixed"),
|
||
fit_stretch=group.get("fit_stretch", True))
|
||
for name in group["alpha"]:
|
||
line = fit["lines"][name]
|
||
out[name] = dict(shift_meV=line["shift"] * 1000, sigma=line["sigma"],
|
||
fwhm=line["fwhm"], chi2_ndf=fit["chi2_ndf"])
|
||
except Exception as error: # noqa: BLE001
|
||
errors.append(f"xrf {run}: {error!r}")
|
||
return out
|
||
|
||
|
||
def classical_lines(config, errors):
|
||
"""Line fits of the classical 3x3 spectra (model independent, cached)."""
|
||
out = {}
|
||
edges = np.arange(0.0, 40.025, 0.05)
|
||
centres = 0.5 * (edges[1:] + edges[:-1])
|
||
for run, groups in XRF_FITS.items():
|
||
try:
|
||
reference = classical_and_model(config, run, n_frames=10_000)
|
||
counts = np.histogram(reference["classical"][:, 3], edges)[0]
|
||
for group in groups:
|
||
fit = fit_lines(centres, counts, group["alpha"], group.get("beta"),
|
||
window=group["window"],
|
||
beta_ratio_fixed=group.get("ratio_fixed"),
|
||
fit_stretch=group.get("fit_stretch", True), tails=False)
|
||
for name in group["alpha"]:
|
||
line = fit["lines"][name]
|
||
out[name] = dict(shift_meV=line["shift"] * 1000, sigma=line["sigma"],
|
||
fwhm=line["fwhm"])
|
||
except Exception as error: # noqa: BLE001
|
||
errors.append(f"classical {run}: {error!r}")
|
||
return out
|
||
|
||
|
||
def efficiency(config, checkpoint, tag, errors):
|
||
out = {}
|
||
for run in XRF_FITS:
|
||
try:
|
||
reference = classical_and_model(config, run, n_frames=10_000,
|
||
checkpoint=checkpoint, tag=tag)
|
||
classical, grid = reference["classical"], reference["grid"]
|
||
found = reference["classical_match"] >= 0
|
||
isolated = reference["isolation_classical"] > 4.0
|
||
energy = classical[:, 3]
|
||
deposited = reference["balance"][:, 1].sum()
|
||
out[run] = dict(
|
||
isolated_above5=float(found[isolated & (energy >= 5)].mean()),
|
||
isolated_3to5=float(found[isolated & (energy >= 3) & (energy < 5)].mean()),
|
||
count_ratio=float(len(grid) / len(classical)),
|
||
energy_balance=float(grid[:, 3].sum() / deposited))
|
||
except Exception as error: # noqa: BLE001
|
||
errors.append(f"efficiency {run}: {error!r}")
|
||
return out
|
||
|
||
|
||
def run_summaries(folder_of, errors):
|
||
out = {}
|
||
for run in list(XRF_FITS) + [DARK_RUN]:
|
||
try:
|
||
summary = json.load(open(folder_of(run) / "summary.json"))
|
||
out[run] = dict(n_photons=summary["n_photons"], n_frames=summary["n_frames"],
|
||
subpixel=summary.get("subpixel_nonuniformity_before"),
|
||
fps=summary.get("frames_per_second"),
|
||
precision=summary.get("precision", "fp32"))
|
||
except Exception as error: # noqa: BLE001
|
||
errors.append(f"summary {run}: {error!r}")
|
||
return out
|
||
|
||
|
||
def write_report(models, classical, out_dir):
|
||
labels = list(models)
|
||
lines = [f"# Poly model comparison ({time.strftime('%Y-%m-%d %H:%M')})", ""]
|
||
|
||
lines += ["## Models", "",
|
||
"| model | channels | R | params | best epoch / epochs | FP16 frame/s (400×400, model only) |",
|
||
"|---|---|---|---|---|---|"]
|
||
for label in labels:
|
||
m = models[label]["meta"]
|
||
lines.append(f"| {label} | {m.get('channels', 'n/a')} | {m.get('R', 'n/a')} | "
|
||
f"{m.get('params', 'n/a')} | {m.get('best_epoch', 'n/a')} / "
|
||
f"{m.get('epochs', 'n/a')} | {fmt(m.get('speed_fp16'), '.0f')} |")
|
||
|
||
lines += ["", "## MC, fixed-occupancy scan (threshold 0.3)", "",
|
||
"| model | occupancy | efficiency | purity | RMS [µm] | σ68 [µm] | E σ68 [keV] | sub-pixel |",
|
||
"|---|---|---|---|---|---|---|---|"]
|
||
for label in labels:
|
||
for occupancy, s in models[label]["mc"].items():
|
||
if occupancy not in ("1%", "5%", "10%"):
|
||
continue
|
||
lines.append(f"| {label} | {occupancy} | {s['efficiency']*100:.2f} % | "
|
||
f"{s['purity']*100:.2f} % | {s['rms_1d_um']:.3f} | "
|
||
f"{s['sigma68_1d_um']:.3f} | {s.get('energy_sigma68_keV', float('nan')):.3f} | "
|
||
f"{s['subpixel_nonuniformity']*100:.1f} % |")
|
||
|
||
lines += ["", "## MC energy response (occupancy ≤ 2.5 %)", "",
|
||
"core σ [keV] / line shift g(E) − E [meV] / tails |ΔE| > 1 keV [%] / efficiency [%]", "",
|
||
"| E [keV] | " + " | ".join(labels) + " |", "|---|" + "---|" * len(labels)]
|
||
for energy in RESPONSE_ENERGIES:
|
||
cells = []
|
||
for label in labels:
|
||
r = models[label]["response"].get(energy)
|
||
cells.append("—" if r is None else
|
||
f"{r['core_sigma']:.3f} / {r['shift_meV']:+.0f} / "
|
||
f"{r['tail_percent']:.1f} / {r['efficiency']*100:.1f}")
|
||
lines.append(f"| {energy} | " + " | ".join(cells) + " |")
|
||
|
||
lines += ["", "## XRF lines (AgMo, CuCrCo slit factor 10)", "",
|
||
"fitted − tabulated [meV] / core σ [keV] / FWHM [keV]", "",
|
||
"| line | " + " | ".join(labels + ["classical 3×3"]) + " |",
|
||
"|---|" + "---|" * (len(labels) + 1)]
|
||
for name in XRF_LINES:
|
||
cells = []
|
||
for source in [models[label]["xrf"] for label in labels] + [classical]:
|
||
r = source.get(name)
|
||
cells.append("n/a" if r is None else
|
||
f"{r['shift_meV']:+.0f} / {r['sigma']:.3f} / {r['fwhm']:.2f}")
|
||
lines.append(f"| {name} ({line_centroid(name):.2f}) | " + " | ".join(cells) + " |")
|
||
|
||
lines += ["", "## Efficiency against the classical finder, fake rate, sub-pixel", "",
|
||
"| model | run | isolated ≥ 5 keV found | isolated 3–5 keV found | grid / classical | energy balance | sub-pixel excess rms |",
|
||
"|---|---|---|---|---|---|---|"]
|
||
for label in labels:
|
||
for run in XRF_FITS:
|
||
e = models[label]["efficiency"].get(run, {})
|
||
s = models[label]["runs"].get(run, {})
|
||
lines.append(f"| {label} | {run} | {fmt(e.get('isolated_above5', np.nan)*100, '.2f')} % | "
|
||
f"{fmt(e.get('isolated_3to5', np.nan)*100, '.1f')} % | "
|
||
f"{fmt(e.get('count_ratio'), '.3f')} | {fmt(e.get('energy_balance'), '.4f')} | "
|
||
f"{fmt((s.get('subpixel') or np.nan)*100, '.1f')} % |")
|
||
lines += ["", "| model | dark run photons | frames |", "|---|---|---|"]
|
||
for label in labels:
|
||
s = models[label]["runs"].get(DARK_RUN, {})
|
||
lines.append(f"| {label} | {s.get('n_photons', 'n/a')} | {s.get('n_frames', 'n/a')} |")
|
||
|
||
problems = [(label, e) for label in labels for e in models[label]["errors"]]
|
||
if problems:
|
||
lines += ["", "## Missing pieces", ""]
|
||
lines += [f"- {label}: {error}" for label, error in problems]
|
||
lines += ["", "Plots: `report_response.png`."]
|
||
(out_dir / "report.md").write_text("\n".join(lines) + "\n")
|
||
|
||
|
||
def plot_response(models, out_dir):
|
||
figure, axes = plt.subplots(1, 3, figsize=(17, 4.4))
|
||
for label, model in models.items():
|
||
curve = model["response"].get("_curve")
|
||
if not curve:
|
||
continue
|
||
centre = np.asarray(curve["centre"])
|
||
axes[0].plot(centre, np.asarray(curve["shift"]) * 1000, lw=1.2, label=label)
|
||
axes[1].plot(centre, curve["sigma"], lw=1.2, label=label)
|
||
axes[2].plot(centre, np.asarray(curve["efficiency"]) * 100, lw=1.2, label=label)
|
||
axes[0].set_ylim(-200, 200)
|
||
for axis, ylabel, title in zip(axes, ("g(E) − E [meV]", "core σ [keV]", "efficiency [%]"),
|
||
("line shift", "energy resolution", "detection efficiency")):
|
||
axis.set_xlabel("E_true [keV]"); axis.set_ylabel(ylabel); axis.set_title(title)
|
||
axis.grid(alpha=.3); axis.legend(fontsize=7)
|
||
figure.tight_layout()
|
||
figure.savefig(out_dir / "report_response.png", dpi=130)
|
||
plt.close(figure)
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--model", nargs=3, action="append", required=True,
|
||
metavar=("LABEL", "EXPERIMENT", "TAG"))
|
||
parser.add_argument("--config", default=str(HERE.parent / "Configs" /
|
||
"infer_measurement_poly_Moench014.yaml"))
|
||
parser.add_argument("--out", required=True)
|
||
args = parser.parse_args()
|
||
|
||
out_dir = Path(args.out)
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
output_folder = Path(OmegaConf.load(args.config).inference.output_folder)
|
||
models = {}
|
||
for label, experiment, tag in args.model:
|
||
experiment = Path(experiment)
|
||
tag = None if tag in ("-", "") else tag
|
||
errors = []
|
||
print(f"[compare] {label}: {experiment} (tag {tag})", flush=True)
|
||
|
||
def folder_of(run, tag=tag):
|
||
return output_folder / (run + (f"_{tag}" if tag else ""))
|
||
|
||
models[label] = dict(
|
||
meta=model_meta(experiment, errors), mc=mc_metrics(experiment, errors),
|
||
response=response(experiment, errors), xrf=xrf_lines(folder_of, errors),
|
||
efficiency=efficiency(args.config, str(experiment / "Models" / "best.pth"),
|
||
tag, errors),
|
||
runs=run_summaries(folder_of, errors), errors=errors)
|
||
for error in errors:
|
||
print(f" missing: {error}", flush=True)
|
||
classical_errors = []
|
||
classical = classical_lines(args.config, classical_errors)
|
||
if classical_errors:
|
||
models[next(iter(models))]["errors"] += classical_errors
|
||
|
||
write_report(models, classical, out_dir)
|
||
plot_response(models, out_dir)
|
||
|
||
def clean(value):
|
||
if isinstance(value, dict):
|
||
return {str(k): clean(v) for k, v in value.items()}
|
||
if isinstance(value, (list, tuple)):
|
||
return [clean(v) for v in value]
|
||
if isinstance(value, (np.floating, np.integer)):
|
||
return value.item()
|
||
return value
|
||
with open(out_dir / "report.json", "w") as handle:
|
||
json.dump(clean(dict(models=models, classical=classical)), handle, indent=1)
|
||
print(f"[compare] wrote {out_dir/'report.md'}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
main()
|
||
except Exception: # noqa: BLE001
|
||
traceback.print_exc()
|
||
sys.exit(1)
|