def __init__(self, *args, **kwargs):
super(BKJPsiEEAnalyzer, self).__init__(*args, **kwargs)
# stuff I need to instantiate only once
self.vtxfit = VertexFitter()
# create a std::vector<reco::Track> to be passed to the fitter
self.tofit_el = ROOT.std.vector('reco::Track')()
# create a std::vector<pat::PackedCandidate> to be passed to the fitter
self.tofit_pc = ROOT.std.vector('pat::PackedCandidate')()
def beginLoop(self, setup):
super(BKJPsiAnalyzer, self).beginLoop(setup)
self.counters.addCounter('BKJPsiAnalyzer')
count = self.counters.counter('BKJPsiAnalyzer')
count.register('all events')
count.register('>= 2 muons')
count.register('dimuon mass < 6')
count.register('good dimuon vtx')
count.register('B(KLL) mass < 6')
# stuff I need to instantiate only once
self.vtxfit = VertexFitter()
# create a std::vector<reco::RecoChargedCandidate> to be passed to the fitter
self.tofit_cc = ROOT.std.vector('reco::RecoChargedCandidate')()
# create a std::vector<pat::PackedCandidate> to be passed to the fitter
self.tofit_pc = ROOT.std.vector('pat::PackedCandidate')()
def __init__(self, tk1, tk2, requireVtx=False):
self.tk1_ = tk1
self.tk2_ = tk2
self.arbitrate()
self.vtx = None
self.vtxprob = -99.
if requireVtx:
# stuff I need to instantiate only once
self.vtxfit = VertexFitter()
# create a std::vector<reco::Track> to be passed to the fitter
self.tofit = ROOT.std.vector('reco::Track')()
# fill the vector
self.tofit.push_back(self.tk1_.physObj)
self.tofit.push_back(self.tk2_.physObj)
# fit the vertex and save the information
self.makeVtx_()
class BKJPsiEEAnalyzer(Analyzer):
'''
'''
def __init__(self, *args, **kwargs):
super(BKJPsiEEAnalyzer, self).__init__(*args, **kwargs)
# stuff I need to instantiate only once
self.vtxfit = VertexFitter()
# create a std::vector<reco::Track> to be passed to the fitter
self.tofit_el = ROOT.std.vector('reco::Track')()
# create a std::vector<pat::PackedCandidate> to be passed to the fitter
self.tofit_pc = ROOT.std.vector('pat::PackedCandidate')()
def declareHandles(self):
super(BKJPsiEEAnalyzer, self).declareHandles()
# miniAOD collections
self.handles['electrons' ] = AutoHandle('slimmedElectrons' , 'std::vector<pat::Electron>' )
self.handles['muons' ] = AutoHandle('slimmedMuons' , 'std::vector<pat::Muon>' )
self.handles['losttracks'] = AutoHandle('lostTracks' , 'std::vector<pat::PackedCandidate>')
self.handles['pfcands' ] = AutoHandle('packedPFCandidates' , 'std::vector<pat::PackedCandidate>')
self.handles['pvs' ] = AutoHandle('offlineSlimmedPrimaryVertices', 'std::vector<reco::Vertex>' )
# enriched miniAOD collections
self.handles['eletrackMap'] = AutoHandle(('ttk', 'eleTtkMap', 'TTK'), 'std::vector<pair<edm::Ptr<pat::Electron>,reco::Track> >')
# AOD collections
self.handles['beamspot' ] = AutoHandle('offlineBeamSpot', 'reco::BeamSpot')
def beginLoop(self, setup):
super(BKJPsiEEAnalyzer, self).beginLoop(setup)
self.counters.addCounter('BKJPsiEEAnalyzer')
count = self.counters.counter('BKJPsiEEAnalyzer')
count.register('all events')
count.register('>= 2 electrons')
count.register('diele mass < 6')
count.register('diele dz < 1')
count.register('good diele vtx')
count.register('B(KLL) mass < 6')
def process(self, event):
self.readCollections(event.input)
self.counters.counter('BKJPsiEEAnalyzer').inc('all events')
# vertex stuff
event.pvs = self.handles['pvs' ].product()
event.beamspot = self.handles['beamspot'].product()
# get the tracks
allpf = map(PhysicsObject, self.handles['pfcands' ].product())
losttracks = map(PhysicsObject, self.handles['losttracks'].product())
# merge the track collections
event.alltracks = sorted([tt for tt in allpf + losttracks if tt.charge() != 0 and abs(tt.pdgId()) not in (11,13)], key = lambda x : x.pt(), reverse = True)
# get the offline electrons and muons
event.electrons = map(Electron, self.handles['electrons'].product())
event.muons = map(Muon , self.handles['muons' ].product())
# get the electron-track map BEFORE selections
eletrks = self.handles['eletrackMap'].product()
# check consistency between electron-track map and electrons
if len(event.electrons) != len(eletrks):
print 'run %d \tlumi %d \tevent %d' %(event.run, event.lumi, event.eventId)
print 'different number of electrons and ele tracks, %d and %d respectively' %(len(event.electrons), len(eletrks))
for ii in event.electrons:
print ii, ii.gsfTrack().pt()
for ii in eletrks:
print ii.first.pt(), ii.first.eta(), ii.first.phi(), ii.second.pt()
set_trace()
for jj in range(len(event.electrons)):
event.electrons[jj].trackFromGsfTrack = eletrks[jj].second
event.electrons[jj].ptr = eletrks[jj].first
if event.electrons[jj].gsfTrack().pt() != eletrks[jj].second.pt():
set_trace()
# preselect electrons
event.electrons = [ele for ele in event.electrons if self.testEle(ele)]
# at least two electrons
if len(event.electrons)<2: return False
self.counters.counter('BKJPsiEEAnalyzer').inc('>= 2 electrons')
# build all di-ele pairs
dieles = [(ele1, ele2) for ele1, ele2 in combinations(event.electrons, 2)]
# opposite sign di-muon
dieles = [(ele1, ele2) for ele1, ele2 in dieles if ele1.charge() != ele2.charge()]
# SAVE THEM ALL, ALSO NON RESONANT!
# invariant mass http://cmslxr.fnal.gov/source/DataFormats/EgammaCandidates/interface/GsfElectron.h#0795
# dieles = [(ele1, ele2) for ele1, ele2 in dieles if (ele1.p4(1) + ele2.p4(1)).mass() < 6.]
dieles = [(ele1, ele2) for ele1, ele2 in dieles if (ele1.p4() + ele2.p4()).mass() < 6.]
if not len(dieles): return False
self.counters.counter('BKJPsiEEAnalyzer').inc('diele mass < 6')
# select by dz
dieles = [(ele1, ele2) for ele1, ele2 in dieles if abs(ele1.vz() - ele2.vz()) < 1.]
if not len(dieles): return False
self.counters.counter('BKJPsiEEAnalyzer').inc('diele dz < 1')
selDieles = []
# dieles with a good vertex
for diele in dieles:
# clear the vectors
self.tofit_el.clear()
self.tofit_pc.clear()
# create the vector of bla bla FIXME!
for il in diele:
# set_trace()
self.tofit_el.push_back(il.trackFromGsfTrack)
# fit it!
# set_trace()
svtree = self.vtxfit.Fit(self.tofit_el, self.tofit_pc) # actual vertex fitting
# check that the vertex is good
if not svtree.get().isEmpty() and svtree.get().isValid():
svtree.movePointerToTheTop()
sv = svtree.currentDecayVertex().get()
recoSv = makeRecoVertex(sv, kinVtxTrkSize=2) # need to do some gymastics
selDieles.append(diele)
if len(selDieles)==0: return False
self.counters.counter('BKJPsiEEAnalyzer').inc('good diele vtx')
cands = []
# set_trace()
# add a track
for diele, tk in product(selDieles, event.alltracks):
# try to fit a vertex only if the track's close to the di-ele cand
if max([abs(ee.vz() - tk.vz()) for ee in diele]) > 1.5:
continue
# remove double countings
if deltaR(tk, diele[0])<0.01 or deltaR(tk, diele[1])<0.01:
continue
# pt and sanity checks on the track
# set_trace()
if not tk.bestTrack():
continue
if tk.pt()<0.8:
continue
# kaon mass
tk.setMass(0.493677)
# clear the vectors
self.tofit_el.clear()
self.tofit_pc.clear()
# create a RecoChargedCandidate for each reconstructed lepton and flush it into the vector
for il in diele:
self.tofit_el.push_back(il.trackFromGsfTrack)
# push the track into the vector
self.tofit_pc.push_back(tk.physObj)
# fit it!
svtree = self.vtxfit.Fit(self.tofit_el, self.tofit_pc) # actual vertex fitting
# check that the vertex is good
if not svtree.get().isEmpty() and svtree.get().isValid():
svtree.movePointerToTheTop()
sv = svtree.currentDecayVertex().get()
recoSv = makeRecoVertex(sv, kinVtxTrkSize=3) # need to do some gymastics
cands.append(BKLL(diele[0], diele[1], tk, recoSv, event.beamspot))
if not len(cands): return False
self.counters.counter('BKJPsiEEAnalyzer').inc('good diele vtx')
cands = [cand for cand in cands if cand.b().mass()>4 and cand.b().mass()<6]
if not len(cands): return False
self.counters.counter('BKJPsiEEAnalyzer').inc('B(KLL) mass < 6')
# set_trace()
# for jj in cands: print jj.b().mass(), jj.b().pt(), jj.b().eta(), jj.b().phi(), jj.ll().mass(), jj.ls2d(), jj.vtxprob(), jj.llcone(), jj.bcone()
event.myB = sorted(cands, key = lambda x : x.vtxprob(), reverse=True)[0]
return True
def testEle(self, ele):
return ele.pt()>2. and \
abs(ele.eta())<2.5
class Kst(object):
'''
Builds a K*(892) candidate out of two tracks.
Assigns \pi and K mass hypothesis, find the assignment that returns the
invariant mass closer to the nominal K*(892) mass.
As an option, it can be required that the two tracks make a vertex.
'''
def __init__(self, tk1, tk2, requireVtx=False):
self.tk1_ = tk1
self.tk2_ = tk2
self.arbitrate()
self.vtx = None
self.vtxprob = -99.
if requireVtx:
# stuff I need to instantiate only once
self.vtxfit = VertexFitter()
# create a std::vector<reco::Track> to be passed to the fitter
self.tofit = ROOT.std.vector('reco::Track')()
# fill the vector
self.tofit.push_back(self.tk1_.physObj)
self.tofit.push_back(self.tk2_.physObj)
# fit the vertex and save the information
self.makeVtx_()
def makeVtx_(self):
# fit it!
svtree = self.vtxfit.Fit(self.tofit) # actual vertex fitting
# check that the vertex is good
if not svtree.get().isEmpty() and svtree.get().isValid():
svtree.movePointerToTheTop()
sv = svtree.currentDecayVertex().get()
recoSv = makeRecoVertex(
sv, kinVtxTrkSize=2) # need to do some gymastics
self.vtx = recoSv
self.vtxprob = ROOT.TMath.Prob(self.vtx.chi2(), int(self.vtx().ndof()))
def arbitrate(self):
pi1 = ROOT.Math.LorentzVector('<ROOT::Math::PxPyPzM4D<double> >')(
self.tk1_.px(), self.tk1_.py(), self.tk1_.pz(), m_pi_ch)
pi2 = ROOT.Math.LorentzVector('<ROOT::Math::PxPyPzM4D<double> >')(
self.tk2_.px(), self.tk2_.py(), self.tk2_.pz(), m_pi_ch)
k1 = ROOT.Math.LorentzVector('<ROOT::Math::PxPyPzM4D<double> >')(
self.tk1_.px(), self.tk1_.py(), self.tk1_.pz(), m_k_ch)
k2 = ROOT.Math.LorentzVector('<ROOT::Math::PxPyPzM4D<double> >')(
self.tk2_.px(), self.tk2_.py(), self.tk2_.pz(), m_k_ch)
# assign the mass hypotheses
if abs((pi1 + k2).mass() - m_k_st_zero) < abs((pi2 + k1).mass() -
m_k_st_zero):
self.pi_ = PhysicsObject(
ROOT.pat.PackedCandidate(self.tk1_.physObj))
self.k_ = PhysicsObject(ROOT.pat.PackedCandidate(
self.tk2_.physObj))
else:
self.pi_ = PhysicsObject(
ROOT.pat.PackedCandidate(self.tk2_.physObj))
self.k_ = PhysicsObject(ROOT.pat.PackedCandidate(
self.tk1_.physObj))
self.pi_.setMass(m_pi_ch)
self.k_.setMass(m_k_ch)
self.pi_.setPdgId(self.pi_.charge() * 211)
self.k_.setPdgId(self.k_.charge() * 321)
def charge(self):
return self.tk1_.charge() + self.tk2_.charge()
def pi(self):
return self.pi_
def k(self):
return self.k_
def p4(self):
return self.pi().p4() + self.k().p4()
def mass(self):
return self.p4().mass()
def eta(self):
return self.p4().eta()
def phi(self):
return self.p4().phi()
def e(self):
return self.p4().e()
def pt(self):
return self.p4().pt()
def p(self):
return math.sqrt(self.p4().px()**2 + self.p4().py()**2 +
self.p4().pz()**2)
def px(self):
return self.p4().px()
def py(self):
return self.p4().py()
def pz(self):
return self.p4().pz()
def pdgId():
# FIXME! get the correct pdgId based on the kaon an pion charges
return 313 if self.k().charge() > 0 else -313
class BKJPsiAnalyzer(Analyzer):
'''
'''
def declareHandles(self):
super(BKJPsiAnalyzer, self).declareHandles()
self.handles['electrons'] = AutoHandle('slimmedElectrons',
'std::vector<pat::Electron>')
self.handles['muons'] = AutoHandle('slimmedMuons',
'std::vector<pat::Muon>')
self.handles['losttracks'] = AutoHandle(
'lostTracks', 'std::vector<pat::PackedCandidate>')
self.handles['pfcands'] = AutoHandle(
'packedPFCandidates', 'std::vector<pat::PackedCandidate>')
self.handles['pvs'] = AutoHandle(
('offlineSlimmedPrimaryVertices', '', 'PAT'),
'std::vector<reco::Vertex>')
self.handles['beamspot'] = AutoHandle(('offlineBeamSpot', '', 'RECO'),
'reco::BeamSpot')
def beginLoop(self, setup):
super(BKJPsiAnalyzer, self).beginLoop(setup)
self.counters.addCounter('BKJPsiAnalyzer')
count = self.counters.counter('BKJPsiAnalyzer')
count.register('all events')
count.register('>= 2 muons')
count.register('dimuon mass < 6')
count.register('good dimuon vtx')
count.register('B(KLL) mass < 6')
# stuff I need to instantiate only once
self.vtxfit = VertexFitter()
# create a std::vector<reco::RecoChargedCandidate> to be passed to the fitter
self.tofit_cc = ROOT.std.vector('reco::RecoChargedCandidate')()
# create a std::vector<pat::PackedCandidate> to be passed to the fitter
self.tofit_pc = ROOT.std.vector('pat::PackedCandidate')()
def process(self, event):
self.readCollections(event.input)
self.counters.counter('BKJPsiAnalyzer').inc('all events')
# vertex stuff
event.pvs = self.handles['pvs'].product()
event.beamspot = self.handles['beamspot'].product()
# get the tracks
allpf = map(PhysicsObject, self.handles['pfcands'].product())
losttracks = map(PhysicsObject, self.handles['losttracks'].product())
# merge the track collections
event.alltracks = sorted([
tt for tt in allpf + losttracks
if tt.charge() != 0 and abs(tt.pdgId()) not in (11, 13)
],
key=lambda x: x.pt(),
reverse=True)
# get the offline electrons and muons
event.electrons = map(Electron, self.handles['electrons'].product())
event.muons = map(Muon, self.handles['muons'].product())
# preselect muons
event.muons = [muon for muon in event.muons if self.testMuon(muon)]
if len(event.muons) < 2: return False
self.counters.counter('BKJPsiAnalyzer').inc('>= 2 muons')
# build all di-muon pairs
dimuons = [(mu1, mu2) for mu1, mu2 in combinations(event.muons, 2)]
# opposite sign di-muon
dimuons = [(mu1, mu2) for mu1, mu2 in dimuons
if mu1.charge() != mu2.charge()]
# SAVE THEM ALL, ALSO NON RESONANT!
# invariant mass
dimuons = [(mu1, mu2) for mu1, mu2 in dimuons
if (mu1.p4() + mu2.p4()).mass() < 6.]
if not len(dimuons): return False
self.counters.counter('BKJPsiAnalyzer').inc('dimuon mass < 6')
selDimuons = []
# dimuons with a good vertex
for dimuon in dimuons:
# clear the vectors
self.tofit_cc.clear()
self.tofit_pc.clear()
# create a RecoChargedCandidate for each reconstructed lepton and flush it into the vector
for il in dimuon:
# if the reco particle is a displaced thing, it does not have the p4() method, so let's build it
myp4 = ROOT.Math.LorentzVector(
'<ROOT::Math::PxPyPzE4D<double> >')(
il.px(), il.py(), il.pz(),
math.sqrt(il.mass()**2 + il.px()**2 + il.py()**2 +
il.pz()**2))
ic = ROOT.reco.RecoChargedCandidate(
) # instantiate a dummy RecoChargedCandidate
ic.setCharge(il.charge()) # assign the correct charge
ic.setP4(myp4) # assign the correct p4
ic.setTrack(il.track()) # set the correct TrackRef
if ic.track().isNonnull(): # check that the track is valid
self.tofit_cc.push_back(ic)
# further sanity check: two *distinct* tracks
if self.tofit_cc.size(
) == 2 and self.tofit_cc[0].track() != self.tofit_cc[1].track():
# fit it!
svtree = self.vtxfit.Fit(
self.tofit_cc, self.tofit_pc) # actual vertex fitting
# check that the vertex is good
if not svtree.get().isEmpty() and svtree.get().isValid():
svtree.movePointerToTheTop()
sv = svtree.currentDecayVertex().get()
recoSv = makeRecoVertex(
sv, kinVtxTrkSize=2) # need to do some gymastics
selDimuons.append(dimuon)
if len(selDimuons) == 0: return False
self.counters.counter('BKJPsiAnalyzer').inc('good dimuon vtx')
cands = []
# add a track
for dimu, tk in product(selDimuons, event.alltracks):
# remove double countings
if deltaR(tk, dimu[0]) < 0.01 or deltaR(tk, dimu[1]) < 0.01:
continue
# pt and sanity checks on the track
# set_trace()
if not tk.bestTrack():
continue
if tk.pt() < 0.8:
continue
# clear the vectors
self.tofit_cc.clear()
self.tofit_pc.clear()
# create a RecoChargedCandidate for each reconstructed lepton and flush it into the vector
for il in dimu:
# if the reco particle is a displaced thing, it does not have the p4() method, so let's build it
myp4 = ROOT.Math.LorentzVector(
'<ROOT::Math::PxPyPzE4D<double> >')(
il.px(), il.py(), il.pz(),
math.sqrt(il.mass()**2 + il.px()**2 + il.py()**2 +
il.pz()**2))
ic = ROOT.reco.RecoChargedCandidate(
) # instantiate a dummy RecoChargedCandidate
ic.setCharge(il.charge()) # assign the correct charge
ic.setP4(myp4) # assign the correct p4
ic.setTrack(il.track()) # set the correct TrackRef
if ic.track().isNonnull(): # check that the track is valid
self.tofit_cc.push_back(ic)
set_trace()
m_k = 0.493677
# push the track into the vector
self.tofit_pc.push_back(tk.physObj)
# fit it!
try:
svtree = self.vtxfit.Fit(
self.tofit_cc, self.tofit_pc) # actual vertex fitting
except:
set_trace()
# check that the vertex is good
if not svtree.get().isEmpty() and svtree.get().isValid():
svtree.movePointerToTheTop()
sv = svtree.currentDecayVertex().get()
recoSv = makeRecoVertex(
sv, kinVtxTrkSize=3) # need to do some gymastics
cands.append(BKLL(dimu[0], dimu[1], tk, recoSv,
event.beamspot))
if not len(cands): return False
self.counters.counter('BKJPsiAnalyzer').inc('good dimuon vtx')
cands = [cand for cand in cands if cand.b().mass() < 6]
if not len(cands): return False
self.counters.counter('BKJPsiAnalyzer').inc('B(KLL) mass < 6')
# set_trace()
# for jj in cands: print jj.b().mass(), jj.b().pt(), jj.b().eta(), jj.b().phi(), jj.ll().mass(), jj.ls2d(), jj.vtxprob(), jj.llcone(), jj.bcone()
event.myB = sorted(cands, key=lambda x: x.vtxprob(), reverse=True)[0]
return True
def testMuon(self, muon):
return muon.pt()>1. and \
abs(muon.eta())<2.5 and \
muon.isGlobalMuon()