def costheta_CS(pt1, eta1, phi1, pt2, eta2, phi2, l1id):
mu1 = TLorentzVector()
mu2 = TLorentzVector()
Q = TLorentzVector()
if (l1id < 0):
mu1.SetPtEtaPhiM(pt1, eta1, phi1, 0)
mu2.SetPtEtaPhiM(pt2, eta2, phi2, 0)
else:
mu1.SetPtEtaPhiM(pt2, eta2, phi2, 0)
mu2.SetPtEtaPhiM(pt1, eta1, phi1, 0)
Q = mu1 + mu2
mu1plus = ((2.0)**(-0.5)) * (mu1.E() + mu1.Pz())
mu1minus = ((2.0)**(-0.5)) * (mu1.E() - mu1.Pz())
mu2plus = ((2.0)**(-0.5)) * (mu2.E() + mu2.Pz())
mu2minus = ((2.0)**(-0.5)) * (mu2.E() - mu2.Pz())
costheta = ((2.0 / Q.Mag()) /
((Q.Mag()**2.0 + Q.Pt()**2.0)**(0.5))) * (mu1plus * mu2minus -
mu1minus * mu2plus)
return costheta
def insideLoop(self,t):
if not self.isData:
for i in range(t.nLHEPdfWeight):
self.obj['SumOfPDFweights'].Fill(i+1, t.LHEPdfWeight[i])
for i in range(t.nLHEScaleWeight):
self.obj['SumOfScaleWeights'].Fill(i+1, t.LHEScaleWeight[i])
# Three leptons
looseLeptons = []
mll = []
for imu in range(t.nMuon):
if t.Muon_pt[imu] >= 10:
l = TLorentzVector()
l.SetPtEtaPhiM(t.Muon_pt[imu], t.Muon_eta[imu], t.Muon_phi[imu], t.Muon_mass[imu])
looseLeptons.append(l)
for iel in range(t.nElectron):
if t.Electron_pt >= 10:
l = TLorentzVector()
l.SetPtEtaPhiM(t.Electron_pt[iel], t.Electron_eta[iel], t.Electron_phi[iel], t.Electron_mass[iel])
looseLeptons.append(l)
if len(looseLeptons) < 2: return
mll = (looseLeptons[0]+looseLeptons[1]).M()
#print mll
if mll<12: return
self.obj['EventTree'].Fill()
def getGenVpt(event):
GenVpt = 0.
VPt = -1.
idx_V = -1
LepP = TLorentzVector()
LepM = TLorentzVector()
particle_masses = {11 : 0.000511,
12 : 0.,
13 : 0.10566,
14 : 0.,
15 : 1.77686,
16 : 0.}
for idx_gen in range(event.nGenPart):
if event.GenPart_pdgId[idx_gen] == 23 or event.GenPart_pdgId[idx_gen] == 24:
if VPt<=0.:
VPt = event.GenPart_pt[idx_gen]
idx_V = idx_gen
else:
if event.GenPart_genPartIdxMother[idx_gen]==idx_V:
VPt = event.GenPart_pt[idx_gen]
idx_V = idx_gen
elif (event.GenPart_status[idx_gen]==1 and event.GenPart_pdgId[idx_gen] >= +11 and event.GenPart_pdgId[idx_gen] <= +16) or (event.GenPart_pdgId[idx_gen]== +15 and event.GenPart_status[idx_gen]==2):
if event.GenPart_pt[idx_gen]>LepP.Pt():
particle_mass = particle_masses[abs(event.GenPart_pdgId[idx_gen])]
LepP.SetPtEtaPhiM(event.GenPart_pt[idx_gen],event.GenPart_eta[idx_gen],event.GenPart_phi[idx_gen],particle_mass)
elif (event.GenPart_status[idx_gen]==1 and event.GenPart_pdgId[idx_gen] >= -16 and event.GenPart_pdgId[idx_gen] <= -11) or (event.GenPart_pdgId[idx_gen]== -15 and event.GenPart_status[idx_gen]==2):
if event.GenPart_pt[idx_gen]>LepM.Pt():
particle_mass = particle_masses[abs(event.GenPart_pdgId[idx_gen])]
LepM.SetPtEtaPhiM(event.GenPart_pt[idx_gen],event.GenPart_eta[idx_gen],event.GenPart_phi[idx_gen],particle_mass)
if VPt > 0.:
GenVpt = VPt
elif LepP.Pt() > 0. and LepM.Pt() > 0.:
GenVpt = (LepP+LepM).Pt()
return GenVpt
def test_deltaR2(self):
"""Test that the deltaR2 method properly uses either eta or
theta depending on the collider configuration
"""
Collider.BEAMS = 'pp'
tlv1 = TLorentzVector()
tlv1.SetPtEtaPhiM(10, 1.1, 0, 0)
tlv2 = TLorentzVector()
tlv2.SetPtEtaPhiM(10, 1.2, 0, 0)
ptc1 = Particle(1, 1, tlv1)
ptc2 = Particle(1, 1, tlv2)
dR = math.sqrt( deltaR2(ptc1, ptc2))
self.assertAlmostEqual(dR, 0.1)
Collider.BEAMS = 'ee'
tlv1 = TLorentzVector()
tlv1.SetPtEtaPhiM(10, 1.1, 0, 0)
tlv1.SetTheta(1.1)
tlv2 = TLorentzVector()
tlv2.SetPtEtaPhiM(10, 1.2, 0, 0)
tlv2.SetTheta(1.2)
ptc1 = Particle(1, 1, tlv1)
ptc2 = Particle(1, 1, tlv2)
dR = math.sqrt( deltaR2(ptc1, ptc2))
self.assertAlmostEqual(dR, 0.1)
def test_iso1(self):
p4 = TLorentzVector()
p4.SetPtEtaPhiM(10, 0, 0, 0.105)
lepton = Particle(13, 1, p4)
p4 = TLorentzVector()
p4.SetPtEtaPhiM(1, 0, 0, 0.105)
ptc = Particle(211, 1, p4)
# test iso calc
computer = IsolationComputer([EtaPhiCircle(0.4)])
iso = computer.compute(lepton, [ptc,ptc])
self.assertEqual(iso.sumpt, 2*ptc.pt())
self.assertEqual(iso.sume, 2*ptc.e())
self.assertEqual(iso.num, 2)
# test IsolationInfo addition
iso2 = copy.copy(iso)
iso2 += iso
self.assertEqual(iso2.sumpt, 4*ptc.pt())
self.assertEqual(iso2.sume, 4*ptc.e())
self.assertEqual(iso2.num, 4)
# test veto
computer = IsolationComputer([EtaPhiCircle(0.4)], [EtaPhiCircle(0.1)])
iso = computer.compute(lepton, [ptc])
self.assertEqual(iso.sumpt, 0.)
def computeProcessedFeatures(muon1_p4, muon2_p4, unpairedMuon_p4, jpsi_p4, bc_p4):
if ( bc_p4.M() != 0):
bcPtCorrected = (bcPdgMass * bc_p4.Pt())/ bc_p4.M()
else:
bcPtCorrected = bc_p4.Pt()
muonsSystem_p4 = muon1_p4 + muon2_p4 + unpairedMuon_p4
bcCorrected_p4 = TLorentzVector()
bcCorrected_p4.SetPtEtaPhiM(bcPtCorrected,
muonsSystem_p4.Eta(),
muonsSystem_p4.Phi(),
#bc_p4.M())
bcPdgMass)
boostToBcCorrectedRestFrame = -bcCorrected_p4.BoostVector()
#boostToJpsiRestFrame = -(muon1_p4+muon2_p4).BoostVector()
boostToJpsiRestFrame = -jpsi_p4.BoostVector()
unpairedMuonBoostedToBcCorrectedRestFrame_p4 = TLorentzVector()
unpairedMuonBoostedToBcCorrectedRestFrame_p4.SetPtEtaPhiM(unpairedMuon_p4.Pt(), unpairedMuon_p4.Eta(), unpairedMuon_p4.Phi(), unpairedMuon_p4.M())
unpairedMuonBoostedToBcCorrectedRestFrame_p4.Boost(boostToBcCorrectedRestFrame)
unpairedMuonBoostedToJpsiRestFrame_p4 = TLorentzVector()
unpairedMuonBoostedToJpsiRestFrame_p4.SetPtEtaPhiM(unpairedMuon_p4.Pt(), unpairedMuon_p4.Eta(), unpairedMuon_p4.Phi(), unpairedMuon_p4.M())
unpairedMuonBoostedToJpsiRestFrame_p4.Boost(boostToJpsiRestFrame)
nn_energyBcRestFrame = unpairedMuonBoostedToBcCorrectedRestFrame_p4.E()
#nn_missMass2 = (bcCorrected_p4 - muon1_p4 - muon2_p4 - unpairedMuon_p4).M2()
nn_missMass2 = (bcCorrected_p4 - jpsi_p4 - unpairedMuon_p4).M2()
#nn_q2 = (bc_p4 - muon1_p4 - muon2_p4).M2()
nn_q2 = (bcCorrected_p4 - jpsi_p4).M2()
#nn_missPt = bcCorrected_p4.Pt() - muon1_p4.Pt() - muon2_p4.Pt() - unpairedMuon_p4.Pt()
nn_missPt = bcCorrected_p4.Pt() - jpsi_p4.Pt() - unpairedMuon_p4.Pt()
nn_energyJpsiRestFrame = unpairedMuonBoostedToJpsiRestFrame_p4.E()
#nn_varPt = (muon1_p4 + muon2_p4).Pt() - unpairedMuon_p4.Pt()
nn_varPt = jpsi_p4.Pt() - unpairedMuon_p4.Pt()
nn_deltaRMu1Mu2 = muon1_p4.DeltaR(muon2_p4)
nn_unpairedMuPhi = unpairedMuon_p4.Phi()
nn_unpairedMuPt = unpairedMuon_p4.Pt()
nn_unpairedMuEta = unpairedMuon_p4.Eta()
featuresEntry = np.array([
[bcCorrected_p4.Pt()],
[bcCorrected_p4.Px()],
[bcCorrected_p4.Py()],
[bcCorrected_p4.Pz()],
[bcCorrected_p4.E()],
[nn_energyBcRestFrame],
[nn_missMass2],
[nn_q2],
[nn_missPt],
[nn_energyJpsiRestFrame],
[nn_varPt],
[nn_deltaRMu1Mu2],
[nn_unpairedMuPhi],
[nn_unpairedMuPt],
[nn_unpairedMuEta]],
dtype=np.double)
return featuresEntry
def invariantMassPt(p1_pt, p1_eta, p1_phi, p1_mass, p2_pt, p2_eta, p2_phi,
p2_mass):
if p1_pt < 0 or p2_pt < 0: return -1
p1 = TLorentzVector()
p2 = TLorentzVector()
p1.SetPtEtaPhiM(p1_pt, p1_eta, p1_phi, p1_mass)
p2.SetPtEtaPhiM(p2_pt, p2_eta, p2_phi, p2_mass)
return (p1 + p2).Pt()
def filterDR(obj, collection):
"""Returns the given object filtered from the given collection."""
objVec = TLorentzVector()
objVec.SetPtEtaPhiM(obj.pt(), obj.eta(), obj.phi(), obj.mass())
for p in collection:
pVec = TLorentzVector()
pVec.SetPtEtaPhiM(p.pt(), p.eta(), p.phi(), p.mass())
if objVec.DeltaR(pVec) > 0.3:
return obj
def sintheta_CM(pt1, eta1, phi1, ptz, etaz, phiz, mass):
mu1 = TLorentzVector()
zb = TLorentzVector()
mu1.SetPtEtaPhiM(pt1, eta1, phi1, 0)
zb.SetPtEtaPhiM(ptz, etaz, phiz, mass)
Sintheta = 2.0 * (pt1 / mass) * math.sin(zb.Angle(mu1.Vect()))
#if ( zb.Angle(mu1.Vect()) < 0.0 ):
#print zb.Angle(mu1.Vect()), "******%%%%%%%%%%*"
return Sintheta
def process(self, event):
self.tree.reset()
zprimes_ele = getattr(event, self.cfg_ana.zprime_ele)
zprimes_muon = getattr(event, self.cfg_ana.zprime_muon)
if len(zprimes_ele) + len(zprimes_muon) == 0: return
jets = getattr(event, self.cfg_ana.jets)
for ijet, jet in enumerate(jets):
if ijet == 3:
break
fillParticle(self.tree, 'jet{ijet}'.format(ijet=ijet + 1), jet)
self.tree.fill('weight', sign(event.weight))
met = getattr(event, self.cfg_ana.met)
fillMet(self.tree, 'met', met)
zprimes_ele.sort(key=lambda x: x.m(), reverse=True)
zprimes_muon.sort(key=lambda x: x.m(), reverse=True)
Zp = TLorentzVector()
if len(zprimes_ele) > 0 and len(zprimes_muon) == 0:
fillParticle(self.tree, 'zprime_ele', zprimes_ele[0])
fillLepton(self.tree, 'lep1', zprimes_ele[0].legs[0])
fillLepton(self.tree, 'lep2', zprimes_ele[0].legs[1])
Zp.SetPtEtaPhiM(zprimes_ele[0].pt(), zprimes_ele[0].eta(),
zprimes_ele[0].phi(), zprimes_ele[0].m())
self.tree.fill('zprime_y', Zp.Rapidity())
elif len(zprimes_muon) > 0 and len(zprimes_ele) == 0:
fillParticle(self.tree, 'zprime_muon', zprimes_muon[0])
fillLepton(self.tree, 'lep1', zprimes_muon[0].legs[0])
fillLepton(self.tree, 'lep2', zprimes_muon[0].legs[1])
Zp.SetPtEtaPhiM(zprimes_muon[0].pt(), zprimes_muon[0].eta(),
zprimes_muon[0].phi(), zprimes_muon[0].m())
self.tree.fill('zprime_y', Zp.Rapidity())
else:
if zprimes_ele[0].m() > zprimes_muon[0].m():
fillParticle(self.tree, 'zprime_ele', zprimes_ele[0])
fillLepton(self.tree, 'lep1', zprimes_ele[0].legs[0])
fillLepton(self.tree, 'lep2', zprimes_ele[0].legs[1])
Zp.SetPtEtaPhiM(zprimes_ele[0].pt(), zprimes_ele[0].eta(),
zprimes_ele[0].phi(), zprimes_ele[0].m())
self.tree.fill('zprime_y', Zp.Rapidity())
else:
fillParticle(self.tree, 'zprime_muon', zprimes_muon[0])
fillLepton(self.tree, 'lep1', zprimes_muon[0].legs[0])
fillLepton(self.tree, 'lep2', zprimes_muon[0].legs[1])
Zp.SetPtEtaPhiM(zprimes_muon[0].pt(), zprimes_muon[0].eta(),
zprimes_muon[0].phi(), zprimes_muon[0].m())
self.tree.fill('zprime_y', Zp.Rapidity())
self.tree.tree.Fill()
def findZ(goodElectronList, goodMuonList, entry):
selpair, pairList, mZ, bestDiff = [], [], 91.19, 99999.
nElectron = len(goodElectronList)
if nElectron > 1:
for i in range(nElectron):
ii = goodElectronList[i]
e1 = TLorentzVector()
e1.SetPtEtaPhiM(entry.Electron_pt[ii], entry.Electron_eta[ii],
entry.Electron_phi[ii], 0.0005)
for j in range(i + 1, nElectron):
jj = goodElectronList[j]
if entry.Electron_charge[ii] != entry.Electron_charge[jj]:
e2 = TLorentzVector()
e2.SetPtEtaPhiM(entry.Electron_pt[jj],
entry.Electron_eta[jj],
entry.Electron_phi[jj], 0.0005)
cand = e1 + e2
mass = cand.M()
if abs(mass - mZ) < bestDiff:
bestDiff = abs(mass - mZ)
if entry.Electron_charge[ii] > 0.:
pairList = [e1, e2]
selpair = [ii, jj]
else:
pairList = [e2, e1]
selpair = [jj, ii]
nMuon = len(goodMuonList)
if nMuon > 1:
# find mass pairings
for i in range(nMuon):
ii = goodMuonList[i]
mu1 = TLorentzVector()
mu1.SetPtEtaPhiM(entry.Muon_pt[ii], entry.Muon_eta[ii],
entry.Muon_phi[ii], 0.105)
for j in range(i + 1, nMuon):
jj = goodMuonList[j]
if entry.Muon_charge[ii] != entry.Muon_charge[jj]:
mu2 = TLorentzVector()
mu2.SetPtEtaPhiM(entry.Muon_pt[jj], entry.Muon_eta[jj],
entry.Muon_phi[jj], 0.105)
cand = mu1 + mu2
mass = cand.M()
if abs(mass - mZ) < bestDiff:
bestDiff = abs(mass - mZ)
if entry.Muon_charge[ii] > 0.:
pairList = [mu1, mu2]
selpair = [ii, jj]
else:
pairList = [mu2, mu1]
selpair = [jj, ii]
# first particle of pair is positive
#print selpair
return pairList, selpair
def invariantDoubleMass(Z_pt, Z_eta, Z_phi, jet1_pt, jet1_eta, jet1_phi,
jet2_pt, jet2_eta, jet2_phi):
if jet1_pt < 0 or jet2_pt < 0: return -1
j1 = TLorentzVector()
j2 = TLorentzVector()
Z = TLorentzVector()
Z.SetPtEtaPhiM(Z_pt, Z_eta, Z_phi, 91.)
j1.SetPtEtaPhiM(jet1_pt, jet1_eta, jet1_phi, 5.)
j2.SetPtEtaPhiM(jet2_pt, jet2_eta, jet2_phi, 5.)
return (j1 + j2 + Z).M()
def p4(self, correction=''): # pt, eta, phi, m
corrp4 = TLorentzVector()
corrp4.SetPtEtaPhiM(self.pt('corr'), self.eta(), self.phi(),
self.mass())
uncorrp4 = TLorentzVector()
uncorrp4.SetPtEtaPhiM(self.pt('uncorr'), self.eta(), self.phi(),
self.mass())
# decide which value to return
if correction == 'corr': return corrp4
elif correction == 'uncorr': return uncorrp4
elif self.corrected: return corrp4
else: return uncorrp4
def masses(evt, i, j):
"""
Return the i'th jet mass, the j'th jet mass, and their invariant mass
"""
i, j = leading_subleading(evt)
lead_vector = TLorentzVector()
lead_vector.SetPtEtaPhiM(evt.JetPt[i], evt.JetEta[i], evt.JetPhi[i],
evt.JetM[i])
subl_vector = TLorentzVector()
subl_vector.SetPtEtaPhiM(evt.JetPt[j], evt.JetEta[j], evt.JetPhi[j],
evt.JetM[j])
return evt.JetM[i], evt.JetM[j], (lead_vector + subl_vector).Mag()
def invariantMassTracksAsPions(track1, track2):
trkTlv1 = TLorentzVector()
trkTlv2 = TLorentzVector()
trkTlv1.SetPtEtaPhiM(track1.pt(), track1.eta(), track1.phi(), 0.13957018)
trkTlv2.SetPtEtaPhiM(track2.pt(), track2.eta(), track2.phi(), 0.13957018)
invMass = (trkTlv1 + trkTlv2).M()
return invMass
def mass(tau1, tau2, METpx, METpy):
"""
Calculate and return the collinear mass and momentum fractions
of tau1 and tau2
TODO: set visible mass of taus. 1.2 GeV for 3p and 0.8 GeV for 1p
"""
recTau1 = TLorentzVector()
recTau2 = TLorentzVector()
# tau 4-vector; synchronize for MMC calculation
if tau1.numTrack < 3:
recTau1.SetPtEtaPhiM(tau1.pt, tau1.eta, tau1.phi, 800.) # MeV
else:
recTau1.SetPtEtaPhiM(tau1.pt, tau1.eta, tau1.phi, 1200.) # MeV
if tau2.numTrack < 3:
recTau2.SetPtEtaPhiM(tau2.pt, tau2.eta, tau2.phi, 800.) # MeV
else:
recTau2.SetPtEtaPhiM(tau2.pt, tau2.eta, tau2.phi, 1200.) # MeV
K = ROOT.TMatrixD(2, 2)
K[0][0] = recTau1.Px()
K[0][1] = recTau2.Px()
K[1][0] = recTau1.Py()
K[1][1] = recTau2.Py()
if K.Determinant() == 0:
return -1., -1111., -1111.
M = ROOT.TMatrixD(2, 1)
M[0][0] = METpx
M[1][0] = METpy
Kinv = K.Invert()
X = Kinv * M
X1 = X(0, 0)
X2 = X(1, 0)
x1 = 1. / (1. + X1)
x2 = 1. / (1. + X2)
p1 = recTau1 * (1. / x1)
p2 = recTau2 * (1. / x2)
m_col = (p1 + p2).M()
m_vis = (recTau1 + recTau2).M()
return m_vis, m_col, x1, x2
def jet_Lorentz_4v(jet):
for n in np.arange(len(jet)):
if np.sum(jet[n,:]) != 0:
v = TLorentzVector(0, 0, 0, 0)
v.SetPtEtaPhiM(jet[n,0], jet[n,1], jet[n,2], jet[n,3])
jet[n,:] = v.E(), v.Px(), v.Py(), v.Pz()
return jet
class Met(BaseMET):
def __init__(self, fccmet):
self.fccmet = fccmet
self._sum_et = fccmet.scalarSum()
self._tlv = TLorentzVector()
self._tlv.SetPtEtaPhiM(fccmet.magnitude(), 0., fccmet.phi(), 0.)
self._charge = 0.
class SimTrack(BaseParticle):
def __init__(self, simtrack, simvertex=None):
self.simtrack = simtrack
self._charge = simtrack.charge()
self._pid = simtrack.type()
self._status = 1
self._tlv = TLorentzVector()
p4 = simtrack.momentum()
self._tlv.SetPtEtaPhiM(p4.pt(), p4.eta(), p4.phi(), p4.mass())
if simvertex:
self._vertex = Vertex(simvertex)
else:
self._vertex = None
def vertex(self):
return self._vertex
def __getattr__(self, attr):
return getattr(self.simtrack, attr)
def __str__(self):
tmp = '{className} : pdgid = {pdgid:5}, status = {status:3}, q = {q:2} {p4}, {vertex}'
return tmp.format(className=self.__class__.__name__,
pdgid=self.pdgid(),
status=self.status(),
q=self.q(),
p4=super(BaseParticle, self).__str__(),
vertex=self._vertex.__str__() if self.vertex else '')
def __repr__(self):
return str(self)
def get_jets_vectors(self):
'''
Returns a list of 4-momenta for jets looking only at jets
that are cleaned from FatJets.
A list of indexes in the collection of CleanJet is returned as a reference.
'''
jets = []
coll_ids = []
for ijnf in range(len(self.JetNotFat_coll)):
jetindex = self.JetNotFat_coll[ijnf].jetIdx
# index in the original Jet collection
rawjetid = self.Jet_coll[jetindex].jetIdx
pt, eta, phi, mass = self.Jet_coll[jetindex].pt, \
self.Jet_coll[jetindex].eta,\
self.Jet_coll[jetindex].phi, \
self.rawJet_coll[rawjetid].mass
if abs(eta) > 10: continue
vec = TLorentzVector()
vec.SetPtEtaPhiM(pt, eta, phi, mass)
# check if different from the previous one
if self.debug:
print "Jet index: ", jetindex, "> pt:", pt, " eta:", eta, " phi:", phi, " mass:", mass
jets.append(vec)
coll_ids.append(jetindex)
return jets, coll_ids
def invariantMassTracks(track1, track2):
trkTlv1 = TLorentzVector()
trkTlv2 = TLorentzVector()
#trkTlv1.SetPxPyPzE(track1.px(), track1.py(), track1.pz(), track1.pt()*np.cosh(track1.eta()))
#trkTlv2.SetPxPyPzE(track2.px(), track2.py(), track2.pz(), track2.pt()*np.cosh(track2.eta()))
trkTlv1.SetPtEtaPhiM(track1.pt(), track1.eta(), track1.phi(), 0.13957018)
trkTlv2.SetPtEtaPhiM(track2.pt(), track2.eta(), track2.phi(), 0.13957018)
invMass = (trkTlv1 + trkTlv2).M()
return invMass
def generateRandomVector( ):
x = TLorentzVector( )
x.SetPtEtaPhiM( random.uniform( 0, 50 ),
random.uniform( -5, 5 ),
random.uniform( -math.pi, math.pi ),
0. )
return x
def TransBoosted_Angle(pt1,eta1,phi1,pt2,eta2,phi2,ptz,etaz,phiz,mass,Met,Metphi): mu1 = TLorentzVector() mu2 = TLorentzVector() zb = TLorentzVector() met = TLorentzVector() mu1.SetPtEtaPhiM(pt1,0,phi1,0) mu2.SetPtEtaPhiM(pt2,0,phi2,0) zb.SetPtEtaPhiM(ptz,0,phiz,mass) MET.SetPtEtaPhiM(Met,0,Metphi,0) MET.Boost(-zb.Px()/zb.Et(),-zb.Py()/zb.Et(),0) mu1.Boost(-zb.Px()/zb.Et(),-zb.Py()/zb.Et(),0) mu2.Boost(-zb.Px()/zb.Et(),-zb.Py()/zb.Et(),0) angleBoost_1 = MET.DeltaPhi(mu1) angleBoost_2 = MET.DeltaPhi(mu2) angleBoost_Z = MET.DeltaPhi(zb) return [angleBoost_Z,angleBoost_1,angleBoost_2]
def isBoostedTop(event, hadTop_p4=None, minPt=200., maxDR=0.4):
# categorize the event as boosted if there is a large R jet matched to top
# LargeR jet pT > 200 GeV from Delphes
for ljet in event.FatJet:
# pT cut on large R jet
if ljet.PT < minPt:
continue
# large R jet mass window: [120, 220] GeV
if ljet.Mass < 120. or ljet.Mass > 220.:
continue
# match the large jet with the hadronic top parton by dR
if hadTop_p4 is not None:
ljet_p4 = TLorentzVector()
ljet_p4.SetPtEtaPhiM(ljet.PT, ljet.Eta, ljet.Phi, ljet.Mass)
dR = ljet_p4.DeltaR(hadTop_p4)
if dR > maxDR:
continue
# ljet is matched to the hadronic top
return True
# None of the large R jets meet the requirements
return False
def make_LHEparticle(p, id=None):
if id != None:
fv = TLorentzVector()
try:
m = p.Mass
except AttributeError:
m = 0
try:
pT = p.PT
except AttributeError:
pT = p.MET
fv.SetPtEtaPhiM(pT, p.Eta, p.Phi, m)
try:
status = p.Status
except AttributeError:
status = 1
_dict={'id':id,'status':status,\
'mother1':0,'mother2':0,'color1':0,'color2':0,\
'px':fv.Px(),'py':fv.Py(),'pz':fv.Pz(),'e':fv.Energy(),'m':fv.M(),'lifetime':0,'spin':0}
lhepart = lhef.LHEParticle(**_dict)
return lhepart
else:
print("id must be specified")
pass
class Jet:
def __init__(self, jet_obj, subjet_array):
self.pT = jet_obj.PT
self.eta = jet_obj.Eta
self.phi = jet_obj.Phi
self.mass = jet_obj.Mass
self.momentum = TLorentzVector()
self.momentum.SetPtEtaPhiM(self.pT, self.eta, self.phi, self.mass)
self.full_obj = jet_obj
momenta = [array2lorentz(a) for a in subjet_array]
self.trimmed_momentum = momenta[0]
self.sj_momentum = momenta[1:]
self.cells = self.read_constituents(self.full_obj.Constituents)
def read_constituents(self, constituents):
cells = []
for const in constituents:
if isinstance(const, ROOT.Tower):
c = Cell(const)
if c.notempty():
cells.append(c)
return cells
def __array__(self):
ar = np.zeros((len(self.cells), 3))
for i, cell in enumerate(self.cells):
ar[i, :] = np.array(cell)
return ar
def __len__(self):
return len(self.cells)
def __getitem__(self, item):
return self.cells[item]
def Boosted_Angle(pt1, eta1, phi1, pt2, eta2, phi2, ptz, etaz, phiz, mass):
mu1 = TLorentzVector()
mu2 = TLorentzVector()
zb = TLorentzVector()
mu1.SetPtEtaPhiM(pt1, eta1, phi1, 0)
mu2.SetPtEtaPhiM(pt2, eta2, phi2, 0)
angle = mu1.Angle(mu2.Vect())
zb.SetPtEtaPhiM(ptz, etaz, phiz, mass)
angle_Z1 = zb.Angle(mu1.Vect())
angle_Z2 = zb.Angle(mu2.Vect())
mu1.Boost(-zb.Px() / zb.E(), -zb.Py() / zb.E(), -zb.Pz() / zb.E())
mu2.Boost(-zb.Px() / zb.E(), -zb.Py() / zb.E(), -zb.Pz() / zb.E())
angleBoost = mu1.Angle(mu2.Vect())
angleBoost_Z1 = zb.Angle(mu1.Vect())
angleBoost_Z2 = zb.Angle(mu2.Vect())
#print "******&&&&******", angle, angleBoost
return [angleBoost, angle, angleBoost_Z1, angle_Z1, angle_Z2]
class particle:
def __init__(self, tree=None, i=None, original=None):
if tree: #Build particle instance from TTree
self.PID = tree.Particle[i].PID
self.mom = tree.Particle[i].Mother1
self.daughters = []
self.m = tree.Particle[i].M
self.e = tree.Particle[i].E
self.p = TLorentzVector()
self.p.SetPxPyPzE(tree.Particle[i].Px, tree.Particle[i].Py,
tree.Particle[i].Pz, tree.Particle[i].E)
self.pt = self.p.Pt()
self.eta = self.p.Eta()
self.phi = self.p.Phi()
self.SF = 1
self.smear = None
self.res = None
elif original: #Make a resolution smeared version of the original
self.PID = copy(original.PID)
self.mom = copy(original.mom)
self.daughters = copy(original.daughters)
self.res = jetEnergyResolution(original.p.E())
self.smearE = -1
while self.smearE < 0:
self.smearE = np.random.normal(0.985, self.res)
# self.offsetM = -1
# while self.offsetM+original.p.M() < 5: self.offsetM = np.random.normal(5, 5)
# self.smearM = -1
# while self.smearM < 0.5: self.smearM = np.random.normal(1, 1)
self.p = TLorentzVector()
self.p.SetPtEtaPhiM(original.p.Pt() * self.smearE,
original.p.Eta(), original.p.Phi(), 0)
# self.p.SetPtEtaPhiM( original.p.Pt() * self.smearE,
# original.p.Eta(),
# original.p.Phi(),
# (original.p.M()+self.offsetM) * self.smearM)
self.pt = self.p.Pt()
self.eta = self.p.Eta()
self.phi = self.p.Phi()
self.m = self.p.M()
self.e = self.p.E()
self.SF = 1
def getDump(self):
out = "PID " + str(self.PID).rjust(3) + " | mom " + str(
self.mom).rjust(3) + " | mass " + str(self.m).ljust(
12) + " | pt " + str(self.pt).ljust(20) + " | eta " + str(
self.eta).ljust(20) + " | phi " + str(self.phi).ljust(20)
if self.res:
out += " | res " + str(self.res).ljust(20) + " | smear " + str(
self.smear).ljust(20)
return out
def dump(self):
print(self.getDump())
def sumTLVs():
"""Test function test behavior of TLorentzVectors."""
print ">>> create TLVs..."
pt = [40.0, 30.0]
eta = [1.5, -2.2]
phi = [0.0, 0.0]
m = [4.0, 4.0]
tlv1 = TLorentzVector()
tlv1.SetPtEtaPhiM(pt[0], eta[0], phi[0], m[0])
tlv2 = TLorentzVector()
tlv2.SetPtEtaPhiM(pt[1], eta[1], phi[1], m[1])
print ">>> tlv1.SetPtEtaPhiM(%.3g,%.3g,%.3g,%.3g)" % (pt[0], eta[0],
phi[0], m[0])
print ">>> tlv2.SetPtEtaPhiM(%.3g,%.3g,%.3g,%.3g)" % (pt[1], eta[1],
phi[1], m[1])
def TagVars(collections):
results=[]; Et_ratio=[]; Dphi=[]; projB=[]; DzTagMuK=[]; DzTagMuL1=[];
DzTagMuL2=[];
tracks=collections[0]
Bcands=collections[2]
trgmuons=(collections[1])
trgmuon_vec=TLorentzVector()
trgmuon_vec.SetPtEtaPhiM(0,0,0,0)
trgmuon_vz=-99
for trgmuon in trgmuons:
if not getattr(trgmuon,"isTriggering"):
continue;
trgmuon_vec.SetPtEtaPhiM(getattr(trgmuon,"pt"),getattr(trgmuon,"eta"),getattr(trgmuon,"phi"),0.105)
trgmuon_vz=getattr(trgmuon,"vz")
break
if trgmuon_vec.M()==0:
result=[[0], [0], [0], [0], [0], [0]]
return result
sum_track_vec=trgmuon_vec;
sum_track=trgmuon_vec.Pt();
for track in tracks:
track_vec=TLorentzVector();
track_vec.SetPtEtaPhiM(getattr(track,"pt"),getattr(track,"eta"),getattr(track,"phi"),0.139)
if trgmuon_vec.DrEtaPhi(track_vec)>0.4:
continue
sum_track_vec=sum_track_vec+track_vec;
sum_track+=track_vec.Pt()
for Bcand in Bcands:
Bcand_vec=TLorentzVector();
Bcand_vec.SetPtEtaPhiM(getattr(Bcand,"fit_pt"),getattr(Bcand,"fit_eta"),getattr(Bcand,"fit_phi"),getattr(Bcand,"fit_mass"))
if sum_track>0:
Et_ratio.append(Bcand_vec.Et()/sum_track)
else:
Et_ratio.append(0)
Dphi.append(Bcand_vec.DeltaPhi(sum_track_vec))
projB.append(Bcand_vec*sum_track_vec)
DzTagMuK.append( abs(trgmuon_vz-getattr(Bcand,"kVz")) )
DzTagMuL1.append( abs(trgmuon_vz-getattr(Bcand,"l1Vz")) )
DzTagMuL2.append( abs(trgmuon_vz-getattr(Bcand,"l2Vz")) )
result=[Et_ratio, Dphi, projB, DzTagMuL1, DzTagMuL2, DzTagMuK]
return result
