Skip to content
Snippets Groups Projects
rootable.py 19.4 KiB
Newer Older
  • Learn to ignore specific revisions
  • 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404
    import numpy as np
    from numpy.typing import ArrayLike
    import uproot as ur
    
    
    class Rootable:
        """
        this class uses uproot to load pxd data from root files and converts them into
        native python data structures.
        it can load the cluster information, uses the digits to generate the adc matrices,
        coordinates, layer and ladders and finally also 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])
    
            # every line in this corresponds to one entry in the array above, this is used
            # to put the projected uv plane in the right 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]])
    
            # every entry here corresponds to the entries in the array above, these are
            # used for rotating the projected uv plane
            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])
    
            # the layer and ladder arrays, for finding them from sensor id
            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])
    
            # all transpormaations are stored in a dict, with the sensor id as a keyword
            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])] = [self.panelLayer[i], self.panelLadder[i]]
    
            # these are the branch names for cluster info in the root file
            self.gotClusters = False
            self.clusters = ['PXDClusters/PXDClusters.m_clsCharge',
                             'PXDClusters/PXDClusters.m_seedCharge',
                             'PXDClusters/PXDClusters.m_clsSize',
                             'PXDClusters/PXDClusters.m_uSize',
                             'PXDClusters/PXDClusters.m_vSize',
                             'PXDClusters/PXDClusters.m_uPosition',
                             'PXDClusters/PXDClusters.m_vPosition',
                             'PXDClusters/PXDClusters.m_sensorID']
    
            # these are the branch names for cluster digits in the root file
            self.digits = ['PXDDigits/PXDDigits.m_uCellID',
                           'PXDDigits/PXDDigits.m_vCellID',
                           'PXDDigits/PXDDigits.m_charge']
    
            # this establishes the relationship between clusters and digits
            # because for some reaseon the branch for digits has a different
            # size than the cluster branch
            self.clusterToDigis = 'PXDClustersToPXDDigits/m_elements/m_elements.m_to'
    
            # these are the branch names for monte carlo data in the root file
            self.mcData = ['MCParticles/MCParticles.m_pdg',
                           'MCParticles/MCParticles.m_momentum_x',
                           'MCParticles/MCParticles.m_momentum_y',
                           'MCParticles/MCParticles.m_momentum_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.clusterToMC = 'PXDClustersToMCParticles/m_elements/m_elements.m_to'
            self.mcToCluster = 'PXDClustersToMCParticles/m_elements/m_elements.m_from'
    
            # this dict stores the data
            self.data = {}
    
        def __getitem__(self, index: str | int | ArrayLike) -> np.ndarray | dict:
            """
            this makes the class subscriptable, one can retrieve one coloumn by using
            strings as keywords, or get a row by using integer indices or arrays
            """
            if isinstance(index, str):
                return self.data[index]
            return {key: value[index] for key, value in self.data.items()}
    
        def loadData(self, file: str) -> None:
            """
            reads the file off of the harddrive, it automatically creates event numbers
            file: str = it's the whole file path + .root ending
            """
            self.eventTree = ur.open(f'{file}:tree')
            self._genEventNumbers()
    
        def _genEventNumbers(self) -> None:
            """
            a private method that gets called on file import
            it generates the event numbers from the jagged arrays
            coming from the branches
            """
            eventNumbers = []
            clusters = self.eventTree.arrays('PXDClusters/PXDClusters.m_clsCharge', library='np')['PXDClusters/PXDClusters.m_clsCharge']
            for i in range(len(clusters)):
                eventNumbers.append(np.array([i]*len(clusters[i])))
            self.data['eventNumbers'] = self._flatten(eventNumbers)
    
        def _getData(self, keyword: str, library: str = 'np') -> np.ndarray:
            """
            a private method for converting branches into something useful, namely
            into numpy arrays, if the keyward library is set to np.
            keyword: str = the full branch name
            library: str = can be 'np' (numpy), 'pd' (pandas) or 'ak' (akward)
                           see uproot documentation for more info
            """
            try:
                data = self.eventTree.arrays(keyword, library=library)[keyword]
                return self._flatten(data)
            except:
                return KeyError
    
        def _flatten(self, structure: ArrayLike, maxDepth: int = None, currentDepth: int = 0) -> np.ndarray:
            """
            this is a private function, that gets called during loading branches
            it flattens ragged array, one can set the depths to which one wants to flatten
            structure: the list/array to flatten
            maxDepth: int = the amount of flattening
            currentDepth: int = don't touch this, it's used for recursively calling
            """
            flat_list = []
    
            for element in structure:
                if isinstance(element, (list, np.ndarray)) and (maxDepth is None or currentDepth < maxDepth):
                    flat_list.extend(self._flatten(element, maxDepth, currentDepth + 1))
                else:
                    flat_list.append(element)
    
            return np.array(flat_list, dtype=object)
    
        def getClusters(self) -> None:
            """
            this uses the array from __init__ to load different branches into the data dict
            """
            self.gotClusters = True
            for branch in self.clusters:
                data = self._getData(branch)
                keyword = branch.split('_')[-1]
                self.data[keyword] = data
    
        def getMatrices(self) -> None:
            """
            loads the digit branches into arrays and converts them into adc matrices
            """
            uCellIDs = self.eventTree.arrays(self.digits[0], library='np')[self.digits[0]]
            vCellIDs = self.eventTree.arrays(self.digits[1], library='np')[self.digits[1]]
            cellCharges = self.eventTree.arrays(self.digits[2], library='np')[self.digits[2]]
    
            # this establishes the relation between digits and clusters, it's still
            # shocking to me, that this is necessary, why aren't digits stored in the
            # same way as clusters, than one wouldn't need to jump through hoops just
            # to have the data in a usable und sensible manner
            # root is such a retarded file format
            clusterDigits = self.eventTree.arrays(self.clusterToDigis, library='np')[self.clusterToDigis]
    
            self.data['cluster'] = self._genMatrices(uCellIDs, vCellIDs, cellCharges, clusterDigits).astype('int')
    
        def _genMatrices(self, uCellIDs: ArrayLike, vCellIDs: ArrayLike, cellCharges: ArrayLike, clusterDigits: ArrayLike, matrixSize: tuple = (9, 9)) -> np.ndarray:
            """
            this takes the ragged/jagged digit arrays and converts them into 9x9 matrices
            it's a rather slow process because of all the looping
            """
            plotRange = np.array(matrixSize) // 2
            events = []
    
            for event in range(len(cellCharges)):
                adcValues = []
                digitsU = np.array(uCellIDs[event])
                digitsV = np.array(vCellIDs[event])
                digitsCharge = np.array(cellCharges[event])
                digitIndices = clusterDigits[event]
    
                for indices in digitIndices:
                    cacheImg = np.zeros(matrixSize)
                    maxChargeIndex = digitsCharge[indices].argmax()
                    uMax, vMax = digitsU[indices[maxChargeIndex]], digitsV[indices[maxChargeIndex]]
                    uPos, vPos = digitsU[indices] - uMax + plotRange[0], digitsV[indices] - vMax + plotRange[1]
    
                    valid_indices = (uPos >= 0) & (uPos < matrixSize[0]) & (vPos >= 0) & (vPos < matrixSize[1])
                    cacheImg[uPos[valid_indices].astype(int), vPos[valid_indices].astype(int)] = digitsCharge[indices][valid_indices]
                    adcValues.append(cacheImg)
    
                events.extend(adcValues)
    
            return np.array(events, dtype=object)
    
        def genCoordisnate(self) -> None:
            """
            converting the uv coordinates, together with sensor ids, into xyz coordinates
            """
            if self.gotClusters is False:
                self.getClusters()
            xcoords, ycoords, zcoords = self._getCartesian(self.data['uPosition'], self.data['vPosition'], self.data['sensorID'])
            self.data['xPosition'] = xcoords
            self.data['yPosition'] = ycoords
            self.data['zPosition'] = zcoords
    
        def _getCartesian(self, uPositions: ArrayLike, vPositions: ArrayLike, sensorIDs: ArrayLike) -> tuple[np.ndarray]:
            """
            a private method for transposing/converting 2d uv coords into 3d xyz coordinates
            """
            length = len(sensorIDs)
            xArr, yArr, zArr = np.zeros(length), np.zeros(length), np.zeros(length)
    
            # iterting over the cluster arrays
            for index, (u, v, sensor_id) in enumerate(zip(uPositions, vPositions, sensorIDs)):
                # grabbing the shift vector and rotation angle
                shift, angle = self.transformation[str(sensor_id)]
    
                # setting up rotation matrix
                theta = np.deg2rad(angle)
                rotMatrix = np.array([[np.cos(theta), -np.sin(theta), 0], [np.sin(theta), np.cos(theta), 0], [0, 0, 1]])
    
                # projecting uv coordinates into 3d space
                point = np.array([u, 0, v])
    
                # shifting and rotating the projected vector
                shifted = rotMatrix.dot(point) + shift
                xArr[index], yArr[index], zArr[index] = shifted
    
            return xArr, yArr, zArr
    
        def getLayers(self) -> None:
            """
            looks up the corresponding layers and ladders for every cluster
            """
            if self.gotClusters is False:
                self.getClusters()
            layers, ladders = [], []
            for id in self.data['sensorID']:
                layer, ladder = self.layersLadders[str(id)]
                layers.append(layer)
                ladders.append(ladder)
            self.data['layers'] = np.array(layers)
            self.data['ladder'] = np.array(ladders)
    
        def getMCData(self) -> None:
            """
            this loads the monte carlo from the root file
            """
    
            # the monte carlo data, they are longer than the cluster data
            pdg = self.eventTree.arrays(self.mcData[0], library='np')[self.mcData[0]]
            momentumX = self.eventTree.arrays(self.mcData[1], library='np')[self.mcData[1]]
            momentumY = self.eventTree.arrays(self.mcData[2], library='np')[self.mcData[2]]
            momentumZ = self.eventTree.arrays(self.mcData[3], library='np')[self.mcData[3]]
    
            # 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 = self.eventTree.arrays(self.clusterToMC, library='np')[self.clusterToMC]
            mcToCluster = self.eventTree.arrays(self.mcToCluster, library='np')[self.mcToCluster]
    
            # 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 = self.eventTree.arrays('PXDClusters/PXDClusters.m_clsCharge', library='np')['PXDClusters/PXDClusters.m_clsCharge']
    
            # reorganizing MC data
            momentumXList = []
            momentumYList = []
            momentumZList = []
            pdgList = []
            clusterNumbersList = []
            for i in range(len(clusterToMC)):
                # _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 = self._fillMCList(mcToCluster[i], clusterToMC[i], len(clsCharge[i]))
                clusterNumbersList.append(fullClusterReferences)
                pdgs, xmom, ymom, zmom = self._getMCData(fullClusterReferences, pdg[i], momentumX[i], momentumY[i], momentumZ[i])
                momentumXList.append(xmom)
                momentumYList.append(ymom)
                momentumZList.append(zmom)
                pdgList.append(pdgs)
    
            self.data['momentumX'] = self._flatten(momentumXList)
            self.data['momentumY'] = self._flatten(momentumYList)
            self.data['momentumZ'] = self._flatten(momentumZList)
            self.data['pdg'] = self._flatten(pdgList)
            self.data['clsNumber'] = self._flatten(clusterNumbersList)
    
        def _findMissing(self, lst: list, length: int) -> list:
            """
            a private method for finding missing elements in mc data arrays
            """
            return sorted(set(range(0, length)) - set(lst))
    
        def _fillMCList(self, fromClusters: ArrayLike, toClusters: ArrayLike, length: ArrayLike) -> list:
            """
            a private method for filling MC data arrays where clusters don't have
            any information
            """
            missingIndex = self._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
    
        def _getMCData(self, toClusters: ArrayLike, pdgs: ArrayLike, xMom: ArrayLike, yMom: ArrayLike, zMom: ArrayLike) -> tuple[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
            """
            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)
    
        def getStructuredArray(self) -> np.ndarray:
            """
            this converts the data dict of this class into a structured numpy array
            """
            # Create a list to hold the dtype specifications
            dtype = []
    
            # Iterate through the dictionary keys and values
            for key, value in self.data.items():
                # Determine the data type of the first value in the list
                sampleValue = value[0]
    
                if isinstance(sampleValue, np.ndarray):
                    # If the value is an array, use its shape and dtype
                    fieldDtype = (sampleValue.dtype, sampleValue.shape)
                else:
                    # Otherwise, use the type of the value itself
                    fieldDtype = type(sampleValue)
    
                # Append the key and data type to the dtype list
                dtype.append((key, fieldDtype))
    
            # Convert the dictionary to a list of tuples
            keys = list(self.data.keys())
            dataList = [tuple(self.data[key][i] for key in keys) for i in range(len(self.data[keys[0]]))]
    
            # Create the structured array
            structuredArray = np.array(dataList, dtype=dtype)
    
            return structuredArray