Files
DataProcess/UsefulFuncs.py

283 lines
13 KiB
Python

import os
import numpy as np
from ROOT import TH1D, TH2D, TCanvas, TF1
from multiprocessing import Pool
from array import array
import h5py
from multiprocessing import Pool
import aare
_cfg = {} ### global config
_aduToKev3DMap = None ### global caliFile
_pedestalAduFrame = None ### global pedestalAduFrame
_noiseEneFrame = None ### global noiseEneFrame
nChunks = 16
def init(cfg):
global _cfg
global _aduToKev3DMap
_cfg = cfg
_aduToKev3DMap = cfg['caliFile']
Roi = cfg['Roi']
global nChunks
nChunks = _cfg.get('NChunks', 16)
def convertAduFrameToEnergyFrame(aduFrame):
if _aduToKev3DMap is None:
return aduFrame
NX, NY = _cfg['NX'], _cfg['NY']
if _aduToKev3DMap.shape[0] == 1640:
idxFrame = (np.abs(aduFrame//10)).astype(np.int32)
idxFrame[idxFrame > _aduToKev3DMap.shape[0]-2] = _aduToKev3DMap.shape[0]-2
NY, NX = _aduToKev3DMap.shape[1:3]
_energyFrame0 = _aduToKev3DMap[idxFrame, np.arange(NY).reshape(-1, 1), np.arange(NX)]
_energyFrame1 = _aduToKev3DMap[idxFrame+1, np.arange(NY).reshape(-1, 1), np.arange(NX)]
energyFrame = _energyFrame0 + (_energyFrame1 - _energyFrame0)/10 * (np.abs(aduFrame) - idxFrame*10)
energyFrame *= np.sign(aduFrame)
else:
aduPerBin = 10
nExtraBins = 1000 // aduPerBin
idxFrame = (aduFrame//aduPerBin).astype(np.int32) + nExtraBins
idxFrame[idxFrame > _aduToKev3DMap.shape[0]-2] = _aduToKev3DMap.shape[0]-2
NY, NX = _aduToKev3DMap.shape[1:3]
_energyFrame0 = _aduToKev3DMap[idxFrame, np.arange(NY).reshape(-1, 1), np.arange(NX)]
_energyFrame1 = _aduToKev3DMap[idxFrame+1, np.arange(NY).reshape(-1, 1), np.arange(NX)]
energyFrame = _energyFrame0 + (_energyFrame1 - _energyFrame0)/aduPerBin * (aduFrame - (idxFrame-nExtraBins)*aduPerBin)
if 'energyScalingFactor' in _cfg:
energyFrame *= _cfg['energyScalingFactor']
return energyFrame
def bookHistograms(energy, suffix='', energyBinWidth=0.1, isMC=False):
histMaxEnergy = energy * 5.5
nEnergyBins = int(histMaxEnergy // energyBinWidth)
roi_x0, roi_x1, roi_y0, roi_y1 = _cfg['Roi'][0], _cfg['Roi'][1], _cfg['Roi'][2], _cfg['Roi'][3]
roi_width = roi_x1 - roi_x0
roi_height = roi_y1 - roi_y0
interpolationBins = 10
histograms = {}
histograms['h1_ClusterEnergy'] = TH1D(
f'h1_ClusterEnergy{suffix}', f'h1_ClusterEnergy{suffix}', nEnergyBins, -1, histMaxEnergy
)
histograms['h1_ClusterEnergy'].SetTitle('Cluster Energy;Energy [keV];Counts')
histograms['h1_PixelEnergy'] = TH1D(
f'h1_PixelEnergy{suffix}', f'h1_PixelEnergy{suffix}', nEnergyBins, -1, histMaxEnergy
)
histograms['h1_PixelEnergy'].SetTitle('Pixel Energy;Energy [keV];Counts')
histograms['h2_HitsSumFrame'] = TH2D(
f'h2_HitsSumFrame{suffix}', f'h2_HitsSumFrame{suffix}', roi_width, roi_x0, roi_x1, roi_height, roi_y0, roi_y1
)
histograms['h2_HitsSumFrame'].SetTitle('Hits Sum Frame;X [pixel];Y [pixel];Energy Sum [keV]')
histograms['h2_DoubleHitsSumFrame'] = TH2D(
f'h2_DoubleHitsSumFrame{suffix}', f'h2_DoubleHitsSumFrame{suffix}', roi_width, roi_x0, roi_x1, roi_height, roi_y0, roi_y1
)
histograms['h2_DoubleHitsSumFrame'].SetTitle('Double Hits Sum Frame;X [pixel];Y [pixel];Energy Sum [keV]')
histograms['h2_ClusterSizeVsEnergy'] = TH2D(
f'h2_ClusterSizeVsEnergy{suffix}', f'h2_ClusterSizeVsEnergy{suffix}', nEnergyBins, -1, histMaxEnergy, 30, 0, 30
)
histograms['h2_ClusterSizeVsEnergy'].SetTitle('Cluster Size vs Energy;Energy [keV];Cluster Size [pixel]')
histograms['h2_DiameterVsEnergy'] = TH2D(
f'h2_DiameterVsEnergy{suffix}', f'h2_DiameterVsEnergy{suffix}', nEnergyBins, -1, histMaxEnergy, 10, 0, 10
)
histograms['h2_DiameterVsEnergy'].SetTitle('Cluster Diameter vs Energy;Energy [keV];Cluster Diameter [pixel]')
for hist in histograms.values():
hist.SetDirectory(0)
return histograms
def _processFrames(idxChunk): ### for both single and double photon events, using
startFrame, endFrame = _cfg['NFrame'] * idxChunk // nChunks, min(_cfg['NFrame'] * (idxChunk + 1) // nChunks, _cfg['NFrame'])
print(f'Processing frames from {startFrame} to {endFrame}...')
Energy = _cfg['energy']
selectionRange = _cfg['selectionRange']
signalFileNames = _cfg['signalFileNames']
NX = _cfg['NX']
NY = _cfg['NY']
headerSize = _cfg['headerSize']
Roi = _cfg['Roi']
nFramePerFile = os.path.getsize(f'{signalFileNames[0]}') // (NX * NY + 56) // 2
_hists = bookHistograms(Energy, f'_{idxChunk}')
### create a cluster finder
nxROI = Roi[1] - Roi[0]
nyROI = Roi[3] - Roi[2]
CF = aare.VarClusterFinder((nxROI, nyROI), 5) ### arg1: frame size, arg2: threshold
CF.set_peripheralThresholdFactor(3)
CF.set_noiseMap(_noiseEneFrame)
if 'writeClusters' in _cfg and _cfg['writeClusters'] == True:
h5_1Photon_file = h5py.File(f'{_cfg["outputFolder"]}/Clusters_1Photon_chunk{idxChunk}.h5', 'w')
dset_1Photon_clusters = h5_1Photon_file.create_dataset(
'clusters', (0, 3, 3), maxshape=(None, 3, 3), dtype='f4',
chunks=True, compression='gzip'
)
dset_1photon_refs = h5_1Photon_file.create_dataset(
'referencePoint', (0, 2), maxshape=(None, 2), dtype='i4',
chunks=True
)
cluster_1Photon_list = []
refpoint_1Photon_list = []
h5_2Photon_file = h5py.File(f'{_cfg["outputFolder"]}/Clusters_2Photon_chunk{idxChunk}.h5', 'w')
dset_2Photon_clusters = h5_2Photon_file.create_dataset(
'clusters', (0, 6, 6), maxshape=(None, 6, 6), dtype='f4',
chunks=True, compression='gzip'
)
dset_2Photon_refs = h5_2Photon_file.create_dataset(
'referencePoint', (0, 2), maxshape=(None, 2), dtype='i4',
chunks=True
)
cluster_2Photon_list = []
refpoint_2Photon_list = []
def flush_h5_buffer():
old_len = dset_1Photon_clusters.shape[0]
new_len = old_len + len(cluster_1Photon_list)
dset_1Photon_clusters.resize((new_len, 3, 3))
dset_1photon_refs.resize((new_len, 2))
dset_1Photon_clusters[old_len:new_len] = cluster_1Photon_list
dset_1photon_refs[old_len:new_len] = refpoint_1Photon_list
cluster_1Photon_list.clear()
refpoint_1Photon_list.clear()
old_len = dset_2Photon_clusters.shape[0]
new_len = old_len + len(cluster_2Photon_list)
dset_2Photon_clusters.resize((new_len, 6, 6))
dset_2Photon_refs.resize((new_len, 2))
dset_2Photon_clusters[old_len:new_len] = cluster_2Photon_list
dset_2Photon_refs[old_len:new_len] = refpoint_2Photon_list
cluster_2Photon_list.clear()
refpoint_2Photon_list.clear()
for idxFrame in range(startFrame, endFrame):
if idxFrame % 10000 == 0:
print(f'Processing frame {idxFrame}...')
idxFile, idxFrame = divmod(idxFrame, nFramePerFile)
try:
if 'nFiber' in _cfg and _cfg['nFiber'] == 2:
offset = (headerSize + idxFrame * (NX * NY//2 + headerSize)) * 2
fileName1 = signalFileNames[idxFile]
fileName2 = fileName1.replace('d0', 'd1')
_signalAduFrame1 = np.fromfile(f'{fileName1}', dtype=np.uint16, offset=offset, count=NX * NY //2).astype(np.int32).reshape(NY//2, NX)
_signalAduFrame2 = np.fromfile(f'{fileName2}', dtype=np.uint16, offset=offset, count=NX * NY //2).astype(np.int32).reshape(NY//2, NX)
signalAduFrame = np.concatenate([_signalAduFrame1, _signalAduFrame2], axis=0)
del _signalAduFrame1, _signalAduFrame2
else:
offset = (headerSize + idxFrame * (NX * NY + headerSize)) * 2
signalAduFrame = np.fromfile(f'{signalFileNames[idxFile]}', dtype=np.uint16, offset=offset, count=NX * NY).astype(np.int32).reshape(NX, NY)
except Exception as e:
print(f'Error reading frame {idxFrame} from file {signalFileNames[idxFile]}: {e}, stop processing this chunk ({startFrame}-{endFrame})')
return _hists
signalAduFrame = signalAduFrame - _pedestalAduFrame
signalEneFrame = convertAduFrameToEnergyFrame(signalAduFrame)
CF.find_clusters_X(signalEneFrame[Roi[2]:Roi[3], Roi[0]:Roi[1]])
clusters = CF.hits()
for idxCluster in range(len(clusters)):
clusterSize = clusters['size'][idxCluster]
energy = clusters['energy'][idxCluster]
_hists['h1_ClusterEnergy'].Fill(energy)
xs = clusters['cols'][idxCluster][:clusterSize]
ys = clusters['rows'][idxCluster][:clusterSize]
enes = clusters['enes'][idxCluster][:clusterSize]
if energy < 5:
continue
if 2 * Energy - selectionRange < energy < 2 * Energy + selectionRange:
for i in range(len(xs)):
_hists['h2_DoubleHitsSumFrame'].Fill(xs[i]+Roi[0], ys[i]+Roi[2], enes[i])
if 'writeClusters' in _cfg and _cfg['writeClusters'] == True:
ref_x = (min(xs) + max(xs)) // 2
ref_y = (min(ys) + max(ys)) // 2
cluster_2Photon = np.zeros((6, 6), dtype=np.float32)
for i in range(len(xs)):
x_rel = xs[i] - int(ref_x) + 3
y_rel = ys[i] - int(ref_y) + 3
if 0 <= x_rel < 6 and 0 <= y_rel < 6:
cluster_2Photon[y_rel, x_rel] = enes[i]
cluster_2Photon_list.append(cluster_2Photon)
refpoint_2Photon_list.append([int(ref_x)+Roi[0], int(ref_y)+Roi[2]])
_hists['h2_ClusterSizeVsEnergy'].Fill(energy, clusterSize)
_diameterX = max(xs) - min(xs) + 1
_diameterY = max(ys) - min(ys) + 1
_diameter = max(_diameterX, _diameterY)
_hists['h2_DiameterVsEnergy'].Fill(energy, _diameter)
for i in range(len(xs)):
_hists['h1_PixelEnergy'].Fill(enes[i])
_hists['h2_HitsSumFrame'].Fill(xs[i]+Roi[0], ys[i]+Roi[2], enes[i])
if 'writeClusters' in _cfg and _cfg['writeClusters'] == True:
flush_h5_buffer()
h5_1Photon_file.close()
h5_2Photon_file.close()
return _hists
def process():
with Pool(_cfg['NThread'], maxtasksperchild=1) as p:
results = p.map(_processFrames, range(nChunks))
global hists
hists = bookHistograms(_cfg['energy'])
for h in hists.values():
h.SetDirectory(0)
for _hists in results:
for key in hists.keys():
hists[key].Add(_hists[key])
del _hists
def getPedestalAndNoise():
NX, NY = _cfg['NX'], _cfg['NY']
pedestalFileName = _cfg['pedestalFileName']
if 'nFiber' in _cfg and _cfg['nFiber'] == 2:
pedestalFileName2 = pedestalFileName.replace('d0', 'd1')
_pedestalAduFrames1 = np.fromfile(f'{pedestalFileName}', dtype=np.uint16).astype(np.int32)
_pedestalAduFrames1 = _pedestalAduFrames1.reshape(-1, NX * NY //2 + 56)
_pedestalAduFrames1 = _pedestalAduFrames1[:, 56:].reshape(-1, NY//2, NX)
_pedestalAduFrames2 = np.fromfile(f'{pedestalFileName2}', dtype=np.uint16).astype(np.int32)
_pedestalAduFrames2 = _pedestalAduFrames2.reshape(-1, NX * NY //2 + 56)
_pedestalAduFrames2 = _pedestalAduFrames2[:, 56:].reshape(-1, NY//2, NX)
pedestalAduFrames = np.concatenate([_pedestalAduFrames1, _pedestalAduFrames2], axis=1)
del _pedestalAduFrames1, _pedestalAduFrames2
else:
pedestalAduFrames = np.fromfile(f'{pedestalFileName}', dtype=np.uint16).astype(np.int32)
pedestalAduFrames = pedestalAduFrames.reshape(-1, NX * NY + 56)
pedestalAduFrames = pedestalAduFrames[:, 56:].reshape(-1, NX, NY)
pedestalAduFrames = pedestalAduFrames[pedestalAduFrames.shape[0]//10:] # skip the first 10% frames
global _pedestalAduFrame, _noiseEneFrame
_pedestalAduFrame = np.mean(pedestalAduFrames, axis=0)
noiseAduFrame = np.std(pedestalAduFrames, axis=0, ddof=1)
with Pool(16) as p:
pedestalEneFrames = p.map(convertAduFrameToEnergyFrame, pedestalAduFrames)
_noiseEneFrame = np.std(pedestalEneFrames, axis=0, ddof=1)
print(f'Average noise = {np.mean(_noiseEneFrame):.3f} keV; {np.mean(noiseAduFrame):.3f} ADU')
del pedestalAduFrames, pedestalEneFrames
def getPedestalAndNoise_simplified():
NX, NY = _cfg['NX'], _cfg['NY']
pedestalFileName = _cfg['pedestalFileName']
pedestalAduFrames = np.fromfile(f'{pedestalFileName}', dtype=np.uint16).astype(np.int32)
pedestalAduFrames = pedestalAduFrames.reshape(-1, NX * NY + 56)
pedestalAduFrames = pedestalAduFrames[:, 56:].reshape(-1, NX, NY)
pedestalAduFrames = pedestalAduFrames[pedestalAduFrames.shape[0]//10:] # skip the first 10% frames
global _pedestalAduFrame, _noiseEneFrame
_pedestalAduFrame = np.mean(pedestalAduFrames, axis=0)
noiseAduFrame = np.std(pedestalAduFrames, axis=0, ddof=1)
_noiseEneFrame = convertAduFrameToEnergyFrame(noiseAduFrame)
print(f'Average noise = {np.mean(_noiseEneFrame):.3f} keV; {np.mean(noiseAduFrame):.3f} ADU')
del pedestalAduFrames