SiStripMonitorClient/scripts/PhaseITreeProducer.py

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