--- UserCode/RootMacros/tree2hists.py 2010/05/20 19:44:30 1.3 +++ UserCode/RootMacros/tree2hists.py 2010/05/23 20:56:57 1.4 @@ -1,33 +1,34 @@ #!/usr/bin/env python +""" +Create ROOT Histograms from one or more ROOT TTrees or TNtuples. + +Options are specified in the given configuration file. +""" +## Created by Michael Anderson (mbanderson@wisc.edu), May 2010 # -# Run by typing: -# ./tree2hists.py -# -# Creates histograms from TTree in one or more files. -# Different cuts & scales can be set for each file. -# -# Details are specified in a configruation file, and run with: -# ./tree2hists.py config.py -# -# Michael Anderson -# Original: Nov 4, 2009 -# Updated: May 20, 2010 +# Create configuration file: +# tree2hists.py +# Edit, then run with config file: +# tree2hists.py config.py + +######## Import python libraries ############################################# import sys # For exiting program if sys.version_info < (2, 6): print "Go to a CMSSW dir and type 'cmsenv' first. (this loads a modern version of python)" sys.exit(0) -from ROOT import TFile, TH1F, TH2F, TH3F, TTree, gROOT # Import any ROOT class you want -from copy import deepcopy # For copying histograms -from math import pi -from array import array # For making Float_t array ROOT wants (for hists) -from datetime import datetime # For output filename -from os import path # For finding file +from ROOT import TFile, TTree, TH1F, TH2F, TH3F, gROOT # Import any ROOT class you want +from copy import deepcopy # For copying histograms +from math import pi # For use in histogram bounds +from array import array # For making Float_t array ROOT wants (for hists) +from datetime import datetime # For output filename +from os import path # For finding file + +######## Define classes and generators ####################################### -######################################## -# Class for root files to store -# file's name, scale, cuts, etc... class rootFile: + """Wrapper for TTrees and TNtuples, allowing association with + a scale and cuts.""" def __init__(self, fileName, scale=1.0, cuts="", tree=""): self.name = fileName self.scale = scale @@ -45,36 +46,31 @@ class rootFile: sys.exit(0) print "\t %s contains %.0f entries before cuts, %.0f after." % (tree, self.ttree.GetEntries(), self.ttree.GetEntries(cuts)) -# Class to store varibles, histograms to plot them into, and cuts class Plot: - def __init__(self, treeVariable, histogram, cuts=""): + """Wrapper for TH1 objects, associating TTree variables with a histogram""" + def __init__(self, treeVariable, histogram, cuts="", storeErrors=True): self.treeVariable = treeVariable self.histogram = histogram self.cuts = cuts - self.histogram.Sumw2() # Store errors + if storeErrors: self.histogram.Sumw2() -# This joins a list of cuts (strings) -# into something ROOT can handle -# example: given several strings, like "1<2","","5>4" -# this returns the sting "1<2&&5>4" -def joinNonEmpty(*listOfCuts): +def joinCuts(*listOfCuts): + """Joines list of cuts (strings) into something ROOT can handle. + Example: given ("1<2","","5>4") returns '1<2&&5>4'""" listOfNonEmptyCuts=[] for cut in listOfCuts: if cut: listOfNonEmptyCuts.append(cut) return '&&'.join(listOfNonEmptyCuts) -def writeDefaultConfig(): - if path.exists('t2h_config.py'): - print "Specify a config file, like:" - print "./tree2hists.py t2h_config.py" - sys.exit(0) +def writeDefaultT2HConfig(): + """Writes configuration file for tree2hists""" defaultConfig = '''# Configuration file for tree2hists.py # Created %s. from tree2hists import * -listOfFiles = [rootFile("MultiPhotonAnalyzer_SDEG.root", scale=1.0, tree="NTuples/Analysis"), - rootFile("MultiPhotonAnalyzer_SD_EG_May7.root", scale=1.0, tree="NTuples/Analysis")] +listOfFiles = [rootFile("MultiPhotonAnalyzer_SDEG.root", tree="NTuples/Analysis", scale=1.0, cuts=""), + rootFile("MultiPhotonAnalyzer_SD_EG_May7.root", tree="NTuples/Analysis", scale=1.0 cuts="")] outputFilename = "Hists_Data_%%s.root" %% datetime.now().strftime("%%b%%d_%%I%%p") @@ -82,14 +78,14 @@ cutForAllFiles = "(!TTBit[36] && !TTBit[ "&&((isEB[0] && (seedSeverity[0]!=3 && seedSeverity[0]!=4 ) && (seedRecoFlag[0] != 2) ) || isEE[0])" # All plots are made for each "cutSet". -# A cut set is 3 things: folder name to store hists in, string to add to hist titles, and actual cuts -cutSets = ( +# A "cutSet" is 3 things: folder name to store hists in, string to add to hist titles, and cuts for these hists. +# Set cutSet = [] to make all plots. +cutSets = [ ("barrel15to20", "(|#eta|<1.45, 1515&&et[0]<20&&abs(eta[0])<1.45"), ("barrel20to30", "(|#eta|<1.45, 2020&&et[0]<30&&abs(eta[0])<1.45"), ("endcap15to20", "(1.7<|#eta|<2.5, 1515&&et[0]<20&&abs(eta[0])>1.7&&abs(eta[0])<2.5"), ("endcap20to30", "(1.7<|#eta|<2.5, 2020&&et[0]<30&&abs(eta[0])>1.7&&abs(eta[0])<2.5") -# ("","",""), # all plots with no special cuts - ) + ] # Define histograms to plot bins_et = array("f", [15.0, 20.0, 30.0, 50.0, 80.0, 120.0]) @@ -108,33 +104,30 @@ listOfPlots = [ f.close() print "Created default configuration file: t2h_config.py" print "Edit it, and run with:" - print " tree2hists t2h_config.py" + print " tree2hists.py t2h_config.py" ######################################## -######################################## -if __name__ == '__main__': - if len(sys.argv) > 1: - if path.isfile(sys.argv[1]): - config_file = sys.argv[1].split('.')[0] - try: - exec("from " + config_file + " import *") - except Exception, e: - print e - sys.exit(1) - else: - print "%s not found." % sys.argv[1] - print "Create default config file by running tree2hists.py with no argument." - sys.exit(1) - else: - writeDefaultConfig() - sys.exit(0) +######## Define the main program ############################################# +def tree2hists_main(config_file): + try: + # Import only certain variables + _temp = __import__(config_file, globals(), locals(), ['listOfFiles','outputFilename','cutForAllFiles','cutSets','listOfPlots'], -1) + listOfFiles = _temp.listOfFiles + outputFilename = _temp.outputFilename + cutForAllFiles = _temp.cutForAllFiles + cutSets = _temp.cutSets + listOfPlots = _temp.listOfPlots + if not cutSets: cutSets = [("","","")] # Make all plots, no extra cuts + except Exception, e: + print e + sys.exit(1) if path.isfile('rootlogon.C'): print "Loading rootlogon.C" - gROOT.Macro('rootlogon.C') # Run ROOT logon script + gROOT.Macro('rootlogon.C') # Run ROOT logon script (for loading of functions) - outputFile = TFile(outputFilename, "recreate") # Open output file + outputFile = TFile(outputFilename, "recreate") if not outputFile.IsZombie(): print "Opened %s for output." % outputFilename else: @@ -146,28 +139,25 @@ if __name__ == '__main__': print "\nCuts applied to all plots from all files:\n %s" % cutForAllFiles numberOfPlots = len(listOfPlots) print "\nCreating %i plots for each of %i cut sets..." % (numberOfPlots, len(cutSets)) - for setOfCuts in cutSets: + for setOfCuts in cutSets: # Make all plots for each cutSet histNamePostfix, titlePostfix, currentCutSet = setOfCuts print '\n Cut set "%s": %s' % (histNamePostfix, currentCutSet) - # Loop over all things to plot - for i, plot in enumerate(listOfPlots): + for i, plot in enumerate(listOfPlots): # Loop over plots newTitle = ' '.join((plot.histogram.GetTitle(), titlePostfix)) newPlot = deepcopy(plot) dir = histNamePostfix listOfPlotsToWrite.append((dir, newPlot)) newPlot.histogram.SetTitle(newTitle) - # Print plot being made - print " %i %s >> %s/%s" % (i, newPlot.treeVariable, dir, newPlot.histogram.GetName()), + for aFile in listOfFiles: # Loop over all TTrees + tempHist = newPlot.histogram.Clone("temp") # Create temp histogram + cuts = joinCuts(cutForAllFiles, aFile.cuts, currentCutSet, newPlot.cuts) # Set cuts for temp + aFile.ttree.Draw( "%s >> temp" % newPlot.treeVariable, cuts, "goff") # Draw into temp; graphics off + tempHist.Scale(aFile.scale) # Scale temp + newPlot.histogram.Add(tempHist) # Add temp to total histogram + print " %3i %7i %s >> %s/%s" % (i, newPlot.histogram.GetEntries(), newPlot.treeVariable, dir, newPlot.histogram.GetName()), if newPlot.cuts: - print "\textra cuts: %s" % newPlot.cuts, # Plot-specific cuts - # Loop over all TTrees (from the different files) - for aFile in listOfFiles: - tempHist = newPlot.histogram.Clone("temp") # Create temp histogram - cuts = joinNonEmpty(cutForAllFiles, aFile.cuts, currentCutSet, newPlot.cuts) # Set cuts for temp - aFile.ttree.Draw( "%s >> temp" % newPlot.treeVariable, cuts, "goff") # Draw into temp; graphics off - tempHist.Scale(aFile.scale) # Scale temp - newPlot.histogram.Add(tempHist) # Add temp to total histogram - print "%i" % newPlot.histogram.GetEntries() + print "\textra cuts: %s" % newPlot.cuts, # plot-specific cuts + print print "done." # Store and save/close files @@ -178,9 +168,30 @@ if __name__ == '__main__': outputFile.cd(dir) plot.histogram.Write() - print "Closing files...", outputFile.Close() for aFile in listOfFiles: aFile.file.Close() - print "done.\n\nHistograms stored in\n %s" % outputFilename + print "\n\nHistograms stored in\n %s" % outputFilename ######################################## + +if __name__ == "__main__": + if len(sys.argv) > 1: + if path.isfile(sys.argv[1]): + config_file = sys.argv[1].split('.')[0] + tree2hists_main(config_file) + else: + print "%s not found." % sys.argv[1] + print "Create default config file by running tree2hists.py with no argument." + sys.exit(1) + else: + if path.exists('t2h_config.py'): + print "Specify a config file, like:" + print "tree2hists.py t2h_config.py" + sys.exit(1) + writeDefaultT2HConfig() + sys.exit(0) + #import cProfile + #cProfile.run('main()', 'fooprof') + #import pstats + #p = pstats.Stats('fooprof') + #p.sort_stats('cumulative').print_stats(15)