forked from mstamenk/hhh-analysis-framework
-
Notifications
You must be signed in to change notification settings - Fork 0
/
hhh_variables.py
137 lines (103 loc) · 4.14 KB
/
hhh_variables.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
# Script to add variables for spanet
import os, ROOT
import itertools
ROOT.ROOT.EnableImplicitMT()
computem = '''
float computemH(int type, float j1_pt, float j1_eta, float j1_phi, float j1_mass, float j1_breg ,float j2_pt, float j2_eta, float j2_phi, float j2_mass, float j2_breg){
TLorentzVector j1;
TLorentzVector j2;
j1.SetPtEtaPhiM(j1_pt*j1_breg, j1_eta, j1_phi, j1_mass);
j2.SetPtEtaPhiM(j2_pt*j2_breg, j2_eta, j2_phi, j2_mass);
if (type == 0) return (j1+j2).M();
else if (type == 1) return (j1+j2).Pt();
else if (type == 2) return (j1+j2).Eta();
else if (type == 3) return (j1+j2).Phi();
else if (type == 4) return j1.DeltaR(j2);
else return 0;
}
'''
ROOT.gInterpreter.Declare(computem)
# 12 34 56 78 910
# 13
unique = []
for j in ['jet1','jet2','jet3','jet4','jet5','jet6','jet7','jet8','jet9','jet10']:
for k in ['jet1','jet2','jet3','jet4','jet5','jet6','jet7','jet8','jet9','jet10']:
if j in k and k in j: continue
perm = [j,k]
perm_inv = [k,j]
if perm_inv in unique: continue
else: unique.append(perm)
#print(unique)
#print(len(unique))
def add_h1_mass_corrected(df):
j1 = 'bcand1'
j2 = 'bcand2'
variables = ['%sPt'%j1,'%sEta'%j1,'%sPhi'%j1,'%sMass'%j1,'%sbRegCorr'%j1,'%sPt'%j2,'%sEta'%j2,'%sPhi'%j2,'% sMass'%j2,'%sbRegCorr'%j2]
df = df.Define('h1_breg_mass','computemH(0,%s)'%','.join(variables))
return df
def add_hhh_variables(df):
masses = []
pts = []
etas = []
phis = []
drs = []
for i in range(len(unique)):
perm = unique[i]
j1,j2 = perm
variables = ['%sPt'%j1,'%sEta'%j1,'%sPhi'%j1,'%sMass'%j1,'%sbRegCorr'%j1,'%sPt'%j2,'%sEta'%j2,'%sPhi'%j2,'%sMass'%j2,'%sbRegCorr'%j2]
mass = 'mass%s%s'%(j1,j2)
pt = 'pt%s%s'%(j1,j2)
eta = 'eta%s%s'%(j1,j2)
phi = 'phi%s%s'%(j1,j2)
dr = 'dr%s%s'%(j1,j2)
df = df.Define(mass, 'computemH(0,%s)'%','.join(variables))
df = df.Define(pt, 'computemH(1,%s)'%','.join(variables))
df = df.Define(eta, 'computemH(2,%s)'%','.join(variables))
df = df.Define(phi, 'computemH(3,%s)'%','.join(variables))
df = df.Define(dr, 'computemH(4,%s)'%','.join(variables))
masses.append(mass)
pts.append(pt)
etas.append(eta)
phis.append(phi)
drs.append(dr)
return df,masses,pts,etas,phis,drs
def add_hhh_variables_resolved(df):
masses = []
pts = []
etas = []
phis = []
drs = []
for i in range(len(unique)):
perm = unique[i]
j1,j2 = perm
variables = ['%sPt'%j1,'%sEta'%j1,'%sPhi'%j1,'%sMass'%j1,'%sbRegCorr'%j1,'%sPt'%j2,'%sEta'%j2,'%sPhi'%j2,'%sMass'%j2,'%sbRegCorr'%j2]
mass = 'mass%s%s'%(j1,j2)
pt = 'pt%s%s'%(j1,j2)
eta = 'eta%s%s'%(j1,j2)
phi = 'phi%s%s'%(j1,j2)
dr = 'dr%s%s'%(j1,j2)
#df = df.Define(mass, 'computemH(0,%s)/h1_t3_mass'%','.join(variables))
df = df.Define(mass, 'computemH(0,%s)'%','.join(variables))
#df = df.Define(pt, 'computemH(1,%s)/h1_t3_pt'%','.join(variables))
df = df.Define(pt, 'computemH(1,%s)'%','.join(variables))
df = df.Define(eta, 'computemH(2,%s)'%','.join(variables))
df = df.Define(phi, 'computemH(3,%s)'%','.join(variables))
df = df.Define(dr, 'computemH(4,%s)'%','.join(variables))
masses.append(mass)
pts.append(pt)
etas.append(eta)
phis.append(phi)
drs.append(dr)
return df,masses,pts,etas,phis,drs
#df.Snapshot('Events', 'GluGluToHHHTo6B_SM_spanet.root')
if __name__ == '__main__':
f_in = 'GluGluToHHHTo6B_SM'
path = '/isilon/data/users/mstamenk/eos-triple-h/v25/mva-inputs-HLT-boosted-bdt-v25-inclusive-loose-wp-0ptag-2018/inclusive/'
df = ROOT.RDataFrame('Events',path + '/' + f_in + '.root')
#j1 = 'jet1'
#j2 = 'jet2'
#variables = ['%sPt'%j1,'%sEta'%j1,'%sPhi'%j1,'%sMass'%j1,'%sPt'%j2,'%sEta'%j2,'%sPhi'%j2,'%sMass'%j2]
#df = df.Define('mH12', 'computemH(0,%s)'%','.join(variables))
df,masses,pts,etas,phis = add_hhh_variables(df)
h = df.Histo1D('massjet1jet2')
h.Draw()