DTCalibration/scripts/mergeVDriftHistosByStation.py

0001 #! /usr/bin/env python3
0002
0003 import ROOT
0004 import sys
0005
0006 def getHistoName(wheel,station,sector):
0007     wheelStr = 'W' + str(wheel)
0008     stationStr = 'St' + str(station)
0009     sectorStr = 'Sec' + str(sector)
0010     name = "hRPhiVDriftCorr_" + wheelStr + "_" + stationStr + "_" + sectorStr
0011
0012     return name
0013
0014 def mergeHistosWheelSector(file, wheel, sector):
0015
0016     histWheelSector = None
0017     for station in range(1,5):
0018         if sector in (13,14) and station != 4: continue
0019         name = getHistoName(wheel,station,sector)
0020         hist = file.Get(name)
0021         if hist:
0022             print("Adding",hist.GetName())
0023             if not histWheelSector: histWheelSector = hist.Clone( "h_W%d_Sec%d" % (wheel,sector) )
0024             else: histWheelSector.Add(hist)
0025
0026     return histWheelSector
0027
0028 def mergeHistosWheelStation(file, wheel, station):
0029
0030     sectors = range(1,13)
0031     if station == 4: sectors.extend([13,14])
0032     histWheelStation = None
0033     for sector in sectors:
0034         name = getHistoName(wheel,station,sector)
0035         hist = file.Get(name)
0036         if hist:
0037             print("Adding",hist.GetName())
0038             if not histWheelStation: histWheelStation = hist.Clone( "h_W%d_St%d" % (wheel,station) )
0039             else: histWheelStation.Add(hist)
0040
0041     return histWheelStation
0042
0043 if __name__ == '__main__':
0044     import optparse
0045     parser = optparse.OptionParser ("Usage: %prog [--options]")
0046     # Options
0047     parser.add_option("-f","--file", dest="file", help="Input file name")
0048     parser.add_option("-o","--out", dest="out", default="merged.root", help="Output file name")
0049     (options, args) = parser.parse_args()
0050
0051     if not options.file:
0052         parser.error('must set an input file')
0053
0054     file = ROOT.TFile(options.file,"READ")
0055     ROOT.gROOT.cd()
0056
0057     wheels = range(-2,3)
0058     stations = range(1,5)
0059     sectors = range(1,15)
0060     histos = {}
0061     for wheel in wheels:
0062         for station in stations:
0063             print("Merging histos from Wheel %d, Station %d" % (wheel,station))
0064             histos[(wheel,station)] = mergeHistosWheelStation(file,wheel,station)
0065
0066     file.Close()
0067
0068     outputFile = ROOT.TFile(options.out,"RECREATE")
0069     outputFile.cd()
0070     for wheel in wheels:
0071         wheelStr = 'W' + str(wheel)
0072         for station in stations:
0073             stationStr = 'St' + str(station)
0074             for sector in sectors:
0075                 if sector in (13,14) and station != 4: continue
0076                 sectorStr = 'Sec' + str(sector)
0077                 name = "hRPhiVDriftCorr_" + wheelStr + "_" + stationStr + "_" + sectorStr
0078                 print("Writing",name)
0079                 histos[(wheel,station)].Clone(name).Write()
0080
0081     outputFile.Close()
0082
0083     sys.exit(0)