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,138 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Self-tests for the grid formulation. Run after touching src/.
|
||||
|
||||
python scripts/00_selftest.py
|
||||
|
||||
1. encode -> decode round trip: the coordinate convention is self-consistent.
|
||||
2. D4 augmentation keeps frame and truth coordinates aligned.
|
||||
3. Translation equivariance: a model trained on SxS frames gives *identical*
|
||||
predictions when the same content is embedded in a larger frame -- the claim
|
||||
that D1/D2 rest on. If this test fails, train-small/infer-large is invalid.
|
||||
4. Position-space merging removes duplicates but keeps genuinely distinct
|
||||
photons in adjacent cells -- the property the 3x3 max-pool NMS lacked.
|
||||
"""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(HERE.parent / "src"))
|
||||
from datasets import apply_d4, encode_targets # noqa: E402
|
||||
from decode import decode # noqa: E402
|
||||
from models import GridNet, valid_mask # noqa: E402
|
||||
|
||||
all_passed = True
|
||||
|
||||
|
||||
def check(name, condition, detail=""):
|
||||
global all_passed
|
||||
all_passed &= bool(condition)
|
||||
print(f" [{'PASS' if condition else 'FAIL'}] {name}"
|
||||
f"{' ' + detail if detail else ''}")
|
||||
|
||||
|
||||
print("1. encode -> decode round trip")
|
||||
frame_size, energy_scale = 32, 12.0
|
||||
rng = np.random.default_rng(0)
|
||||
n_photons = 40
|
||||
photons = np.zeros((n_photons, 4), np.float32)
|
||||
photons[:, 0] = rng.uniform(0, frame_size, n_photons)
|
||||
photons[:, 1] = rng.uniform(0, frame_size, n_photons)
|
||||
photons[:, 2] = rng.uniform(5, 30, n_photons)
|
||||
cells = (photons[:, 1].astype(int) * frame_size + photons[:, 0].astype(int))
|
||||
_, unique_index = np.unique(cells, return_index=True) # drop same-cell duplicates
|
||||
photons = photons[np.sort(unique_index)]
|
||||
|
||||
objectness, offset, energy = encode_targets(photons, frame_size, energy_scale)
|
||||
prediction = torch.zeros(1, 4, frame_size, frame_size)
|
||||
prediction[0, 0] = torch.from_numpy(
|
||||
np.where(objectness > 0, 10.0, -10.0).astype(np.float32))
|
||||
prediction[0, 1:3] = torch.from_numpy(offset)
|
||||
prediction[0, 3:4] = torch.from_numpy(energy)
|
||||
|
||||
points = decode(prediction, threshold=0.5, merge_distance=0.0, halo=0,
|
||||
energy_scale=energy_scale)[0].numpy()
|
||||
check("photon count", len(points) == len(photons), f"{len(points)} vs {len(photons)}")
|
||||
order_points = np.lexsort((points[:, 0], points[:, 1]))
|
||||
order_truth = np.lexsort((photons[:, 0], photons[:, 1]))
|
||||
position_error = np.abs(points[order_points, :2] - photons[order_truth, :2]).max()
|
||||
energy_error = np.abs(points[order_points, 2] - photons[order_truth, 2]).max()
|
||||
check("position round trip", position_error < 1e-4,
|
||||
f"max |dx| = {position_error:.2e} px")
|
||||
check("energy round trip", energy_error < 1e-3,
|
||||
f"max |dE| = {energy_error:.2e} keV")
|
||||
|
||||
print("2. D4 augmentation consistency")
|
||||
frame = np.zeros((frame_size, frame_size), np.float32)
|
||||
for x, y, photon_energy, _ in photons: # put a delta at each photon
|
||||
frame[int(y), int(x)] += photon_energy
|
||||
worst_mismatch = 0.0
|
||||
for operation in range(8):
|
||||
transformed_frame, transformed_positions = apply_d4(frame, photons[:, :2], operation)
|
||||
for (x, y), (x0, y0) in zip(transformed_positions, photons[:, :2]):
|
||||
original = frame[int(y0), int(x0)]
|
||||
worst_mismatch = max(worst_mismatch,
|
||||
abs(transformed_frame[int(y), int(x)] - original))
|
||||
check("charge follows the truth under all 8 ops", worst_mismatch < 1e-5,
|
||||
f"max mismatch {worst_mismatch:.2e}")
|
||||
|
||||
print("3. translation equivariance (train small -> infer large)")
|
||||
torch.manual_seed(0)
|
||||
network = GridNet(channels=(32, 64, 64, 64)).eval()
|
||||
radius = network.rf_radius
|
||||
small_frame = torch.randn(1, 1, frame_size, frame_size) * 0.05
|
||||
small_frame[0, 0, 12:15, 9:12] += torch.tensor([[.1, .3, .1],
|
||||
[.4, 3.0, .5],
|
||||
[.1, .4, .1]])
|
||||
large_frame = torch.randn(1, 1, 3 * frame_size, 3 * frame_size) * 0.05
|
||||
row_offset, column_offset = 40, 27
|
||||
large_frame[0, 0, row_offset:row_offset + frame_size,
|
||||
column_offset:column_offset + frame_size] = small_frame[0, 0]
|
||||
with torch.no_grad():
|
||||
small_prediction = network(small_frame)
|
||||
large_prediction = network(large_frame)
|
||||
embedded = large_prediction[:, :, row_offset:row_offset + frame_size,
|
||||
column_offset:column_offset + frame_size]
|
||||
mask = valid_mask(frame_size, frame_size, radius)
|
||||
inside_difference = (small_prediction - embedded)[:, :, mask].abs().max().item()
|
||||
check(f"identical inside the halo (R={radius})", inside_difference < 1e-4,
|
||||
f"max |diff| = {inside_difference:.2e}")
|
||||
# R must also be tight: the ring one pixel *inside* the halo must already differ
|
||||
inner_ring = valid_mask(frame_size, frame_size, radius - 1) & ~mask
|
||||
ring_difference = (small_prediction - embedded)[:, :, inner_ring].abs().max().item()
|
||||
check(f"ring at distance R-1={radius-1} already differs "
|
||||
f"(R is tight, not over-conservative)",
|
||||
ring_difference > 1e-6, f"|diff| = {ring_difference:.2e}")
|
||||
|
||||
print("4. position-space merging")
|
||||
# two adjacent cells pointing at the same physical place -> one photon
|
||||
duplicate = torch.full((1, 4, 8, 8), -10.0)
|
||||
duplicate[0, 1:3] = 0.0
|
||||
duplicate[0, 3] = 1.0
|
||||
duplicate[0, 0, 4, 4] = 6.0
|
||||
duplicate[0, 1, 4, 4] = 0.45 # x = 4.95
|
||||
duplicate[0, 0, 4, 5] = 5.0
|
||||
duplicate[0, 1, 4, 5] = -0.45 # x = 5.05, same point
|
||||
merged = decode(duplicate, threshold=0.5, merge_distance=0.3, halo=0,
|
||||
energy_scale=12.0)[0]
|
||||
check("duplicate in adjacent cells merged to one", len(merged) == 1,
|
||||
f"got {len(merged)} points")
|
||||
|
||||
# two adjacent cells pointing at genuinely different places -> two photons
|
||||
distinct = torch.full((1, 4, 8, 8), -10.0)
|
||||
distinct[0, 1:3] = 0.0
|
||||
distinct[0, 3] = 1.0
|
||||
distinct[0, 0, 4, 4] = 6.0
|
||||
distinct[0, 1, 4, 4] = -0.3 # x = 4.2
|
||||
distinct[0, 0, 4, 5] = 5.0
|
||||
distinct[0, 1, 4, 5] = 0.3 # x = 5.8, 1.6 px away
|
||||
kept = decode(distinct, threshold=0.5, merge_distance=0.3, halo=0,
|
||||
energy_scale=12.0)[0]
|
||||
check("distinct photons in adjacent cells both kept", len(kept) == 2,
|
||||
f"got {len(kept)} points (3x3 max-pool NMS would have given 1)")
|
||||
|
||||
print("\n" + ("ALL TESTS PASSED" if all_passed else "SOME TESTS FAILED"))
|
||||
sys.exit(0 if all_passed else 1)
|
||||
@@ -0,0 +1,178 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
WP1.1 -- build the alpha(E, t) / beta(E, t) look-up table for continuous-energy
|
||||
charge-cloud sampling.
|
||||
|
||||
python scripts/01_build_charge_lut.py --config Configs/charge_lut_Moench014.yaml
|
||||
|
||||
For every energy the charge-transport simulation gives the transverse carrier
|
||||
distribution in 512 depth bins. Each bin is fitted with a generalised normal
|
||||
distribution exactly as ``McGeneration/McGenerator.py::parameterization`` does,
|
||||
so the shapes stay identical to the ones validated against measurement.
|
||||
|
||||
What is new is what happens next. ``parameterization`` compresses each energy
|
||||
into 9 numbers by fitting alpha(t) with a 4-parameter polynomial and beta(t)
|
||||
with a 5-parameter form; those 9 numbers are strongly correlated and
|
||||
individually unstable, so interpolating them in energy produces unphysical
|
||||
curves in between. Here the fitted *curves* are tabulated instead: the drift
|
||||
time depends only on the sensor, bias and temperature, so the depth grid is the
|
||||
same at every energy and (energy, drift time) is already a regular grid --
|
||||
alpha and beta can simply be looked up and interpolated bilinearly, with no
|
||||
functional form in between.
|
||||
|
||||
Output ``data/charge_lut_<sensor>_<hv>V.npz``:
|
||||
energies (nE,) keV
|
||||
z0 (nZ,) um, conversion depth measured from the entrance
|
||||
drift_time (nZ,) ns
|
||||
alpha, beta (nE, nZ) per-depth GGD fit
|
||||
alpha_err, beta_err fit errors
|
||||
n_carriers (nE, nZ) entries behind each fit
|
||||
"""
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from multiprocessing import Pool
|
||||
|
||||
import numpy as np
|
||||
from array import array
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from ROOT import TF1, TH1F, TH1, gROOT
|
||||
gROOT.SetBatch(True)
|
||||
TH1.AddDirectory(False)
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
sys.path.insert(0, os.path.join(HERE, "..", "src"))
|
||||
from charge_model import drift_time_allpix2 # noqa: E402
|
||||
|
||||
_CONF = None
|
||||
|
||||
|
||||
def fit_generalised_normal(coords, index):
|
||||
"""The per-depth GGD fit of McGenerator.parameterization, unchanged."""
|
||||
rms = np.std(coords)
|
||||
if not np.isfinite(rms) or rms <= 0:
|
||||
return np.nan, np.nan, np.nan, np.nan
|
||||
bin_width = rms / 20
|
||||
n_bins = int(100 / bin_width)
|
||||
histogram = TH1F(f"_h1_{index}", "", n_bins, -rms * 5, rms * 5)
|
||||
histogram.FillN(len(coords), array("d", coords), array("d", [1] * len(coords)))
|
||||
function = TF1(
|
||||
f"f1_{index}",
|
||||
"[0]/(2*[1]*TMath::Gamma(1/[0])) * exp(-pow(TMath::Abs(x)/[1], [0]))",
|
||||
-5 * rms, 5 * rms) # par[0]: beta, par[1]: alpha
|
||||
function.SetParLimits(0, 2, 5)
|
||||
function.SetParLimits(1, .5, 30)
|
||||
function.SetParameters(2., rms * 1.5)
|
||||
histogram.Scale(histogram.GetNbinsX() / (10 * rms) / histogram.Integral())
|
||||
histogram.Fit(function, "Q")
|
||||
beta, alpha = function.GetParameter(0), function.GetParameter(1)
|
||||
beta_err, alpha_err = function.GetParError(0), function.GetParError(1)
|
||||
del histogram, function
|
||||
return alpha, beta, alpha_err, beta_err
|
||||
|
||||
|
||||
def process_energy(energy_keV):
|
||||
conf = _CONF
|
||||
folder = conf.simulation.folder.format(sensor=conf.sensor.code,
|
||||
hv=conf.sensor.hv, energy=energy_keV)
|
||||
n_bins = int(conf.simulation.z_bins)
|
||||
method = conf.simulation.method
|
||||
started = time.time()
|
||||
alpha = np.full(n_bins, np.nan)
|
||||
beta = np.full(n_bins, np.nan)
|
||||
alpha_err = np.full(n_bins, np.nan)
|
||||
beta_err = np.full(n_bins, np.nan)
|
||||
n_carriers = np.zeros(n_bins, np.int64)
|
||||
for index in range(n_bins):
|
||||
stem = (f"{conf.sensor.code}_{conf.sensor.hv}V_{{axis}}_{energy_keV}keV_"
|
||||
f"{index}_of_{n_bins}_{method}.npy")
|
||||
try:
|
||||
xs = np.load(os.path.join(folder, stem.format(axis="xs")))
|
||||
ys = np.load(os.path.join(folder, stem.format(axis="ys")))
|
||||
except FileNotFoundError:
|
||||
continue
|
||||
coords = np.concatenate((xs, ys))
|
||||
n_carriers[index] = len(coords)
|
||||
(alpha[index], beta[index],
|
||||
alpha_err[index], beta_err[index]) = fit_generalised_normal(coords, index)
|
||||
print(f" {energy_keV:2d} keV: {np.isfinite(alpha).sum()}/{n_bins} depth bins fitted "
|
||||
f"({time.time()-started:.0f}s)", flush=True)
|
||||
return energy_keV, alpha, beta, alpha_err, beta_err, n_carriers
|
||||
|
||||
|
||||
def initialise(conf):
|
||||
global _CONF
|
||||
_CONF = conf
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--config", required=True)
|
||||
parser.add_argument("--workers", type=int, default=8)
|
||||
args = parser.parse_args()
|
||||
conf = OmegaConf.load(args.config)
|
||||
|
||||
energies = list(range(int(conf.energies.first), int(conf.energies.last) + 1,
|
||||
int(conf.energies.step)))
|
||||
n_bins = int(conf.simulation.z_bins)
|
||||
thickness = float(conf.sensor.thickness_um)
|
||||
edges = np.linspace(0, thickness, n_bins + 1)
|
||||
z0 = 0.5 * (edges[:-1] + edges[1:])
|
||||
drift_time = np.array([drift_time_allpix2(
|
||||
depth, thickness_um=thickness, hv=float(conf.sensor.hv),
|
||||
depletion_voltage=float(conf.sensor.depletion_voltage),
|
||||
temperature=float(conf.sensor.temperature)) for depth in z0])
|
||||
|
||||
print(f"[lut] {conf.sensor.code} @{conf.sensor.hv} V, {len(energies)} energies "
|
||||
f"{energies[0]}-{energies[-1]} keV, {n_bins} depth bins")
|
||||
print(f"[lut] drift time {drift_time.min():.2f} - {drift_time.max():.2f} ns")
|
||||
|
||||
with Pool(processes=args.workers, initializer=initialise, initargs=(conf,)) as pool:
|
||||
results = pool.map(process_energy, energies)
|
||||
|
||||
order = {e: i for i, e in enumerate(energies)}
|
||||
shape = (len(energies), n_bins)
|
||||
alpha, beta = np.full(shape, np.nan), np.full(shape, np.nan)
|
||||
alpha_err, beta_err = np.full(shape, np.nan), np.full(shape, np.nan)
|
||||
n_carriers = np.zeros(shape, np.int64)
|
||||
for energy, a, b, ae, be, nc in results:
|
||||
i = order[energy]
|
||||
alpha[i], beta[i], alpha_err[i], beta_err[i], n_carriers[i] = a, b, ae, be, nc
|
||||
|
||||
# Depth bins the simulation never produced (e.g. the deepest bins of an
|
||||
# energy whose run was cut short) can be filled from the neighbouring
|
||||
# energies at the *same* depth bin, i.e. the same drift time. The cloud size
|
||||
# is set mainly by the drift time, so this is far closer than
|
||||
# ChargeCloudLUT._fill_gaps, which holds the last fitted depth constant.
|
||||
filled = ~np.isfinite(alpha) | ~np.isfinite(beta)
|
||||
if conf.get("fill_missing", None) == "energy" and filled.any():
|
||||
grid = np.array(energies, float)
|
||||
for j in np.flatnonzero(filled.any(axis=0)):
|
||||
good = ~filled[:, j]
|
||||
if good.sum() < 2:
|
||||
continue
|
||||
for surface in (alpha, beta):
|
||||
surface[filled[:, j], j] = np.interp(grid[filled[:, j]],
|
||||
grid[good], surface[good, j])
|
||||
for i in np.flatnonzero(filled.any(axis=1)):
|
||||
bins = np.flatnonzero(filled[i])
|
||||
print(f"[lut] {energies[i]:2d} keV: {len(bins)} depth bins filled from "
|
||||
f"neighbouring energies (bins {bins.min()}-{bins.max()})")
|
||||
|
||||
os.makedirs(os.path.dirname(conf.output), exist_ok=True)
|
||||
np.savez(conf.output, energies=np.array(energies, float), z0=z0,
|
||||
drift_time=drift_time, alpha=alpha, beta=beta,
|
||||
alpha_err=alpha_err, beta_err=beta_err, n_carriers=n_carriers,
|
||||
filled_from_energy=filled if conf.get("fill_missing", None) == "energy"
|
||||
else np.zeros(shape, bool),
|
||||
sensor=conf.sensor.code, hv=float(conf.sensor.hv),
|
||||
thickness_um=thickness)
|
||||
print(f"[lut] wrote {conf.output} "
|
||||
f"alpha {np.nanmin(alpha):.2f}-{np.nanmax(alpha):.2f} um, "
|
||||
f"beta {np.nanmin(beta):.2f}-{np.nanmax(beta):.2f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,158 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
WP1.1 validation -- is the alpha(E,t) / beta(E,t) table good enough to sample
|
||||
continuous energies from?
|
||||
|
||||
Four questions:
|
||||
1. how noisy are the tabulated surfaces, and are they monotone in depth?
|
||||
2. leave-one-out: drop an energy, interpolate it back from its neighbours,
|
||||
and see how far off the interpolation is. With +-1 keV neighbours this
|
||||
tests a 2 keV gap, twice the spacing the table actually uses.
|
||||
3. does the table reproduce the 9-parameter fits that McGenerator uses at the
|
||||
energies where both exist?
|
||||
4. what happens if the 9 fit parameters are interpolated in energy instead
|
||||
of the curves -- the thing this table exists to avoid.
|
||||
"""
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(HERE.parent / "src"))
|
||||
from charge_model import ChargeCloudLUT # noqa: E402
|
||||
|
||||
GGD_PAR_DIR = Path('/home/xie_x1/MLXID/McGeneration/Notebooks')
|
||||
|
||||
|
||||
def parametric_alpha(pars, t):
|
||||
return pars[0] + pars[1] * np.sqrt(t) + pars[2] * t + pars[3] * t ** 2
|
||||
|
||||
|
||||
def parametric_beta(pars, t):
|
||||
return pars[4] * (t - pars[5]) ** pars[6] + pars[7] * np.exp(pars[8] * t) + 2
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--lut", default="data/charge_lut_Moench014_150V.npz")
|
||||
parser.add_argument("--sensor", default="Moench014")
|
||||
parser.add_argument("--hv", type=int, default=150)
|
||||
parser.add_argument("--outdir", default="Results/charge_lut")
|
||||
args = parser.parse_args()
|
||||
|
||||
lut = ChargeCloudLUT(args.lut)
|
||||
energies, time = lut.energies, lut._time
|
||||
alpha, beta = lut._alpha, lut._beta
|
||||
outdir = Path(args.outdir); outdir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# ---- 1. noise and monotonicity ------------------------------------
|
||||
print("1. surfaces")
|
||||
print(f" grid {alpha.shape[0]} energies x {alpha.shape[1]} depth bins, "
|
||||
f"drift time {time.min():.2f}-{time.max():.2f} ns")
|
||||
# bin-to-bin scatter after removing the smooth trend
|
||||
from scipy.signal import savgol_filter
|
||||
jitter_a, jitter_b = [], []
|
||||
for row in range(alpha.shape[0]):
|
||||
jitter_a.append(np.std(alpha[row] - savgol_filter(alpha[row], 21, 3)) /
|
||||
np.mean(alpha[row]))
|
||||
jitter_b.append(np.std(beta[row] - savgol_filter(beta[row], 21, 3)) /
|
||||
np.mean(beta[row]))
|
||||
print(f" bin-to-bin scatter about a local fit: alpha {np.mean(jitter_a)*100:.2f} %, "
|
||||
f"beta {np.mean(jitter_b)*100:.2f} %")
|
||||
rising = np.mean([np.all(np.diff(alpha[r]) > -0.02) for r in range(alpha.shape[0])])
|
||||
print(f" alpha monotone in drift time for {rising*100:.0f} % of energies")
|
||||
at_limit = ((alpha <= 0.501).mean(), (beta <= 2.001).mean())
|
||||
print(f" fits sitting on the parameter limits: alpha {at_limit[0]*100:.1f} %, "
|
||||
f"beta {at_limit[1]*100:.1f} % (shortest drifts, cloud smaller than a micron)")
|
||||
|
||||
# ---- 2. leave-one-out in energy -----------------------------------
|
||||
print("\n2. leave-one-out in energy (interpolating across a 2 keV gap)")
|
||||
rows = []
|
||||
for i in range(1, len(energies) - 1):
|
||||
a_hat = 0.5 * (alpha[i - 1] + alpha[i + 1])
|
||||
b_hat = 0.5 * (beta[i - 1] + beta[i + 1])
|
||||
rows.append((energies[i],
|
||||
np.max(np.abs(a_hat - alpha[i]) / alpha[i]) * 100,
|
||||
np.mean(np.abs(a_hat - alpha[i]) / alpha[i]) * 100,
|
||||
np.max(np.abs(b_hat - beta[i]) / beta[i]) * 100))
|
||||
rows = np.array(rows)
|
||||
print(f" alpha: mean |error| {rows[:,2].mean():.3f} %, worst energy "
|
||||
f"{rows[rows[:,1].argmax(),0]:.0f} keV at {rows[:,1].max():.2f} % (max over depth)")
|
||||
print(f" beta : worst {rows[:,3].max():.2f} %")
|
||||
print(" per energy (keV, mean%, max%):",
|
||||
", ".join(f"{r[0]:.0f}:{r[2]:.2f}/{r[1]:.2f}" for r in rows[::5]))
|
||||
|
||||
# ---- 3. against the 9-parameter fits ------------------------------
|
||||
print("\n3. table vs the 9-parameter fits McGenerator uses")
|
||||
valid = (time > 0.3) & (time < time.max() - 0.2)
|
||||
comparison = []
|
||||
for i, energy in enumerate(energies):
|
||||
path = GGD_PAR_DIR / f'ggdPar_{args.sensor}_{args.hv}V_{int(energy)}keV.npy'
|
||||
if not path.exists():
|
||||
continue
|
||||
pars = np.load(path)
|
||||
da = (parametric_alpha(pars, time[valid]) - alpha[i][valid]) / alpha[i][valid]
|
||||
db = (parametric_beta(pars, time[valid]) - beta[i][valid]) / beta[i][valid]
|
||||
comparison.append((energy, np.mean(np.abs(da)) * 100, np.mean(np.abs(db)) * 100))
|
||||
comparison = np.array(comparison)
|
||||
print(f" {len(comparison)} energies compared over {valid.sum()} depth bins "
|
||||
f"inside the parametric fit's validity range")
|
||||
print(f" mean |alpha difference| {comparison[:,1].mean():.2f} % "
|
||||
f"(worst {comparison[:,1].max():.2f} % at {comparison[comparison[:,1].argmax(),0]:.0f} keV)")
|
||||
print(f" mean |beta difference| {comparison[:,2].mean():.2f} %")
|
||||
|
||||
# ---- 4. interpolating the 9 parameters instead of the curves -------
|
||||
print("\n4. interpolating the 9 fit parameters instead (what this table avoids)")
|
||||
worst = 0.0
|
||||
for target in range(6, 30):
|
||||
low = GGD_PAR_DIR / f'ggdPar_{args.sensor}_{args.hv}V_{target-1}keV.npy'
|
||||
high = GGD_PAR_DIR / f'ggdPar_{args.sensor}_{args.hv}V_{target+1}keV.npy'
|
||||
if not (low.exists() and high.exists()):
|
||||
continue
|
||||
blended = 0.5 * (np.load(low) + np.load(high))
|
||||
i = int(np.argmin(np.abs(energies - target)))
|
||||
da = np.abs(parametric_alpha(blended, time[valid]) - alpha[i][valid]) / alpha[i][valid]
|
||||
worst = max(worst, da.max() * 100)
|
||||
if target in (8, 12, 20, 28):
|
||||
print(f" {target:2d} keV: parameter blend is off by up to {da.max()*100:6.2f} %"
|
||||
f" (curve interpolation: {rows[rows[:,0]==target][0,1]:.2f} %)")
|
||||
print(f" worst case over all energies: {worst:.1f} %")
|
||||
|
||||
# ---- figure --------------------------------------------------------
|
||||
fig, ax = plt.subplots(2, 3, figsize=(16, 8))
|
||||
for i in range(0, len(energies), 5):
|
||||
ax[0, 0].plot(time, alpha[i], lw=1, label=f'{energies[i]:.0f} keV')
|
||||
ax[0, 1].plot(time, beta[i], lw=1)
|
||||
ax[0, 2].plot(time, lut.cloud_sigma(energies[i], time), lw=1)
|
||||
ax[0, 0].set_ylabel(r'$\alpha$ [$\mu$m]'); ax[0, 0].legend(fontsize=8)
|
||||
ax[0, 1].set_ylabel(r'$\beta$')
|
||||
ax[0, 2].set_ylabel(r'cloud RMS [$\mu$m]')
|
||||
for a in ax[0]:
|
||||
a.set_xlabel('drift time [ns]'); a.grid(alpha=.3)
|
||||
|
||||
im = ax[1, 0].imshow(alpha, aspect='auto', origin='lower', cmap='viridis',
|
||||
extent=[time.min(), time.max(), energies[0], energies[-1]])
|
||||
ax[1, 0].set_title(r'$\alpha(E, t)$ [$\mu$m]'); plt.colorbar(im, ax=ax[1, 0])
|
||||
im = ax[1, 1].imshow(beta, aspect='auto', origin='lower', cmap='viridis',
|
||||
extent=[time.min(), time.max(), energies[0], energies[-1]])
|
||||
ax[1, 1].set_title(r'$\beta(E, t)$'); plt.colorbar(im, ax=ax[1, 1])
|
||||
for a in ax[1, :2]:
|
||||
a.set_xlabel('drift time [ns]'); a.set_ylabel('energy [keV]')
|
||||
|
||||
ax[1, 2].plot(rows[:, 0], rows[:, 1], 'o-', label=r'curve interp, max over depth')
|
||||
ax[1, 2].plot(rows[:, 0], rows[:, 2], 's-', label=r'curve interp, mean')
|
||||
ax[1, 2].set_xlabel('energy left out [keV]'); ax[1, 2].set_ylabel(r'$\alpha$ error [%]')
|
||||
ax[1, 2].set_yscale('log'); ax[1, 2].legend(fontsize=8); ax[1, 2].grid(alpha=.3)
|
||||
ax[1, 2].set_title('leave-one-out across a 2 keV gap')
|
||||
fig.tight_layout()
|
||||
fig.savefig(outdir / 'charge_lut_validation.png', dpi=130)
|
||||
print(f"\nwrote {outdir / 'charge_lut_validation.png'}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,176 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
WP1.2 -- generate grid-based MC frames from the existing single-photon stamps.
|
||||
|
||||
python scripts/02_generate_frames.py --config Configs/generate_frames_12keV.yaml
|
||||
python scripts/02_generate_frames.py --config ... --splits train val
|
||||
python scripts/02_generate_frames.py --config ... --splits eval_scan
|
||||
python scripts/02_generate_frames.py --config ... --smoke # tiny, fast
|
||||
|
||||
Each shard is one HDF5 file:
|
||||
frames (N, S, S) float32 keV
|
||||
photons (M, 4) float32 x, y, E[keV], z[um] (frame coordinates)
|
||||
frame_ptr (N+1,) int64 photons of frame i = photons[ptr[i]:ptr[i+1]]
|
||||
occupancy (N,) float32 the occupancy each frame was thrown at
|
||||
"""
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from multiprocessing import Pool
|
||||
|
||||
import h5py
|
||||
import numpy as np
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "src"))
|
||||
from frame_generator import generate_shard # noqa: E402
|
||||
|
||||
SMOKE_FRAMES = 400
|
||||
|
||||
|
||||
def stamp_path(conf, index):
|
||||
return os.path.join(conf.stamps.folder, conf.stamps.pattern.format(index))
|
||||
|
||||
|
||||
def run_job(job):
|
||||
started = time.time()
|
||||
frames, photons, frame_pointer, occupancy, stats = generate_shard(
|
||||
stamp_files=job["stamp_files"],
|
||||
n_frames=job["n_frames"],
|
||||
frame_size=job["frame_size"],
|
||||
occupancy_spec=job["occupancy"],
|
||||
seed=job["seed"],
|
||||
margin=job["margin"],
|
||||
noise_keV=job["noise_keV"],
|
||||
cm_row_keV=job["cm_row_keV"],
|
||||
cm_col_keV=job["cm_col_keV"],
|
||||
edge_fraction=job["edge_fraction"],
|
||||
mode=job["mode"], lut_path=job.get("lut_path"),
|
||||
energy_spec=job.get("energy_spec"),
|
||||
pixel_size_um=job["pixel_size_um"],
|
||||
depletion_voltage=job.get("depletion_voltage"),
|
||||
temperature=job.get("temperature"),
|
||||
)
|
||||
os.makedirs(os.path.dirname(job["output"]), exist_ok=True)
|
||||
with h5py.File(job["output"], "w") as handle:
|
||||
handle.create_dataset("frames", data=frames, dtype="f4")
|
||||
handle.create_dataset("photons", data=photons, dtype="f4")
|
||||
handle.create_dataset("frame_ptr", data=frame_pointer, dtype="i8")
|
||||
handle.create_dataset("occupancy", data=occupancy, dtype="f4")
|
||||
for key, value in job["attrs"].items():
|
||||
handle.attrs[key] = value
|
||||
for key, value in stats.items():
|
||||
handle.attrs[key] = value
|
||||
return dict(output=job["output"], seconds=time.time() - started, **stats)
|
||||
|
||||
|
||||
def build_jobs(conf, splits, smoke=False):
|
||||
jobs = []
|
||||
base_seed = int(conf.runtime.seed)
|
||||
generation = conf.get("generation", {})
|
||||
mode = str(generation.get("mode", "stamps"))
|
||||
common = dict(
|
||||
frame_size=int(conf.frame.size),
|
||||
margin=int(conf.frame.margin),
|
||||
noise_keV=float(conf.noise.pixel_keV),
|
||||
cm_row_keV=float(conf.noise.common_mode_row_keV),
|
||||
cm_col_keV=float(conf.noise.common_mode_col_keV),
|
||||
mode=mode,
|
||||
pixel_size_um=float(conf.meta.pixel_size_um),
|
||||
)
|
||||
if mode == "lut":
|
||||
common.update(lut_path=str(generation.lut),
|
||||
depletion_voltage=float(conf.meta.depletion_voltage),
|
||||
temperature=float(conf.meta.temperature))
|
||||
attrs = dict(energy_keV=float(conf.meta.get("energy", 0.0)),
|
||||
sensor=str(conf.meta.sensor),
|
||||
hv=int(conf.meta.hv), pixel_size_um=float(conf.meta.pixel_size_um),
|
||||
frame_size=int(conf.frame.size), margin=int(conf.frame.margin),
|
||||
noise_keV=float(conf.noise.pixel_keV))
|
||||
|
||||
for split in splits:
|
||||
if split == "eval_scan":
|
||||
scan = conf.eval_scan
|
||||
stamp_files = ([] if mode == "lut"
|
||||
else [stamp_path(conf, index) for index in scan.stamp_files])
|
||||
n_frames = SMOKE_FRAMES if smoke else int(scan.frames_each)
|
||||
for point_index, occupancy in enumerate(scan.occupancies):
|
||||
tag = f"occ{float(occupancy) * 100:g}pct".replace(".", "p")
|
||||
jobs.append(dict(
|
||||
output=os.path.join(conf.output.folder, "eval_scan", f"{tag}.h5"),
|
||||
stamp_files=([] if mode == "lut"
|
||||
else [stamp_files[point_index % len(stamp_files)]]),
|
||||
energy_spec=(OmegaConf.to_container(conf.spectrum["eval"],
|
||||
resolve=True)
|
||||
if mode == "lut" else None),
|
||||
n_frames=n_frames,
|
||||
occupancy={"mode": "fixed", "value": float(occupancy)},
|
||||
seed=base_seed + 9000 + point_index,
|
||||
edge_fraction=float(scan.edge_fraction),
|
||||
attrs=dict(attrs, split="eval_scan",
|
||||
occupancy_nominal=float(occupancy)),
|
||||
**common))
|
||||
continue
|
||||
|
||||
split_conf = conf.splits[split]
|
||||
if mode == "lut":
|
||||
first_file, last_file = 0, int(split_conf.n_shards) - 1
|
||||
else:
|
||||
first_file, last_file = conf.stamps[f"{split}_files"]
|
||||
n_frames = SMOKE_FRAMES if smoke else int(split_conf.frames_per_shard)
|
||||
for shard_index, file_index in enumerate(range(int(first_file),
|
||||
int(last_file) + 1)):
|
||||
jobs.append(dict(
|
||||
output=os.path.join(conf.output.folder, split,
|
||||
f"shard_{shard_index:02d}.h5"),
|
||||
stamp_files=([] if mode == "lut"
|
||||
else [stamp_path(conf, file_index)]),
|
||||
energy_spec=(OmegaConf.to_container(conf.spectrum[split], resolve=True)
|
||||
if mode == "lut" else None),
|
||||
n_frames=n_frames,
|
||||
occupancy=OmegaConf.to_container(split_conf.occupancy, resolve=True),
|
||||
seed=base_seed + 1000 * (splits.index(split) + 1) + shard_index,
|
||||
edge_fraction=float(conf.edge_augmentation.fraction),
|
||||
attrs=dict(attrs, split=split, mode=mode),
|
||||
**common))
|
||||
return jobs
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--config", required=True)
|
||||
parser.add_argument("--splits", nargs="+",
|
||||
default=["train", "val", "test", "eval_scan"])
|
||||
parser.add_argument("--smoke", action="store_true",
|
||||
help=f"{SMOKE_FRAMES} frames per shard, for a quick check")
|
||||
parser.add_argument("--workers", type=int, default=None)
|
||||
args = parser.parse_args()
|
||||
|
||||
conf = OmegaConf.load(args.config)
|
||||
jobs = build_jobs(conf, list(args.splits), smoke=args.smoke)
|
||||
n_workers = args.workers or int(conf.runtime.n_workers)
|
||||
print(f"[gen] {len(jobs)} shards, {n_workers} workers, "
|
||||
f"frame {conf.frame.size}x{conf.frame.size}")
|
||||
print(f"[gen] output -> {conf.output.folder}")
|
||||
|
||||
started = time.time()
|
||||
with Pool(processes=min(n_workers, len(jobs))) as pool:
|
||||
results = pool.map(run_job, jobs)
|
||||
|
||||
total_photons = sum(r["n_photons"] for r in results)
|
||||
total_frames = sum(r["n_frames"] for r in results)
|
||||
total_same_cell = sum(r["n_same_cell"] for r in results)
|
||||
for result in sorted(results, key=lambda r: r["output"]):
|
||||
print(f" {os.path.relpath(result['output'], conf.output.folder):32s} "
|
||||
f"{result['n_frames']:7d} frames {result['n_photons']:9d} photons "
|
||||
f"<occ>={result['mean_occupancy']*100:5.2f}% "
|
||||
f"same-cell={result['n_same_cell']:6d} ({result['seconds']:.1f}s)")
|
||||
print(f"[gen] total {total_frames} frames, {total_photons} photons, "
|
||||
f"same-cell loss {total_same_cell} "
|
||||
f"({100.0*total_same_cell/max(total_photons,1):.2f}%), "
|
||||
f"{time.time()-started:.1f}s")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,250 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
WP3 -- train the grid-based model.
|
||||
|
||||
python scripts/03_train.py --config Configs/train_grid_12keV.yaml
|
||||
python scripts/03_train.py --config ... --epochs 3 --tag smoke
|
||||
|
||||
Every epoch the validation pass reports both the loss and the *detection*
|
||||
metrics (efficiency / purity / position RMS), because the loss alone says
|
||||
nothing about the working point.
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
from omegaconf import OmegaConf
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(HERE.parent / "src"))
|
||||
from datasets import GridFrameDataset, collate_with_photons # noqa: E402
|
||||
from decode import decode # noqa: E402
|
||||
from losses import GridLoss # noqa: E402
|
||||
from metrics import DetectionStats # noqa: E402
|
||||
from models import build_model, valid_mask # noqa: E402
|
||||
|
||||
LOSS_PARTS = ("objectness", "offset", "energy")
|
||||
|
||||
|
||||
def shard_files(folder, split, shard_range):
|
||||
first, last = shard_range
|
||||
return [os.path.join(folder, split, f"shard_{index:02d}.h5")
|
||||
for index in range(int(first), int(last) + 1)]
|
||||
|
||||
|
||||
def prepare_output_dir(conf, tag=None):
|
||||
date = datetime.now().strftime("%y%m%d")
|
||||
base = Path(conf.experiment.results_dir)
|
||||
run_index = 0
|
||||
while True:
|
||||
name = (f"{date}_{conf.experiment.name}"
|
||||
+ (f"_{tag}" if tag else "") + f"_{run_index:02d}")
|
||||
output_dir = base / name
|
||||
if not output_dir.exists():
|
||||
break
|
||||
run_index += 1
|
||||
(output_dir / "Models").mkdir(parents=True)
|
||||
(output_dir / "Plots").mkdir(parents=True)
|
||||
OmegaConf.save(conf, output_dir / "config.yaml")
|
||||
return output_dir
|
||||
|
||||
|
||||
def run_epoch(model, loader, loss_function, mask, optimizer=None,
|
||||
grad_scaler=None, use_amp=False):
|
||||
is_training = optimizer is not None
|
||||
model.train(is_training)
|
||||
totals = dict(loss=0.0, objectness=0.0, offset=0.0, energy=0.0)
|
||||
n_samples = 0
|
||||
for batch in loader:
|
||||
frames, objectness, offset, energy = [t.cuda(non_blocking=True)
|
||||
for t in batch[:4]]
|
||||
with torch.autocast("cuda", enabled=use_amp):
|
||||
prediction = model(frames)
|
||||
loss, parts = loss_function(prediction.float(), objectness, offset,
|
||||
energy, mask)
|
||||
if is_training:
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
if grad_scaler is not None:
|
||||
grad_scaler.scale(loss).backward()
|
||||
grad_scaler.step(optimizer)
|
||||
grad_scaler.update()
|
||||
else:
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
batch_size = frames.shape[0]
|
||||
totals["loss"] += loss.item() * batch_size
|
||||
for part in LOSS_PARTS:
|
||||
totals[part] += parts[part].item() * batch_size
|
||||
n_samples += batch_size
|
||||
return {key: value / max(n_samples, 1) for key, value in totals.items()}
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def run_detection_metrics(model, loader, conf, halo, frame_size, use_amp=False):
|
||||
model.eval()
|
||||
stats = DetectionStats(pixel_size_um=float(conf.eval.pixel_size_um),
|
||||
match_radius=float(conf.eval.match_radius),
|
||||
halo=halo, frame_size=frame_size)
|
||||
for frames, objectness, offset, energy, photons in loader:
|
||||
frames = frames.cuda(non_blocking=True)
|
||||
with torch.autocast("cuda", enabled=use_amp):
|
||||
prediction = model(frames)
|
||||
points = decode(prediction.float(),
|
||||
threshold=float(conf.eval.threshold),
|
||||
merge_distance=float(conf.eval.merge_distance),
|
||||
halo=halo,
|
||||
predict_energy=bool(conf.model.predict_energy),
|
||||
energy_scale=float(conf.data.energy_scale))
|
||||
stats.update([frame_points.numpy() for frame_points in points],
|
||||
[frame_truth.numpy() for frame_truth in photons])
|
||||
return stats.summary()
|
||||
|
||||
|
||||
def plot_curves(history, output_dir):
|
||||
epochs = np.arange(1, len(history["train_loss"]) + 1)
|
||||
figure, axes = plt.subplots(1, 3, figsize=(15, 4.2))
|
||||
axes[0].plot(epochs, history["train_loss"], label="train")
|
||||
axes[0].plot(epochs, history["val_loss"], label="val")
|
||||
axes[0].set_yscale("log")
|
||||
axes[0].set_xlabel("epoch")
|
||||
axes[0].set_ylabel("total loss")
|
||||
axes[0].legend()
|
||||
axes[0].grid(alpha=.3)
|
||||
|
||||
axes[1].plot(epochs, history["efficiency"], label="efficiency")
|
||||
axes[1].plot(epochs, history["purity"], label="purity")
|
||||
axes[1].set_xlabel("epoch")
|
||||
axes[1].set_ylim(0, 1.02)
|
||||
axes[1].legend()
|
||||
axes[1].grid(alpha=.3)
|
||||
|
||||
axes[2].plot(epochs, history["rms_1d_um"], label="RMS 1D")
|
||||
axes[2].plot(epochs, history["sigma68_1d_um"], label=r"$\sigma_{68}$")
|
||||
axes[2].set_xlabel("epoch")
|
||||
axes[2].set_ylabel(r"residual [$\mu$m]")
|
||||
axes[2].set_yscale("log")
|
||||
axes[2].legend()
|
||||
axes[2].grid(alpha=.3)
|
||||
|
||||
figure.tight_layout()
|
||||
figure.savefig(output_dir / "Plots" / "training_curves.png", dpi=130)
|
||||
plt.close(figure)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--config", required=True)
|
||||
parser.add_argument("--epochs", type=int, default=None)
|
||||
parser.add_argument("--tag", default=None)
|
||||
args = parser.parse_args()
|
||||
|
||||
conf = OmegaConf.load(args.config)
|
||||
if args.epochs:
|
||||
conf.training.epochs = args.epochs
|
||||
|
||||
torch.manual_seed(conf.training.seed)
|
||||
np.random.seed(conf.training.seed)
|
||||
|
||||
output_dir = prepare_output_dir(conf, args.tag)
|
||||
print(f"[train] output -> {output_dir}")
|
||||
|
||||
dataset_kwargs = dict(input_scale=float(conf.data.input_scale),
|
||||
energy_scale=float(conf.data.energy_scale))
|
||||
train_dataset = GridFrameDataset(
|
||||
shard_files(conf.data.folder, "train", conf.data.train_shards),
|
||||
augment=bool(conf.data.augment), name="train",
|
||||
max_frames=conf.data.max_train_frames, **dataset_kwargs)
|
||||
val_dataset = GridFrameDataset(
|
||||
shard_files(conf.data.folder, "val", conf.data.val_shards),
|
||||
augment=False, name="val", return_photons=True,
|
||||
max_frames=conf.data.max_val_frames, **dataset_kwargs)
|
||||
|
||||
n_workers = int(conf.data.num_workers)
|
||||
train_loader = DataLoader(train_dataset, batch_size=int(conf.data.batch_size_train),
|
||||
shuffle=True, num_workers=n_workers, pin_memory=True,
|
||||
drop_last=True, persistent_workers=n_workers > 0)
|
||||
val_loader = DataLoader(val_dataset, batch_size=int(conf.data.batch_size_val),
|
||||
shuffle=False, num_workers=n_workers, pin_memory=True,
|
||||
collate_fn=collate_with_photons,
|
||||
persistent_workers=n_workers > 0)
|
||||
|
||||
model = build_model(conf).cuda()
|
||||
halo = model.rf_radius
|
||||
frame_size = train_dataset.frame_size
|
||||
mask = valid_mask(frame_size, frame_size, halo, device="cuda")
|
||||
n_parameters = sum(p.numel() for p in model.parameters())
|
||||
print(f"[train] GridNet: {n_parameters} parameters, receptive-field radius "
|
||||
f"R={halo} px (halo), valid region "
|
||||
f"{frame_size-2*halo}x{frame_size-2*halo} of {frame_size}x{frame_size}")
|
||||
|
||||
loss_function = GridLoss(conf).cuda()
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=float(conf.training.lr),
|
||||
weight_decay=float(conf.training.weight_decay))
|
||||
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
|
||||
optimizer, "min", factor=float(conf.training.scheduler_factor),
|
||||
patience=int(conf.training.scheduler_patience))
|
||||
use_amp = bool(conf.training.amp)
|
||||
grad_scaler = torch.amp.GradScaler("cuda") if use_amp else None
|
||||
|
||||
history = {key: [] for key in ("train_loss", "val_loss", "efficiency",
|
||||
"purity", "rms_1d_um", "sigma68_1d_um")}
|
||||
best_loss = float("inf")
|
||||
for epoch in range(1, int(conf.training.epochs) + 1):
|
||||
started = time.time()
|
||||
train_metrics = run_epoch(model, train_loader, loss_function, mask,
|
||||
optimizer, grad_scaler, use_amp)
|
||||
val_metrics = run_epoch(model, val_loader, loss_function, mask,
|
||||
None, None, use_amp)
|
||||
detection = run_detection_metrics(model, val_loader, conf, halo,
|
||||
frame_size, use_amp)
|
||||
scheduler.step(val_metrics["loss"])
|
||||
|
||||
history["train_loss"].append(train_metrics["loss"])
|
||||
history["val_loss"].append(val_metrics["loss"])
|
||||
history["efficiency"].append(detection["efficiency"])
|
||||
history["purity"].append(detection["purity"])
|
||||
history["rms_1d_um"].append(detection["rms_1d_um"])
|
||||
history["sigma68_1d_um"].append(detection["sigma68_1d_um"])
|
||||
|
||||
print(f"[{epoch:3d}/{conf.training.epochs}] "
|
||||
f"train {train_metrics['loss']:.4f} "
|
||||
f"(obj {train_metrics['objectness']:.4f} "
|
||||
f"off {train_metrics['offset']:.5f} "
|
||||
f"ene {train_metrics['energy']:.5f}) "
|
||||
f"val {val_metrics['loss']:.4f} | "
|
||||
f"eff {detection['efficiency']*100:5.2f}% "
|
||||
f"pur {detection['purity']*100:5.2f}% "
|
||||
f"| RMS {detection['rms_1d_um']:.3f} um "
|
||||
f"s68 {detection['sigma68_1d_um']:.3f} um "
|
||||
f"| lr {optimizer.param_groups[0]['lr']:.1e} "
|
||||
f"| {time.time()-started:.0f}s", flush=True)
|
||||
|
||||
checkpoint = dict(model=model.state_dict(),
|
||||
conf=OmegaConf.to_container(conf),
|
||||
epoch=epoch, metrics=detection)
|
||||
if val_metrics["loss"] < best_loss:
|
||||
best_loss = val_metrics["loss"]
|
||||
torch.save(checkpoint, output_dir / "Models" / "best.pth")
|
||||
if epoch in list(conf.training.checkpoint_epochs) or epoch == conf.training.epochs:
|
||||
torch.save(checkpoint, output_dir / "Models" / f"E{epoch}.pth")
|
||||
torch.save(checkpoint, output_dir / "Models" / "last.pth")
|
||||
plot_curves(history, output_dir)
|
||||
with open(output_dir / "history.json", "w") as handle:
|
||||
json.dump(history, handle, indent=1)
|
||||
|
||||
print(f"[train] done. best val loss {best_loss:.4f}. "
|
||||
f"models in {output_dir/'Models'}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,344 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
WP4 -- evaluate a grid model on the fixed-occupancy MC scan.
|
||||
|
||||
python scripts/04_evaluate.py --checkpoint Results/<exp>/Models/best.pth
|
||||
python scripts/04_evaluate.py --checkpoint ... --data <folder>/eval_scan
|
||||
|
||||
Produces, per occupancy point:
|
||||
efficiency / purity at the working threshold, position RMS and sigma68,
|
||||
energy bias/resolution, sub-pixel non-uniformity (risk R1),
|
||||
plus the threshold ROC, the pair-resolution curve (efficiency vs distance to the
|
||||
nearest other photon -- the figure of merit that replaces "efficiency vs
|
||||
occupancy"), and residual vs conversion depth z.
|
||||
"""
|
||||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
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
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
sys.path.insert(0, str(HERE.parent / "src"))
|
||||
from datasets import GridFrameDataset, collate_with_photons # noqa: E402
|
||||
from decode import decode # noqa: E402
|
||||
from metrics import (DetectionStats, greedy_match, # noqa: E402
|
||||
nearest_neighbour_distance)
|
||||
from models import build_model # noqa: E402
|
||||
|
||||
THRESHOLDS = np.array([.05, .1, .2, .3, .4, .5, .6, .7, .8, .9, .95, .98])
|
||||
NEIGHBOUR_BIN_EDGES = np.array([0, .5, 1, 1.5, 2, 2.5, 3, 4, 6, 1e9])
|
||||
ENERGY_BIN_WIDTH = 2.5 # keV; one bin only for a
|
||||
# monochromatic set, which is fine
|
||||
|
||||
|
||||
def energy_bin_edges(energies, width=ENERGY_BIN_WIDTH):
|
||||
"""2.5 keV bins on a fixed grid, spanning the energies of the set."""
|
||||
low = np.floor(energies.min() / width) * width
|
||||
high = np.ceil(energies.max() / width) * width
|
||||
return np.arange(low, high + width / 2, width)
|
||||
|
||||
|
||||
def occupancy_of_filename(path):
|
||||
match = re.search(r"occ([0-9p]+)pct", os.path.basename(path))
|
||||
return float(match.group(1).replace("p", ".")) / 100 if match else float("nan")
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def run_file(model, conf, path, halo, batch_size=512, workers=8, max_frames=None,
|
||||
roc_frames=2000):
|
||||
dataset = GridFrameDataset([path], augment=False, return_photons=True,
|
||||
input_scale=float(conf.data.input_scale),
|
||||
energy_scale=float(conf.data.energy_scale),
|
||||
max_frames=max_frames, name=os.path.basename(path))
|
||||
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False,
|
||||
num_workers=workers, collate_fn=collate_with_photons,
|
||||
pin_memory=True)
|
||||
frame_size = dataset.frame_size
|
||||
working_threshold = float(conf.eval.threshold)
|
||||
lowest_threshold = float(THRESHOLDS.min())
|
||||
|
||||
stats = DetectionStats(pixel_size_um=float(conf.eval.pixel_size_um),
|
||||
match_radius=float(conf.eval.match_radius),
|
||||
halo=halo, frame_size=frame_size)
|
||||
roc_predicted, roc_true = [], []
|
||||
|
||||
for frames, objectness, offset, energy, photons in loader:
|
||||
frames = frames.cuda(non_blocking=True)
|
||||
prediction = model(frames).float()
|
||||
points = decode(prediction, threshold=lowest_threshold,
|
||||
merge_distance=float(conf.eval.merge_distance), halo=halo,
|
||||
predict_energy=bool(conf.model.predict_energy),
|
||||
energy_scale=float(conf.data.energy_scale))
|
||||
points = [frame_points.numpy() for frame_points in points]
|
||||
truth = [frame_truth.numpy() for frame_truth in photons]
|
||||
stats.update([p[p[:, 3] > working_threshold] for p in points], truth)
|
||||
if len(roc_predicted) < roc_frames:
|
||||
roc_predicted.extend(points)
|
||||
roc_true.extend(truth)
|
||||
|
||||
summary = stats.summary()
|
||||
differential = stats.differential()
|
||||
|
||||
# ---- threshold sweep on a subset -------------------------------------
|
||||
roc = []
|
||||
for threshold in THRESHOLDS:
|
||||
subset_stats = DetectionStats(pixel_size_um=float(conf.eval.pixel_size_um),
|
||||
match_radius=float(conf.eval.match_radius),
|
||||
halo=halo, frame_size=frame_size)
|
||||
subset_stats.update([p[p[:, 3] > threshold] for p in roc_predicted[:roc_frames]],
|
||||
roc_true[:roc_frames])
|
||||
subset_summary = subset_stats.summary()
|
||||
roc.append((float(threshold), subset_summary["efficiency"],
|
||||
subset_summary["purity"], subset_summary["rms_1d_um"]))
|
||||
|
||||
# ---- pair resolution: efficiency vs truth nearest-neighbour distance --
|
||||
n_bins = len(NEIGHBOUR_BIN_EDGES) - 1
|
||||
n_true_binned = np.zeros(n_bins)
|
||||
n_matched_binned = np.zeros(n_bins)
|
||||
for predicted, truth in zip(roc_predicted[:roc_frames], roc_true[:roc_frames]):
|
||||
predicted = predicted[predicted[:, 3] > working_threshold]
|
||||
if len(truth) == 0:
|
||||
continue
|
||||
inside = ((truth[:, 0] >= halo) & (truth[:, 0] < frame_size - halo) &
|
||||
(truth[:, 1] >= halo) & (truth[:, 1] < frame_size - halo))
|
||||
truth = truth[inside]
|
||||
if len(truth) == 0:
|
||||
continue
|
||||
neighbour_distance = nearest_neighbour_distance(truth[:, :2])
|
||||
bin_index = np.digitize(neighbour_distance, NEIGHBOUR_BIN_EDGES) - 1
|
||||
for index in range(n_bins):
|
||||
n_true_binned[index] += (bin_index == index).sum()
|
||||
if len(predicted):
|
||||
_, index_true = greedy_match(predicted[:, :2], truth[:, :2],
|
||||
float(conf.eval.match_radius))
|
||||
for index in range(n_bins):
|
||||
n_matched_binned[index] += (bin_index[index_true] == index).sum()
|
||||
pair_resolution = dict(
|
||||
edges=NEIGHBOUR_BIN_EDGES[:-1].tolist(),
|
||||
n_true=n_true_binned.tolist(),
|
||||
n_match=n_matched_binned.tolist(),
|
||||
efficiency=(n_matched_binned / np.maximum(n_true_binned, 1)).tolist())
|
||||
|
||||
return (summary, differential, roc, pair_resolution,
|
||||
stats.versus_energy(energy_bin_edges(dataset.photons[:, 2])))
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--checkpoint", required=True)
|
||||
parser.add_argument("--data", default=None,
|
||||
help="folder with the eval_scan h5 files")
|
||||
parser.add_argument("--max-frames", type=int, default=None)
|
||||
parser.add_argument("--workers", type=int, default=8)
|
||||
parser.add_argument("--outdir", default=None)
|
||||
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
|
||||
print(f"[eval] checkpoint epoch {checkpoint['epoch']}, R={halo}, "
|
||||
f"{sum(p.numel() for p in model.parameters())} params, "
|
||||
f"threshold {conf.eval.threshold}, merge {conf.eval.merge_distance} px")
|
||||
|
||||
data_folder = args.data or os.path.join(conf.data.folder, "eval_scan")
|
||||
files = sorted(glob.glob(os.path.join(data_folder, "*.h5")),
|
||||
key=occupancy_of_filename)
|
||||
if not files:
|
||||
sys.exit(f"no eval files in {data_folder}")
|
||||
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 = occupancy_of_filename(path)
|
||||
summary, differential, roc, pair_resolution, per_energy = run_file(
|
||||
model, conf, path, halo, workers=args.workers, max_frames=args.max_frames)
|
||||
results[f"{occupancy*100:g}%"] = dict(occupancy=occupancy, summary=summary,
|
||||
roc=roc, pair=pair_resolution,
|
||||
per_energy=per_energy,
|
||||
_differential=differential)
|
||||
print(f" occ {occupancy*100:5.2f}% eff {summary['efficiency']*100:6.2f}% "
|
||||
f"pur {summary['purity']*100:6.2f}% RMS {summary['rms_1d_um']:.3f} um "
|
||||
f"s68 {summary['sigma68_1d_um']:.3f} um "
|
||||
f"dE {summary.get('energy_rms_keV', float('nan')):.3f} keV "
|
||||
f"subpix-nonunif {summary['subpixel_nonuniformity']*100:.1f}%",
|
||||
flush=True)
|
||||
|
||||
populated = [row for row in list(results.values())[len(results)//2]["per_energy"]
|
||||
if row["n_true"] > 100]
|
||||
if len(populated) > 1:
|
||||
print(f"\n per energy bin at occupancy "
|
||||
f"{list(results.values())[len(results)//2]['occupancy']*100:g}%:")
|
||||
print(f" {'keV':>11} {'n':>9} {'eff%':>7} {'RMS um':>8} {'s68 um':>8} "
|
||||
f"{'dE bias':>8} {'dE rms':>8}")
|
||||
for row in populated:
|
||||
print(f" {row['energy_low']:4.1f}-{row['energy_high']:4.1f} "
|
||||
f"{row['n_true']:9d} {row['efficiency']*100:7.2f} "
|
||||
f"{row['rms_1d_um']:8.3f} {row['sigma68_1d_um']:8.3f} "
|
||||
f"{row['energy_bias_keV']:8.3f} {row['energy_rms_keV']:8.3f}")
|
||||
plot_versus_energy(results, output_dir)
|
||||
|
||||
plot_all(results, output_dir, float(conf.eval.threshold))
|
||||
serialisable = {label: {key: value for key, value in result.items()
|
||||
if key != "_differential"}
|
||||
for label, result in results.items()}
|
||||
with open(output_dir / "eval_summary.json", "w") as handle:
|
||||
json.dump(serialisable, handle, indent=1)
|
||||
print(f"[eval] wrote {output_dir/'eval_summary.json'} and {output_dir/'Plots'}")
|
||||
|
||||
|
||||
def plot_versus_energy(results, output_dir):
|
||||
"""Efficiency, position resolution and energy resolution against photon energy."""
|
||||
figure, axes = plt.subplots(1, 3, figsize=(15, 4.4))
|
||||
for label, result in results.items():
|
||||
rows = [row for row in result["per_energy"] if row["n_true"] > 100]
|
||||
if len(rows) < 2:
|
||||
continue
|
||||
centre = np.array([0.5 * (r["energy_low"] + r["energy_high"]) for r in rows])
|
||||
axes[0].plot(centre, [r["efficiency"] * 100 for r in rows], "o-", label=f"occ {label}")
|
||||
axes[1].plot(centre, [r["sigma68_1d_um"] for r in rows], "o-", label=f"occ {label}")
|
||||
axes[2].plot(centre, [r["energy_rms_keV"] for r in rows], "o-", label=f"occ {label}")
|
||||
for axis, ylabel in zip(axes, ("efficiency [%]", r"position $\sigma_{68}$ [$\mu$m]",
|
||||
r"energy RMS [keV]")):
|
||||
axis.set_xlabel("photon energy [keV]")
|
||||
axis.set_ylabel(ylabel)
|
||||
axis.grid(alpha=.3)
|
||||
axis.legend(fontsize=8)
|
||||
figure.tight_layout()
|
||||
figure.savefig(output_dir / "Plots" / "eval_vs_energy.png", dpi=130)
|
||||
plt.close(figure)
|
||||
|
||||
|
||||
def plot_all(results, output_dir, working_threshold):
|
||||
occupancy = np.array([r["occupancy"] for r in results.values()]) * 100
|
||||
efficiency = np.array([r["summary"]["efficiency"] for r in results.values()]) * 100
|
||||
purity = np.array([r["summary"]["purity"] for r in results.values()]) * 100
|
||||
rms = np.array([r["summary"]["rms_1d_um"] for r in results.values()])
|
||||
sigma68 = np.array([r["summary"]["sigma68_1d_um"] for r in results.values()])
|
||||
|
||||
figure, axes = plt.subplots(2, 3, figsize=(16, 9))
|
||||
|
||||
axes[0, 0].plot(occupancy, efficiency, "o-", label="efficiency")
|
||||
axes[0, 0].plot(occupancy, purity, "s-", label="purity")
|
||||
axes[0, 0].set_xlabel("occupancy [%]")
|
||||
axes[0, 0].set_ylabel("[%]")
|
||||
axes[0, 0].set_title(f"working point p>{working_threshold}")
|
||||
axes[0, 0].legend()
|
||||
axes[0, 0].grid(alpha=.3)
|
||||
|
||||
axes[0, 1].plot(occupancy, rms, "o-", label="RMS")
|
||||
axes[0, 1].plot(occupancy, sigma68, "s-", label=r"$\sigma_{68}$")
|
||||
axes[0, 1].set_xlabel("occupancy [%]")
|
||||
axes[0, 1].set_ylabel(r"residual [$\mu$m]")
|
||||
axes[0, 1].legend()
|
||||
axes[0, 1].grid(alpha=.3)
|
||||
|
||||
for label, result in results.items():
|
||||
roc = np.array(result["roc"])
|
||||
axes[0, 2].plot(roc[:, 2] * 100, roc[:, 1] * 100, ".-", label=f"occ {label}")
|
||||
axes[0, 2].set_xlabel("purity [%]")
|
||||
axes[0, 2].set_ylabel("efficiency [%]")
|
||||
axes[0, 2].set_title("threshold sweep")
|
||||
axes[0, 2].legend(fontsize=8)
|
||||
axes[0, 2].grid(alpha=.3)
|
||||
|
||||
for label, result in results.items():
|
||||
residual = result["_differential"]["residual_um"]
|
||||
if len(residual):
|
||||
axes[1, 0].hist(residual.ravel(), bins=200, range=(-6, 6),
|
||||
histtype="step", density=True, label=f"occ {label}")
|
||||
axes[1, 0].set_xlabel(r"residual [$\mu$m]")
|
||||
axes[1, 0].set_yscale("log")
|
||||
axes[1, 0].legend(fontsize=8)
|
||||
axes[1, 0].grid(alpha=.3)
|
||||
|
||||
for label, result in results.items():
|
||||
pair = result["pair"]
|
||||
edges = np.array(pair["edges"])
|
||||
pair_efficiency = np.array(pair["efficiency"]) * 100
|
||||
populated = np.array(pair["n_true"]) > 50
|
||||
axes[1, 1].plot(edges[populated], pair_efficiency[populated], "o-",
|
||||
label=f"occ {label}")
|
||||
axes[1, 1].set_xlabel("distance to nearest other photon [px]")
|
||||
axes[1, 1].set_ylabel("efficiency [%]")
|
||||
axes[1, 1].set_xlim(0, 6)
|
||||
axes[1, 1].set_title("pair resolution")
|
||||
axes[1, 1].legend(fontsize=8)
|
||||
axes[1, 1].grid(alpha=.3)
|
||||
|
||||
middle = list(results.values())[len(results) // 2]
|
||||
subpixel = middle["_differential"]["subpixel"]
|
||||
if len(subpixel):
|
||||
histogram, _, _ = np.histogram2d(subpixel[:, 0], subpixel[:, 1], bins=25,
|
||||
range=[[0, 1], [0, 1]])
|
||||
image = axes[1, 2].imshow(histogram.T / histogram.mean(), origin="lower",
|
||||
extent=[0, 1, 0, 1], vmin=0.5, vmax=1.5,
|
||||
cmap="viridis")
|
||||
plt.colorbar(image, ax=axes[1, 2])
|
||||
axes[1, 2].set_title(f"sub-pixel uniformity (occ {middle['occupancy']*100:g}%)\n"
|
||||
"flat = no grid artefact (risk R1)")
|
||||
axes[1, 2].set_xlabel("sub-pixel x")
|
||||
axes[1, 2].set_ylabel("sub-pixel y")
|
||||
|
||||
figure.tight_layout()
|
||||
figure.savefig(output_dir / "Plots" / "eval_scan.png", dpi=130)
|
||||
plt.close(figure)
|
||||
|
||||
# ---- residual vs depth, vs isolation, and energy ---------------------
|
||||
figure, axes = plt.subplots(1, 3, figsize=(16, 4.4))
|
||||
depth_edges = np.linspace(0, 650, 14)
|
||||
neighbour_edges = np.array([0, .5, 1, 1.5, 2, 2.5, 3, 4, 6, 10])
|
||||
for label, result in results.items():
|
||||
differential = result["_differential"]
|
||||
if not len(differential["z_um"]):
|
||||
continue
|
||||
depth_bin = np.digitize(differential["z_um"], depth_edges) - 1
|
||||
depth_rms = [differential["residual_um"][depth_bin == index].std()
|
||||
if (depth_bin == index).sum() > 50 else np.nan
|
||||
for index in range(len(depth_edges) - 1)]
|
||||
axes[0].plot(0.5 * (depth_edges[:-1] + depth_edges[1:]), depth_rms, "o-",
|
||||
label=f"occ {label}")
|
||||
|
||||
neighbour_bin = np.digitize(differential["nn_px"], neighbour_edges) - 1
|
||||
neighbour_rms = [differential["residual_um"][neighbour_bin == index].std()
|
||||
if (neighbour_bin == index).sum() > 50 else np.nan
|
||||
for index in range(len(neighbour_edges) - 1)]
|
||||
axes[1].plot(neighbour_edges[:-1], neighbour_rms, "o-", label=f"occ {label}")
|
||||
|
||||
if np.isfinite(differential["e_pred"]).any():
|
||||
axes[2].hist(differential["e_pred"] - differential["e_true"], bins=150,
|
||||
range=(-4, 4), histtype="step", density=True,
|
||||
label=f"occ {label}")
|
||||
axes[0].set_xlabel(r"conversion depth z [$\mu$m]")
|
||||
axes[0].set_ylabel(r"RMS [$\mu$m]")
|
||||
axes[0].set_title("resolution vs charge-cloud size")
|
||||
axes[0].legend(fontsize=8)
|
||||
axes[0].grid(alpha=.3)
|
||||
axes[1].set_xlabel("distance to nearest other photon [px]")
|
||||
axes[1].set_ylabel(r"RMS [$\mu$m]")
|
||||
axes[1].legend(fontsize=8)
|
||||
axes[1].grid(alpha=.3)
|
||||
axes[2].set_xlabel(r"$E_{pred}-E_{true}$ [keV]")
|
||||
axes[2].set_yscale("log")
|
||||
axes[2].legend(fontsize=8)
|
||||
axes[2].grid(alpha=.3)
|
||||
figure.tight_layout()
|
||||
figure.savefig(output_dir / "Plots" / "eval_differential.png", dpi=130)
|
||||
plt.close(figure)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,492 @@
|
||||
#!/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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -0,0 +1,363 @@
|
||||
#!/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)
|
||||
@@ -0,0 +1,148 @@
|
||||
#!/usr/bin/env bash
|
||||
# Overnight run (2026-09-27/28): poly 2-20 keV model for MOENCH040 @150 V, then
|
||||
# the 12 keV MaxIV knife edge at 0-3 filters and the 15 keV SOLEIL Siemens star.
|
||||
#
|
||||
# setsid nohup bash scripts/overnight_poly2to20keV_Moench040.sh \
|
||||
# > Results/overnight_poly2to20keV_Moench040/driver.log 2>&1 < /dev/null &
|
||||
# DRY_RUN=1 bash scripts/overnight_poly2to20keV_Moench040.sh # print the plan only
|
||||
#
|
||||
# Part A (this file): charge LUT -> MC frames -> training -> MC evaluation.
|
||||
# Part B (overnight_poly2to20keV_Moench040_partB.sh): measured runs + notebooks.
|
||||
# Part B is started at the end of part A once it has been marked ready
|
||||
# (state/partB.ready), so it can be finished and tested while A trains.
|
||||
#
|
||||
# Every stage logs to $BASE/logs/<stage>.log and leaves state/<stage>.done when
|
||||
# it succeeds; a rerun skips finished stages. Progress with times goes to
|
||||
# $BASE/status.md.
|
||||
|
||||
set -u
|
||||
cd /home/xie_x1/MLXID/GridBased
|
||||
PY=/home/xie_x1/software/miniforge3/envs/pytorch_nightly/bin/python
|
||||
BASE=Results/overnight_poly2to20keV_Moench040
|
||||
LOG=$BASE/logs
|
||||
STATE=$BASE/state
|
||||
STATUS=$BASE/status.md
|
||||
LUT_CONF=Configs/charge_lut_Moench040_2to20keV.yaml
|
||||
LUT=data/charge_lut_Moench040_150V_2to20keV.npz
|
||||
GEN=Configs/generate_frames_poly2to20keV_Moench040.yaml
|
||||
TRAIN=Configs/train_grid_poly2to20keV_Moench040_c96.yaml
|
||||
DATA=/mnt/sls_det_storage/moench_data/MLXID/Samples/GridBased/poly2to20keV_Moench040_150V
|
||||
EPOCHS=200
|
||||
DRY_RUN=${DRY_RUN:-0}
|
||||
mkdir -p "$LOG" "$STATE"
|
||||
|
||||
note() {
|
||||
[ "$DRY_RUN" = 1 ] || echo "- $(date '+%m-%d %H:%M:%S') $*" >> "$STATUS"
|
||||
echo "$(date '+%m-%d %H:%M:%S') $*"
|
||||
}
|
||||
|
||||
# run NAME COMMAND... -- skip if done, log, time, mark
|
||||
run() {
|
||||
local name=$1; shift
|
||||
if [ -f "$STATE/$name.done" ]; then
|
||||
note "$name: already done, skipped"; return 0
|
||||
fi
|
||||
if [ "$DRY_RUN" = 1 ]; then
|
||||
echo "[dry] $name: $*"; return 0
|
||||
fi
|
||||
note "$name: started"
|
||||
local started=$(date +%s)
|
||||
"$@" > "$LOG/$name.log" 2>&1
|
||||
local status=$?
|
||||
local minutes=$(( ($(date +%s) - started) / 60 ))
|
||||
if [ $status -eq 0 ]; then
|
||||
touch "$STATE/$name.done"
|
||||
note "$name: done (${minutes} min)"
|
||||
else
|
||||
note "$name: **FAILED** (exit $status after ${minutes} min), see logs/$name.log"
|
||||
fi
|
||||
return $status
|
||||
}
|
||||
|
||||
check_frames() {
|
||||
$PY - "$DATA" <<'EOF'
|
||||
import glob, sys
|
||||
import h5py
|
||||
folder = sys.argv[1]
|
||||
expected = {"train": (12, 100_000), "val": (2, 40_000), "test": (2, 40_000),
|
||||
"eval_scan": (5, 8_000)}
|
||||
bad = []
|
||||
for split, (n_files, n_frames) in expected.items():
|
||||
files = sorted(glob.glob(f"{folder}/{split}/*.h5"))
|
||||
if len(files) != n_files:
|
||||
bad.append(f"{split}: {len(files)} files, expected {n_files}")
|
||||
for path in files:
|
||||
with h5py.File(path) as f:
|
||||
frames, energy = f["frames"].shape[0], f["photons"][:, 2]
|
||||
if frames != n_frames:
|
||||
bad.append(f"{path}: {frames} frames, expected {n_frames}")
|
||||
if energy.min() < 1.999 or energy.max() > 20.001:
|
||||
bad.append(f"{path}: energy {energy.min():.3f}-{energy.max():.3f} keV")
|
||||
print("\n".join(bad) if bad else "all shards complete, energies within 2-20 keV")
|
||||
sys.exit(1 if bad else 0)
|
||||
EOF
|
||||
}
|
||||
|
||||
check_training() { # experiment dir
|
||||
[ -f "$1/Models/best.pth" ] && $PY -c "
|
||||
import json, sys
|
||||
h = json.load(open('$1/history.json'))
|
||||
print(len(h['val_loss']), 'epochs, best val loss', min(h['val_loss']))
|
||||
sys.exit(0 if len(h['val_loss']) == $EPOCHS else 1)"
|
||||
}
|
||||
|
||||
note "=== driver started (pid $$) ==="
|
||||
|
||||
# --- 1. charge-cloud table ---------------------------------------------------
|
||||
if ! run lut $PY scripts/01_build_charge_lut.py --config $LUT_CONF --workers 19; then
|
||||
note "charge LUT failed: stopping"; exit 1
|
||||
fi
|
||||
# diagnostic only: its comparison with the 9-parameter fits needs MOENCH040
|
||||
# ggdPar files that exist for a few energies only, so a failure does not stop
|
||||
# the chain
|
||||
run lut_validation $PY scripts/01b_validate_charge_lut.py --lut $LUT \
|
||||
--sensor Moench040 --outdir Results/charge_lut_Moench040_2to20keV
|
||||
|
||||
# --- 2. MC frames, 2-20 keV (~35 min) ----------------------------------------
|
||||
if ! run generate $PY scripts/02_generate_frames.py --config $GEN; then
|
||||
note "frame generation failed: stopping"; exit 1
|
||||
fi
|
||||
if ! run check_frames check_frames; then
|
||||
note "frame check failed: stopping"; exit 1
|
||||
fi
|
||||
|
||||
# --- 3. training, [32, 64, 96, 96], 200 epochs (~1 h 40 min) -----------------
|
||||
if ! run train $PY -u scripts/03_train.py --config $TRAIN; then
|
||||
note "training failed: stopping"; exit 1
|
||||
fi
|
||||
if [ "$DRY_RUN" = 1 ]; then
|
||||
experiment="Results/<date>_poly2to20keV_Moench040_c96_00"
|
||||
else
|
||||
[ -f "$STATE/experiment" ] || \
|
||||
grep -m1 "\[train\] output ->" "$LOG/train.log" | awk '{print $NF}' > "$STATE/experiment"
|
||||
experiment=$(cat "$STATE/experiment")
|
||||
if ! run check_train check_training "$experiment"; then
|
||||
note "training incomplete in $experiment: stopping"; exit 1
|
||||
fi
|
||||
fi
|
||||
checkpoint=$experiment/Models/best.pth
|
||||
|
||||
# --- 4. MC evaluation --------------------------------------------------------
|
||||
run evaluate $PY scripts/04_evaluate.py --checkpoint "$checkpoint"
|
||||
run response $PY scripts/07_energy_response.py --checkpoint "$checkpoint"
|
||||
note "=== part A finished: $experiment ==="
|
||||
|
||||
# --- 5. part B: measured runs and notebooks ----------------------------------
|
||||
PARTB=scripts/overnight_poly2to20keV_Moench040_partB.sh
|
||||
if [ "$DRY_RUN" = 1 ]; then
|
||||
echo "[dry] part B: bash $PARTB $experiment"; exit 0
|
||||
fi
|
||||
waited=0
|
||||
while [ ! -f "$STATE/partB.ready" ] && [ $waited -lt 7200 ]; do
|
||||
[ $waited -eq 0 ] && note "waiting for state/partB.ready"
|
||||
sleep 60; waited=$((waited + 60))
|
||||
done
|
||||
if [ ! -f "$STATE/partB.ready" ]; then
|
||||
note "part B was never marked ready: stopping after part A"; exit 1
|
||||
fi
|
||||
bash "$PARTB" "$experiment"
|
||||
@@ -0,0 +1,111 @@
|
||||
#!/usr/bin/env bash
|
||||
# Overnight run (2026-09-27/28), part B: the measured runs and the notebooks.
|
||||
# Started by overnight_poly2to20keV_Moench040.sh once part A (training) is done
|
||||
# and state/partB.ready exists; can also be run by hand:
|
||||
#
|
||||
# setsid nohup bash scripts/overnight_poly2to20keV_Moench040_partB.sh \
|
||||
# Results/<experiment> >> Results/overnight_poly2to20keV_Moench040/driver.log 2>&1 < /dev/null &
|
||||
#
|
||||
# Order: dark runs (fake rate, 1 min) -> knife edge 3/2/1/0 filters, edge + flat
|
||||
# (63 M frames of 101x101) -> knife-edge notebook -> Siemens star, four parts of
|
||||
# 500 files -> Siemens-star notebook -> HTML copies of both notebooks.
|
||||
# A failed run is logged and skipped; the notebooks report what is missing.
|
||||
|
||||
set -u
|
||||
cd /home/xie_x1/MLXID/GridBased
|
||||
ENV_BIN=/home/xie_x1/software/miniforge3/envs/pytorch_nightly/bin
|
||||
PY=$ENV_BIN/python
|
||||
BASE=Results/overnight_poly2to20keV_Moench040
|
||||
LOG=$BASE/logs
|
||||
STATE=$BASE/state
|
||||
STATUS=$BASE/status.md
|
||||
REPORT=$BASE/report
|
||||
KNIFE=Configs/infer_measurement_poly_Moench040_knife.yaml
|
||||
STAR=Configs/infer_measurement_poly_Moench040_star.yaml
|
||||
TAG=poly2to20
|
||||
mkdir -p "$LOG" "$STATE" "$REPORT"
|
||||
|
||||
experiment=${1:-$(cat "$STATE/experiment" 2>/dev/null)}
|
||||
checkpoint=$experiment/Models/best.pth
|
||||
if [ ! -f "$checkpoint" ]; then
|
||||
echo "- $(date '+%m-%d %H:%M:%S') part B: no checkpoint at '$checkpoint'" >> "$STATUS"
|
||||
exit 1
|
||||
fi
|
||||
# the CPU side of the inference is small per-batch work; the default thread
|
||||
# count (16) thrashes as soon as anything else runs (32 s instead of 0.1 s per
|
||||
# decode batch next to the frame generator)
|
||||
export OMP_NUM_THREADS=4 MKL_NUM_THREADS=4
|
||||
export PATH=$ENV_BIN:$PATH # the notebooks' kernel runs `python`
|
||||
export GRID_POLY_EXPERIMENT=$experiment GRID_POLY_TAG=$TAG
|
||||
|
||||
note() {
|
||||
echo "- $(date '+%m-%d %H:%M:%S') $*" >> "$STATUS"
|
||||
echo "$(date '+%m-%d %H:%M:%S') $*"
|
||||
}
|
||||
|
||||
run() {
|
||||
local name=$1; shift
|
||||
if [ -f "$STATE/$name.done" ]; then
|
||||
note "$name: already done, skipped"; return 0
|
||||
fi
|
||||
note "$name: started"
|
||||
local started=$(date +%s)
|
||||
"$@" > "$LOG/$name.log" 2>&1
|
||||
local status=$?
|
||||
local minutes=$(( ($(date +%s) - started) / 60 ))
|
||||
if [ $status -eq 0 ]; then
|
||||
touch "$STATE/$name.done"
|
||||
note "$name: done (${minutes} min) $(grep -o '([0-9]* frame/s)' "$LOG/$name.log" | tail -1)"
|
||||
else
|
||||
note "$name: **FAILED** (exit $status after ${minutes} min), see logs/$name.log"
|
||||
fi
|
||||
return $status
|
||||
}
|
||||
|
||||
notebook() { # notebook name -> executed in place + HTML copy in report/
|
||||
local name=$1
|
||||
jupyter nbconvert --to notebook --execute --inplace \
|
||||
--ExecutePreprocessor.timeout=3600 --ExecutePreprocessor.kernel_name=python3 \
|
||||
"Notebooks/$name.ipynb" && \
|
||||
jupyter nbconvert --to html --output-dir "$REPORT" "Notebooks/$name.ipynb"
|
||||
}
|
||||
|
||||
note "=== part B started (pid $$), model $checkpoint ==="
|
||||
|
||||
# Two lanes in parallel: the knife edge is bound by the GPU and the Python
|
||||
# loop (101x101 frames, ~10^4 frame/s), the Siemens star by reading the NFS
|
||||
# store and remapping ~820 photons per frame (~3000 frame/s), so they mostly
|
||||
# use different resources.
|
||||
|
||||
knife_lane() {
|
||||
# dark runs first: fake photons on 1000 frames each (a minute)
|
||||
for run_name in Pedestal_1filters_12keV Pedestal_2filters_12keV Pedestal_3filters_pos0_12keV; do
|
||||
run infer_$run_name $PY scripts/05_infer_measurement.py --config $KNIFE \
|
||||
--run $run_name --checkpoint "$checkpoint" --tag $TAG
|
||||
done
|
||||
# the four flux points, edge and flat (63 M frames, ~1.5-2 h)
|
||||
for filters in 3filters_pos0 2filters 1filters 0filters; do
|
||||
for kind in Edge Flat; do
|
||||
run infer_${kind}_${filters} $PY scripts/05_infer_measurement.py --config $KNIFE \
|
||||
--run ${kind}_${filters}_12keV --checkpoint "$checkpoint" --tag $TAG
|
||||
done
|
||||
done
|
||||
rm -f "$STATE/notebook_knife.done" # always re-run on what exists
|
||||
run notebook_knife notebook KnifeEdge_poly2to20keV_Moench040
|
||||
}
|
||||
|
||||
star_lane() {
|
||||
# lower left + lower right in one crop, four parts of 500 files (~2 h)
|
||||
for part in 0 1 2 3; do
|
||||
run infer_star_part$part $PY scripts/05_infer_measurement.py --config $STAR \
|
||||
--run SiemensStar_15keV_part$part --checkpoint "$checkpoint" --tag $TAG
|
||||
done
|
||||
rm -f "$STATE/notebook_star.done"
|
||||
run notebook_star notebook SiemensStar_poly2to20keV_15keV
|
||||
}
|
||||
|
||||
knife_lane &
|
||||
star_lane &
|
||||
wait
|
||||
|
||||
note "=== part B finished; notebooks in Notebooks/, HTML copies in $REPORT ==="
|
||||
@@ -0,0 +1,140 @@
|
||||
#!/usr/bin/env bash
|
||||
# Overnight run (2026-09-18): poly models trained on 3-30 keV, four network sizes.
|
||||
#
|
||||
# setsid nohup bash scripts/overnight_poly3to30keV.sh \
|
||||
# > Results/overnight_poly3to30keV/driver.log 2>&1 < /dev/null &
|
||||
# DRY_RUN=1 bash scripts/overnight_poly3to30keV.sh # print the plan only
|
||||
#
|
||||
# Every stage logs to Results/overnight_poly3to30keV/logs/<stage>.log and leaves
|
||||
# state/<stage>.done when it succeeds; a rerun skips finished stages, so the
|
||||
# script can simply be started again after a failure. Progress, with times, is
|
||||
# appended to Results/overnight_poly3to30keV/status.md. A failure of the frame
|
||||
# generation stops everything; a failure inside one model skips the rest of that
|
||||
# model and goes on with the next. The last stage writes report/report.md.
|
||||
|
||||
set -u
|
||||
cd /home/xie_x1/MLXID/GridBased
|
||||
PY=/home/xie_x1/software/miniforge3/envs/pytorch_nightly/bin/python
|
||||
BASE=Results/overnight_poly3to30keV
|
||||
LOG=$BASE/logs
|
||||
STATE=$BASE/state
|
||||
STATUS=$BASE/status.md
|
||||
INFER=Configs/infer_measurement_poly_Moench014.yaml
|
||||
GEN=Configs/generate_frames_poly3to30keV_Moench014.yaml
|
||||
DATA=/mnt/sls_det_storage/moench_data/MLXID/Samples/GridBased/poly3to30keV_Moench014_150V
|
||||
REFERENCE=Results/260901_poly5to30keV_Moench014_00 # the current 5-30 keV model
|
||||
MODELS="c128 c96 c64 r3"
|
||||
declare -A CHANNELS=([c128]="[32,64,128,128]" [c96]="[32,64,96,96]"
|
||||
[c64]="[32,64,64,64]" [r3]="[32,64,128]")
|
||||
DRY_RUN=${DRY_RUN:-0}
|
||||
mkdir -p "$LOG" "$STATE"
|
||||
|
||||
note() {
|
||||
[ "$DRY_RUN" = 1 ] || echo "- $(date '+%m-%d %H:%M:%S') $*" >> "$STATUS"
|
||||
echo "$(date '+%m-%d %H:%M:%S') $*"
|
||||
}
|
||||
|
||||
# run NAME COMMAND... -- skip if done, log, time, mark
|
||||
run() {
|
||||
local name=$1; shift
|
||||
if [ -f "$STATE/$name.done" ]; then
|
||||
note "$name: already done, skipped"; return 0
|
||||
fi
|
||||
if [ "$DRY_RUN" = 1 ]; then
|
||||
echo "[dry] $name: $*"; return 0
|
||||
fi
|
||||
note "$name: started"
|
||||
local started=$(date +%s)
|
||||
"$@" > "$LOG/$name.log" 2>&1
|
||||
local status=$?
|
||||
local minutes=$(( ($(date +%s) - started) / 60 ))
|
||||
if [ $status -eq 0 ]; then
|
||||
touch "$STATE/$name.done"
|
||||
note "$name: done (${minutes} min)"
|
||||
else
|
||||
note "$name: **FAILED** (exit $status after ${minutes} min), see logs/$name.log"
|
||||
fi
|
||||
return $status
|
||||
}
|
||||
|
||||
check_frames() {
|
||||
$PY - "$DATA" <<'EOF'
|
||||
import glob, sys
|
||||
import h5py
|
||||
folder = sys.argv[1]
|
||||
expected = {"train": (12, 100_000), "val": (2, 40_000), "test": (2, 40_000),
|
||||
"eval_scan": (5, 8_000)}
|
||||
bad = []
|
||||
for split, (n_files, n_frames) in expected.items():
|
||||
files = sorted(glob.glob(f"{folder}/{split}/*.h5"))
|
||||
if len(files) != n_files:
|
||||
bad.append(f"{split}: {len(files)} files, expected {n_files}")
|
||||
for path in files:
|
||||
with h5py.File(path) as f:
|
||||
frames, energy = f["frames"].shape[0], f["photons"][:, 2]
|
||||
if frames != n_frames:
|
||||
bad.append(f"{path}: {frames} frames, expected {n_frames}")
|
||||
if energy.min() < 2.999 or energy.max() > 30.001:
|
||||
bad.append(f"{path}: energy {energy.min():.3f}-{energy.max():.3f} keV")
|
||||
print("\n".join(bad) if bad else "all shards complete, energies within 3-30 keV")
|
||||
sys.exit(1 if bad else 0)
|
||||
EOF
|
||||
}
|
||||
|
||||
check_training() { # experiment dir
|
||||
[ -f "$1/Models/best.pth" ] && $PY -c "
|
||||
import json, sys
|
||||
h = json.load(open('$1/history.json'))
|
||||
print(len(h['val_loss']), 'epochs, best val loss', min(h['val_loss']))
|
||||
sys.exit(0 if len(h['val_loss']) == 300 else 1)"
|
||||
}
|
||||
|
||||
note "=== driver started (pid $$) ==="
|
||||
|
||||
# --- 1. charge-cloud table (done interactively before launch) ---------------
|
||||
run lut $PY scripts/01_build_charge_lut.py --config Configs/charge_lut_Moench014_3to30keV.yaml --workers 14
|
||||
run lut_validation $PY scripts/01b_validate_charge_lut.py \
|
||||
--lut data/charge_lut_Moench014_150V_3to30keV.npz --outdir Results/charge_lut_3to30keV
|
||||
|
||||
# --- 2. MC frames, 3-30 keV (~1 h) ------------------------------------------
|
||||
if ! run generate $PY scripts/02_generate_frames.py --config $GEN; then
|
||||
note "frame generation failed: stopping"; exit 1
|
||||
fi
|
||||
if ! run check_frames check_frames; then
|
||||
note "frame check failed: stopping"; exit 1
|
||||
fi
|
||||
|
||||
# --- 3. four models, one after the other (~3 h each) ------------------------
|
||||
# its XRF runs were moved to <run>_poly5to30 when the c96 model became the default
|
||||
compare_args=(--model "5-30 keV [32,64,128,128] (previous default)" "$REFERENCE" poly5to30)
|
||||
for model in $MODELS; do
|
||||
if ! run train_$model $PY -u scripts/03_train.py \
|
||||
--config Configs/train_grid_poly3to30keV_Moench014_$model.yaml; then
|
||||
note "$model: training failed, skipping its evaluation"; continue
|
||||
fi
|
||||
if [ "$DRY_RUN" = 1 ]; then
|
||||
experiment="Results/<date>_poly3to30keV_Moench014_${model}_00"
|
||||
else
|
||||
[ -f "$STATE/$model.experiment" ] || \
|
||||
grep -m1 "\[train\] output ->" "$LOG/train_$model.log" | awk '{print $NF}' \
|
||||
> "$STATE/$model.experiment"
|
||||
experiment=$(cat "$STATE/$model.experiment")
|
||||
if ! run check_train_$model check_training "$experiment"; then
|
||||
note "$model: training incomplete in $experiment, skipping its evaluation"
|
||||
continue
|
||||
fi
|
||||
fi
|
||||
checkpoint=$experiment/Models/best.pth
|
||||
run evaluate_$model $PY scripts/04_evaluate.py --checkpoint "$checkpoint"
|
||||
run response_$model $PY scripts/07_energy_response.py --checkpoint "$checkpoint"
|
||||
for xrf in AgMo_1_factor_10 CuCrCo_1_factor_10 Pedestal_CuCrCo_1; do
|
||||
run xrf_${xrf}_$model $PY scripts/05_infer_measurement.py --config $INFER \
|
||||
--run $xrf --checkpoint "$checkpoint" --tag $model
|
||||
done
|
||||
compare_args+=(--model "3-30 keV ${CHANNELS[$model]}" "$experiment" $model)
|
||||
done
|
||||
|
||||
# --- 4. comparison report ----------------------------------------------------
|
||||
rm -f "$STATE/compare.done" # always refresh the report
|
||||
run compare $PY scripts/08_compare_models.py --out $BASE/report "${compare_args[@]}"
|
||||
note "=== driver finished; report: $BASE/report/report.md ==="
|
||||
@@ -0,0 +1,83 @@
|
||||
#!/usr/bin/env bash
|
||||
# Siemens star again with the masked column common mode (2026-09-28).
|
||||
# The 09-27 run used the plain column median, which the ~5 % occupancy of the
|
||||
# star pulls up: the 15 keV line came out at 14.47 keV. Outputs carry the tag
|
||||
# _poly2to20_cmmasked, next to the old _poly2to20 ones.
|
||||
#
|
||||
# setsid nohup bash scripts/rerun_star_column_masked.sh \
|
||||
# >> Results/overnight_poly2to20keV_Moench040/driver.log 2>&1 < /dev/null &
|
||||
#
|
||||
# After parts 0-1 (half of the 20 M frames) the notebook is executed on them and
|
||||
# kept as SiemensStar_poly2to20keV_15keV_50pct.ipynb (+ HTML, summary json), so
|
||||
# a result exists even if the rest does not finish; after all four parts the
|
||||
# main notebook is executed in place.
|
||||
|
||||
set -u
|
||||
cd /home/xie_x1/MLXID/GridBased
|
||||
ENV_BIN=/home/xie_x1/software/miniforge3/envs/pytorch_nightly/bin
|
||||
PY=$ENV_BIN/python
|
||||
BASE=Results/overnight_poly2to20keV_Moench040
|
||||
LOG=$BASE/logs
|
||||
STATE=$BASE/state
|
||||
STATUS=$BASE/status.md
|
||||
REPORT=$BASE/report
|
||||
STAR=Configs/infer_measurement_poly_Moench040_star.yaml
|
||||
TAG=poly2to20_cmmasked
|
||||
experiment=$(cat "$STATE/experiment")
|
||||
checkpoint=$experiment/Models/best.pth
|
||||
export OMP_NUM_THREADS=4 MKL_NUM_THREADS=4
|
||||
export PATH=$ENV_BIN:$PATH GRID_POLY_TAG=$TAG
|
||||
NOTEBOOK=Notebooks/SiemensStar_poly2to20keV_15keV
|
||||
|
||||
note() {
|
||||
echo "- $(date '+%m-%d %H:%M:%S') $*" >> "$STATUS"
|
||||
echo "$(date '+%m-%d %H:%M:%S') $*"
|
||||
}
|
||||
|
||||
run() {
|
||||
local name=$1; shift
|
||||
if [ -f "$STATE/$name.done" ]; then
|
||||
note "$name: already done, skipped"; return 0
|
||||
fi
|
||||
note "$name: started"
|
||||
local started=$(date +%s)
|
||||
"$@" > "$LOG/$name.log" 2>&1
|
||||
local status=$?
|
||||
local minutes=$(( ($(date +%s) - started) / 60 ))
|
||||
if [ $status -eq 0 ]; then
|
||||
touch "$STATE/$name.done"
|
||||
note "$name: done (${minutes} min) $(grep -o '([0-9]* frame/s)' "$LOG/$name.log" | tail -1)"
|
||||
else
|
||||
note "$name: **FAILED** (exit $status after ${minutes} min), see logs/$name.log"
|
||||
fi
|
||||
return $status
|
||||
}
|
||||
|
||||
snapshot() { # executed copy of the notebook on the first two parts
|
||||
GRID_STAR_RUNS=SiemensStar_15keV_part0,SiemensStar_15keV_part1 \
|
||||
GRID_STAR_SUMMARY=siemens_star_poly2to20keV_summary_50pct.json \
|
||||
jupyter nbconvert --to notebook --execute --ExecutePreprocessor.timeout=3600 \
|
||||
--ExecutePreprocessor.kernel_name=python3 \
|
||||
--output SiemensStar_poly2to20keV_15keV_50pct.ipynb "$NOTEBOOK.ipynb" && \
|
||||
jupyter nbconvert --to html --output-dir "$REPORT" "${NOTEBOOK}_50pct.ipynb"
|
||||
}
|
||||
|
||||
full() {
|
||||
jupyter nbconvert --to notebook --execute --inplace --ExecutePreprocessor.timeout=3600 \
|
||||
--ExecutePreprocessor.kernel_name=python3 "$NOTEBOOK.ipynb" && \
|
||||
jupyter nbconvert --to html --output-dir "$REPORT" "$NOTEBOOK.ipynb"
|
||||
}
|
||||
|
||||
note "=== Siemens star rerun, column_masked (pid $$), tag $TAG ==="
|
||||
for part in 0 1; do
|
||||
run infer_star_cmmasked_part$part $PY scripts/05_infer_measurement.py --config $STAR \
|
||||
--run SiemensStar_15keV_part$part --checkpoint "$checkpoint" --tag $TAG
|
||||
done
|
||||
run notebook_star_cmmasked_50pct snapshot
|
||||
for part in 2 3; do
|
||||
run infer_star_cmmasked_part$part $PY scripts/05_infer_measurement.py --config $STAR \
|
||||
--run SiemensStar_15keV_part$part --checkpoint "$checkpoint" --tag $TAG
|
||||
done
|
||||
rm -f "$STATE/notebook_star_cmmasked.done"
|
||||
run notebook_star_cmmasked full
|
||||
note "=== Siemens star rerun finished ==="
|
||||
Reference in new issue
Block a user