1#!/usr/bin/env python
2"""Illustration of the assignment of pythonic user information to PseudoJets:
3
4  - the script has a ParticleInfo class, to store information about particles
5
6  - the read_event(...) function creates a ParticleInfo for each
7    particle and assigns it to the corresponding PseudoJet, using the
8    PseudoJet.set_python_info(...) call.
9
10  - the print_jets(...) function gets the jet constituents, examines
11    the ParticleInfo for each one and uses it to determine additional
12    information about each jet. It uses the PseudoJet.python_info(...)
13    call.
14
15For this script to work, make sure that the installation location for
16the fastjet python module (cf. fastjet-config --pythonpath) is
17included in your PYTHONPATH environment variable.
18
19"""
20from __future__ import print_function
21
22import fastjet as fj
23import gzip
24
25def main():
26
27    # get the banner out of the way early on
28    fj.ClusterSequence.print_banner()
29    print()
30
31    # set up our jet definition and a jet selector
32    jet_def = fj.JetDefinition(fj.antikt_algorithm, 0.4)
33    selector = fj.SelectorPtMin(15.0) & fj.SelectorAbsRapMax(4.5)
34    print("jet definition is:",jet_def)
35    print("jet selector is:", selector,"\n")
36
37    filename = '../data/Pythia-dijet-ptmin100-lhc-pileup-1ev.dat'
38    f = open(filename,'r')
39    #filename = '/Users/gsalam/work/fastjet/data/Pythia-PtMin50-LHC-10kev.dat.gz'
40    #f = gzip.GzipFile(filename,'rb')
41
42    # get the event
43    iev = 0
44    while True:
45        event = read_event(f)
46        iev += 1
47        if (len(event) == 0): break
48
49        # cluster it
50        jets = selector(jet_def(event))
51
52        # print some info
53        npileup = 0
54        for p in event:
55            if (p.python_info().subevent_index > 0): npileup += 1
56        print("Event {0} has {1} particles (of which {2} from pileup)".format(iev, len(event), npileup))
57        print_jets(jets)
58
59#----------------------------------------------------------------------
60def print_jets(jets):
61    print("{0:>5s} {1:>10s} {2:>10s} {3:>10s} {4:>12s} {5:>12s} {6:>12s}".format(
62        "jet #", "pt", "rap", "phi", "primary pt", "N particles", "N photons"))
63
64    for ijet in range(len(jets)):
65        jet = jets[ijet]
66
67        # figure out how many particles and how many photons the jet contains
68        # and how much pt comes from the primary vertex
69        constituents = jet.constituents()
70        nphotons = 0
71        primary_pt = 0
72        for c in constituents:
73            if (c.python_info().pdg_id == 22): nphotons += 1
74            if (c.python_info().subevent_index <= 0): primary_pt += c.pt()
75
76        print("{0:5d} {1:10.3f} {2:10.4f} {3:10.4f} {4:10.3f} {5:12d} {6:12d}".format(
77            ijet, jet.pt(), jet.rap(), jet.phi(), primary_pt, len(constituents), nphotons))
78
79#----------------------------------------------------------------------
80class ParticleInfo(object):
81    """illustrative class for use in assigning pythonic user information
82    to a PseudoJet.
83    """
84    def __init__(self, subevent_index, particle_index, pdg_id=0):
85        self.subevent_index = subevent_index
86        self.particle_index = particle_index
87        self.pdg_id = pdg_id
88
89    def set_pdg_id(self, pdg_id):
90        self.pdg_id = pdg_id
91
92    def __str__(self):
93        return "subevent_index={0}, particle_index={1}, pdg_id={2}".format(
94            self.subevent_index, self.particle_index, self.pdg_id)
95
96#----------------------------------------------------------------------
97event_index = 0
98def read_event(file_or_filename):
99    """
100Routine that can take either an existing opened file object, or a
101filename (which it will open itself) and then reads an event from that
102file. An event is deemed to end when the file comes to an end or when
103the reader encounters the string "#END".
104
105The event is converted to a python list of PseudoJets
106    """
107
108    # open the file if necessary
109    if (isinstance(file_or_filename,str)) : f = open(file_or_filename, 'r')
110    else                                  : f = file_or_filename
111
112    # create an empty list
113    event = []
114    particle_index = 0
115    subevent_index = -1
116    while True:
117        line = f.readline()
118        # exit if the file has come to an end
119        if   (not line): break
120        # or if we reach the string "#END"
121        if   (len(line) >=4 and line[0:4] == '#END'): break
122
123        # #SUBSTART means that we start a new subevent
124        if (len(line)>=9 and line[0:9] == '#SUBSTART'):
125            subevent_index += 1
126            continue
127
128        # ignore comment lines or empty lines
129        elif (line[0] == '#' or len(line) <= 1): continue
130
131        # assume we have a good line and split it into px, py, pz, E
132        p = line.split()
133        # create the PseudoJet
134        pj = fj.PseudoJet(float(p[0]),float(p[1]),float(p[2]),float(p[3]))
135
136        # then create its user info
137        info = ParticleInfo(subevent_index, particle_index)
138        # optionally with a pdg_id (if it was in the file)
139        if (len(p) > 4): info.set_pdg_id(int(p[4]))
140
141        # finally assign the user info and append the particle the event
142        pj.set_python_info(info)
143        event.append(pj);
144
145        # don't forget to increment it
146        particle_index += 1
147
148    global event_index
149    event_index += 1
150    return event
151
152if __name__ == '__main__':
153    main()
154