SiStripMonitorClient/scripts/PhaseITreeProducer.py

0001 #!/usr/bin/env python3
0002
0003 from __future__ import print_function
0004 import sys
0005 from ROOT import TTree, TFile, gROOT, TClass
0006 from array import array
0007 from copy import deepcopy
0008
0009 gROOT.SetBatch()        # don't pop up canvases
0010
0011 #get data files
0012
0013 def getFileInPath(rfile):
0014    import os
0015    for dir in os.environ['CMSSW_SEARCH_PATH'].split(":"):
0016      if os.path.exists(os.path.join(dir,rfile)): return os.path.join(dir,rfile)
0017    return None
0018
0019 # Default values
0020 inputFileName = "DQM.root"
0021 outputFileName = "DQMTree.root"
0022 #detIDsFileName = "DATA/detids.dat"
0023 detIDsFileName = getFileInPath('DQM/SiStripMonitorClient/data/detids.dat')
0024
0025 class ModuleLvlValuesReader:
0026
0027   ############################################################################
0028
0029   def __TraverseDirTree(self, dir):
0030
0031     currPath = (dir.GetPath().split(":/"))[1]
0032
0033     for obj in dir.GetListOfKeys():
0034       if not obj.IsFolder():
0035         if obj.ReadObjectAny(TClass.GetClass("TH2")):
0036           th2 = deepcopy(obj.ReadObj())
0037           name = th2.GetName()
0038           if 6 < th2.GetNbinsX() < 10 and name.find("per") != -1 and name.find("Lumisection") == -1: #take only module lvl plots
0039             print(''.join([dir.GetPath(), '/', name]))
0040
0041             # fix when there are plots starting with the same strings in different directories
0042             prefix = ""
0043             for i in self.dirs:
0044               if currPath.startswith(i):
0045                 prefix = self.dirsAliases[i]
0046                 break
0047             # print(currPath, prefix)
0048             th2.SetName(prefix + th2.GetName())
0049             self.listOfNumHistograms.append(th2)
0050       else:
0051         self.__TraverseDirTree(obj.ReadObj())
0052
0053   def __GetPartStr(self, isXlowerThanZero, isYlowerThanZero):
0054     if isXlowerThanZero and isYlowerThanZero:
0055       return "mO"
0056     if isXlowerThanZero and isYlowerThanZero == False:
0057       return "mI"
0058     if isXlowerThanZero == False and isYlowerThanZero:
0059       return "pO"
0060     if isXlowerThanZero == False and isYlowerThanZero == False:
0061       return "pI"
0062
0063   def __GetBarrelSector(self, layer, signedLadder, signedModule): #adapted from PixelBarrelName
0064     theLadder = abs(signedLadder)
0065     theModule = abs(signedModule)
0066
0067     sector = 0
0068
0069     if layer == 1:
0070
0071       if theLadder == 1:
0072         if theModule >= 2:
0073           return 1
0074         else:
0075           return 2
0076       if theLadder == 2:
0077         if theModule >= 3:
0078           return 2
0079         else:
0080           return 3
0081       if theLadder == 3:
0082         if theModule >= 4:
0083           return 3
0084         else:
0085           return 4
0086       if theLadder == 4:
0087         if theModule >= 2:
0088           return 5
0089         else:
0090           return 6
0091       if theLadder == 5:
0092         if theModule >= 3:
0093           return 6
0094         else:
0095           return 7
0096       if theLadder == 6:
0097         if theModule >= 4:
0098           return 7
0099         else:
0100           return 8
0101     # here is used simplified form of assignment, see source file for reference
0102     elif layer == 2:
0103       i = theLadder // 5
0104       sector = i * 3
0105       shortLadder = theLadder - 5 * i
0106       for i in range(0, shortLadder, 2):
0107         sector = sector + 1
0108       return sector
0109     elif layer == 3:
0110       sector = 1
0111       for i in range(2, theLadder, 3):
0112         if (i + 1) % 3 == 0:
0113           sector = sector + 1
0114       return sector
0115     elif layer == 4:
0116       sector = (theLadder + 3) // 4
0117       return sector
0118
0119   def __BuildOnlineBarrelName(self, signedModule, signedLadder, layer): #in Phase1 it is assumed that there are only full modules
0120     thePart = self.__GetPartStr(signedModule < 0, signedLadder < 0)
0121     theSector = str(self.__GetBarrelSector(layer, signedLadder, signedModule))
0122     return "BPix_B" + thePart + "_SEC" + theSector + "_LYR" + str(layer) + "_LDR" + str(abs(signedLadder)) + "F_MOD" + str(abs(signedModule))
0123
0124   def __BuildOnlineDiskName(self, signedDisk, signedBlade, panel, ring):
0125     thePart = self.__GetPartStr(signedDisk < 0, signedBlade < 0)
0126     return "FPix_B" + thePart + "_D" + str(abs(signedDisk)) + "_BLD" + str(abs(signedBlade)) + "_PNL" + str(panel) + "_RNG" + str(ring)
0127
0128   def __GroupHistograms(self):
0129     currentGroupName = ""
0130     groupOfHists = []
0131     self.groupedHistograms = []
0132
0133     ##### GROUP ALL LAYERS/RINGS HAVING SIMILAR INFORMATION
0134     for obj in self.listOfNumHistograms:
0135       objName = obj.GetName()
0136       objNameSplit = objName.split("_")
0137       objNameCollected = ''.join(objNameSplit[0:-1])
0138       if objNameCollected != currentGroupName:
0139         if len(groupOfHists):
0140           self.groupedHistograms.append(groupOfHists)
0141           groupOfHists = []
0142
0143         currentGroupName = objNameCollected
0144       groupOfHists.append(obj)
0145     self.groupedHistograms.append(groupOfHists) #the last group
0146
0147   def __CreateDummyStructAsStr(self, dicData):
0148     str = "struct MyStruct{"
0149     str = str + "Int_t key;"
0150     leafStr = "key/I"
0151     for k in dicData:
0152       str = str + "Float_t " + k + ";"
0153       leafStr = leafStr + ":" + k + "/F"
0154     str = str + "};"
0155     return str, leafStr
0156
0157   ############################################################################
0158
0159   def __init__(self, inputDQMName, outputDQMName, modDicName):
0160     self.inputFileName = inputDQMName
0161     self.outputFileName = outputDQMName
0162     self.detIDsFileName = modDicName
0163
0164     index = self.inputFileName.find('R000')
0165     runNumber = self.inputFileName[index+4:index+10]
0166
0167     self.dirs = ['DQMData/Run %s/PixelPhase1/Run summary/Phase1_MechanicalView' % (runNumber),
0168                  'DQMData/Run %s/PixelPhase1/Run summary/Tracks' % (runNumber)]
0169     self.dirsAliases = {self.dirs[0]:"", self.dirs[1]: "T"}
0170
0171     self.inputFile = TFile(self.inputFileName)
0172     self.listOfNumHistograms = []
0173     self.availableNames = []
0174
0175     self.maxLadderToLayer = {6:1, 14:2, 22:3, 32:4}
0176     self.maxBladeToRing = {11:1, 17:2}
0177
0178     self.internalData = {}
0179
0180     if self.inputFile.IsOpen():
0181       print("%s opened successfully!" % (self.inputFileName))
0182       #Get all neeeded histograms
0183       for dir in self.dirs:
0184         self.__TraverseDirTree(self.inputFile.Get(dir))
0185       print("Histograms to read %d" % (len(self.listOfNumHistograms)))
0186
0187       self.detDict = {}
0188
0189       with open(self.detIDsFileName, "r") as detIDs:  # create dictionary online -> rawid
0190         for entry in detIDs:
0191           items = entry.replace("\n", " ").split(" ")
0192           self.detDict.update({items[1] : int(items[0])})
0193           # init internal data structure
0194           self.internalData.update({int(items[0]) : {}})
0195
0196       self.__GroupHistograms()
0197
0198     else:
0199       print("Unable to open file %s" % (self.inputFileName))
0200
0201   def ReadHistograms(self):
0202     if self.inputFile.IsOpen():
0203       for group in self.groupedHistograms:
0204         # name = ''.join(group[0].GetName().split("_")[0:-1])
0205         name = ''.join(group[0].GetName().split("_per_")[0])
0206         self.availableNames.append(name)
0207         # print(name)
0208         for obj in group:
0209           nbinsX = obj.GetNbinsX()
0210           nbinsY = obj.GetNbinsY()
0211
0212           if nbinsX == 9: # BARREL
0213             maxX = nbinsX // 2
0214             maxY = nbinsY // 2
0215
0216             for x in range(-maxX, maxX + 1):
0217               if x == 0:
0218                 continue
0219               for y in range(-maxY, maxY + 1, 1):
0220                 if y == 0:
0221                   continue
0222                 onlineName = self.__BuildOnlineBarrelName(x, y, self.maxLadderToLayer[maxY])
0223                 self.internalData[self.detDict[onlineName]].update({name : obj.GetBinContent(x + maxX + 1, y + maxY + 1)})
0224
0225           elif nbinsX == 7: # FORWARD
0226             maxX = nbinsX // 2
0227             maxY = nbinsY // 4
0228
0229             for x in range(-maxX, maxX + 1):
0230               if x == 0:
0231                 continue
0232               for y in range(-maxY, maxY + 1):
0233                 if int(y) == 0:
0234                   continue
0235                 for panel in range(1, 3):
0236                   onlineName = self.__BuildOnlineDiskName(x, y, panel, self.maxBladeToRing[maxY])
0237                   self.internalData[self.detDict[onlineName]].update({name : obj.GetBinContent(x + maxX + 1, (y + maxY) * 2 + panel)})
0238           else:
0239             print("Unrecognized plot")
0240       else:
0241         print("Histograms saved to internal data structure")
0242   def CreateTree(self): # too much complication lvl, use CreateTree2
0243     if len(self.internalData):
0244       # <- TOTAL WORKAROUND: START -> #
0245       # SEE: http://wlav.web.cern.ch/wlav/pyroot/tpytree.html
0246
0247       leafStr = ""
0248       for key in self.internalData:
0249         s, leafStr = self.__CreateDummyStructAsStr(self.internalData[key])
0250         print(s)
0251         print(leafStr)
0252         gROOT.ProcessLine(s)
0253         break #all modules are assumed to have the same set of measured parameters
0254
0255       from ROOT import MyStruct
0256       ms = MyStruct()
0257       # <- TOTAL WORKAROUND: END -> #
0258
0259       self.outputFile = TFile(self.outputFileName, "recreate")
0260       tree = TTree("tree", "readData")
0261       tree.Branch("b", ms, leafStr)
0262
0263       tree.SetBranchAddress("b", tree)
0264
0265       for key in self.internalData:
0266         setattr(ms, "key", key)
0267         for d in self.internalData[key]:
0268           setattr(ms, d, (self.internalData[key])[d])
0269         tree.Fill()
0270
0271         # break
0272       tree.Write()
0273       self.outputFile.Close()
0274
0275       print("Data saved as TTree object")
0276
0277   def CreateTree2(self): # use for TkCommissioner
0278     if len(self.internalData):
0279
0280       self.outputFile = TFile(self.outputFileName, "recreate")
0281       tree = TTree("tree", "readData")
0282
0283       key = array("i", [0])
0284       tree.Branch("detid", key, "detid/i")
0285
0286       dat = {}
0287       for k in self.internalData:
0288         for d in self.internalData[k]:
0289           dat.update({d : array("f", [0])})
0290           tree.Branch(d, dat[d], d + "/F")
0291         break
0292
0293       for k in self.internalData:
0294         key[0] = k
0295         for d in self.internalData[k]:
0296           (dat[d])[0] = (self.internalData[k])[d]
0297         tree.Fill()
0298
0299       tree.Write()
0300       self.outputFile.Close()
0301
0302       print("Data saved as TTree object")
0303
0304   def DumpData(self):
0305     for key in self.internalData:
0306       print("#"*20)
0307       print(key)
0308       module = self.internalData[key]
0309       for d in module:
0310         print((d, module[d]))
0311
0312     for i in self.availableNames:
0313       print(i)
0314     print(len(self.availableNames))
0315
0316   def __del__(self):
0317     if self.inputFile:
0318       if self.inputFile.IsOpen():
0319         self.inputFile.Close()
0320
0321 #--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--
0322 for i in range(1, len(sys.argv), 1):
0323   if i == 1:
0324     inputFileName = sys.argv[i]
0325   elif i == 2:
0326 #    detIDsFileName = sys.argv[i]
0327 #  elif i == 3:
0328     outputFileName = sys.argv[i]
0329
0330 #readerObj = ModuleLvlValuesReader(inputFileName, outputFileName)
0331 readerObj = ModuleLvlValuesReader(inputFileName, outputFileName, detIDsFileName)
0332 readerObj.ReadHistograms()
0333 readerObj.CreateTree2()
0334 # readerObj.DumpData()