""" 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.')