Skip to content
Snippets Groups Projects
Commit abb71a7a authored by johannes bilk's avatar johannes bilk
Browse files

initial commit

parents
No related branches found
No related tags found
No related merge requests found
import numpy as np
import uproot as ur
class fromRoot:
def __init__(self):
# panel ids werden benutzt um layer und lader zu bestimmen und um u/v zu kartesisch um zu rechnen
self.panelIDs = np.array([ 8480, 8512, 8736, 8768, 8992, 9024, 9248, 9280,
9504, 9536, 9760, 9792, 10016, 10048, 10272, 10304,
16672, 16704, 16928, 16960, 17184, 17216, 17440, 17472,
17696, 17728, 17952, 17984, 18208, 18240, 18464, 18496,
18720, 18752, 18976, 19008, 19232, 19264, 19488, 19520])
# die koordinaten verschiebung der u/v koordinaten in korrekte panel position
self.panelShifts = np.array([[1.3985 , 0.2652658 , 3.68255],
[ 1.3985 , 0.23238491, -0.88255],
[ 0.80146531, 1.17631236, 3.68255],
[ 0.82407264, 1.15370502, -0.88255],
[-0.2582769 , 1.3985 , 3.68255],
[-0.2322286 , 1.3985 , -0.88255],
[-1.17531186, 0.80246583, 3.68255 ],
[-1.15510614, 0.82267151, -0.88255],
[-1.3985 , -0.2645974 , 3.68255],
[-1.3985 , -0.23012119, -0.88255],
[-0.80591227, -1.17186534, 3.68255],
[-0.82344228, -1.15433536, -0.88255],
[ 0.26975836, -1.3985 , 3.68255],
[ 0.23326624, -1.3985 , -0.88255],
[ 1.1746111 , -0.80316652, 3.68255],
[ 1.15205703, -0.82572062, -0.88255],
[ 2.2015 , 0.26959865, 5.01305],
[ 2.2015 , 0.2524582 , -1.21305],
[ 1.77559093, 1.32758398, 5.01305],
[ 1.78212569, 1.31626522, -1.21305],
[ 0.87798948, 2.03516717, 5.01305],
[ 0.88478563, 2.03124357, -1.21305],
[-0.26129975, 2.2015 , 5.01305],
[-0.25184137, 2.2015 , -1.21305],
[-1.32416655, 1.77756402, 5.01305],
[-1.31417539, 1.78333226, -1.21305],
[-2.03421133, 0.87964512, 5.01305],
[-2.02960691, 0.88762038, -1.21305],
[-2.2015 , -0.25954151, 5.01305],
[-2.2015 , -0.24969109, -1.21305],
[-1.77636043, -1.32625112, 5.01305],
[-1.78138268, -1.31755219, -1.21305],
[-0.87493138, -2.03693277, 5.01305 ],
[-0.8912978 , -2.02748378, -1.21305],
[ 0.26489725, -2.2015 , 5.01305],
[ 0.25364439, -2.2015 , -1.21305],
[ 1.3269198 , -1.7759744 , 5.01305],
[ 1.32258793, -1.77847528, -1.21305],
[ 2.03616649, -0.87625871, 5.01305],
[ 2.02936825, -0.8880338 , -1.21305]])
# drehwinkel um panels korrekt auszurichten
self.panelRotations = np.array([ 90, 90, 135, 135, 180, 180, 225, 225, 270, 270, 315, 315, 360,
360, 405, 405, 90, 90, 120, 120, 150, 150, 180, 180, 210, 210,
240, 240, 270, 270, 300, 300, 330, 330, 360, 360, 390, 390, 420,
420])
# ladder und layer ids
self.panelLayer = np.array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2])
self.panelLadder = np.array([1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 12, 12, 13, 13, 14, 14, 15, 15, 16, 16, 17, 17, 18, 18, 19, 19, 20, 20, 21, 21])
# generierung der look-up tabels für u/v -> x/y/z transformation
self.transformation = {}
self.layersLadders = {}
for i in range(len(self.panelIDs)):
self.transformation[str(self.panelIDs[i])] = [self.panelShifts[i], self.panelRotations[i]]
self.layersLadders[str(self.panelIDs[i])] = [panelLayer[i], panelLadder[i]]
def loadData(self, file, path='.'):
self.eventTree = ur.open(f'{path}/{file}.root:tree')
# liest den event tree aus, man muss das voll schlüsselwort angeben
def getData(self, keyword: str):
try:
return self.eventTree.arrays(keyword, library='np')[keyword]
except:
return KeyError
# generiert für jeden cluster eine event nummer
def genEventNumbers(self, clusters):
eventNumbers = []
for i in range(len(clusters)):
eventNumbers.append(np.array([i]*len(clusters[i])))
return flatten(eventNumbers)
# organisiert mc und digit daten aller events um, sodass sie zu clustern passen
def getEventData(self, relations, *args):
returnList = []
for i in range(len(args)):
stuffList = []
for item in relations:
stuffList.append([0] * len(item))
returnList.append(stuffList)
for i, references in enumerate(relations):
for k, index in enumerate(references):
for j in range(len(args)):
returnList[j][i][k] = args[j][index]
if len(returnList) == 1:
return returnList[0]
else:
return returnList
def flatten(self, structure, max_depth=None, current_depth=0):
flat_list = []
for element in structure:
if isinstance(element, (list, np.ndarray)) and (max_depth is None or current_depth < max_depth):
flat_list.extend(self.flatten(element, max_depth, current_depth + 1))
else:
flat_list.append(element)
return np.array(flat_list, dtype=object)
# generiert alle pixel matrizen aller events
def getClustersFlattened(self, uCellIDs, vCellIDs, cellCharges, clusterDigits, matrixSize=(9,9)):
length = 0
start = 0
for item in cellCharges:
length += len(item)
events = [0] * length
plotRange = int(np.round(matrixSize[0]/2)), int(np.round(matrixSize[1]/2))
for event in range(len(cellCharges)):
adcValues = []
digitsU = uCellIDs[event]
digitsV = vCellIDs[event]
digitsCharge = cellCharges[event]
digitIndices = clusterDigits[event]
for indices in digitIndices:
cacheImg = np.zeros(matrixSize)
maxChargeIndex = digitsCharge[indices].argmax()
uMax = digitsU[indices[maxChargeIndex]]
vMax = digitsV[indices[maxChargeIndex]]
for index in indices:
uPos = digitsU[index]
vPos = digitsV[index]
uu = int(uPos) - int(uMax) + plotRange[0]
vv = int(vPos) - int(vMax) + plotRange[1]
if uu >= 0 and uu < matrixSize[0] and vv >= 0 and vv < matrixSize[1]:
cacheImg[uu,vv] = digitsCharge[index]
adcValues.append(cacheImg)
stop = len(adcValues)
events[start:start+stop] = adcValues
start += stop
return np.array(events, dtype=object)
# rotiert und verschiebt eine Koordinate
def rotShiftVector(self, vector, angle, shift=[0,0,0], scale=1):
theta = np.deg2rad(angle)
rotMatrix = np.array([[np.cos(theta),-np.sin(theta),0],[np.sin(theta),np.cos(theta),0],[0,0,1]])
scaleMatrix = np.array([[1,0,0],[0,1,0],[0,0,scale]])
return rotMatrix.dot(scaleMatrix.dot(vector)) + shift
# berechnet die koordinaten aller Events
def getCartesianFlattened(self, uPositions, vPositions, sensorIDs, transformations: dict):
length = 0
start = 0
for item in sensorIDs:
length += len(item)
xArr, yArr, zArr = [0] * length, [0] * length, [0] * length
for event in range(len(sensorIDs)):
xyz = []
uPos = uPositions[event]
vPos = vPositions[event]
sensors = sensorIDs[event]
points = np.vstack((uPos, np.zeros(len(uPos)), vPos)).T
for point, id in zip(points, sensors):
shift, angle = transformations[str(id)]
shifted = rotShiftVector(point, angle, shift)
xyz.append(shifted)
if len(xyz) > 0:
stop = len(xyz)
xArr[start:start+stop] = np.array(xyz)[:,0]
yArr[start:start+stop] = np.array(xyz)[:,1]
zArr[start:start+stop] = np.array(xyz)[:,2]
start += stop
return np.array(xArr, dtype=object), np.array(yArr, dtype=object), np.array(zArr, dtype=object)
# bestimmt die layer und ladder nummer eines events
def getLayers(self, sensorIDs, layersLadders: dict):
layers, ladders = [], []
for id in sensorIDs:
layer, ladder = layersLadders[str(id)]
layers.append(layer)
ladders.append(ladder)
return np.array(layers), np.array(ladders)
# findet die fehlenden event referenzen in mc daten, setzt sie gleich -1
def findMissing(self, lst: list, length: int) -> list:
return sorted(set(range(0, length)) - set(lst))
# füllt die mc-cluster beziehungs arrays mit fehlenden werten auf
def fillMCList(self, fromClusters, toClusters, length):
missingIndex = findMissing(fromClusters, length)
testList = [-1] * length
fillIndex = 0
for i in range(len(testList)):
if i in missingIndex:
testList[i] = -1
else:
try:
testList[i] = int(toClusters[fillIndex])
except TypeError:
testList[i] = int(toClusters[fillIndex][0])
fillIndex += 1
return testList
# organisiert mc daten eines events um
def getMCData(self, toClusters, pdgs, xMom, yMom, zMom):
pxList, pyList, pzList = [], [], []
pdgList = []
for references in toClusters:
if references == -1:
pxList.append(0)
pyList.append(0)
pzList.append(0)
pdgList.append(0)
else:
pxList.append(xMom[references])
pyList.append(yMom[references])
pzList.append(zMom[references])
pdgList.append(pdgs[references])
return np.array(pdgList,dtype=list), np.array(pxList,dtype=list), np.array(pyList,dtype=list), np.array(pzList,dtype=list)
events = ur.open('{}/{}.root:tree'.format(settings['inPath'], settings['inFile']))
# cluster daten
print('loading cluster data...')
clsCharge = getData(events, 'PXDClusters/PXDClusters.m_clsCharge')
iterLength = 0
for i in range(len(clsCharge)):
iterLength += len(clsCharge[i])
if settings['clusterCharge'] is True or settings['eventNumbers'] is True:
clsChargeFlat = flatten(clsCharge)
eventNumbers = genEventNumbers(clsCharge)
if settings['seedCharge'] is True:
seedCharge = getData(events, 'PXDClusters/PXDClusters.m_seedCharge')
seedChargeFlat = flatten(seedCharge)
if settings['clusterSize'] is True:
clsSize = getData(events, 'PXDClusters/PXDClusters.m_clsSize')
clsSizeFlat = flatten(clsSize)
if settings['uvSize'] is True:
uSize = getData(events, 'PXDClusters/PXDClusters.m_uSize')
vSize = getData(events, 'PXDClusters/PXDClusters.m_vSize')
uSizeFlat = flatten(uSize)
vSizeFlat = flatten(vSize)
if settings['uvPosition'] is True or settings['coordinates'] is True:
uPosition = getData(events, 'PXDClusters/PXDClusters.m_uPosition')
vPosition = getData(events, 'PXDClusters/PXDClusters.m_vPosition')
uPositionFlat = flatten(uPosition)
vPositionFlat = flatten(vPosition)
if settings['sensorID'] is True or settings['coordinates'] is True or settings['ladder'] or settings['layer']:
sensorID = getData(events, 'PXDClusters/PXDClusters.m_sensorID')
sensorIDFlat = flatten(sensorID)
# koordinaten, layer und ladder
if settings['coordinates'] is True:
print('calculating coordinates...')
xcoords, ycoords, zcoords = getCartesianFlattened(uPosition, vPosition, sensorID, transformation)
if settings['layer'] is True or settings['ladder'] is True:
layers, ladders = getLayers(sensorIDFlat, layersLadders)
# digits laden
if settings['uvCellID'] is True or settings['cellCharge'] is True or settings['pxdClusters'] is True:
print('loading cell data...')
uCellID = getData(events, 'PXDDigits/PXDDigits.m_uCellID')
vCellID = getData(events, 'PXDDigits/PXDDigits.m_vCellID')
charge = getData(events, 'PXDDigits/PXDDigits.m_charge')
digitClusters = getData(events, 'PXDClustersToPXDDigits/m_elements/m_elements.m_to')
# digits umorganisieren
uCellIDFlat = []
vCellIDFlat = []
chargeFlat = []
for i in range(len(digitClusters)):
udigits, vdigits, cdigits = getEventData(digitClusters[i], uCellID[i], vCellID[i], charge[i])
uCellIDFlat.append(udigits)
vCellIDFlat.append(vdigits)
chargeFlat.append(cdigits)
# digits flachen
uCellIDFlat = flatten(uCellIDFlat)
vCellIDFlat = flatten(vCellIDFlat)
chargeFlat = flatten(chargeFlat)
# adc matrix bestimmen
if settings['pxdClusters'] is True:
print('generating pixel cluster')
clustersFlat = getClustersFlattened(uCellID, vCellID, charge, digitClusters)
# mc laden
if settings['pdg'] or settings['momenta'] or settings['clusterNumbers']:
print('loading mc info...')
pdg = getData(events, 'MCParticles/MCParticles.m_pdg')
momentumX = getData(events, 'MCParticles/MCParticles.m_momentum_x')
momentumY = getData(events, 'MCParticles/MCParticles.m_momentum_y')
momentumZ = getData(events, 'MCParticles/MCParticles.m_momentum_z')
mcClusters = getData(events, 'PXDClustersToMCParticles/m_elements/m_elements.m_to')
mcfromClusters = getData(events, 'PXDClustersToMCParticles/m_elements/m_elements.m_from')
# mc umorganisieren
momentumXList = []
momentumYList = []
momentumZList = []
pdgList = []
clusterNumbersList = []
for i in range(len(mcClusters)):
fullClusterReferences = fillMCList(mcfromClusters[i], mcClusters[i], len(clsCharge[i]))
clusterNumbersList.append(fullClusterReferences)
pdgs, xmom, ymom, zmom = getMCData(fullClusterReferences, pdg[i], momentumX[i], momentumY[i], momentumZ[i])
momentumXList.append(xmom)
momentumYList.append(ymom)
momentumZList.append(zmom)
pdgList.append(pdgs)
# mc flachen
momentumXFlat = flatten(momentumXList)
momentumYFlat = flatten(momentumYList)
momentumZFlat = flatten(momentumZList)
pdgFlat = flatten(pdgList)
clusterNumbersFlat = flatten(clusterNumbersList)
if settings['saveFormat'] == 'structured':
print('preparing data array...')
# daten type definieren
dataType = []
if settings['eventNumbers'] is True:
dataType.append(('event', 'i4'))
if settings['clusterNumbers'] is True:
dataType.append(('clsNumber', 'i4'))
if settings['pxdClusters'] is True:
dataType.append(('cluster', 'i4', (9,9)))
if settings['clusterCharge'] is True:
dataType.append(('clsCharge', 'i4'))
if settings['seedCharge'] is True:
dataType.append(('seedCharge', 'i4'))
if settings['clusterSize'] is True:
dataType.append(('clsSize', 'i4'))
if settings['uvSize'] is True:
dataType.append(('uSize', 'i4'))
dataType.append(('vSize', 'i4'))
if settings['uvCellID'] is True:
dataType.append(('uCellID', 'object'))
dataType.append(('vCellID', 'object'))
if settings['cellCharge'] is True:
dataType.append(('charge', 'object'))
if settings['coordinates'] is True:
dataType.append(('xPosition', 'f8'))
dataType.append(('yPosition', 'f8'))
dataType.append(('zPosition', 'f8'))
if settings['uvPosition'] is True:
dataType.append(('uPosition', 'f8'))
dataType.append(('vPosition', 'f8'))
if settings['sensorID'] is True:
dataType.append(('sensorID', 'i8'))
if settings['layer'] is True:
dataType.append(('layer', 'i4'))
if settings['ladder'] is True:
dataType.append(('ladder', 'i4'))
if settings['momenta'] is True:
dataType.append(('xMomentum', 'f8'))
dataType.append(('yMomentum', 'f8'))
dataType.append(('zMomentum', 'f8'))
if settings['pdg'] is True:
dataType.append(('pdg', 'i4'))
# structured array definieren
forArray = []
for i in range(iterLength):
toAppend = []
if settings['eventNumbers'] is True:
toAppend.append(eventNumbers[i])
if settings['clusterNumbers'] is True:
toAppend.append(clusterNumbersFlat[i])
if settings['pxdClusters'] is True:
toAppend.append(clustersFlat[i])
if settings['clusterCharge'] is True:
toAppend.append(clsChargeFlat[i])
if settings['seedCharge'] is True:
toAppend.append(seedChargeFlat[i])
if settings['clusterSize'] is True:
toAppend.append(clsSizeFlat[i])
if settings['uvSize'] is True:
toAppend.append(uSizeFlat[i])
toAppend.append(vSizeFlat[i])
if settings['uvCellID'] is True:
toAppend.append(uCellIDFlat[i])
toAppend.append(vCellIDFlat[i])
if settings['cellCharge'] is True:
toAppend.append(chargeFlat[i])
if settings['coordinates'] is True:
toAppend.append(xcoords[i])
toAppend.append(ycoords[i])
toAppend.append(zcoords[i])
if settings['uvPosition'] is True:
toAppend.append(uPositionFlat[i])
toAppend.append(vPositionFlat[i])
if settings['sensorID'] is True:
toAppend.append(sensorIDFlat[i])
if settings['layer'] is True:
toAppend.append(layers[i])
if settings['ladder'] is True:
toAppend.append(ladders[i])
if settings['momenta'] is True:
toAppend.append(momentumXFlat[i])
toAppend.append(momentumYFlat[i])
toAppend.append(momentumZFlat[i])
if settings['pdg'] is True:
toAppend.append(pdgFlat[i])
toAppend = tuple(toAppend)
forArray.append(toAppend)
saveArr = np.array(forArray, dtype=dataType)
elif settings['saveFormat'] == 'array' or settings['saveFormat'] == 'txt':
print('preparing data array...')
# structured array definieren
forArray = []
for i in range(iterLength):
toAppend = np.array([0])
if settings['eventNumbers'] is True:
toAppend = np.concatenate((toAppend, [eventNumbers[i]]))
if settings['clusterNumbers'] is True:
toAppend = np.concatenate((toAppend, [clusterNumbersFlat[i]]))
if settings['pxdClusters'] is True:
toAppend = np.concatenate((toAppend, clustersFlat[i].flatten()))
if settings['clusterCharge'] is True:
toAppend = np.concatenate((toAppend, [clsChargeFlat[i]]))
if settings['seedCharge'] is True:
toAppend = np.concatenate((toAppend, [seedChargeFlat[i]]))
if settings['clusterSize'] is True:
toAppend = np.concatenate((toAppend, [clsSizeFlat[i]]))
if settings['coordinates'] is True:
toAppend = np.concatenate((toAppend, [xcoords[i]]))
toAppend = np.concatenate((toAppend, [ycoords[i]]))
toAppend = np.concatenate((toAppend, [zcoords[i]]))
if settings['uvPosition'] is True:
toAppend = np.concatenate((toAppend, [uPositionFlat[i]]))
toAppend = np.concatenate((toAppend, [vPositionFlat[i]]))
if settings['sensorID'] is True:
toAppend = np.concatenate((toAppend, [sensorIDFlat[i]]))
if settings['layer'] is True:
toAppend = np.concatenate((toAppend, [layers[i]]))
if settings['ladder'] is True:
toAppend = np.concatenate((toAppend, [ladders[i]]))
if settings['momenta'] is True:
toAppend = np.concatenate((toAppend, [momentumXFlat[i]]))
toAppend = np.concatenate((toAppend, [momentumYFlat[i]]))
toAppend = np.concatenate((toAppend, [momentumZFlat[i]]))
if settings['pdg'] is True:
toAppend = np.concatenate((toAppend, [pdgFlat[i]]))
forArray.append(toAppend[1:])
saveArr = np.array(forArray)
# array speichern
if settings['saveFormat'] == 'structured' or settings['saveFormat'] == 'array':
print('saving data array as {} to {}'.format(settings['outFile'], settings['outPath']))
np.save('{}/{}.npy'.format(settings['outPath'], settings['outFile']), saveArr)
if settings['saveFormat'] == 'txt':
print('saving txt data as {} to {}'.format(settings['outFile'], settings['outPath']))
with open('{}/{}.txt'.format(settings['outPath'], settings['outFile']), 'w') as f:
for line in saveArr:
for item in line:
f.write(str(item))
f.write(' ')
f.write('\n')
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment