def MergeTuples(path, oldTreeName, newTreeName):
myChain = TChain(oldTreeName)
for fname in glob.glob(path):
print(fname)
myChain.Add(fname)
myTree = myChain.CloneTree(-1, "fast")
myTree.SetName(newTreeName)
return myTree
def dump(self, location, ntuples):
if any(ntuples):
self.f.cd()
self.f.mkdir(location)
self.f.cd(location)
for ntuple in ntuples:
t_chain = TChain()
for f in self.input_:
m = self.matchLine.match(f)
if m:
newName = m.group(1)
os.rename(f, newName)
f = newName
self._logger.info('ntuple: %s | file: %s', ntuple, f)
t_chain.Add(f + '/' + location + '/' + ntuple)
t_ntuple = t_chain.CloneTree(-1)
t_ntuple.Write("", TObject.kOverwrite)
def __call__(self, inputDS_list, basepath, trigger, treename):
if not type(inputDS_list) is list: inputDS_list = [inputDS_list]
from ROOT import TChain, TObject
self._file.cd()
self._file.mkdir(basepath + '/' + trigger)
self._file.cd(basepath + '/' + trigger)
cobject = TChain()
for inputDS in inputDS_list:
self._logger.info( ('Copy tree name %s in %s to %s')%(treename,\
inputDS, self._outputDS) )
location = inputDS + '/' + basepath + '/' + trigger + '/' + treename
cobject.Add(location)
if cobject.GetEntries() == 0:
self._logger.warning(
('There is no events into this path: %s') % (location))
del cobject
return False
else:
self._logger.info(('Copy %d events...') % (cobject.GetEntries()))
try: # Copy protection
copy_cobject = cobject.CloneTree(-1)
try: # Write protection
copy_cobject.Write("", TObject.kOverwrite)
del copy_cobject
except: # error
del copy_cobject
self._logger.error('Can not write the tree')
return False
except: # error
del cobject
self._logger.error('Can not copy the tree')
return False
del cobject
# Everything is good
return True
def merge(self, filelist, i):
'''
Takes a python list "filelist" as argument, merges all the root files from that list
into an output root file, whose name is decided by "i"
'''
# String manipulation to build output name
outfile = 'merger_out' + self.subdir + '/'
temp_index = "%05d" % (i, )
if 'reco' in filelist[0]:
temp_int = basename(filelist[0]).find('.reco') - 6
outfile = outfile + basename(
filelist[0])[0:temp_int] + temp_index + '.reco.root'
elif 'mc' in filelist[0]:
temp_int = basename(filelist[0]).find('.mc') - 6
outfile = outfile + basename(
filelist[0])[0:temp_int] + temp_index + '.mc.root'
else:
raise Exception(
"Could not identify file type. Looking for '.reco.root' or '.mc.root'"
)
del temp_int, temp_index
if isfile(outfile):
return
# Open files, build TChain
dmpch = DmpChain("CollectionTree")
metachain = TChain("RunMetadataTree")
for f in filelist:
dmpch.Add(f)
metachain.Add(filelist[0])
# Get Metadata values from "old" files
nrofevents = 0
for f in filelist:
fo = TFile.Open(f)
rmt = fo.Get("RunMetadataTree")
simuheader = DmpRunSimuHeader()
rmt.SetBranchAddress("DmpRunSimuHeader", simuheader)
rmt.GetEntry(0)
nrofevents += simuheader.GetEventNumber()
fo.Close()
fob = TFile.Open(filelist[0])
rmt = fob.Get("RunMetadataTree")
simuheader = DmpRunSimuHeader()
rmt.SetBranchAddress("DmpRunSimuHeader", simuheader)
rmt.GetEntry(0)
runNumber = simuheader.GetRunNumber()
spectrumType = simuheader.GetSpectrumType()
sourceType = simuheader.GetSourceType()
vertexRadius = simuheader.GetVertexRadius()
sourceGen = simuheader.GetSourceGen()
maxEne = simuheader.GetMaxEne()
minEne = simuheader.GetMinEne()
version = simuheader.GetVersion()
fluxFile = simuheader.GetFluxFile()
orbitFile = simuheader.GetOrbitFile()
prescale = simuheader.GetPrescale()
startMJD = simuheader.GetMJDstart()
stopMJD = simuheader.GetMJDstop()
seed = simuheader.GetSeed()
fob.Close()
# Write new file
tf = TFile(outfile, "RECREATE")
ot = dmpch.CloneTree(-1, "fast")
om = metachain.CloneTree(0)
simhdr = DmpRunSimuHeader()
om.SetBranchAddress("DmpRunSimuHeader", simhdr)
for i in xrange(metachain.GetEntries()):
metachain.GetEntry(i)
simhdr.SetEventNumber(nrofevents)
simhdr.SetRunNumber(runNumber)
simhdr.SetSpectrumType(spectrumType)
simhdr.SetSourceType(sourceType)
simhdr.SetVertexRadius(vertexRadius)
simhdr.SetSourceGen(sourceGen)
simhdr.SetMaxEne(maxEne)
simhdr.SetMinEne(minEne)
simhdr.SetVersion(version)
simhdr.SetFluxFile(fluxFile)
simhdr.SetOrbitFile(orbitFile)
simhdr.SetPrescale(prescale)
simhdr.SetMJDstart(startMJD)
simhdr.SetMJDstop(stopMJD)
simhdr.SetSeed(seed)
om.Fill()
tf.Write()
tf.Close()
self.equivalence[basename(outfile)] = filelist
for x in [
dmpch, metachain, nrofevents, runNumber, spectrumType,
sourceType, vertexRadius, sourceGen, maxEne, minEne, version,
fluxFile, orbitFile, prescale, startMJD, stopMJD, seed,
simuheader, rmt, simhdr, outfile
]:
try:
del x
except:
continue
def main(options,args):
gROOT.Reset()
#define trigger unprescaled periods
jpsi_Trig_unP = {#'HLT_Mu0_Track0_Jpsi' :[133446,141882],
#'HLT_Mu0_TkMu0_Jpsi' :[140116,144114],
#'HLT_Mu0_TkMu0_OST_Jpsi' :[146428,148058],
#'HLT_Mu0_TkMu0_OST_Jpsi_Tight_v2':[148819,149182],
#'HLT_Mu0_TkMu0_OST_Jpsi_Tight_v3':[149291,149442]}
'HLT_DoubleMu0' :[133446,147116],
'HLT_DoubleMu0_Quarkonium_v1' :[147196,149442]}
#load calcPol routine
gROOT.ProcessLine('.L calcPol.C+')
inputs = TChain(options.treeName)
for arg in args:
print 'Adding: ',arg
inputs.Add(arg)
out = TFile.Open(options.output,'RECREATE')
outTree = inputs.CloneTree(0)
# polarization stuff
gROOT.ProcessLine('struct theBranches { TLorentzVector* JpsiP; TLorentzVector* muPosP; TLorentzVector* muNegP; Double_t JpsiMass; Double_t JpsiPt; Double_t JpsiRap; Double_t muPosPt; Double_t muPosEta; Double_t muPosPhi; Double_t muNegPt; Double_t muNegEta; Double_t muNegPhi; Double_t costh_CS; Double_t phi_CS; Double_t costh_HX; Double_t phi_HX; Double_t costh_PHX; Double_t phi_PHX; Double_t costh_sGJ; Double_t phi_sGJ; Double_t costh_GJ1; Double_t phi_GJ1; Double_t costh_GJ2; Double_t phi_GJ2; }')
# trigger stuff
gROOT.ProcessLine('struct triggers { Int_t runNb; Int_t HLT_Mu0_Track0_Jpsi; Int_t HLT_Mu0_TkMu0_Jpsi; Int_t HLT_Mu0_TkMu0_OST_Jpsi; Int_t HLT_Mu0_TkMu0_OST_Jpsi_Tight_v2; Int_t HLT_Mu0_TkMu0_OST_Jpsi_Tight_v3; Int_t HLT_DoubleMu0; Int_t HLT_DoubleMu0_Quarkonium_v1; }')
branches = ROOT.theBranches()
thetriggers = ROOT.triggers()
inputs.SetBranchAddress('runNb',AddressOf(thetriggers,'runNb'))
inputs.SetBranchAddress('HLT_Mu0_Track0_Jpsi',AddressOf(thetriggers,'HLT_Mu0_Track0_Jpsi'))
inputs.SetBranchAddress('HLT_Mu0_TkMu0_Jpsi',AddressOf(thetriggers,'HLT_Mu0_TkMu0_Jpsi'))
inputs.SetBranchAddress('HLT_Mu0_TkMu0_OST_Jpsi',AddressOf(thetriggers,'HLT_Mu0_TkMu0_OST_Jpsi'))
inputs.SetBranchAddress('HLT_Mu0_TkMu0_OST_Jpsi_Tight_v2',AddressOf(thetriggers,'HLT_Mu0_TkMu0_OST_Jpsi_Tight_v2'))
inputs.SetBranchAddress('HLT_Mu0_TkMu0_OST_Jpsi_Tight_v3',AddressOf(thetriggers,'HLT_Mu0_TkMu0_OST_Jpsi_Tight_v3'))
inputs.SetBranchAddress('HLT_DoubleMu0',AddressOf(thetriggers,'HLT_DoubleMu0'))
inputs.SetBranchAddress('HLT_DoubleMu0_Quarkonium_v1',AddressOf(thetriggers,'HLT_DoubleMu0_Quarkonium_v1'))
outTree.SetBranchAddress('runNb',AddressOf(thetriggers,'runNb'))
outTree.SetBranchAddress('HLT_Mu0_Track0_Jpsi',AddressOf(thetriggers,'HLT_Mu0_Track0_Jpsi'))
outTree.SetBranchAddress('HLT_Mu0_TkMu0_Jpsi',AddressOf(thetriggers,'HLT_Mu0_TkMu0_Jpsi'))
outTree.SetBranchAddress('HLT_Mu0_TkMu0_OST_Jpsi',AddressOf(thetriggers,'HLT_Mu0_TkMu0_OST_Jpsi'))
outTree.SetBranchAddress('HLT_Mu0_TkMu0_OST_Jpsi_Tight_v2',AddressOf(thetriggers,'HLT_Mu0_TkMu0_OST_Jpsi_Tight_v2'))
outTree.SetBranchAddress('HLT_Mu0_TkMu0_OST_Jpsi_Tight_v3',AddressOf(thetriggers,'HLT_Mu0_TkMu0_OST_Jpsi_Tight_v3'))
outTree.SetBranchAddress('HLT_DoubleMu0',AddressOf(thetriggers,'HLT_DoubleMu0'))
outTree.SetBranchAddress('HLT_DoubleMu0_Quarkonium_v1',AddressOf(thetriggers,'HLT_DoubleMu0_Quarkonium_v1'))
inputs.SetBranchAddress('JpsiP',AddressOf(branches,'JpsiP'))
inputs.SetBranchAddress('muPosP',AddressOf(branches,'muPosP'))
inputs.SetBranchAddress('muNegP',AddressOf(branches,'muNegP'))
phiFolded = [0,0,0,0,0,0]
thetaAdjusted = [0,0,0,0,0,0]
outTree.Branch("JpsiMass",AddressOf(branches,'JpsiMass'),"JpsiMass/D");
outTree.Branch("JpsiPt",AddressOf(branches,'JpsiPt'),"JpsiPt/D");
outTree.Branch("JpsiRap",AddressOf(branches,'JpsiRap'),"JpsiRap/D");
outTree.Branch("muPosPt",AddressOf(branches,'muPosPt'),"muPosPt/D");
outTree.Branch("muPosEta",AddressOf(branches,'muPosEta'),"muPosEta/D");
outTree.Branch("muPosPhi",AddressOf(branches,'muPosPhi'),"muPosPhi/D");
outTree.Branch("muNegPt",AddressOf(branches,'muNegPt'),"muNegPt/D");
outTree.Branch("muNegEta",AddressOf(branches,'muNegEta'),"muNegEta/D");
outTree.Branch("muNegPhi",AddressOf(branches,'muNegPhi'),"muNegPhi/D");
outTree.Branch("costh_CS",AddressOf(branches,'costh_CS'),"costh_CS/D");
outTree.Branch("phi_CS",AddressOf(branches,'phi_CS'),"phi_CS/D");
outTree.Branch("costh_HX",AddressOf(branches,'costh_HX'),"costh_HX/D");
outTree.Branch("phi_HX",AddressOf(branches,'phi_HX'),"phi_HX/D");
outTree.Branch("costh_PHX",AddressOf(branches,'costh_PHX'),"costh_PHX/D");
outTree.Branch("phi_PHX",AddressOf(branches,'phi_PHX'),"phi_PHX/D");
outTree.Branch("costh_sGJ",AddressOf(branches,'costh_sGJ'),"costh_sGJ/D");
outTree.Branch("phi_sGJ",AddressOf(branches,'phi_sGJ'),"phi_sGJ/D");
outTree.Branch("costh_GJ1",AddressOf(branches,'costh_GJ1'),"costh_GJ1/D");
outTree.Branch("phi_GJ1",AddressOf(branches,'phi_GJ1'),"phi_GJ1/D");
outTree.Branch("costh_GJ2",AddressOf(branches,'costh_GJ2'),"costh_GJ2/D");
outTree.Branch("phi_GJ2",AddressOf(branches,'phi_GJ2'),"phi_GJ2/D");
nEvents=inputs.GetEntries()
print 'There are ',nEvents,' events to process!'
FiducialCutCountPos = 0
FiducialCutCountNeg = 0
for j in range(0,nEvents):
nb = inputs.GetEntry(j)
if branches.muPosP.Pt() == branches.muNegP.Pt():
continue
passesTrigger = True
if not options.isMC:
passesTrigger=False
for trig in jpsi_Trig_unP.keys():
if thetriggers.runNb >= jpsi_Trig_unP[trig][0] and thetriggers.runNb <= jpsi_Trig_unP[trig][-1] and getattr(thetriggers,trig) == 1:
passesTrigger = True
if not passesTrigger:
continue
branches.muPosPt = branches.muPosP.Pt()
branches.muPosEta = branches.muPosP.Eta()
branches.muPosPhi = branches.muPosP.Phi()
branches.muNegPt = branches.muNegP.Pt()
branches.muNegEta = branches.muNegP.Eta()
branches.muNegPhi = branches.muNegP.Phi()
#(a) on the positive muon
if((abs(branches.muPosEta) < jpsi.etaPS[0] and branches.muPosPt < jpsi.pTMuMin[0]) or
(abs(branches.muPosEta) > jpsi.etaPS[0] and abs(branches.muPosEta) < jpsi.etaPS[1] and branches.muPosP.P() < jpsi.pMuMin[1]) or
(abs(branches.muPosEta) > jpsi.etaPS[1] and abs(branches.muPosEta) < jpsi.etaPS[2] and branches.muPosPt < jpsi.pTMuMin[2])):
FiducialCutCountPos += 1
continue
#(b) on the negative muon
if((abs(branches.muNegEta) < jpsi.etaPS[0] and branches.muNegPt < jpsi.pTMuMin[0]) or
(abs(branches.muNegEta) > jpsi.etaPS[0] and abs(branches.muNegEta) < jpsi.etaPS[1] and branches.muNegP.P() < jpsi.pMuMin[1]) or
(abs(branches.muNegEta) > jpsi.etaPS[1] and abs(branches.muNegEta) < jpsi.etaPS[2] and branches.muNegPt < jpsi.pTMuMin[2])):
FiducialCutCountNeg += 1
continue
ROOT.calcPol(branches.muPosP,branches.muNegP)
#set up KLUDGE to make DoubleMu0 and DoubleMu0_Quarkonium to be the same as they have the same acceptance map
if thetriggers.HLT_DoubleMu0_Quarkonium_v1 == 1:
thetriggers.HLT_DoubleMu0 = 1
if options.isFolded:
for frame in range(len(jpsi.frameLabel)):
phiFolded[frame] = ROOT.thisPhi[frame]
thetaAdjusted[frame] = ROOT.thisCosTh[frame]
if ROOT.thisPhi[frame] >= -90 and ROOT.thisPhi[frame] < 0.:
phiFolded[frame] *= -1
elif ROOT.thisPhi[frame] > 90 and ROOT.thisPhi[frame] < 180:
phiFolded[frame] = 180 - ROOT.thisPhi[frame]
thetaAdjusted[frame] *= -1
elif ROOT.thisPhi[frame] > -180 and ROOT.thisPhi[frame] < -90:
phiFolded[frame] = 180 + ROOT.thisPhi[frame]
thetaAdjusted[frame] *= -1
branches.costh_CS = thetaAdjusted[0]
branches.phi_CS = phiFolded[0]
branches.costh_HX = thetaAdjusted[1]
branches.phi_HX = phiFolded[1]
branches.costh_PHX = thetaAdjusted[2]
branches.phi_PHX = phiFolded[2]
branches.costh_sGJ = thetaAdjusted[3]
branches.phi_sGJ = phiFolded[3]
branches.costh_GJ1 = thetaAdjusted[4]
branches.phi_GJ1 = phiFolded[4]
branches.costh_GJ2 = thetaAdjusted[5]
branches.phi_GJ2 = phiFolded[5]
else:
branches.costh_CS = ROOT.thisCosTh[0]
branches.phi_CS = ROOT.thisPhi[0]
branches.costh_HX = ROOT.thisCosTh[1]
branches.phi_HX = ROOT.thisPhi[1]
branches.costh_PHX = ROOT.thisCosTh[2]
branches.phi_PHX = ROOT.thisPhi[2]
branches.costh_sGJ = ROOT.thisCosTh[3]
branches.phi_sGJ = ROOT.thisPhi[3]
branches.costh_GJ1 = ROOT.thisCosTh[4]
branches.phi_GJ1 = ROOT.thisPhi[4]
branches.costh_GJ2 = ROOT.thisCosTh[5]
branches.phi_GJ2 = ROOT.thisPhi[5]
branches.JpsiMass = branches.JpsiP.M()
branches.JpsiPt = branches.JpsiP.Pt()
branches.JpsiRap = branches.JpsiP.Rapidity()
outTree.Fill()
out.Write()
out.Close()
'BdToJpsiKstar892': ['/home/ltsai/Data/condor/totTree_dir_BdToJpsiKstar892.root'],
'BdToJpsiKstar1432': ['/home/ltsai/Data/condor/totTree_dir_BdToJpsiKstar1430.root'],
'BdToJpsiKpi': ['/home/ltsai/Data/condor/totTree_dir_BdToJpsiKpi.root'],
'BsToJpsiKK': ['/home/ltsai/Data/condor/totTree_dir_BsToJpsiKK.root'],
'BsToJpsiPhi': [],
'BsToJpsiF': ['/home/ltsai/Data/condor/totTree_dir_BsToJpsiF2K1525.root'],
'LbL0': ['/home/ltsai/Data/CRABdata/finalResult/totTree_dir_posLbToJpsiLam0.root'],
}
for mcName, files in mcfiles.iteritems():
chain=TChain('VertexCompCandAnalyzer/{0}'.format(branchTag))
if len(files) == 0:
print mcName + " is skipped"
continue
for mcfile in files:
chain.Add(mcfile)
mctree=chain.CloneTree(0)
mctree.SetName(mcName)
for event in chain:
nCand=event.candSize
keepEvent=False
for i in range(nCand):
if failMyPreselection(event,i): continue
keepEvent=True
if keepEvent:
mctree.Fill()
mctree.Write()
newFile.Close()
os.system('mv creating.root {0}'.format(OUTPUT_FILENAME))
def doSlimSMQCD(treeName,outputFile,inputFiles,year):
from ROOT import TChain
from ROOT import TFile
from ROOT import ROOT
#import rootlogon
ROOT.gROOT.SetBatch(1)
import sys
print ("sys.argv = ", sys.argv)
if not len(sys.argv)>=2:
raise Exception ("Must specify inputFiles as argument!")
#inputFiles = sys.argv[1].split(',')
inputFiles = inputFiles
print ("inputFiles = ", inputFiles)
#get main tree
ch = TChain(treeName)
for file in inputFiles:
ch.Add(file)
nEntries = ch.GetEntries()
#print ("nEntries = ", nEntries)
#*****set branches*****
#set branch status, at first, all off
ch.SetBranchStatus("*", 0)
from JetN2N.SlimSMQCDBranches import SlimSMQCDBranchesList
from JetN2N.SlimSMQCDBranches2011 import SlimSMQCDBranchesList2011
#event information
if year == '2011':
for branch in SlimSMQCDBranchesList2011:
print (branch)
ch.SetBranchStatus(branch,1)
else:
for branch in SlimSMQCDBranchesList:
print (branch)
ch.SetBranchStatus(branch,1)
#*****set branches end*****
#get trigger tree
chTri = TChain(treeName+"Meta/TrigConfTree")
for file in inputFiles:
chTri.Add(file)
if chTri.LoadTree(0) < 0:
chTri = 0
#get bunch tree
chBunch = TChain(treeName+"Meta/BunchConfTree")
for file in inputFiles:
chBunch.Add(file)
if chBunch.LoadTree(0) < 0:
chBunch = 0
#get cut flow tree
#chCutFlow = TChain(treeName+"Meta/CutFlowTree")
#for file in inputFiles:
# chCutFlow.Add(file)
# Write to new file
outFile = outputFile
newFile = TFile(outFile, "RECREATE")
#new tree
ch_new = ch.CloneTree(0)
# event counter
#cutName=["preD3PD","postD3PD","trigCut"]
#evnum=[0,0,0]
#weight=[0,0,0]
#event selection
for i in range(nEntries):
ch.GetEntry(i)
#evnum[1]+=1
#if hasattr(ch,"mcevt_weight") \
# and ch.mcevt_weight.size() !=0 \
# and ch.mcevt_weight[0].size() !=0:
# w=ch.mcevt_weight[0][0]
#else:
# w=1
#weight[1]+=w
ch_new.Fill()
#evnum[2]+=1
#weight[2]+=w
newFile.cd()
ch_new.Write()
#nEntriesNew = ch_new.GetEntries()
#print ("nEntriesForNewFile = ", nEntriesNew)
#check cut flow at D3PD level
#if chCutFlow.GetEntries() != 0:
# for i in range (chCutFlow.GetEntries()):
# chCutFlow.GetEntry(i)
# cutFlowName=chCutFlow.name
# for j in range(cutFlowName.size()):
# if cutName.at(j) == "AllExecutedEvents":
# evnum[0]+=chCutFlow.nAcceptedEvents.at(j)
# weight[0]+=chCutFlow.nWeightedAcceptedEvents.at(j)
#else:
# evnum[0]=evnum[1]
# weight[0]=weight[1]
#copy trigger meta data
newFile.cd()
newdir = newFile.mkdir( treeName+"Meta", treeName+"Meta" )
newdir.cd()
if chTri != 0:
chTri_new = chTri.CloneTree()
chTri_new.Write()
if chBunch != 0:
chBunch_new = chBunch.CloneTree()
chBunch_new.Write()
def runPostProd(indir, outtreedir):
#####
print "\nGet the lumi JSON file"
#####
outjson = 'allcalibtrees.json'
getLumiJSON.getLumiJSON(indir, outjson)
print 'done'
#####
print "\nGet the luminosity"
#####
parser = argparse.ArgumentParser(description='Compute luminosity of a set of samples.')
options = parser.parse_args()
options.inputfiles = [outjson]
options.local = False
options.username = '******'
alllumi = ''
with open(outjson) as f:
d = json.load(f)
d_ = {}
for k in d:
d_[int(k)] = d[k]
alllumi = compute_sample_luminosity.compute_luminosity(d, options)
with open('allcalibtrees.csv', 'w') as o:
o.write(alllumi)
headerLine = ''
brilInfo = {}
# Get Header to fill the dir
with open('allcalibtrees.csv') as f:
firstCommentedLine = True
for line in f:
if '#' not in line:
continue
if firstCommentedLine:
firstCommentedLine = False
continue
elif len(headerLine) < 1:
#run:fill,ls,time,beamstatus,E(GeV),delivered(/ub),recorded(/ub),avgpu,source
headerLine = line.strip().replace('#', '').replace(':', ',', 1).split(',')
break
# print headerLine
irun = headerLine.index('run')
ifill = headerLine.index('fill')
ils = headerLine.index('ls')
itime = headerLine.index('time')
idelivered = headerLine.index('delivered(/ub)')
irecorded = headerLine.index('recorded(/ub)')
iavgpu = headerLine.index('avgpu')
# Get the info itself
with open('allcalibtrees.csv') as f:
for line in f:
if 'Summary' in line:
break
if '#' in line:
continue
# 247243:3829,1:0,06/06/15 12:55:37,STABLE BEAMS,6500,460.468,0.000,17.2,BCM1F^M
#257490:4420,1:1,09/25/15 18:55:00,STABLE BEAMS,6500,59085.017,57661.731,0.0,PXL^M
l = line.strip().replace(':', ',', 1).split(',')
if len(l) < (max(irun, ifill, ils, itime, idelivered, irecorded, iavgpu) + 1):
continue
brilInfo[l[itime]] = {}
brilInfo[l[itime]]['run'] = int(l[irun])
brilInfo[l[itime]]['fill'] = int(l[ifill])
brilInfo[l[itime]]['ls'] = int(l[ils].split(':')[0])
brilInfo[l[itime]]['delivered'] = float(l[idelivered])
brilInfo[l[itime]]['recorded'] = float(l[irecorded])
brilInfo[l[itime]]['avgpu'] = float(l[iavgpu])
# break
# print brilInfo
print 'done'
#####
print '\nAdd the luminosity to the tree'
#####
if not os.path.exists(outtreedir):
os.makedirs(outtreedir)
outFile = outtreedir + '/output.root'
newFile = TFile(outFile, 'recreate')
for treename in ['eventtree']: #, 'clustertree']: # ["t"]
chain = TChain(treename)
chain.Add(indir + 'output_30*root')
t = chain.CloneTree(0)
nEntries = chain.GetEntries()
nSkipped = 0
print '\ttree=', treename, 'nEntries=', nEntries
delivered = array('f', [0.])
recorded = array('f', [0.])
avgpu = array('f', [0.])
t.Branch('brilcalc_delivered', delivered, 'brilcalc_delivered/F' )
t.Branch('brilcalc_recorded', recorded, 'brilcalc_recorded/F' )
t.Branch('brilcalc_avgpu', avgpu, 'brilcalc_avgpu/F' )
for i in range(chain.GetEntries()):
if (i % 1000 == 0 and i < 10000) or (i % 10000 == 0):
print '\t\t\tTreating event i= %i / %i (%.1f %%)' % (i, nEntries, float(i)/float(nEntries) * 100.)
chain.GetEntry(i)
# print chain.run
hasLumiInfo = False
for k in brilInfo:
# print brilInfo[k]['run']
# break
if (treename == 'eventtree' and brilInfo[k]['run'] == chain.run_) or (treename == 'clustertree' and brilInfo[k]['run'] == chain.run):
if (treename == 'eventtree' and brilInfo[k]['ls'] == chain.lumi_) or (treename == 'clustertree' and brilInfo[k]['run'] == chain.lumi):
hasLumiInfo = True
delivered[0] = brilInfo[k]['delivered']
recorded[0] = brilInfo[k]['recorded']
avgpu[0] = brilInfo[k]['avgpu']
# print 'voila', i, chain.run, chain.lumi
# print delivered
if not hasLumiInfo:
nSkipped += 1
# print '\tSkipping event %i: has no lumi info (run= %i, ls= %i)' % (i, chain.run, chain.lumi)
continue
t.Fill()
# break
print '\tnSkipped= %i (%.1f %%)' % (nSkipped, float(nSkipped) / nEntries * 100)
newFile.Write()
newFile.Close()
print 'done'
