xie_x1andClaude Sonnet 5
2683fd3576
MC generation -> pile-up assembly -> CNN vs. classical eta-interpolation (scored against MC ground truth) -> optional training -> qualitative inference on real measurement data, self-contained and runnable from a single conda env (see README.md). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
142 lines
6.1 KiB
Python
142 lines
6.1 KiB
Python
"""
|
|
Assemble 2-photon and 3-photon pile-up clusters from the single-photon MC
|
|
clusters produced by 01_generate_mc_singlephoton.py, following the same
|
|
algorithm as DataProcess/GeneratePileupSample.py: place photons near each
|
|
other, verify with aare's cluster finder that they merged into exactly one
|
|
connected cluster, then crop around that cluster's energy-weighted
|
|
centroid. Ground truth positions are exact.
|
|
|
|
Differs from DataProcess/GeneratePileupSample.py in one way: this places
|
|
extra photons within a small radius of the first one (pile-up close to
|
|
guaranteed by construction) instead of uniformly anywhere in the frame,
|
|
since the latter's accept rate is too low (<1%) for a quick demo run.
|
|
|
|
A fresh aare.VarClusterFinder is created per frame. Reusing one across many
|
|
find_clusters_X() calls -- as both the original script and an earlier
|
|
version of this one did -- accumulates stale hits internally and starts
|
|
returning dozens of spurious clusters even on pure noise after a few
|
|
thousand calls; this likely also affects DataProcess/GeneratePileupSample.py's
|
|
real accept rate, not just this demo's.
|
|
|
|
find_clusters_X() also zeroes the found cluster's pixels in the input frame
|
|
in place, so the restoring line below (`frame[ys, xs] = enes`) is required
|
|
before cropping -- it's not a no-op, despite looking like one.
|
|
|
|
Output: data/pileup_samples/{2,3}photon_pileup_12keV_demo.npz
|
|
samples: (N, size, size) float32, size=6 for 2ph, 9 for 3ph
|
|
labels: (N, n_photons, 4) float32 [x, y, z, energy] per photon
|
|
|
|
Needs: numpy, aare (`pip install aare`). Use the `mlxid_demo` conda env
|
|
(see README.md for how to build it).
|
|
"""
|
|
from pathlib import Path
|
|
import numpy as np
|
|
import aare
|
|
|
|
HERE = Path(__file__).resolve().parent.parent
|
|
MC_SAMPLE_DIR = HERE / 'data/mc_samples'
|
|
OUTPUT_DIR = HERE / 'data/pileup_samples'
|
|
|
|
NOISE_KEV = 0.13
|
|
N_EVENTS_DEMO = 40_000 ### per photon count; production uses ~1e6/thread
|
|
PLACEMENT_RADIUS = 2 ### extra photons land within +/- this many pixels of photon 0
|
|
|
|
PILEUP_CONFIGS = [
|
|
{'n_photons': 2, 'sample_size': 6}, ### matches DeepLearning train_2photon.yaml n_size
|
|
{'n_photons': 3, 'sample_size': 9}, ### matches DeepLearning train_3photon.yaml n_size
|
|
]
|
|
|
|
|
|
def load_single_photon_clusters():
|
|
all_samples, all_labels = [], []
|
|
for f in sorted(MC_SAMPLE_DIR.glob('12keV_Moench040_150V_*.npz')):
|
|
data = np.load(f)
|
|
all_samples.append(data['samples'])
|
|
all_labels.append(data['labels'])
|
|
samples = np.concatenate(all_samples, axis=0)
|
|
labels = np.concatenate(all_labels, axis=0)
|
|
if samples.shape[-1] == 5:
|
|
### extract the central 3x3 part of the 5x5 cluster, adjusting labels
|
|
samples = samples[:, 1:4, 1:4]
|
|
labels = labels.copy()
|
|
labels[:, :2] -= 1
|
|
return samples, labels
|
|
|
|
|
|
def assemble(n_photons, sample_size, single_samples, single_labels, n_events):
|
|
frame_size = sample_size * 3 ### same margin as DataProcess/GeneratePileupSample.py
|
|
frame_center = frame_size // 2 - 1
|
|
|
|
pileup_samples, pileup_labels = [], []
|
|
n_done, n_tried = 0, 0
|
|
while n_done < n_events:
|
|
n_tried += 1
|
|
frame = np.random.normal(0, NOISE_KEV, size=(frame_size, frame_size))
|
|
labels = []
|
|
for k in range(n_photons):
|
|
idx = np.random.randint(0, single_samples.shape[0])
|
|
s = single_samples[idx]
|
|
lab = single_labels[idx].copy()
|
|
if k == 0:
|
|
ref_x, ref_y = frame_center, frame_center
|
|
else:
|
|
ref_x = frame_center + np.random.randint(-PLACEMENT_RADIUS, PLACEMENT_RADIUS + 1)
|
|
ref_y = frame_center + np.random.randint(-PLACEMENT_RADIUS, PLACEMENT_RADIUS + 1)
|
|
frame[ref_y:ref_y + 3, ref_x:ref_x + 3] += s
|
|
lab[0] += ref_x
|
|
lab[1] += ref_y
|
|
labels.append(lab)
|
|
|
|
CF = aare.VarClusterFinder((frame_size, frame_size), 5)
|
|
CF.set_peripheralThresholdFactor(3)
|
|
CF.set_noiseMap(np.ones((frame_size, frame_size)) * NOISE_KEV)
|
|
CF.set_numberOfNeighbours(4)
|
|
CF.set_empty_surroundingPixels(False)
|
|
CF.find_clusters_X(frame)
|
|
clusters = CF.hits()
|
|
if len(clusters) != 1:
|
|
continue ### photons didn't merge into a single cluster -> discard
|
|
|
|
clusterSize = clusters['size'][0]
|
|
enes = clusters['enes'][0][:clusterSize]
|
|
xs = clusters['cols'][0][:clusterSize]
|
|
ys = clusters['rows'][0][:clusterSize]
|
|
frame[ys, xs] = enes ### find_clusters_X zeroes cluster pixels in-place; restore them before cropping
|
|
x_center = np.sum(xs * enes) / np.sum(enes)
|
|
y_center = np.sum(ys * enes) / np.sum(enes)
|
|
|
|
ref_x = int(x_center - sample_size / 2) + 1
|
|
ref_y = int(y_center - sample_size / 2) + 1
|
|
if ref_x < 0 or ref_y < 0 or ref_x + sample_size > frame_size or ref_y + sample_size > frame_size:
|
|
continue ### centroid landed too close to the frame border
|
|
cropped = frame[ref_y:ref_y + sample_size, ref_x:ref_x + sample_size]
|
|
|
|
for lab in labels:
|
|
lab[0] -= ref_x
|
|
lab[1] -= ref_y
|
|
|
|
pileup_samples.append(cropped)
|
|
pileup_labels.append(np.array(labels))
|
|
n_done += 1
|
|
if n_done % 2000 == 0:
|
|
print(f' [{n_photons}ph] {n_done}/{n_events} events '
|
|
f'(accept rate {n_done / n_tried:.2%})')
|
|
|
|
return np.array(pileup_samples), np.array(pileup_labels)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
single_samples, single_labels = load_single_photon_clusters()
|
|
print(f'[02] Loaded {single_samples.shape[0]} single-photon clusters as building blocks')
|
|
|
|
for cfg in PILEUP_CONFIGS:
|
|
n_photons, sample_size = cfg['n_photons'], cfg['sample_size']
|
|
print(f'[02] Assembling {n_photons}-photon pile-up samples ({sample_size}x{sample_size})...')
|
|
samples, labels = assemble(n_photons, sample_size, single_samples, single_labels, N_EVENTS_DEMO)
|
|
out_file = OUTPUT_DIR / f'{n_photons}photon_pileup_12keV_demo.npz'
|
|
np.savez(out_file, samples=samples, labels=labels)
|
|
print(f'[02] Saved {samples.shape[0]} events -> {out_file}')
|
|
|
|
print('[02] Done.')
|