Add eta interploatioin using clusters from DataProcess

Co-authored-by: Copilot <copilot@github.com>
This commit is contained in:
2026-05-28 16:28:03 +02:00
co-authored by Copilot
parent 65a6b77271
commit 8716f175dd
+81
View File
@@ -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)