Skip to content

Instantly share code, notes, and snippets.

@mverzett

mverzett/tree2better_tree.py

Created June 27, 2016 10:10
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save mverzett/f38f5522f964fb5e9e7b48680d8fd028 to your computer and use it in GitHub Desktop.
Save mverzett/f38f5522f964fb5e9e7b48680d8fd028 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
tree2better_tree.py
from rootpy.io import root_open
from argparse import ArgumentParser
from rootpy.tree import Tree
from rootpy.vector import LorentzVector
import ROOT
from itertools import combinations
ROOT.gROOT.SetBatch(True)
from pdb import set_trace
import rootpy
import logging
rootpy.log['/'].setLevel(logging.FATAL)
import numpy as np
parser = ArgumentParser()
parser.add_argument('infile')
parser.add_argument('out')
parser.add_argument('--limit', type=int, default=-1)
args = parser.parse_args()
infile = root_open(args.infile)
outfile = root_open(args.out, 'w')
itree = infile.tree
otree = Tree('tree')
branch_names = [i.GetName() for i in itree.branches] # get old branches
new_branches = {i : 'F' for i in branch_names} #get new branches
new_branches.update({
'mbb_err' : 'F',
'mjj_err' : 'F',
'mbjj_err' : 'F',
'mlv_err' : 'F',
'j1_btag' : 'F',
'j2_btag' : 'F',
'b1_btag' : 'F',
'b2_btag' : 'F',
'jj_dr' : 'F',
'bb_dr' : 'F',
#TTbar hypothesis
'Wlb_dr' : 'F',
'Whb_dr' : 'F',
'tt_dr' : 'F',
'Wl_pt' : 'F',
'Wh_pt' : 'F',
'th_pt' : 'F',
'tl_pt' : 'F',
'tt_pt' : 'F',
'Wl_y' : 'F',
'Wh_y' : 'F',
'th_y' : 'F',
'tl_y' : 'F',
'tt_y' : 'F',
'lj1_deltaphi' : 'F',
'lj2_deltaphi' : 'F',
'l_costheta_helicity' : 'F',
'j1_costheta_helicity' : 'F',
'j2_costheta_helicity' : 'F',
'bl_costheta_helicity' : 'F',
'bh_costheta_helicity' : 'F',
#higgs hypothesis
'higgs_pt' : 'F',
'higgs_y' : 'F',
'HiggsC_pt' : 'F',
'HiggsC_y' : 'F',
'HiggsC_dr' : 'F',
'HiggsN_pt' : 'F',
'HiggsN_y' : 'F',
'HiggsN_dr' : 'F',
'higgsb_costheta' : 'F',
'HiggsChiggs_costheta' : 'F',
'HiggsNHiggsC_costheta' : 'F',
'HiggsCj1_costheta' : 'F',
'HiggsCj2_costheta' : 'F',
'HiggsCl_costheta' : 'F',
'HiggsCb_costheta' : 'F',
})
otree.create_branches(new_branches)
def costh_hel(prod, gran, ggm):
'cos theta in the helicity frame (final obj, gran-mother, gran-gran-mother)'
p = prod.Clone()
gran_boost = gran.BoostVector()*(-1)
p.Boost(gran_boost)
p3D = p.Vect().Unit()
ttcm = ggm.BoostVector()*(-1)
top = gran.Clone()
top.Boost(ttcm)
return top.Vect().Unit().Dot(p3D)
for idx, entry in enumerate(itree):
if idx % 1000 == 0: print idx
if args.limit > 0 and args.limit < idx: break
#copy old entries
for name in branch_names:
setattr(
otree, name,
getattr(entry, name)
)
#build LVs
lep = LorentzVector()
nu_2= LorentzVector()
nu_1= LorentzVector()
#start with lepton (easy)
lep.SetPtEtaPhiM(entry.lepton_pt, entry.lepton_eta, entry.lepton_phi, 0.)
#use mass_lv to get nu eta (two-fold ambiguity)
deta = ROOT.TMath.ACosH(
entry.m_lv**2/(2*entry.lepton_pt*entry.mem_pt)+
ROOT.TMath.Cos(entry.lepton_phi-entry.mem_phi)
)
if np.isnan(deta):
nu_1.SetPtEtaPhiM(entry.mem_pt, entry.lepton_eta, entry.mem_phi, 0.)
nu_2.SetPtEtaPhiM(entry.mem_pt, entry.lepton_eta, entry.mem_phi, 0.)
otree.mlv_err = (entry.m_lv-(lep+nu).M())/entry.m_lv
else:
nu_1.SetPtEtaPhiM(entry.mem_pt, entry.lepton_eta+deta, entry.mem_phi, 0.)
nu_2.SetPtEtaPhiM(entry.mem_pt, entry.lepton_eta-deta, entry.mem_phi, 0.)
otree.mlv_err = 0.
#find bjets
jets = [LorentzVector() for i in range(4)]
jets[0].SetPtEtaPhiM(entry.jet1_pt, entry.jet1_eta, entry.jet1_phi, 0.)
jets[1].SetPtEtaPhiM(entry.jet2_pt, entry.jet2_eta, entry.jet2_phi, 0.)
jets[2].SetPtEtaPhiM(entry.jet3_pt, entry.jet3_eta, entry.jet3_phi, 0.)
jets[3].SetPtEtaPhiM(entry.jet4_pt, entry.jet4_eta, entry.jet4_phi, 0.)
btags = [entry.jet1_btag, entry.jet2_btag, entry.jet3_btag, entry.jet4_btag]
idx_combo = list(combinations(range(4), 2))
bmax = max(btags[i]+btags[j] for i, j in idx_combo)
best = 9999999999
b1 = None
b2 = None
b_idx = None
for i, j in combinations(range(4), 2):
#print btags[i], btags[j], bmax
if int(btags[i]+btags[j]) < int(bmax): continue
mass = (jets[i]+jets[j]).M()
delta = abs(mass - entry.m_bb)
if best > delta:
best = delta
if jets[i].Pt() > jets[j].Pt():
b1 = jets[i]
b2 = jets[j]
otree.b1_btag = btags[i]
otree.b2_btag = btags[j]
else:
b1 = jets[j]
b2 = jets[i]
otree.b1_btag = btags[j]
otree.b2_btag = btags[i]
b_idx = (i,j)
#find other two
idxs = set(range(4))
idxs.remove(b_idx[0])
idxs.remove(b_idx[1])
i, j = tuple(idxs)
j1, j2 = None, None
if jets[i].Pt() > jets[j].Pt():
j1 = jets[i]
j2 = jets[j]
otree.j1_btag = btags[i]
otree.j2_btag = btags[j]
else:
j1 = jets[j]
j2 = jets[i]
otree.j1_btag = btags[j]
otree.j2_btag = btags[i]
#OK, time for science!
#Errors
otree.mbb_err = (entry.m_bb - (b1+b2).M())/entry.m_bb
otree.mjj_err = (entry.m_jj - (j1+j2).M())/entry.m_jj
#mbjj_err =
#openings
otree.jj_dr = j1.DeltaR(j2)
otree.bb_dr = b1.DeltaR(b2)
#pairs
all_the_stuff_I_know = lep+j1+j2+b1+b2
d1 = abs(entry.m_wwbb - (all_the_stuff_I_know+nu_1).M())
d2 = abs(entry.m_wwbb - (all_the_stuff_I_know+nu_2).M())
if d1 < d2:
Wlep = lep+nu_1
nu = nu_1
else:
Wlep = lep+nu_2
nu = nu_2
Whad = j1+j2
## tt hypothesis
#choose thad hypothesis
th1 = (j1+j2+b1)
th2 = (j1+j2+b2)
d1 = abs(173 - th1.M())
d2 = abs(173 - th2.M())
thad = th1 if abs(d1) < abs(d2) else th2
bhad = b1 if abs(d1) < abs(d2) else b2
blep = b2 if abs(d1) < abs(d2) else b1
tlep = Wlep+blep
otree.mbjj_err = (entry.m_jjj - thad.M())/entry.m_jjj
otree.Wlb_dr = Wlep.DeltaR(blep)
otree.Whb_dr = Whad.DeltaR(bhad)
otree.tt_dr = thad.DeltaR(tlep)
tt = thad+tlep
otree.Wl_pt = Wlep.Pt()
otree.Wh_pt = Whad.Pt()
otree.th_pt = thad.Pt()
otree.tl_pt = tlep.Pt()
otree.tt_pt = tt.Pt()
otree.Wl_y = Wlep.Rapidity()
otree.Wh_y = Whad.Rapidity()
otree.th_y = thad.Rapidity()
otree.tl_y = tlep.Rapidity()
otree.tt_y = tt.Rapidity()
otree.lj1_deltaphi = ROOT.TMath.Cos(lep.DeltaPhi(j1))
otree.lj2_deltaphi = ROOT.TMath.Cos(lep.DeltaPhi(j1))
otree.l_costheta_helicity = costh_hel(lep, tlep, tt)
otree.j1_costheta_helicity = costh_hel(j1 , thad, tt)
otree.j2_costheta_helicity = costh_hel(j2 , thad, tt)
otree.bl_costheta_helicity = costh_hel(blep, tlep, tt)
otree.bh_costheta_helicity = costh_hel(bhad, thad, tt)
## H hypothesis
higgs = b1+b2
if abs((higgs+Wlep).M()-325)<abs((higgs+Whad).M()-325):
HiggsC = higgs+Wlep
HiggsN = HiggsC+Whad
otree.HiggsC_dr = higgs.DeltaR(Wlep)
otree.HiggsN_dr = HiggsC.DeltaR(Whad)
otree.HiggsCj1_costheta = -1.1
otree.HiggsCj2_costheta = -1.1
otree.HiggsCl_costheta = costh_hel(lep, HiggsC, HiggsN)
else:
HiggsC = higgs+Whad
HiggsN = HiggsC+Wlep
otree.HiggsC_dr = higgs.DeltaR(Whad)
otree.HiggsN_dr = HiggsC.DeltaR(lep)
otree.HiggsCj1_costheta
otree.HiggsCj2_costheta
otree.HiggsCl_costheta
otree.HiggsCb_costheta
otree.HiggsCj1_costheta = costh_hel(j1, HiggsC, HiggsN)
otree.HiggsCj2_costheta = costh_hel(j1, HiggsC, HiggsN)
otree.HiggsCl_costheta = -1.1
otree.higgs_pt = higgs.Pt()
otree.higgs_y = higgs.Rapidity()
otree.HiggsC_pt = HiggsC.Pt()
otree.HiggsC_y = HiggsC.Rapidity()
otree.HiggsN_pt = HiggsN.Pt()
otree.HiggsN_y = HiggsN.Rapidity()
#angles
otree.higgsb_costheta = costh_hel(b1, higgs, HiggsC)
otree.HiggsChiggs_costheta = costh_hel(higgs, HiggsC, HiggsN)
#otree.HiggsNHiggsC_costheta = costh_hel(HiggsC, HiggsN, )
otree.HiggsCb_costheta = costh_hel(b1, HiggsC, HiggsN)
#special case
p = HiggsC.Clone()
gran_boost = HiggsN.BoostVector()*(-1)
p.Boost(gran_boost)
p3D = p.Vect().Unit()
otree.HiggsNHiggsC_costheta = HiggsN.Vect().Unit().Dot(p3D)
otree.fill()
otree.Write()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment