example.py

#!/usr/bin/env python

import ROOT as r
import numpy as np

class FatJet:
    def __init__(self, tree, idx):
        # make an object representing the jet at index `idx` by pulling
        # interesting attributes out of the tree.
        self.m = tree.fjet_m[idx]
        self.pt = tree.fjet_pt[idx]
        self.pt_untrim = tree.ut_asso_fjet_pt[idx]
        self.tau21 = tree.fjet_tau21[idx]
        self.tau21_untrim = tree.ut_asso_fjet_tau21[idx]

if __name__ == "__main__":
    from argparse import ArgumentParser

    # read the command line arguments
    parser = ArgumentParser(description="Loop through some ntuple and read stuff")
    parser.add_argument("--tree-name", default="zpj", help="The name of the tree to read.")
    parser.add_argument("input_file", help="the input ntuple")
    args = parser.parse_args()

    # open the file and get the tree
    print "Loading tree '%s' from %s" % (args.tree_name, args.input_file)
    input_file = r.TFile(args.input_file)
    t = input_file.Get(args.tree_name)

    total_events = t.GetEntries()
    print "There are %d events in the tree." % total_events

    # loop thru the events and do stuff
    for i,evt in enumerate(t):
        if i%10==0:
            print "processing event %d" % i
        
        # here's one way to do it...
        # simultaneously loop thru the list of fatjets pT/mass/tau21:
        for (pt, mass, tau21) in zip(evt.fjet_pt, evt.fjet_m, evt.fjet_tau21):
            if pt<400e3: continue
            if mass < 20e3: continue

            print "Event %d: found some jet w/ like high pt and stuff: (%g, %g, %g)" % (i, pt, mass, tau21)

        # but IMO a clearer way is to build up a list of jet "objects".
        # see the very simple class definition at the top of this file.
        n_fatjets = len(evt.fjet_m)
        fatjets_all = [FatJet(t, idx) for idx in xrange(n_fatjets)]

        # calculate the tau21' values for each fatjet
        for jet in fatjets_all:
            # note that in python we can arbitrarilly assign new attributes to an object
            jet.tau21_prime = jet.tau21_untrim #+ 0.04 * np.log(jet.m**2 / jet.pt_untrim / 1e3)

        # now slim down the list of fatjets by making some selection requirements on pT and mass
        selected = fatjets_all
        selected = filter(lambda j: j.m>20e3, selected)
        selected = filter(lambda j: j.pt>200e3, selected)

        print "(Event %d) selected %d of %d fatjets!" % (i, len(selected), n_fatjets)
        for iselected, jet in enumerate(selected):
            # print the selected jets and compare their tau21 and tau21' values
            print "\t%d: (pT, tau21, tau21') = (%d, %.2f, %.2f)" % (iselected, jet.pt/1e3, jet.tau21, jet.tau21_prime)