PyROOT macro for the PyROOT-xAOD hands-on session of the LBL ATLAS Software Tutorial with xAOD

xAODPythonMacro.py

#!/usr/bin/env python

# Set up ROOT and RootCore
import ROOT
ROOT.gROOT.Macro('$ROOTCOREDIR/scripts/load_packages.C')

# Initialize the xAOD infrastructure
ROOT.xAOD.Init()

# Set up the input files (PDSF)
fileName = '/eliza18/atlas/atlasdata/atlaslocalgroupdisk/rucio/valid2/4c/15/AOD.01482225._000107.pool.root.1'
treeName = 'CollectionTree'
f = ROOT.TFile.Open(fileName)

# Make the "transient tree"
t = ROOT.xAOD.MakeTransientTree(f, treeName)

# Setup the output histogram file
f_out = ROOT.TFile('MyHistograms.root', 'recreate')
h_el_pt = ROOT.TH1F('el_pt', 'el_pt', 20, 0, 400)
h_jet_pt = ROOT.TH1F('jet_pt', 'jet_pt', 20, 0, 400)

# For convenience
GeV = 1000.

# Print some information
print('Number of input events: %s' % t.GetEntries())
for entry in xrange(t.GetEntries()):
    t.GetEntry(entry)
    s = 'Processing run #%i, event #%i' % (t.EventInfo.runNumber(),
                                           t.EventInfo.eventNumber())
    print(s)

    # loop over electron collection
    for el in t.ElectronCollection:
        s = '  Electron trackParticle eta = %g, phi = %g' 
        print(s % (el.trackParticle().eta(), el.trackParticle().phi()))

        # Print the difference between the electron pt and the electron's
        # trackParticle pt for all electrons with |eta| < 2.47
        if abs(el.eta()) < 2.47:
            diff = el.pt() - el.trackParticle().pt()
            print('    (pt-trkPt) = %g GeV' % (diff/GeV))

            # Fill histogram with electron pt
            h_el_pt.Fill(el.pt()/GeV)

    # Count and print the number of jets with pt > 20 GeV and |eta| < 2.5
    numJets = 0
    for jet in t.AntiKt4LCTopoJets:
        if jet.pt() > 20*GeV and abs(jet.eta()) < 2.5:
            print('  Jet pt = %g, eta = %g' % (jet.pt()/GeV, jet.eta()))
            numJets += 1

            # Fill histogram with jet pt
            h_jet_pt.Fill(jet.pt()/GeV)

    print('Number of jets found: %i' % numJets)

f_out.Write()
f_out.Close()