import numpy as np from numpy.typing import ArrayLike from uproot import TTree from ..common import fillMCList, extractMatrix class MCtoClusters: """ This class establishes the relationship between clusters and monte carlo data """ def __init__(self) -> None: # behind these keys are the monte carlo info on the simulated data self.mcKeys = { 'pdg': 'MCParticles/MCParticles.m_pdg', 'mass': 'MCParticles/MCParticles.m_mass', 'energy': 'MCParticles/MCParticles.m_energy', 'momentumX': 'MCParticles/MCParticles.m_momentum_x', 'momentumY': 'MCParticles/MCParticles.m_momentum_y', 'momentumZ': 'MCParticles/MCParticles.m_momentum_z', 'validVertex': 'MCParticles/MCParticles.m_validVertex', 'productionTime': 'MCParticles/MCParticles.m_productionTime', 'productionVertexX': 'MCParticles/MCParticles.m_productionVertex_x', 'productionVertexY': 'MCParticles/MCParticles.m_productionVertex_y', 'productionVertexZ': 'MCParticles/MCParticles.m_productionVertex_z', 'decayTime': 'MCParticles/MCParticles.m_decayTime', 'decayVertexX': 'MCParticles/MCParticles.m_decayVertex_x', 'decayVertexY': 'MCParticles/MCParticles.m_decayVertex_y', 'decayVertexZ': 'MCParticles/MCParticles.m_decayVertex_z'} # these two establish the relation ship to an from clusters and monte carlo # there more entries than in the cluster data, but there still mc data missing # for some cluster files self.mcClusterRelations = {'from': 'PXDClustersToMCParticles/m_elements/m_elements.m_from', 'to': 'PXDClustersToMCParticles/m_elements/m_elements.m_to'} def branches(self, *, includeUnselected: bool = False) -> list: return list((self.mcKeys | self.mcClusterRelations).values()) def get(self, eventTree: TTree) -> dict: """ this loads the monte carlo from the root file """ # the monte carlo data, they are longer than the cluster data mcData = eventTree.arrays(self.mcKeys.values(), library='np') pdg = mcData[self.mcKeys['pdg']] momentumX = mcData[self.mcKeys['momentumX']] momentumY = mcData[self.mcKeys['momentumY']] momentumZ = mcData[self.mcKeys['momentumZ']] # this loads the relation ships to and from clusters and mc data # this is the same level of retardedness as with the cluster digits clusterToMC = eventTree.arrays(self.mcClusterRelations['to'], library='np')[self.mcClusterRelations['to']] mcToCluster = eventTree.arrays(self.mcClusterRelations['from'], library='np')[self.mcClusterRelations['from']] # it need the cluster charge as a jagged/ragged array, maybe I could simply # use the event numbers, but I am too tired to fix this shitty file format clsCharge = eventTree.arrays('PXDClusters/PXDClusters.m_clsCharge', library='np')['PXDClusters/PXDClusters.m_clsCharge'] # reorganizing MC data n = len(clusterToMC) momentumXList = np.zeros(n, dtype=object) momentumYList = np.zeros(n, dtype=object) momentumZList = np.zeros(n, dtype=object) pdgList = np.zeros(n, dtype=object) clusterNumbersList = np.zeros(n, dtype=object) for i in range(n): # _fillMCList fills in the missing spots, because there are not mc data for # every cluster, even though there are more entries in this branch than # in the cluster branch... as I said, the root format is retarded fullClusterReferences = fillMCList(mcToCluster[i], clusterToMC[i], len(clsCharge[i])) clusterNumbersList[i] = fullClusterReferences pdgs, xmom, ymom, zmom = self._getMCData(fullClusterReferences, pdg[i], momentumX[i], momentumY[i], momentumZ[i]) momentumXList[i] = xmom momentumYList[i] = ymom momentumZList[i] = zmom pdgList[i] = pdgs return { 'momentumX': np.hstack(momentumXList).astype(float), 'momentumY': np.hstack(momentumYList).astype(float), 'momentumZ': np.hstack(momentumZList).astype(float), 'pdg': np.hstack(pdgList).astype(int), 'clsNumber': np.hstack(clusterNumbersList).astype(int) } @staticmethod def _getMCData(toClusters: ArrayLike, pdgs: ArrayLike, xMom: ArrayLike, yMom: ArrayLike, zMom: ArrayLike) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ after filling and reorganizing MC data arrays one can finally collect the actual MC data, where there's data missing I will with zeros """ n = len(toClusters) pxList = np.zeros(n) pyList = np.zeros(n) pzList = np.zeros(n) pdgList = np.zeros(n, dtype=int) for i, references in enumerate(toClusters): if references == -1: continue # Arrays were initialized to zero else: pxList[i] = xMom[references] pyList[i] = yMom[references] pzList[i] = zMom[references] pdgList[i] = pdgs[references] return pdgList, pxList, pyList, pzList class MCtoDigits: """ This class establishes the relationship between digits and monte carlo data """ def __init__(self) -> None: # these are the sensor IDs of the pxd modules/panels from the root file, they are # use to identify on which panels a cluster event happened 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]) # behind these keys are the monte carlo info on the simulated data self.mcKeys = { 'pdg': 'MCParticles/MCParticles.m_pdg', 'mass': 'MCParticles/MCParticles.m_mass', 'energy': 'MCParticles/MCParticles.m_energy', 'momentumX': 'MCParticles/MCParticles.m_momentum_x', 'momentumY': 'MCParticles/MCParticles.m_momentum_y', 'momentumZ': 'MCParticles/MCParticles.m_momentum_z', 'validVertex': 'MCParticles/MCParticles.m_validVertex', 'productionTime': 'MCParticles/MCParticles.m_productionTime', 'productionVertexX': 'MCParticles/MCParticles.m_productionVertex_x', 'productionVertexY': 'MCParticles/MCParticles.m_productionVertex_y', 'productionVertexZ': 'MCParticles/MCParticles.m_productionVertex_z', 'decayTime': 'MCParticles/MCParticles.m_decayTime', 'decayVertexX': 'MCParticles/MCParticles.m_decayVertex_x', 'decayVertexY': 'MCParticles/MCParticles.m_decayVertex_y', 'decayVertexZ': 'MCParticles/MCParticles.m_decayVertex_z'} # if a file is missing cluster data, I need to establish mc-cluster relation # through the digits... who came up with this insanity? self.mcDigitsInRelations = {'from': 'PXDDigitsToMCParticles/m_elements/m_elements.m_from', 'to': 'PXDDigitsToMCParticles/m_elements/m_elements.m_to'} self.mcDigitsOutRelations = {'from': 'PXDDigitsOUTToMCParticles/m_elements/m_elements.m_from', 'to': 'PXDDigitsOUTToMCParticles/m_elements/m_elements.m_to'} # in order to establish the relation ship between mc and digits and ultematly # between mc and clusters I need this additional branch... # as I said many times, root is the most retarded file format on earth self.digitsInKeys = { 'sensorID': 'PXDDigits/PXDDigits.m_sensorID', 'uCellID': 'PXDDigits/PXDDigits.m_uCellID', 'vCellID': 'PXDDigits/PXDDigits.m_vCellID', 'cellCharge': 'PXDDigits/PXDDigits.m_charge'} self.digitsOutKeys = { 'sensorID': 'PXDDigitsOUT/PXDDigitsOUT.m_sensorID', 'uCellID': 'PXDDigitsOUT/PXDDigitsOUT.m_uCellID', 'vCellID': 'PXDDigitsOUT/PXDDigitsOUT.m_vCellID', 'cellCharge': 'PXDDigitsOUT/PXDDigitsOUT.m_charge'} def branches(self, *, includeUnselected: bool = False) -> list: if includeUnselected is True: return list((self.mcKeys | self.mcDigitsInRelations | self.mcDigitsOutRelations).values()) return list((self.mcKeys | self.mcDigitsInRelations).values()) def get(self, eventTree: TTree, inOut: str = 'inROI') -> dict: pdg, momentumX, momentumY, momentumZ, fromDigits, toDigits, uCellIDs, vCellIDs, cellCharges, clusterSensorIDs = self._selectKeys(eventTree, inOut=inOut) return self._process(pdg, momentumX, momentumY, momentumZ, fromDigits, toDigits, uCellIDs, vCellIDs, cellCharges, clusterSensorIDs) def _selectKeys(self, eventTree: TTree, inOut: str = 'inROI') -> tuple: mcData = eventTree.arrays(self.mcKeys.values(), library='np') pdg = mcData[self.mcKeys['pdg']] momentumX = mcData[self.mcKeys['momentumX']] momentumY = mcData[self.mcKeys['momentumY']] momentumZ = mcData[self.mcKeys['momentumZ']] if inOut == 'inROI': relations = eventTree.arrays(self.mcDigitsInRelations.values(), library='np') fromDigits = relations[self.mcDigitsInRelations['from']] # I need to convert data type of each entry... fucking root toDigits = np.array([np.array([item[0] for item in relations[self.mcDigitsInRelations['to']][i]]) for i in range(len(relations[self.mcDigitsInRelations['to']]))], dtype=object) digits = eventTree.arrays(self.digitsInKeys.values(), library='np') uCellIDs = digits[self.digitsInKeys['uCellID']] vCellIDs = digits[self.digitsInKeys['vCellID']] cellCharges = digits[self.digitsInKeys['cellCharge']] clusterSensorIDs = digits[self.digitsInKeys['sensorID']] else: relations = eventTree.arrays(self.mcDigitsOutRelations.values(), library='np') fromDigits = relations[self.mcDigitsOutRelations['from']] toDigits = np.array([np.array([item[0] for item in relations[self.mcDigitsOutRelations['to']][i]]) for i in range(len(relations[self.mcDigitsOutRelations['to']]))], dtype=object) digits = eventTree.arrays(self.digitsOutKeys.values(), library='np') uCellIDs = digits[self.digitsOutKeys['uCellID']] vCellIDs = digits[self.digitsOutKeys['vCellID']] cellCharges = digits[self.digitsOutKeys['cellCharge']] clusterSensorIDs = digits[self.digitsOutKeys['sensorID']] return pdg, momentumX, momentumY, momentumZ, fromDigits, toDigits, uCellIDs, vCellIDs, cellCharges, clusterSensorIDs def _process(self, pdg: ArrayLike, momentumX: ArrayLike, momentumY: ArrayLike, momentumZ: ArrayLike, fromDigits: ArrayLike, toDigits: ArrayLike, uCellIDs: ArrayLike, vCellIDs: ArrayLike, cellCharges: ArrayLike, clusterSensorIDs: ArrayLike) -> dict: # Loop through each cell charge to populate matrices and process data pdgCodes = [] momentumXList = [] momentumYList = [] momentumZList = [] clsNumbers = [] for i in range(len(clusterSensorIDs)): # Initialize and populate the matrix matrixLadder = np.zeros((250, 768)) matrixLadder[uCellIDs[i], vCellIDs[i]] = cellCharges[i] sensorID = clusterSensorIDs[i] xx, yy = uCellIDs[i], vCellIDs[i] # retrieving mc data and references mcDigits = fillMCList(fromDigits[i], toDigits[i], len(sensorID)) pdgs = pdg[i] momentaX = momentumX[i] momentaY = momentumY[i] momentaZ = momentumZ[i] assert len(mcDigits) == len(sensorID), f'event {i}, mcDigits: {len(mcDigits)} and sensorID: {len(sensorID)}' # here I store all pixels, that have already been visited knownPixels = {id: set() for id in self.panelIDs} for x, y, id, relation in zip(xx, yy, sensorID, mcDigits): if (x, y) in knownPixels[id]: continue matrix, xLower, yLower = extractMatrix(matrixLadder, x, y) vv, uu = np.nonzero(matrix) # Convert local indices of non-zero pixels to global indices globalNonzeroX = vv + xLower globalNonzeroY = uu + yLower # Update knownPixels with the global coordinates of the non-zero pixels knownPixels[id].update(zip(globalNonzeroX, globalNonzeroY)) if relation == -1: pdgCodes.append(0) momentumXList.append(0.0) momentumYList.append(0.0) momentumZList.append(0.0) else: pdgCodes.append(pdgs[relation]) momentumXList.append(momentaX[relation]) momentumYList.append(momentaY[relation]) momentumZList.append(momentaZ[relation]) clsNumbers.append(relation) return { 'pdg': np.array(pdgCodes).astype(int), 'momentumX': np.array(momentumXList), 'momentumY': np.array(momentumYList), 'momentumZ': np.array(momentumZList), 'clsNumber': np.array(clsNumbers).astype(int) }