Different ways to access TTree elements, some slow, some fast. Note cProfile doesn't work great here.

performance_test.py

#!/usr/bin/env python

"""
Example ways to access tree elements, to test relative performance
"""


import ROOT
from array import array
# import cProfile


ROOT.PyConfig.IgnoreCommandLineOptions = True
ROOT.gStyle.SetOptStat(0)
ROOT.gROOT.SetBatch(1)
ROOT.gStyle.SetOptFit(1111)
ROOT.TH1.SetDefaultSumw2(True)
ROOT.gStyle.SetPalette(55)
ROOT.gStyle.SetNumberContours(100)
ROOT.gErrorIgnoreLevel = 1 # turn off the printing output


def setup_branches(tree, branch_names):
    return map(lambda b: tree.GetBranch(b), branch_names)


def activate_branches(tree, branch_names):
    tree.SetBranchStatus("*", 0)
    map(lambda b: tree.SetBranchStatus(b, 1), branch_names)


def tree_plotter_branches(tree, hist, num_entries):
    # pr = cProfile.Profile()
    # print 'branches'
    # pr.enable()

    branch_names = ["pt", "ptRef", "rsp"]
    # branches = map(lambda b: tree.GetBranch(b), branch_names)
    branches = [tree.GetBranch(b) for b in branch_names]

    for i in range(num_entries):
        bytes_read = [b.GetEntry(i) for b in branches]
        # bytes_read = map(lambda b: b.GetEntry(i), branches)
        if tree.rsp < 2:
            hist.Fill(tree.pt, tree.ptRef)

    # pr.disable()
    # pr.dump_stats('profile_branches')

def tree_plotter_branches_activate(tree, hist, num_entries):
    # pr = cProfile.Profile()
    # print 'branches + activate'
    # pr.enable()

    branch_names = ["pt", "ptRef", "rsp"]
    # branches = map(lambda b: tree.GetBranch(b), branch_names)
    branches = [tree.GetBranch(b) for b in branch_names]

    tree.SetBranchStatus("*", 0)
    for b in branch_names:
        tree.SetBranchStatus(b, 1)
    # map(lambda b: tree.SetBranchStatus(b, 1), branch_names)

    for i in range(num_entries):
        bytes_read = [b.GetEntry(i) for b in branches]
        if tree.rsp < 2:
            hist.Fill(tree.pt, tree.ptRef)

    # pr.disable()
    # pr.dump_stats('profile_activate')


def tree_plotter_iter(tree, hist, num_entries):
    # pr = cProfile.Profile()
    # print 'iter'
    # pr.enable()

    for i, evt in enumerate(tree):
        if i >= num_entries:
            break
        if evt.rsp < 2:
            hist.Fill(evt.pt, evt.ptRef)

    # pr.disable()
    # pr.dump_stats('profile_iter')


def tree_plotter_array(tree, hist, num_entries):
    # pr = cProfile.Profile()
    # print 'array'
    # pr.enable()

    pt = array('f', [0.])
    ptRef = array('f', [0.])
    rsp = array('f', [0.])
    
    tree.SetBranchAddress("pt", pt)
    tree.SetBranchAddress("ptRef", ptRef)
    tree.SetBranchAddress("rsp", rsp)
    
    for i in range(num_entries):
        tree.GetEntry(i)
        if rsp[0] < 2:
            hist.Fill(pt[0], ptRef[0])

    # pr.disable()
    # pr.dump_stats('profile_array')


def tree_plotter_draw(tree, hist, num_entries):
    # pr = cProfile.Profile()
    # print 'draw'
    # pr.enable()

    tree.Draw("ptRef:pt>>%s" % hist.GetName(), "rsp<2 && Entry$<%d" % num_entries)

    # pr.disable()
    # pr.dump_stats('profile_array')


# if __name__ == "__main__":
filename = 'tree.root'
f = ROOT.TFile(filename)
tree = f.Get("valid")
if not tree:
    raise RuntimeError("No tree object")

NENTRIES = 1000000
print 'Num entries:', NENTRIES

# These 3 methods should have different performance, but don't appear to...
# tree_plotter_branches(tree, ROOT.TH2D("h2d1", "", 200, 0, 400, 200, 0, 400), NENTRIES)

# tree_plotter_branches_activate(tree, ROOT.TH2D("h2d1activate", "", 200, 0, 400, 200, 0, 400), NENTRIES)

# tree_plotter_iter(tree, ROOT.TH2D("h2d2", "", 200, 0, 400, 200, 0, 400), NENTRIES)

# PRETTY DAMN FAST
tree_plotter_array(tree, ROOT.TH2D("h2d2", "", 200, 0, 400, 200, 0, 400), NENTRIES)

# FASTEST
# tree_plotter_draw(tree, ROOT.TH2D("h2ddraw", "", 200, 0, 400, 200, 0, 400), NENTRIES)

f.Close()