def setUp(self):
# goes along x
self.p4 = TLorentzVector(1, 0, 0, 1.1)
# starts at y = 0.1
self.true_IP = 0.1
self.vertex = TVector3(0, self.true_IP, 0)
self.helix = Helix(1, 1, self.p4, self.vertex)
# global origin
self.origin = TVector3(0, 0, 0)
def test_printout(self):
'''Test that the particle printout is adapted to the collider
beams.'''
ptc = TlvParticle(1, 1, TLorentzVector())
Collider.BEAMS = 'pp'
self.assertIn('pt', ptc.__repr__())
Collider.BEAMS = 'ee'
self.assertIn('theta', ptc.__repr__())
Collider.BEAMS = 'pp'
def rotateFourMomentum(self, pFour, rotMatrix):
pArray = [pFour.X(), pFour.Y(), pFour.Z()]
pRot = np.dot(rotMatrix, np.asarray(pArray))
p = TLorentzVector(0, 0, 0, me)
p.SetPx(float(pRot[0]))
p.SetPy(float(pRot[1]))
p.SetPz(float(pRot[2]))
p.SetE(float(pFour.E()))
return p
def pointingAngle(track1, track2, pv_pos, sv_pos):
trk1Tlv = TLorentzVector()
trk1Tlv.SetPxPyPzE(track1.px(), track1.py(), track1.pz(),
track1.pt() * np.cosh(track1.eta()))
trk2Tlv = TLorentzVector()
trk2Tlv.SetPxPyPzE(track2.px(), track2.py(), track2.pz(),
track2.pt() * np.cosh(track2.eta()))
combinedMomentum = trk2Tlv - trk1Tlv
PVtoSV = [
sv_pos[0] - pv_pos.x(), sv_pos[1] - pv_pos.y(), sv_pos[2] - pv_pos.z()
]
pointingAngle = np.arctan2(PVtoSV[0] - combinedMomentum.X(),
PVtoSV[1] - combinedMomentum.Y())
return pointingAngle
def isCleanJet(Chain, index):
if Chain._jetPt[index] < 25.: return False
if not Chain._jetIsTight[index]: return False
jetVec = TLorentzVector()
jetVec.SetPtEtaPhiE(Chain._jetPt[index], Chain._jetEta[index],
Chain._jetPhi[index], Chain._jetE[index])
for l in xrange(Chain._nL):
lVec = TLorentzVector()
lVec.SetPtEtaPhiE(Chain._lPt[l], Chain._lEta[l], Chain._lPhi[l],
Chain._lE[l])
if jetVec.DeltaR(lVec) < 0.4:
return False
return True
def get_p4(self, ptc_type):
'''Build the particle four vector'''
energy = eval("self.evt." + self.get_ptc(ptc_type) + "_e")
p_x = eval("self.evt." + self.get_ptc(ptc_type) + "_px")
p_y = eval("self.evt." + self.get_ptc(ptc_type) + "_py")
p_z = eval("self.evt." + self.get_ptc(ptc_type) + "_pz")
return TLorentzVector(p_x, p_y, p_z, energy)
def passedInvariantMassCuts(self, lIndices):
vec = []
for l in lIndices:
v = TLorentzVector()
v.SetPtEtaPhiE(self.Chain._lPt[l], self.Chain._lEta[l], self.Chain._lPhi[l], self.Chain._lE[l])
vec.append(v)
Mlll = (vec[0] + vec[1] + vec[2]).M()
if abs(Mlll-91.19) < 15 : return False
return True
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
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 CostCS(p1, charge1, p2): #Cosine of the theta decay angle (top (Q=+2/3)) in the Collins-Soper frame pTop1CM = TLorentzVector(0,0,-1,1) # In the CM. frame pTop2CM = TLorentzVector(0,0,-1,1) # In the CM. frame pProjCM = TLorentzVector(0,0,-1,1) # In the CM. frame pTargCM = TLorentzVector(0,0,-1,1) # In the CM. frame pDitopCM = TLorentzVector(0,0,-1,1) # In the CM. frame pTop1Ditop = TLorentzVector(0,0,-1,1) # In the Ditop rest frame pTop2Ditop = TLorentzVector(0,0,-1,1) # In the Ditop rest frame pProjDitop = TLorentzVector(0,0,-1,1) # In the Ditop rest frame pTargDitop = TLorentzVector(0,0,-1,1) # In the Ditop rest frame beta = TVector3(0,0,0) zaxisCS = TVector3(0,0,0) mp = 0.93827231 ep = 6500. # Fill the Lorentz vector for projectile and target in the CM frame pProjCM.SetPxPyPzE(0.,0.,-ep,TMath.Sqrt(ep*ep+mp*mp)) pTargCM.SetPxPyPzE(0.,0.,+ep,TMath.Sqrt(ep*ep+mp*mp)) # Get the Topons parameters in the CM frame pTop1CM.SetPxPyPzE(p1.Px(),p1.Py(),p1.Pz(),p1.E()) pTop2CM.SetPxPyPzE(p2.Px(),p2.Py(),p2.Pz(),p2.E()) # Obtain the Ditop parameters in the CM frame pDitopCM=pTop1CM+pTop2CM # Translate the Ditop parameters in the Ditop rest frame beta=(-1./pDitopCM.E())*pDitopCM.Vect() if(beta.Mag()>=1): return 666. pTop1Ditop=pTop1CM pTop2Ditop=pTop2CM pProjDitop=pProjCM pTargDitop=pTargCM pTop1Ditop.Boost(beta) pTop2Ditop.Boost(beta) pProjDitop.Boost(beta) pTargDitop.Boost(beta) # Determine the z axis for the CS angle zaxisCS=(((pProjDitop.Vect()).Unit())-((pTargDitop.Vect()).Unit())).Unit(); # Determine the CS angle (angle between Top+ and the z axis defined above) cost = -999 if(charge1>0): cost = zaxisCS.Dot((pTop1Ditop.Vect()).Unit()) else: cost = zaxisCS.Dot((pTop2Ditop.Vect()).Unit()) return cost
def monojet(pdgids, theta, phi, pstar, jetenergy, vertex=None):
particles = []
if vertex is None:
vertex = TVector3(0., 0., 0.)
jetp4star = TLorentzVector()
for pdgid in pdgids[:-1]:
mass, charge = particle_data[pdgid]
phistar = random.uniform(-math.pi, math.pi)
thetastar = random.uniform(-math.pi, math.pi)
sint = math.sin(thetastar)
cost = math.cos(thetastar)
sinp = math.sin(phistar)
cosp = math.cos(phistar)
pz = pstar * cost
px = pstar * sint * cosp
py = pstar * sint * sinp
p4 = TLorentzVector()
p4.SetXYZM(px, py, pz, mass)
jetp4star += p4
particles.append(Particle(p4, vertex, charge, pdgid))
pdgid = pdgids[-1]
mass, charge = particle_data[pdgid]
p4 = TLorentzVector()
p4.SetVectM(-jetp4star.Vect(), mass)
particles.append(Particle(p4, vertex, charge, pdgid))
jetp4star += p4
#boosting to lab
gamma = jetenergy / jetp4star.M()
beta = math.sqrt(1 - 1 / gamma**2)
boostvec = TVector3(
math.sin(theta) * math.cos(phi),
math.sin(theta) * math.sin(phi), math.cos(theta))
boostvec *= beta
boosted_particles = []
jetp4 = LorentzVector()
for ptc in particles:
bp4 = LorentzVector(ptc.p4())
bp4.Boost(boostvec)
jetp4 += bp4
boosted_particles.append(
Particle(bp4, ptc.vertex, ptc.q(), ptc.pdgid()))
# print jetp4.M(), jetp4.E()
return boosted_particles
def Return_Top(jet, id_csv, metPhi):
masses = [-1000, -1000, -1000, -1000, -1000, -1000]
b_tmp = TLorentzVector()
b_tmp.SetPtEtaPhiE(jet[id_csv][0], jet[id_csv][1], jet[id_csv][2],
jet[id_csv][3])
jet_nob = list(jet)
jet_nob.pop(id_csv)
j1dR_tmp = TLorentzVector()
j1dR_tmp.SetPtEtaPhiE(0, 0, 0, 0)
j2dR_tmp = TLorentzVector()
j2dR_tmp.SetPtEtaPhiE(0, 0, 0, 0)
a = list(combinations(jet_nob, 2))
min_DeltaR = 1000
for x in a:
if x[0][0] < -5 or x[1][0] < -5: continue
if x[0][1] < -5 or x[1][1] < -5: continue
if (math.sqrt(
math.pow((x[0][0] - x[1][0]), 2) +
math.pow((x[0][1] - x[1][1]), 2)) < min_DeltaR):
j1dR_tmp.SetPtEtaPhiE(x[0][0], x[0][1], x[0][2], x[0][3])
j2dR_tmp.SetPtEtaPhiE(x[1][0], x[1][1], x[1][2], x[1][3])
topdR_tmp = TLorentzVector()
topdR_tmp = j1dR_tmp + j2dR_tmp + b_tmp
WdR_tmp = TLorentzVector()
WdR_tmp = j1dR_tmp + j2dR_tmp
masses[0] = topdR_tmp.M()
masses[1] = WdR_tmp.M()
masses[2] = math.fabs(DeltaPhi(WdR_tmp.Phi(), b_tmp.Phi()))
masses[3] = math.fabs(DeltaPhi(topdR_tmp.Phi(), b_tmp.Phi()))
masses[4] = math.fabs(DeltaPhi(WdR_tmp.Phi(), metPhi))
masses[5] = math.fabs(DeltaPhi(topdR_tmp.Phi(), metPhi))
return masses
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 __init__(self, leg1, leg2):
if leg2.energy() > leg1.energy():
leg2, leg1 = leg1, leg2
lv1 = TLorentzVector(leg1.px(), leg1.py(), leg1.pz(), leg1.energy())
lv2 = TLorentzVector(leg2.px(), leg2.py(), leg2.pz(), leg2.energy())
lv1 += lv2
super(DiObject, self).__init__(lv1)
self.p4_ = lv1
c1 = ( leg1.px()*leg2.px() + \
leg1.py()*leg2.py() + \
leg1.pz()*leg2.pz() ) / ( leg1.p() * leg2.p() )
#print c1
if c1 >= 1.0: c1 = 1. - 1E-12
if c1 <= -1.0: c1 = -1. + 1E-12
#print leg1, leg2, c1
self.angle_ = acos(c1) * 180. / pi
self.leg1 = leg1
self.leg2 = leg2
self.btag_ = []
self.component_ = []
# protect in case non standard object (as taus...) checking with hasattr
for i in range(8):
if hasattr(leg1, "btag"):
self.btag_.append(max(self.leg1.btag(i), self.leg2.btag(i)))
else:
self.btag_.append(-999.)
if hasattr(leg1, "component"):
self.component_.append(Component(self, i))
else:
self.component_.append(-999.)
self.pdgId_ = 0
if self.leg1.pdgId() != 0: self.pdgId_ = self.leg1.pdgId()
if self.leg2.pdgId() != 0: self.pdgId_ = self.leg2.pdgId()
if hasattr(leg1, "nConstituents"):
self.nConstituents_ = self.leg1.nConstituents(
) + self.leg2.nConstituents()
else:
self.nConstituents_ = -999.
def add_obj(mem, typ, **kwargs):
if kwargs.has_key("p4s"):
pt, eta, phi, mass = kwargs.pop("p4s")
v = TLorentzVector()
v.SetPtEtaPhiM(pt, eta, phi, mass)
elif kwargs.has_key("p4c"):
v = TLorentzVector(*kwargs.pop("p4c"))
obsdict = kwargs.pop("obsdict", {})
o = MEM.Object(v, typ)
if typ == MEM.ObjectType.Jet:
tb, tl = attach_jet_transfer_function(v.Pt(), v.Eta())
o.addTransferFunction(MEM.TFType.qReco, tl)
o.addTransferFunction(MEM.TFType.bReco, tb)
for k, v in obsdict.items():
o.addObs(k, v)
mem.push_back_object(o)
def RecoP4(self, regIdx):
"""
Sums the energy collected in a given signal region
and returns the 4-vector.
"""
p4 = TLorentzVector(0, 0, 0, 0)
en = np.sum(self.recEn[:regIdx], axis=None)
pt = en / TMath.CosH(self.genP4.Eta())
p4.SetPtEtaPhiM(pt, self.genP4.Eta(), self.genP4.Phi(), 0)
return p4
def M_jj(parts): """ find pair of jets with invar mass closest to MW """ MW=80.385 p=TLorentzVector(0,0,0,0) # px,py,pz,E Mjets=100000 for j1,j2 in combinations(parts,2): if (abs((j1.P4()+j2.P4()).M()-MW) < abs(Mjets-MW)): Mjets=(j1.P4()+j2.P4()).M() return Mjets
def findZee(goodElectronList, entry) :
pairList, mZ, bestDiff = [], 91.19, 99999.
nElectron = len(goodElectronList)
if nElectron < 2 : return pairList
# find mass pairings
for i in range(nElectron) :
e1 = TLorentzVector()
e1.SetPtEtaPhiM(entry.Electron_pt[i],entry.Electron_eta[i],entry.Electron_phi[i],0.0005)
for j in range(i+1,nElectron) :
if entry.Electron_charge[i] != entry.Electron_charge[j] :
e2 = TLorentzVector()
e2.SetPtEtaPhiM(entry.Electron_pt[j],entry.Electron_eta[j],entry.Electron_phi[j],0.0005)
cand = e1 + e2
mass = cand.M()
if abs(mass-mZ) < bestDiff :
bestDiff = abs(mass-mZ)
pairList.append([e1,e2])
return pairList
def findZmumu(goodMuonList, entry) :
pairList, mZ, bestDiff = [], 91.19, 99999.
nMuon = len(goodMuonList)
if nMuon < 2 : return pairList
# find mass pairings
for i in range(nMuon) :
mu1 = TLorentzVector()
mu1.SetPtEtaPhiM(entry.Muon_pt[i],entry.Muon_eta[i],entry.Muon_phi[i],0.105)
for j in range(i+1,nMuon) :
if entry.Muon_charge[i] != entry.Muon_charge[j] :
mu2 = TLorentzVector()
mu2.SetPtEtaPhiM(entry.Muon_pt[j],entry.Muon_eta[j],entry.Muon_phi[j],0.105)
cand = mu1 + mu2
mass = cand.M()
if abs(mass-mZ) < bestDiff :
bestDiff = abs(mass-mZ)
pairList.append([mu1,mu2])
return pairList
def mt_with(self, *cands):
myP4 = self.p4()
lepP4 = TLorentzVector()
for lep in cands:
lepP4 += lep.p4()
# squared transverse mass of system =
# (met.Et + lep.Et)^2 - (met + lep).Pt)^2
mt = math.sqrt(
abs((lepP4.et() + myP4.et())**2 - ((lepP4 + myP4).pt())**2))
return mt
def process(self, event):
particles = getattr(event, self.cfg_ana.particles)
missingp4 = TLorentzVector()
sumpt = 0
for ptc in particles:
missingp4 += ptc.p4()
sumpt += ptc.pt()
missingp4 *= -1
met = MET(missingp4, sumpt)
setattr(event, self.instance_label, met)
def setUp(self):
"""maps the space with particles"""
self.ptcs = {}
for eta in range(-30, 30, 2):
eta /= 10.
for phi in range(-30, 30, 2):
phi /= 10.
tlv = TLorentzVector()
tlv.SetPtEtaPhiM(10, eta, phi, 0)
self.ptcs[(eta, phi)] = Particle(1, 0, tlv)
def process(self, event):
#sqrts = self.cfg_ana.sqrts
sqrts = 240.
#sqrts = random.gauss(240.,240.*0.0012)
#sqrts = random.gauss(240.,10.)
jets = getattr(event, self.cfg_ana.input_jets)
misenergy = getattr(event, self.cfg_ana.misenergy)
# solving the equation
vis_p4 = TLorentzVector(0, 0, 0, 0)
for jet in jets:
vis_p4 += jet.p4()
vis = Recoil(0, 0, vis_p4, 1)
# m=m_Z constrain
a = vis.e() * sqrts / (vis.m()**2)
b = a**2 - (sqrts**2 - 91.**2) / (vis.m()**2)
# sf stands for scaling factor
if b < 0:
sf = 1
setattr(event, self.cfg_ana.det, 1)
else:
#sf2 corresponds to the solution where the missing energy is negative therefore sf is the scaling factor that makes sense from a physics pov.
sf = a - math.sqrt(b)
sf2 = a + math.sqrt(b)
setattr(event, self.cfg_ana.det, 2)
######test
#visscaled1=TLorentzVector(0,0,0,0)
#visscaled2=TLorentzVector(0,0,0,0)
#for jet in jets:
# visscaled1+=jet.p4()*sf
# visscaled2+=jet.p4()*sf2
#cms=TLorentzVector(0,0,0,240)
#v1=cms-visscaled1
#v2=cms-visscaled2
#print v1.E(),v2.E()
setattr(event, self.cfg_ana.scalingfac, sf)
setattr(event, self.cfg_ana.scalingfac2, sf2)
scale_factors = [sf] * 2
output = []
for jet, factor in zip(jets, scale_factors):
# the jets should not be deepcopied
# as they are heavy objects containing
# in particular a list of consistuent particles
scaled_jet = copy.copy(jet)
scaled_jet._tlv = copy.deepcopy(jet._tlv)
scaled_jet._tlv *= factor
output.append(scaled_jet)
setattr(event, self.cfg_ana.output_jets, output)
def generateM_dd(p_signal, n_files=100, n_entries=10000):
mMs = 1.0
mMs_resolution = 0.01
mMlow = 0.5
mMhigh = 1.5
masses = array.array('d', [0.1, 0.1])
features = ['p1x', 'p1y', 'p1z', 'E1', 'p2x', 'p2y', 'p2z', 'E2']
for n_file in range(n_files + 1):
if n_file % 10 == 0: print("n_file : ", n_file)
mass, p1x, p1y, p1z, E1, p2x, p2y, p2z, E2 = [], [], [], [], [], [], [], [], []
for n in range(n_entries):
choice = random.random()
if choice < p_signal:
mM = numpy.random.normal(loc=mMs, scale=mMs_resolution)
else:
mM = numpy.random.uniform(mMlow, mMhigh)
MotherMom = TLorentzVector(0.0, 0.0, 0.0, mM)
event = TGenPhaseSpace()
event.SetDecay(MotherMom, 2, masses)
while True:
weight = event.Generate()
weightmax = event.GetWtMax()
assert (0. < weight / weightmax) & (weight / weightmax < 1.)
if random.random() < weight / (weightmax): break
pd1 = event.GetDecay(0)
pd2 = event.GetDecay(1)
p1x.append(pd1.Px())
p1y.append(pd1.Py())
p1z.append(pd1.Pz())
E1.append(pd1.E())
p2x.append(pd2.Px())
p2y.append(pd2.Py())
p2z.append(pd2.Pz())
E2.append(pd2.E())
mass.append(mM)
df = pd.DataFrame(np.array([p1x, p1y, p1z, E1, p2x, p2y, p2z,
E2]).transpose(),
columns=features)
df2 = pd.DataFrame(np.array([mass]).transpose(), columns=['mM'])
#print('df : ', df)
if n_file == n_files: n_file = "optimisation_0"
#df.to_csv("../data/M-dd/M-dd_{}_res_{}_{}.txt".format(p_signal, mMs_resolution, n_file), sep="\t", header = None, index=False)
df2.to_csv("../data/M-dd/M-dd_{}_res_{}_1Dmass_{}.txt".format(
p_signal, mMs_resolution, n_file),
sep="\t",
header=None,
index=False)
def TagVarsMC(collections):
results=[]; Et_ratio=-99.; Dphi=-99.; projB=-99.
tracks=collections[0]
trgmuons=collections[1]
recoB_pt=collections[2]
recoB_eta=collections[3]
recoB_phi=collections[4]
recoB_mass=collections[5]
recoE1_vz=collections[6]
recoE2_vz=collections[7]
recoK_vz=collections[8]
trgmuon_vec=TLorentzVector()
trgmuon_vec.SetPtEtaPhiM(0,0,0,0)
for trgmuon in trgmuons:
if getattr(trgmuon,"isTriggering")==0:
continue
trgmuon_vec.SetPtEtaPhiM(getattr(trgmuon,"pt"),getattr(trgmuon,"eta"),getattr(trgmuon,"phi"),0.105)
break
if trgmuon_vec.M()==0:
result=[-99.,-99.,-99.,-99.,-99.,-99.]
return result
sum_track_vec=trgmuon_vec;
sum_track=trgmuon_vec.Pt();
trgmuon_vz=getattr(trgmuon,"vz")
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()
recoB_vec=TLorentzVector();
recoB_vec.SetPtEtaPhiM(recoB_pt,recoB_eta,recoB_phi,recoB_mass)
if sum_track>0:
Et_ratio=recoB_vec.Et()/sum_track
else:
Et_ratio=0
Dphi=recoB_vec.DeltaPhi(sum_track_vec)
projB=recoB_vec*sum_track_vec
result=[Et_ratio,Dphi,projB,abs(trgmuon_vz-recoE1_vz),abs(trgmuon_vz-recoE2_vz),abs(trgmuon_vz-recoK_vz)]
return result
def JpsiKst_Angles(kaon, pion, mu1, mu2):
"""
paula
"""
P11p = kaon[1]
P12p = pion[1]
P21p = mu1[1]
P22p = mu2[1]
from ROOT import TLorentzVector, TVector3
def NProductV(alpha, Vect):
Vect1 = [Vect.x(), Vect.y(), Vect.z()]
for i in range(0, 3):
Vect1[i] = alpha * Vect[i]
return TVector3(Vect1[0], Vect1[1], Vect1[2])
p1 = TLorentzVector(P11p[0], P11p[1], P11p[2], kaon[0])
p2 = TLorentzVector(P12p[0], P12p[1], P12p[2], pion[0])
p3 = TLorentzVector(P21p[0], P21p[1], P21p[2], mu1[0])
p4 = TLorentzVector(P22p[0], P22p[1], P22p[2], mu2[0])
p1Psi = TLorentzVector(p1)
p12 = TLorentzVector(p1 + p2)
p34 = TLorentzVector(p3 + p4)
BK0S = TLorentzVector(p1 + p2 + p3 + p4)
BPsi = TLorentzVector(p1 + p2 + p3 + p4)
p1.Boost(NProductV(-1. / (p12.E()), p12.Vect()))
BK0S.Boost(NProductV(-1. / (p12.E()), p12.Vect()))
ThK = (p1.Vect()).Angle(-BK0S.Vect())
p1Psi.Boost(NProductV(-1. / (p34.E()), p34.Vect()))
BPsi.Boost(NProductV(-1. / (p34.E()), p34.Vect()))
xtr = BPsi.Vect().Unit()
ytr = (p1Psi.Vect().Unit() - NProductV(
(p1Psi.Vect().Unit()).Dot(xtr), xtr)).Unit()
ztr = xtr.Cross(ytr)
p3.Boost(NProductV(-1. / (p34.E()), p34.Vect()))
Thtr = ztr.Angle(p3.Vect())
Phitr = xtr.Angle(p3.Vect() - NProductV(ztr.Dot(p3.Vect()), ztr))
if (ytr.Angle(p3.Vect() - NProductV(ztr.Dot(p3.Vect()), ztr)) > pi / 2.):
Phitr = -Phitr
return ThK, Thtr, Phitr
def selection(event):
vect = TLorentzVector()
vect.SetPxPyPzE(event.GenJet[0], event.GenJet[1], event.GenJet[2],
event.GenJet[3])
if abs(vect.Eta()) > 0.5 or vect.Pt() < 50 or vect.Pt() > 70:
return False
vect.SetPxPyPzE(event.Jet[0], event.Jet[1], event.Jet[2], event.Jet[3])
if abs(vect.Eta()) > 0.5 or vect.Pt() < 50:
return False
else:
return True
def InvMass(obj1, obj2=0, obj3=0):
''' Invariant mass of two objects '''
if isinstance(obj1, list):
ob = TLorentzVector(obj1[0].p)
for o in obj1[1:]:
ob += o.p
return ob.M()
elif obj3 == 0:
return (obj1.p + obj2.p).M()
else:
return (obj1.p + obj2.p + obj3.p).M()
def returnPhi(self, theH, theL1, theL2, theL3, theL4):
h, l1, l2, l3, l4 = TLorentzVector(theH), TLorentzVector(
theL1), TLorentzVector(theL2), TLorentzVector(
theL3), TLorentzVector(theL4)
# Boost objects to the A rest frame
l1.Boost(-h.BoostVector())
l2.Boost(-h.BoostVector())
l3.Boost(-h.BoostVector())
l4.Boost(-h.BoostVector())
# Build unit vectors orthogonal to the decay planes
Zplane = l1.Vect().Cross(l2.Vect()) # L1 x L2
Hplane = l3.Vect().Cross(l4.Vect()) # B1 x B2
Zplane.SetMag(1.)
Hplane.SetMag(1.)
# Sign of Phi
z1 = l1 + l2
sgn = 1 if z1.Vect().Dot(Zplane.Cross(Hplane)) > 0 else -1
value = sgn * math.acos(Zplane.Dot(Hplane))
if value != value or math.isinf(value): return -5.
return value
def corrected_p4(self, mu, run):
p4 = TLorentzVector(mu.px(), mu.py(), mu.pz(), mu.energy())
if self.isMC:
self.corr.applyPtCorrection(p4, mu.charge())
self.corr.applyPtSmearing(p4, mu.charge(), self.isSync)
else:
corr = self.corrD if run >= 203773 else self.corrABC
corr.applyPtCorrection(p4, mu.charge())
## convert to the proper C++ class (but preserve the mass!)
return ROOT.reco.Muon.PolarLorentzVector(p4.Pt(), p4.Eta(), p4.Phi(),
mu.mass())
