113 lines
4.4 KiB
Python
113 lines
4.4 KiB
Python
from EtaInterpolationFunctions import *
|
|
import matplotlib
|
|
matplotlib.use('Agg')
|
|
|
|
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',
|
|
'2603MaxIV_Edge2Filters_12keV', '2603MaxIV_Flat2Filters_12keV',
|
|
'2603MaxIV_Edge1Filters_12keV', '2603MaxIV_Flat1Filters_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), X, Y, according build_xy_lut_Rosenblatt's convention
|
|
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]) + 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
|
|
|
|
etaX = np.clip(etaX, 0, 1)
|
|
etaY = np.clip(etaY, 0, 1)
|
|
|
|
_eta_hist2D += np.histogram2d(etaX, etaY, bins=NEtaBins, range=[[0, 1], [0, 1]])[0] ### X, Y convention according to build_xy_lut_Rosenblatt's convention
|
|
return _eta_hist2D
|
|
|
|
print('Building eta histogram from clusters...')
|
|
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'][:] ### x, y
|
|
|
|
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
|
|
|
|
etaX = np.clip(etaX, 0, 1)
|
|
etaY = np.clip(etaY, 0, 1)
|
|
|
|
x_frac, y_frac = bilinear_xy_lookup(etaX, etaY, U_tab, V_tab)
|
|
|
|
x_ref = referencePoints[:, 0]
|
|
y_ref = referencePoints[:, 1]
|
|
|
|
x = x_ref + x_frac + clusters.shape[1]//2
|
|
y = y_ref + 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(
|
|
y%1, x%1,
|
|
bins=interpolationBins,
|
|
range=[[0, 1], [0, 1]]
|
|
)[0]
|
|
return (_image, _subpixel_positions)
|
|
|
|
print('Reconstructing positions from clusters and accumulating hits...')
|
|
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) |