def run(self):
# if FILES is a string, treat it as a single file and get the tree
if type(self.FILES) == str:
f = R.TFile.Open(self.FILES)
if not f:
raise IOError
t = f.Get(self.TREE)
if not t:
raise ReferenceError
# if FILES is a list, create a TChain and add all the files
elif type(self.FILES) == list:
t = R.TChain(self.TREE)
for f in self.FILES:
t.Add(f)
if not t:
raise ReferenceError
self.declareHistograms(self.PARAMS)
self.releaseHistograms()
self.begin(self.PARAMS)
# only turn on the branches required
Primitives.SelectBranches(t, DecList=self.BRANCHKEYS)
# either run over the whole tree,
# tree gets addressed with the for loop, or
# split the output into several files,
# and tree gets addressed with GetEntry
# either way, end at MAX if TEST, declare ETree, and analyze()
if self.TREELOOP:
if self.SPLITTING is not None:
CHUNK, JOB = self.SPLITTING
ELOW, EHIGH = JOB * CHUNK, min((JOB + 1) * CHUNK,
t.GetEntries())
for INDEX in xrange(ELOW, EHIGH):
self.INDEX = INDEX
if self.TEST:
if INDEX == self.MAX:
break
t.GetEntry(INDEX)
E = Primitives.ETree(t, self.BRANCHKEYS)
self.analyze(E, self.PARAMS)
else:
for INDEX, EVENT in enumerate(t):
self.INDEX = INDEX
if self.TEST:
if INDEX == self.MAX:
break
E = Primitives.ETree(t, self.BRANCHKEYS)
self.analyze(E, self.PARAMS)
self.end(self.PARAMS)
f.Close()
def analyze(self, E, PARAMS=None):
#selections
if self.TRIGGER:
if not Selections.passedTrigger(E): return
dsaMuons = E.getPrimitives('DSAMUON')
diMuons = E.getPrimitives('DIMUON')
Primitives.CopyExtraRecoMuonInfo(diMuons, dsaMuons)
if '4Mu' in self.NAME:
mu11, mu12, mu21, mu22, X1, X2, H, P, extramu = E.getPrimitives('GEN')
genMuons = (mu11, mu12, mu21, mu22)
genMuonPairs = ((mu11, mu12), (mu21, mu22))
elif '2Mu2J' in self.NAME:
mu1, mu2, j1, j2, X, XP, H, P, extramu = E.getPrimitives('GEN')
genMuons = (mu1, mu2)
genMuonPairs = ((mu1, mu2), )
else:
print "Haven't implemented these samples"
return
selectedDimuons = [dim for dim in diMuons if dim.Lxy() < 330]
for genMuonPair in genMuonPairs:
dimuonMatches, muonMatches, exitcode = AT.matchedDimuons(
genMuonPair, selectedDimuons)
if len(muonMatches[0]):
fillStats(self, genMuonPair[0], muonMatches[0][0]['muon'])
if len(muonMatches[1]):
fillStats(self, genMuonPair[1], muonMatches[1][0]['muon'])
return
eprint('Failed to get tree; exiting')
exit()
# containers for storing error information
KEYS = {
'ACCEPTANCE': False,
'MUON': False,
'DIMUON': False,
}
ErrorMessages = []
# test collections
for i, event in enumerate(t):
if i == 10: break
E = Primitives.ETree(t, DecList=('GEN', 'DSAMUON', 'DIMUON'))
Gens = E.getPrimitives('GEN')
DSAMuons = E.getPrimitives('DSAMUON')
Dimuons = E.getPrimitives('DIMUON')
mu11, mu12 = Gens[0], Gens[1]
try:
accSel = Selections.AcceptanceSelection(mu11)
accSel = Selections.AcceptanceSelection((mu11, mu12))
except Exception as e:
if e.message not in ErrorMessages:
ErrorMessages.append(e.message)
KEYS['ACCEPTANCE'] = True
PREFIX = Constants.PREFIX_CERN
else:
PREFIX = ''
F_NTUPLE = PREFIX + Constants.DIR_EOS + 'NTuples/ntuple_HTo2XTo4Mu_125_20_13.root'
def tprint(msg):
print '\033[32mPRIMITIVES TEST: ' + msg + '\033[m'
def eprint(msg):
print '\033[31mPRIMITIVES TEST: ' + msg + '\033[m'
f = R.TFile.Open(F_NTUPLE)
t = f.Get('SimpleNTupler/DDTree')
# test tree
try:
t.GetEntries()
print('Successfully got tree...')
except:
print('Failed to get tree; exiting')
exit()
# test collections
for i, event in enumerate(t):
if i == 3: break
E = Primitives.ETree(t)
print E
def analyze(self, E, PARAMS=None):
if self.TRIGGER and self.SP is not None:
if not Selections.passedTrigger(E): return
Event = E.getPrimitives('EVENT')
DSAmuons = E.getPrimitives('DSAMUON')
Dimuons = E.getPrimitives('DIMUON')
Primitives.CopyExtraRecoMuonInfo(Dimuons, DSAmuons)
eventWeight = 1.
try:
eventWeight = 1. if Event.weight > 0. else -1.
except:
pass
# decide what set of cuts to apply based on self.CUTS cut string
ALL = 'All' in self.CUTS
PROMPT = '_Prompt' in self.CUTS
NOPROMPT = '_NoPrompt' in self.CUTS
NSTATIONS = '_NS' in self.CUTS
NMUONHITS = '_NH' in self.CUTS
FPTERR = '_FPTE' in self.CUTS
PT = '_PT' in self.CUTS
HLT = '_HLT' in self.CUTS
PC = '_PC' in self.CUTS
LXYERR = '_LXYE' in self.CUTS
MASS = '_M' in self.CUTS
def boolsToMuonCutList(NSTATIONS, NMUONHITS, FPTERR, PT):
cutList = []
if NSTATIONS:
cutList.append('b_nStations')
if NMUONHITS:
cutList.append('b_nMuonHits')
if FPTERR:
cutList.append('b_FPTE')
if PT:
cutList.append('b_pT')
return cutList
def boolsToDimuonCutList(LXYERR, MASS):
cutList = []
if LXYERR:
cutList.append('b_LxyErr')
if MASS:
cutList.append('b_mass')
return cutList
# require muons to pass all selections
if ALL:
DSASelections = [Selections.MuonSelection(muon) for muon in DSAmuons]
selectedDSAmuons = [
mu for idx, mu in enumerate(DSAmuons) if DSASelections[idx]
]
selectedDimuons = Dimuons
# don't require reco muons to pass all selections
else:
selectedDSAmuons = DSAmuons
selectedDimuons = Dimuons
# for PROMPT and NOPROMPT event selections
if PROMPT or NOPROMPT:
highLxySigExists = False
for dimuon in Dimuons:
if dimuon.LxySig() > 3.:
highLxySigExists = True
break
# return if there are LxySig > 3
if PROMPT:
if highLxySigExists:
return
# return if there are NO LxySig > 3 -- that's category 1
elif NOPROMPT:
if not highLxySigExists:
return
if PROMPT or NOPROMPT:
# compute all the baseline selection booleans
DSASelections = [
Selections.MuonSelection(muon, cutList='BaselineMuonCutList')
for muon in DSAmuons
]
# figure out which cuts we actually care about
cutList = boolsToMuonCutList(NSTATIONS, NMUONHITS, FPTERR, PT)
# no selection
if len(cutList) == 0:
selectedDSAmuons = DSAmuons
selectedDimuons = Dimuons
# cutList is some nonzero list, meaning keep only the muons that pass the cut keys in cutList
else:
selectedDSAmuons = [
mu for i, mu in enumerate(DSAmuons)
if DSASelections[i].allOf(*cutList)
]
selectedOIndices = [mu.idx for mu in selectedDSAmuons]
selectedDimuons = [
dim for dim in Dimuons if dim.idx1 in selectedOIndices
and dim.idx2 in selectedOIndices
]
# apply HLT RECO matching
if HLT:
HLTPaths, HLTMuons, L1TMuons = E.getPrimitives('TRIGGER')
DSAMuonsForHLTMatch = [
mu for mu in selectedDSAmuons if abs(mu.eta) < 2.
]
HLTMuonMatches = matchedTrigger(HLTMuons, DSAMuonsForHLTMatch)
if not any([HLTMuonMatches[ij]['matchFound']
for ij in HLTMuonMatches]):
return
# apply pairing criteria and transform selectedDimuons
if PC:
selectedDimuons = applyPairingCriteria(selectedDSAmuons,
selectedDimuons)
if PROMPT or NOPROMPT:
# compute all the baseline selection booleans
DimuonSelections = {
dim.ID: Selections.DimuonSelection(dim,
cutList='BaselineDimuonCutList')
for dim in selectedDimuons
}
# figure out which cuts we actually care about
cutList = boolsToDimuonCutList(LXYERR, MASS)
# cutList is some nonzero list, meaning keep only the muons that pass the cut keys in cutList
if len(cutList) > 0:
selectedDimuons = [
dim for dim in selectedDimuons
if DimuonSelections[dim.ID].allOf(*cutList)
]
# for the MC/Data events, skip events with no dimuons, but not for "no selection"
if (PROMPT or NOPROMPT) and NSTATIONS:
if len(selectedDimuons) == 0: return
# also filter selectedDSAmuons to only be of those indices that are in the final dimuons
if PROMPT or NOPROMPT:
selectedOIndices = []
for dim in selectedDimuons:
selectedOIndices.append(dim.idx1)
selectedOIndices.append(dim.idx2)
selectedOIndices = list(set(selectedOIndices))
selectedDSAmuons = [
mu for mu in selectedDSAmuons if mu.idx in selectedOIndices
]
# all refitted muons
allRefittedMuons = []
for dimuon in selectedDimuons:
allRefittedMuons.append(dimuon.mu1)
allRefittedMuons.append(dimuon.mu2)
# fill histograms for every reco muon
for MUON, recoMuons in (('DSA', selectedDSAmuons), ('REF',
allRefittedMuons)):
self.HISTS[MUON + '_nMuon'].Fill(len(recoMuons), eventWeight)
for muon in recoMuons:
for KEY in CONFIG:
self.HISTS[MUON + '_' + KEY].Fill(CONFIG[KEY]['LAMBDA'](muon),
eventWeight)
for KEY in EXTRACONFIG:
F1 = EXTRACONFIG[KEY]['LAMBDA'][0]
F2 = EXTRACONFIG[KEY]['LAMBDA'][1]
self.HISTS[MUON + '_' + KEY].Fill(F1(muon), F2(muon),
eventWeight)
# get gen particles if this is a signal sample
if self.SP is not None:
if '4Mu' in self.NAME:
mu11, mu12, mu21, mu22, X1, X2, H, P, extramu = E.getPrimitives(
'GEN')
genMuons = (mu11, mu12, mu21, mu22)
genMuonPairs = ((mu11, mu12), (mu21, mu22))
elif '2Mu2J' in self.NAME:
mu1, mu2, j1, j2, X, XP, H, P, extramu = E.getPrimitives('GEN')
genMuons = (mu1, mu2)
genMuonPairs = ((mu1, mu2), )
MuonMatches = {'DSA': [], 'REF': []}
# get matched reco muons
for genMuon in genMuons:
# cut genMuons outside the detector acceptance
# don't do it for now
#genMuonSelection = Selections.AcceptanceSelection(genMuon)
for MUON, recoMuons in (('DSA', selectedDSAmuons), ):
MuonMatches[MUON].append(
matchedMuons(genMuon, recoMuons, vertex='BS'))
# and for refitted muons for matched dimuons
for genMuonPair in genMuonPairs:
for MUON in ('REF', ):
dimuonMatches, muonMatches, exitcode = matchedDimuons(
genMuonPair, selectedDimuons)
MuonMatches[MUON].append(muonMatches[0])
MuonMatches[MUON].append(muonMatches[1])
# fill histograms
# for each major muon type,
# MuonMatches contains 2 lists of matches corresponding to each gen muon
# Each match in each of those 2 lists is a list of individual muon matches
for MUON in ('DSA', 'REF'):
for matches in MuonMatches[MUON]:
for match in matches:
muon = match['muon']
deltaR = match['deltaR']
for KEY in CONFIG:
self.HISTS[MUON + '_' + KEY + '_Matched'].Fill(
CONFIG[KEY]['LAMBDA'](muon), eventWeight)
for KEY in EXTRACONFIG:
F1 = EXTRACONFIG[KEY]['LAMBDA'][0]
F2 = EXTRACONFIG[KEY]['LAMBDA'][1]
self.HISTS[MUON + '_' + KEY + '_Matched'].Fill(
F1(muon), F2(muon), eventWeight)
self.HISTS[MUON + '_deltaRGR_Matched'].Fill(
deltaR, eventWeight)
self.HISTS[MUON + '_nMuon_Matched'].Fill(
len(matches), eventWeight)
if len(matches) > 0:
self.HISTS[MUON + '_deltaRGR_Closest'].Fill(
matches[0]['deltaR'], eventWeight)