-
Notifications
You must be signed in to change notification settings - Fork 1
/
Delphes2SA.py
executable file
·245 lines (200 loc) · 7.89 KB
/
Delphes2SA.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
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
#!/usr/bin/env python3
# Example for converting Delphes output to SimpleAnalysis slimmed format
# lepton id/isolation and flavour tagging levels are not implemented
#
# Assumes Delphes and ROOT is setup already
#
from __future__ import annotations
import argparse
import os
import sys
from array import array
import ROOT
if len(sys.argv) < 3:
print(" Usage: Delphes2SA.py <input Delphes file> <Output SA file>")
sys.exit(1)
# Make sure that the interpreter points to the DELPHES classes in order to read through DELPHES events.
# may need to run something like: export ROOT_INCLUDE_PATH=$ROOT_INCLUDE_PATH:/home/mhance/Delphes/Delphes-3.4.1/:/home/mhance/Delphes/Delphes-3.4.1/external/
delphespath = os.environ.get("DELPHES_PATH")
ROOT.gInterpreter.Declare('#include "%s/classes/DelphesClasses.h"' % delphespath)
ROOT.gInterpreter.Declare(
'#include "%s/external/ExRootAnalysis/ExRootTreeReader.h"' % delphespath
)
ROOT.gSystem.Load("%s/libDelphes.so" % delphespath)
# ROOT.gSystem.Load("libDelphes")
try:
ROOT.gInterpreter.Declare('#include "classes/DelphesClasses.h"')
ROOT.gInterpreter.Declare('#include "external/ExRootAnalysis/ExRootTreeReader.h"')
except Exception:
pass
class NtupleVar:
def __init__(self, name, tree, floatVar=False):
if floatVar:
rootVar = array("f", [0])
tree.Branch(name, rootVar, name + "/F")
else:
rootVar = array("i", [0])
tree.Branch(name, rootVar, name + "/I")
self.var = rootVar
def Set(self, value):
self.var[0] = value
outVectors = []
class NtupleVector:
def __init__(self, name, tree, floatVar=False):
rootVar = ROOT.std.vector("float")() if floatVar else ROOT.std.vector("int")()
tree.Branch(name, rootVar)
self.var = rootVar
outVectors.append(rootVar)
def Add(self, value):
self.var.push_back(value)
class ObjectVector:
def __init__(self, name, tree, storeMass=False):
self.pt = NtupleVector(name + "_pt", tree, True)
self.eta = NtupleVector(name + "_eta", tree, True)
self.phi = NtupleVector(name + "_phi", tree, True)
self.charge = NtupleVector(name + "_charge", tree)
self.objID = NtupleVector(name + "_id", tree)
self.motherID = NtupleVector(name + "_motherID", tree)
self.storeMass = storeMass
if storeMass:
self.mass = NtupleVector(name + "_m", tree, True)
def Add(self, obj, objID=0x7FFFFFFF, charge=999):
self.pt.Add(obj.PT)
self.eta.Add(obj.Eta)
self.phi.Add(obj.Phi)
if charge == 999:
self.charge.Add(obj.Charge)
else:
self.charge.Add(charge)
self.objID.Add(objID)
self.motherID.Add(0) # FIXME truth pointing not implemented
if self.storeMass:
self.mass.Add(obj.Mass)
inputFile = sys.argv[1]
outputFile = sys.argv[2]
parser = argparse.ArgumentParser()
parser.add_argument("--input", action="store")
parser.add_argument("--output", action="store")
parser.add_argument(
"--lumi", action="store", default=1000.0
) # 1 fb-1. not actually used.
parser.add_argument("--XS", action="store", default=1.0) # fb
parser.add_argument("--debug", action="store_true") # not actually used.
args = parser.parse_args()
if args.input is not None:
inputFile = args.input
if args.output is not None:
outputFile = args.output
# Create chain of root trees
chain = ROOT.TChain("Delphes")
chain.Add(inputFile)
# Create object of class ExRootTreeReader
treeReader = ROOT.ExRootTreeReader(chain)
numberOfEntries = treeReader.GetEntries()
# Get pointers to branches used in this analysis
branchEvent = treeReader.UseBranch("Event")
branchMET = treeReader.UseBranch("MissingET")
branchHT = treeReader.UseBranch("ScalarHT")
branchPhoton = treeReader.UseBranch("Photon")
branchElectron = treeReader.UseBranch("Electron")
branchMuon = treeReader.UseBranch("Muon")
branchJet = treeReader.UseBranch("Jet")
branchFatJet = treeReader.UseBranch("FatJet")
# Output file and tree
outFH = ROOT.TFile(outputFile, "RECREATE")
outTree = ROOT.TTree("ntuple", "Simple Analysis slim format from delphes")
outTree.SetDirectory(outFH)
EventNumber = NtupleVar("Event", outTree)
mcChannel = NtupleVar("mcChannel", outTree)
mcVetoCode = NtupleVar("mcVetoCode", outTree)
susyChannel = NtupleVar("susyChannel", outTree)
mcWeights = NtupleVector("mcWeights", outTree, True)
genMET = NtupleVar("genMET", outTree, True)
genHT = NtupleVar("genHT", outTree, True)
pdf_id1 = NtupleVar("pdf_id1", outTree)
pdf_x1 = NtupleVar("pdf_x1", outTree, True)
pdf_pdf1 = NtupleVar("pdf_pdf1", outTree, True)
pdf_id2 = NtupleVar("pdf_id2", outTree)
pdf_x2 = NtupleVar("pdf_x2", outTree, True)
pdf_pdf2 = NtupleVar("pdf_pdf2", outTree, True)
pdf_scale = NtupleVar("pdf_scale", outTree, True)
sumet = NtupleVar("sumet", outTree, True)
met_pt = NtupleVar("met_pt", outTree, True)
met_phi = NtupleVar("met_phi", outTree, True)
electrons = ObjectVector("el", outTree)
muons = ObjectVector("mu", outTree)
taus = ObjectVector("tau", outTree)
photons = ObjectVector("ph", outTree)
jets = ObjectVector("jet", outTree, True)
fatjets = ObjectVector("fatjet", outTree, True)
# not supporting hard-scatter truth record for now
# bit of a hack, but should work:
sumweights = 0.0
for entry in range(numberOfEntries):
treeReader.ReadEntry(entry)
sumweights += float(branchEvent.At(0).Weight)
# will normalize everything to XS
print(f"Using cross section {args.XS}")
weightscale = float(args.XS) / sumweights
# Loop over all events
print(f"Looping over {numberOfEntries} events")
for entry in range(numberOfEntries):
for vec in outVectors:
vec.clear()
# Load selected branches with data from specified event
treeReader.ReadEntry(entry)
# Fill in event info (some is left at default 0)
EventNumber.Set(branchEvent.At(0).Number)
mcChannel.Set(branchEvent.At(0).ProcessID)
mcWeights.Add(branchEvent.At(0).Weight * weightscale)
# FIXME: add PDF info etc. if available
metvec = ROOT.TLorentzVector()
metvec.SetPtEtaPhiM(branchMET.At(0).MET, 0, branchMET.At(0).Phi, 0)
# object id not implemented for now
for idx in range(branchElectron.GetEntries()):
electrons.Add(branchElectron.At(idx))
for idx in range(branchMuon.GetEntries()):
muons.Add(branchMuon.At(idx))
muonvec = ROOT.TLorentzVector()
muonvec.SetPtEtaPhiM(
branchMuon.At(idx).PT, branchMuon.At(idx).Eta, branchMuon.At(idx).Phi, 0.1
)
# remove muons from MET. MET is only built from calorimeter info in Delphes,
# so events with muons will have the muon showing up as MET, which artificially
# increases the MET.
#
# see https://arxiv.org/abs/0903.2225 for an explicit statement, and
# http://arxiv.org/abs/1307.6346 for a more recent reference that doesn't contradict this.
#
metvec = metvec - muonvec
for idx in range(branchPhoton.GetEntries()):
photons.Add(branchPhoton.At(idx), charge=0)
# MET/HT
sumet.Set(branchHT.At(0).HT)
met_pt.Set(metvec.Pt())
met_phi.Set(metvec.Phi())
# If event contains at least 1 jet
for idx in range(branchJet.GetEntries()):
jet = branchJet.At(idx)
jetID = 0x000FDF00 # flag as good jet
if jet.BTag:
jetID |= 0x00F000FF # flags as bjet
jets.Add(jet, jetID)
if jet.TauTag:
tauID = 0xFF
if jet.NCharged == 1:
tauID |= 1 << 10
if jet.NCharged == 3:
tauID |= 1 << 11
taus.Add(jet, tauID)
if branchFatJet:
for idx in range(branchFatJet.GetEntries()):
jet = branchFatJet.At(idx)
jetID = 0x000FDF00 # flag as good jet
if jet.BTag:
jetID |= 0x00F000FF # flags as bjet
fatjets.Add(jet, jetID)
outTree.Fill()
print(f"wrote {outTree.GetEntriesFast()} entries to the tree.")
outFH.Write()
outFH.Close()