snowmass-qstar/genParticleStudy.py at master · andersonjacob/snowmass-qstar · GitHub

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
from optparse import OptionParser

parser = OptionParser()
(opts, args) = parser.parse_args()


import pyroot_logon

from ROOT import gSystem

from DataFormats.FWLite import Events, Handle

from ROOT import *
from array import array

events = Events (args)
handle = Handle('vector<reco::GenParticle>')

label = ('genParticles')

def printParticleDecayTree(mom, prefix = '', statusCutoff = 1):
    if mom.status() >= statusCutoff:
        print prefix+'id:',mom.pdgId(),'status:',mom.status(),\
            'nDau:',mom.numberOfDaughters(),'(pt,eta,phi,m):',\
            '(%.1f,%.1f,%.1f,%.0f)' % (mom.pt(), mom.eta(), mom.phi(),
                                       mom.mass())
        for dau in range(0, mom.numberOfDaughters()):
            daughter = mom.daughter(dau)
            if daughter.status() >= statusCutoff:
                printParticleDecayTree(daughter, '  ' + prefix, statusCutoff)

cuts = [
    'Jet[0].PT > 10000',
    'Jet[1].PT > 10000',
    'Jet[0].Eta < 1',
    'Jet[1].Eta < 1',
    ]

def passSelection(event):
    for cut in cuts:
        passing = eval("event.%s" % cut)
        # print cut,passing
        if not passing:
            return False
    return True


f = TFile.Open(args[0])

Delphes = f.Get("Delphes")

fout = TFile(args[0].replace('.root', '_histograms.root'), "recreate")

bins = [20.]
while (bins[-1] < 100.):
    bins.append(bins[-1] + int(bins[-1]*0.05))

binEdges = array('d', bins)
print bins

massSpectrum = TH1F("massSpectrum", "dijet mass", len(bins)-1, binEdges)

for (eventN, event) in enumerate(events):

    if eventN % 5000 == 0:
        print "processing event %i" % eventN

    event.getByLabel(label, handle)
    genParts = handle.product()

    # printParticleDecayTree(genParts[0], "", 3)

    foundParticle = False
    for genParticle in genParts:
        if abs(genParticle.pdgId()) in [4000001, 4000002]:
            # printParticleDecayTree(genParticle, "", 3)
            massSpectrum.Fill(genParticle.mass()/1000.)
            foundParticle = True
            break

    if not foundParticle:
        printParticleDecayTree(genParts[0], "", 3)
    # massSpectrum.Fill(dijet.M()/1000.)

    # if eventN > 9:
    #     break

fout.Write()

massSpectrum.Scale(1., "width")
massSpectrum.Draw()