An example PyROOT analysis which applies several different NHIT cuts, calculates the fraction of events failing, and produces a plot.

nhit_cut.py

'''An example PyROOT analysis which applies several different NHIT cuts,
calculates the fraction of events failing, and produces a plot.
'''

import sys
import array
from rat import ROOT


def fast_dsreader(filename):
    '''The dsreader built into RAT is a little broken right now.
    
    :param filename: Name of the RAT ROOT file to open
    :returns: Generator of ROOT.RAT.DS.Root objects
    '''
    r = ROOT.RAT.DSReader(filename)
    ds = r.NextEvent()
    while ds:
        yield ROOT.RAT.DS.Root(ds)
        ds = r.NextEvent()


def get_cut_fractions(filename, cuts):
    '''Determine the fraction of events failing the cut.
    
    :param filename: RAT ROOT file to analyze
    :param cuts: List of cuts at which to evaluate
    :returns: A {cut parameter: fraction cut} dictionary
    '''
    dsr = fast_dsreader(filename)

    count_cut = {cut: 0 for cut in cuts}
    count_total = 0

    # loop over events in the file
    for i, ds in enumerate(dsr):
        print 'Processing event', i
        count_total += 1

        # loop over cuts
        for cut in cuts:
            # apply the cut, incrementing counter if it fails
            if ds.GetEV(0).GetNhits() < cut:
                count_cut[cut] += 1

    # convert to fractions
    for cut in count_cut:
        count_cut[cut] = 1.0 * count_cut[cut] / count_total

    return count_cut


def make_cut_plot(cut_fractions):
    '''Take a set of cut results as a {cut value: fraction cut} dictionary and
    turn it into a ROOT plot.

    :param cut_fractions: Dictionary of cut results
    :returns: TGraph plotting the results
    '''
    # create a sorted list and split into x and y arrays via Python magic
    sorted_fractions = sorted(cut_fractions.iteritems())
    x, y = map(lambda x: array.array('d', x), zip(*sorted_fractions))  # fun!

    # make the graph
    g = ROOT.TGraph(len(x), x, y)
    g.SetTitle('')
    g.SetMarkerStyle(20)
    g.SetMarkerColor(ROOT.kBlue)
    g.SetLineColor(ROOT.kBlue)
    g.SetLineWidth(2)
    g.GetXaxis().SetTitle('NHIT cut')
    g.GetYaxis().SetTitle('Fraction of events cut')

    return g


if __name__ == '__main__':
    if len(sys.argv) != 2:
        print 'Usage:', sys.argv[0], 'filename.root'
        sys.exit(1)

    filename = sys.argv[1]

    # define parameter values where we want to try the cut
    nhit_cut_values = range(0, 1000, 10)

    # the get_cut_fractions function does all the work
    cut_fractions = get_cut_fractions(filename, nhit_cut_values)

    # make and draw a nice plot
    plot = make_cut_plot(cut_fractions)
    plot.Draw('alp')

    # wait to exit, so the plot stays open
    raw_input()