From 8716f175ddd9a41d51bc37614a7207fbeadbb118 Mon Sep 17 00:00:00 2001
From: "xiangyu.xie" <xiangyu.xie@psi.ch>
Date: Thu, 28 May 2026 16:28:03 +0200
Subject: [PATCH] Add eta interploatioin using clusters from DataProcess

Co-authored-by: Copilot <copilot@github.com>
---
 etaInterpolationFromClusters.py | 81 +++++++++++++++++++++++++++++++++
 1 file changed, 81 insertions(+)
 create mode 100644 etaInterpolationFromClusters.py

diff --git a/etaInterpolationFromClusters.py b/etaInterpolationFromClusters.py
new file mode 100644
index 0000000..ed751dc
--- /dev/null
+++ b/etaInterpolationFromClusters.py
@@ -0,0 +1,81 @@
+from EtaInterpolationFunctions import *
+import numpy as np
+from array import array
+import matplotlib.pyplot as plt
+import h5py
+from multiprocessing import Pool
+from argparse import ArgumentParser
+
+parser = ArgumentParser()
+parser.add_argument('--runname', '-r', type=str,  default='2603MaxIV_Edge3Filters_pos0_12keV', choices=['2603MaxIV_Edge3Filters_pos0_12keV', '2603MaxIV_FlatField3Filters_pos0_12keV'], help='Name of the run, used to locate the cluster files and save the results')
+args = parser.parse_args()
+RUNNAME = args.runname
+
+
+### cluster files
+clusterFiles = [f'/home/xie_x1/MLXID/DataProcess/Samples/{RUNNAME}/1Photon_CS3_chunk{i}.h5' for i in range(16)]
+
+NEtaBins = 201
+### generate eta map
+def get_eta_hist2D(clusterFile):
+    with h5py.File(clusterFile, 'r') as f:
+        _eta_hist2D = np.zeros((NEtaBins, NEtaBins)) ### shape (201, 201), Y, X
+        clusters = f['clusters'][:] ### shape (N, 3, 3), N, Y, X
+        sum_clusters = np.sum(clusters, axis=(1, 2)) ### shape (N,)
+        position_weights = np.array([-1, 0, 1]) 
+        etaX =  ((np.sum(clusters, axis=1) * position_weights).sum(axis=-1) / sum_clusters + 0.5) % 1
+        etaY =  ((np.sum(clusters, axis=2) * position_weights).sum(axis=-1) / sum_clusters + 0.5) % 1
+
+        _eta_hist2D += np.histogram2d(etaY, etaX, bins=NEtaBins, range=[[0, 1], [0, 1]])[0]
+    return _eta_hist2D
+
+eta_hist2D = np.zeros((NEtaBins, NEtaBins))
+with Pool(16) as pool:
+    results = pool.map(get_eta_hist2D, clusterFiles)
+    for res in results:
+        eta_hist2D += res
+
+# generate U, V lookup tables
+U_tab, V_tab = build_xy_lut_Rosenblatt(eta_hist2D) ### or using build_xy_lut_DoubleCDF
+
+Roi = [1, 100, 1, 100]
+interpolationBins = 10
+
+### slightly better LUTs from strictly selected clusters in /home/xie_x1/MLXID/EtaInterpolation/Examples/etaInterpolation_edge_12keV.ipynb
+# U_tab = np.load('/home/xi/
+clusterFiles = [f'/home/xie_x1/MLXID/DataProcess/Samples/{RUNNAME}/1Photon_CS3_chunk{i}.h5' for i in range(16)]
+
+def reconstruct_position(clusterFile):
+    _image = np.zeros(((Roi[1]-Roi[0])*interpolationBins, (Roi[3]-Roi[2])*interpolationBins))
+    _subpixel_positions = np.zeros((interpolationBins, interpolationBins))
+    with h5py.File(clusterFile, 'r') as f:
+        clusters = f['clusters'][:] ### shape (N, 3, 3), N, Y, X
+        referencePoints = f['referencePoint'][:] ### shape (N, 2), N, (y, x)
+        sum_clusters = np.sum(clusters, axis=(1, 2)) ### shape (N,)
+        position_weights = np.array([-1, 0, 1]) + 0.5 
+        etaX =  (np.sum(clusters, axis=1) * position_weights).sum(axis=-1) / sum_clusters
+        etaY =  (np.sum(clusters, axis=2) * position_weights).sum(axis=-1) / sum_clusters
+        x_frac, y_frac = bilinear_xy_lookup(etaX, etaY, U_tab, V_tab)
+        x = referencePoints[:, 0] + x_frac + clusters.shape[1]//2 
+        y = referencePoints[:, 1] + y_frac + clusters.shape[2]//2
+        _image += np.histogram2d(y, x, bins=(Roi[1]-Roi[0])*interpolationBins, range=[[Roi[0], Roi[1]], [Roi[2], Roi[3]]])[0]
+        _subpixel_positions += np.histogram2d(x%1, y%1, bins=interpolationBins, range=[[0, 1], [0, 1]])[0]
+    return (_image, _subpixel_positions)
+
+reconstructed_image = np.zeros(((Roi[1]-Roi[0])*interpolationBins, (Roi[3]-Roi[2])*interpolationBins))
+subpixel_positions = np.zeros((interpolationBins, interpolationBins))
+with Pool(16) as pool:
+    results = pool.map(reconstruct_position, clusterFiles)
+    for res in results:
+        reconstructed_image += res[0]
+        subpixel_positions += res[1]
+
+plt.figure(figsize=(16, 8), dpi=300)
+plt.subplot(1, 2, 1)
+plt.imshow(reconstructed_image, cmap='viridis', origin='lower')
+plt.colorbar()
+plt.subplot(1, 2, 2)
+plt.imshow(subpixel_positions, cmap='viridis', origin='lower')
+plt.colorbar()
+
+np.save(f'Results/{RUNNAME}_fromClusters.npy', reconstructed_image)
\ No newline at end of file