In This Package:

figures.py

Go to the documentation of this file.

#!/usr/bin/env python
#
# Generate a figure directory based on the contents of a root file
#
# Usage:
#   % figures.py -o /output/path -r /rootfile/directory histograms.root


class FigureGenerator:
    """Figure Generator Class"""
    def __init__(self, filename, fileRoot, diskRoot):
        "Constructor"
        self.filename = filename
        self.fileRoot = fileRoot
        self.diskRoot = diskRoot
        self.file = None
        self.format = "png"
        return

    def process(self):
        """Generate directories and figures"""
        if not self.openFile():
            return
        self.processPath("")
        return
        
    def openFile(self):
        "Open root file"
        import ROOT
        gROOT = ROOT.gROOT
        gROOT.SetBatch(True)
        # Avoid verbose output
        ROOT.gErrorIgnoreLevel = ROOT.kInfo
        gROOT.SetStyle("Plain")
        ROOT.gStyle.SetTitleFillColor(0)
        ROOT.gStyle.SetTitleBorderSize(0)
        ROOT.gStyle.SetPalette(1)
        # I don't like ROOT.  I have to set China timezone here
        # to get proper time axis in figure generation
        ROOT.gSystem.Setenv("TZ","PRC")
        rootFile = ROOT.TFile(self.filename)
        if not rootFile:
            print "Failed to open root file:",filename
            return False
        self.file = rootFile
        self.canvas = ROOT.TCanvas()
        return True

    def processPath(self, itemPath):
        """Process a path in the current root file"""
        #print "Processing:",itemPath
        filePath = self.fileRoot
        diskPath = self.diskRoot
        if len(itemPath)>0:
            filePath += "/"+itemPath
            diskPath += "/"+itemPath
        content = self.file.Get(filePath)
        if not content and itemPath=="":
            print "FIXME: Sim data has no run number"
            self.fileRoot = self.fileRoot[:-7] + "0000000"
            filePath = self.fileRoot
            content = self.file.Get(filePath)
        if not content:
            print "Failed to get item at: ", filePath
            return
        if content.IsA().InheritsFrom("TH1"):
            # Found a histogram.  Draw it.
            if content.IsA().InheritsFrom("TH2F"):
                # Draw 2-D histograms with colz option
                content.Draw("colz");
            elif content.IsA().InheritsFrom("TH1F"):
                # Draw 1-D histograms with log scale
                self.canvas.SetLogy(1)
                content.Draw();
            else:
                content.Draw();
            self.canvas.SaveAs(diskPath+"."+self.format)
            self.canvas.SetLogy(0)
        elif content.IsA().InheritsFrom("TGraph"):
            # Found a graph.  Draw it.
            content.Draw("APL");
            self.canvas.SaveAs(diskPath+"."+self.format)
        elif content.IsA().InheritsFrom("TDirectory"):
            self.makeDirectory(diskPath)
            keylist = content.GetListOfKeys()
            for key in keylist:
                dirItem = key.ReadObj()
                dirItemPath = None
                if len(itemPath)>0:
                    dirItemPath = itemPath+"/"+dirItem.GetName()
                else:
                    dirItemPath = dirItem.GetName()
                self.processPath(dirItemPath)
        else:
            # Won't handle this object
            print "Not handing %s of type %s" % (filepath,content.ClassName())
        return

    def makeDirectory(self, diskPath):
        """Make sure the output directory exists"""
        import os
        os.system("mkdir -p "+diskPath)
        return

if __name__ == "__main__":

    fileRoot = "stats"
    diskRoot = "."
    filename = None
    import sys, getopt
    opts, args = getopt.getopt(sys.argv[1:],"o:r:")
    for opt, arg in opts:
        if opt == "-o":
            diskRoot = arg
        elif opt == "-r":
            fileRoot = arg
    print args
    filename = args[0]
    print "Filename:  ",filename
    print "File Root: ",fileRoot
    print "Disk Root: ",diskRoot
    figGen = FigureGenerator(filename, fileRoot, diskRoot)
    figGen.process()
    print "Done with figure generation"
    
    

---