pyroot_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')


def print_hist(hist):
    c = ROOT.TCanvas("c", "", 800, 800)
    hist.Draw("COLZ")
    c.SaveAs("c.pdf")


# 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

# branch_names = ["pt", "ptRef", "rsp"]
# # branches = map(lambda b: tree.GetBranch(b), branch_names)
# branches = [tree.GetBranch(b) for b in branch_names]
# hist = ROOT.TH2D("h2d1", "", 200, 0, 400, 200, 0, 400)
# num_entries = NENTRIES
# 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)

# 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)

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

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

f.Close()