[Rivet] ATLAS analysis to be included into Rivet

[Roman Lysak]

Hi Hendrik,

thanks a lot for the inclusion of the analysis.
You are right, I modified some older version of root2flat in order to 
work for me.
I'm attaching now the modified version of root2flat trunk version.
There is really just two line change.

Cheers,
   Roman


On 12/06/2012 11:46 AM, Hendrik Hoeth wrote:
> Hi Roman,
>
> thanks for the analysis and your changes to aidamerge and root2flat.
> I've included the analysis and your aidamerge changes into the Rivet svn
> repository, but I'd like to ask you to send me a patch for root2flat.
> Looking at the differences between the current root2flat and your
> version, I'm overwhelmed and don't know what differences are important.
>
> Cheers,
>
>      Hendrik
>

-------------- next part --------------
#! /usr/bin/env python

"""%prog

Read in .root files and write out the histograms in FLAT
format suitable for make-plots.

Use e.g. 'root2flat file1.root -o flathistos' to produce the flat histogram
files which need to be converted to AIDA files using Rivet's flat2aida tool

Usage:
    %prog [options] file1.root

Use --help to get information for all options.
"""

import sys
if sys.version_info[:3] < (2,4,0):
    print "rivet scripts require Python version >= 2.4.0... exiting"
    sys.exit(1)


import os, optparse, logging, ROOT


## Parse options
parser = optparse.OptionParser(usage = __doc__)
parser.add_option("-a", dest="ANALYSIS", default=None,
                  help = "Optionally add analysis name to histogram keys")
parser.add_option("-e", dest="ENUM", default=False, action='store_true',
                  help = "Enumerate histos hepdata style")
parser.add_option("-f", dest="FILENAME", default=None,
                  help = "Force output of all histos in single file")
parser.add_option("-x", dest="XCLUDE", default=None,
                  help = "Exclude histos from conversion based on string pattern")
parser.add_option("-o", dest="OUTDIR", default=".",
                  help = "specify directory in which to write out converted files")
(opts, args) = parser.parse_args()


def readROOT(rootfile):
    """ This is the main function that opens a ROOT file, browses its contents
        for histograms and tries to write the converted histos to files.
    """
    global nobs
    # Open the ROOT file
    f = ROOT.TFile(rootfile)

    # Initial browse to see the structure
    subdirs, histonames, tgraphnames, tcanvases = browse(f)

    # Keep browsing if there are subdirectories TODO: Make this work for
    # arbitrarily deep directory structures
    if len(subdirs) > 0:
        for sd in subdirs:
            t_s, t_h = browse(f, sd)
            histonames.extend(t_h)

    # This will convert and write the histos
    for num, histoname in enumerate(histonames):
        writeHisto(histoname, f.Get(histoname), num+nobs)
    for num, tgraphname in enumerate(tgraphnames):
        writeHisto(tgraphname, f.Get(tgraphname), num + len(histonames)+nobs, True)
    # TCanvas items
    for num, tc in enumerate(tcanvases):
        writeHisto(tc[0], tc[-1], num + len(histonames) + len(tgraphnames) + nobs, tc[1])

    nobs += len(histonames)
    nobs += len(tgraphnames)




def browse(f, branch=None):
    """ This function browses a file/branch, trying to find objects that
        either inherit from TH1 or from TProfile.
    """
    # Prepare return values
    histos  = []
    tgraphs = []
    subdirs = []
    tcanvases = []

    # Get Iterator
    if branch:
        f.Cd(branch)

    primary_keys = ROOT.gDirectory.GetListOfKeys()
    iter_primaries = ROOT.TListIter(primary_keys)

    if branch:
        f.Cd('..')

    # Iterate over iterator
    for i in xrange(len(primary_keys)):
        if branch:
            t_n = branch + '/' + iter_primaries.Next().GetName()
        else:
            t_n = iter_primaries.Next().GetName()
        # Make sure we don't have a NoneType object here
        if f.Get(t_n):
            if opts.XCLUDE is not None and opts.XCLUDE in t_n:
                continue
            # Check if the curent object is a directory
            if type(f.Get(t_n)) == ROOT.TDirectoryFile:
                subdirs.append(t_n)
            # Check if the curent object is a histogram
            elif f.Get(t_n).InheritsFrom("TH1") or f.Get(t_n).InheritsFrom("TProfile"):
                histos.append(t_n)
            elif f.Get(t_n).InheritsFrom("TGraphAsymmErrors") or f.Get(t_n).InheritsFrom("TGraphErrors"):
                tgraphs.append(t_n)
            # Also support graphs, histos stored in TCanvases
            elif f.Get(t_n).InheritsFrom("TCanvas"):
                for ci in f.Get(t_n).GetListOfPrimitives():
                    if ci.InheritsFrom("TH1") or ci.InheritsFrom("TProfile"):
                        tcanvases.append((t_n,False,ci))
                    elif ci.InheritsFrom("TGraphAsymmErrors") or ci.InheritsFrom("TGraphErrors"):
                        tcanvases.append((t_n,True,ci))

    return subdirs, histos, tgraphs, tcanvases


def convertHisto(R_histo):
    """ This function reads a single ROOT histo and converts it into the
        FLAT format suitable for plotting with make-plots.
    """
    title = R_histo.GetTitle().replace("#","\\")
    xtitle= R_histo.GetXaxis().GetTitle().replace("#","\\")
    ytitle= R_histo.GetYaxis().GetTitle().replace("#","\\")
    bins = getBinsFromTH1F(R_histo)
    for bin in bins:
        try:
            binstr = ""
            for key in ["xlow", "xhigh", "y", "y_err_low", "y_err_high"]:
                binstr += "%s: %e " % (key, bin[key])
            logging.info(binstr)
        except:
            pass
    return title, xtitle, ytitle, bins


def convertTGraph(TGraph):
    title = TGraph.GetTitle().replace("#","\\")
    xtitle= TGraph.GetXaxis().GetTitle().replace("#","\\")
    ytitle= TGraph.GetYaxis().GetTitle().replace("#","\\")
    bins = getBinsFromTGraph(TGraph)
    return title, xtitle, ytitle, bins


def getBinsFromTH1F(R_histo):
    """ A little helper function that returns a list of bin-dictionaries).
    """
    allbins=[]
    for ii in xrange(R_histo.GetNbinsX()):
        i = ii + 1
        xlow  = R_histo.GetBinLowEdge(i)
        xhigh = xlow + R_histo.GetBinWidth(i)
        y     = R_histo.GetBinContent(i)
        y_err = R_histo.GetBinError(i)
        bin = {"xlow":xlow, "xhigh":xhigh, "y":y, "y_err_low":y_err, "y_err_high":y_err}
        allbins.append(bin)
    return allbins


def getBinsFromTGraph(TGraph):
    allbins=[]
    X = TGraph.GetX()
    Y = TGraph.GetY()
    for ii in xrange(TGraph.GetN()):
        i = ii #+ 1
        xlow  = X[i] - TGraph.GetErrorXlow(i)
        xhigh = X[i] + TGraph.GetErrorXhigh(i)
        y     = Y[i]
        y_err_low = TGraph.GetErrorYlow(i)
        y_err_high = TGraph.GetErrorYhigh(i)
        bin = {"xlow":xlow, "xhigh":xhigh, "y":y, "y_err_low":y_err_low, "y_err_high":y_err_high}
        allbins.append(bin)
    return allbins


def writeHisto(name, R_histo, rivetid, tgraph=False):
    """ This writes the histogram into a single file, ready to plot with
    make-plots.
    """
    if tgraph:
        title, xlabel, ylabel, bins = convertTGraph(R_histo)
    else:
        title, xlabel, ylabel, bins = convertHisto(R_histo)
    head = "# BEGIN PLOT\nTitle=%s\nLegend=1\nLogY=1\nDrawOnly=%s\n" % (title, name)
    head += "XLabel=%s\nYLabel=%s\n# END PLOT\n" % (xlabel, ylabel)

    if opts.ANALYSIS and opts.ENUM:
        name = os.path.join(opts.ANALYSIS, "d"+str(rivetid).zfill(2)+"-x01-y01")
    elif opts.ANALYSIS and not opts.ENUM:
        name = os.path.join(opts.ANALYSIS, name)

    histo = getFlatHisto(bins, name, title)
    flatname = name.replace("/","_") + ".dat"

    if flatname.startswith("_"):
        flatname = flatname[1:]
    flatfile = os.path.join(opts.OUTDIR, flatname)

    if not opts.FILENAME:
        f = open(flatfile, "w")
    else:
        global nobs
        if os.path.exists(opts.FILENAME) and rivetid == 1:
            print "Error, outputfile '%s' exists! Exiting..."%opts.FILENAME
            sys.exit(1)
        else:
            f = open(opts.FILENAME, "a")

    f.write(head)
    f.write("\n")
    f.write(histo)
    f.close()


def getFlatHisto(bins, name, title):
    """ This returns a histo in the FLAT format. """
    histo= "# BEGIN HISTOGRAM %s\n" % name
    histo += "LineColor=black\n"
    histo += "ErrorBars=1\n"
    histo += "PolyMarker=*\n"
    histo += "Title=%s\n" % title
    for bin in bins:
        histo += "%.8e\t%.8e\t%.8e\t%.8e\t%.8e\n" % (bin["xlow"], bin["xhigh"],
                                                     bin["y"], bin["y_err_low"], bin["y_err_high"])
    histo += "# END HISTOGRAM\n\n"
    return histo


if __name__ == "__main__":
    nobs = 1
    for infile in args:
        if not os.path.exists(opts.OUTDIR):
            os.mkdir(opts.OUTDIR)
        readROOT(infile)
    dest = opts.OUTDIR
    if opts.FILENAME:
        dest=opts.FILENAME
    print  "Done. Written all converted histos to %s" % dest