python/rootplot/tree2hists.py

0001 #!/usr/bin/env python3
0002 """
0003 Create ROOT Histograms from one or more ROOT TTrees or TNtuples.
0004
0005 Options are specified in the given configuration file.
0006 """
0007 from __future__ import print_function
0008
0009 # Create configuration file:
0010 #   tree2hists.py
0011 # Edit, then run with config file:
0012 #   tree2hists.py config.py
0013
0014 __license__ = '''\
0015 Copyright (c) 2010 Michael Anderson <mbanderson@wisc.edu>
0016
0017 Permission is hereby granted, free of charge, to any person obtaining a copy
0018 of this software and associated documentation files (the "Software"), to deal
0019 in the Software without restriction, including without limitation the rights
0020 to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
0021 copies of the Software, and to permit persons to whom the Software is
0022 furnished to do so, subject to the following conditions:
0023
0024 The above copyright notice and this permission notice shall be included in
0025 all copies or substantial portions of the Software.
0026
0027 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
0028 IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
0029 FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
0030 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
0031 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
0032 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
0033 THE SOFTWARE.
0034 '''
0035
0036 ######## Import python libraries #############################################
0037
0038 import sys                              # For exiting program
0039 if '-h' in sys.argv or '--help' in sys.argv:
0040     print('''\
0041 Create ROOT Histograms from one or more ROOT TTrees or TNtuples.
0042
0043 Run by specifying configuration file:
0044   tree2hists config.py
0045
0046 Create default config file by running with no arguments:
0047   tree2hists''')
0048     sys.exit(0)
0049 try:
0050     from ROOT import TFile, TTree, TH1F, TH2F, TH3F, gROOT
0051 except Exception as e:
0052     print(e)
0053     print ("Use a python that has PyROOT installed.")
0054     sys.exit(0)
0055 from copy import deepcopy     # For copying histograms
0056 from math import pi           # For use in histogram bounds
0057 from array import array       # For making Float_t array ROOT wants (for hists)
0058 from datetime import datetime # For output filename
0059 from os import path           # For finding file
0060
0061 ######## Define classes and generators #######################################
0062
0063 class RootTree:
0064     """Wrapper for TTrees and TNtuples, allowing association with
0065     a scale and cuts."""
0066     def __init__(self, treeName, fileName, scale=1.0, cuts=""):
0067         self.fileName = fileName
0068         self.treeName = treeName
0069         self.scale    = scale
0070         self.cuts     = cuts
0071         self.tfile    = TFile()
0072         self.ttree    = TTree()
0073
0074 class Plot:
0075     """Wrapper for TH1 objects, associating TTree variables with a histogram"""
0076     def __init__(self, treeVariable, histogram, cuts="", storeErrors=True):
0077         self.treeVariable = treeVariable
0078         self.histogram    = histogram
0079         self.name         = histogram.GetName()
0080         self.cuts         = cuts
0081         if storeErrors: self.histogram.Sumw2()
0082
0083 def join_cuts(*list_of_cuts):
0084     """Joins list of cuts (strings) into something ROOT can handle.
0085     Example:  given ('1<2','','5>4') returns '1<2&&5>4'"""
0086     list_of_nonempty_cuts = []
0087     for cut in list_of_cuts:
0088         if cut:
0089             list_of_nonempty_cuts.append(cut)
0090     return '&&'.join(list_of_nonempty_cuts)
0091
0092 def duration_to_string(start, end):
0093     timeTaken = end - start
0094     hours, remainder = divmod(timeTaken.seconds, 3600)
0095     minutes, seconds = divmod(remainder, 60)
0096     if hours>0:
0097         return "%i hours, %i minutes" % (hours, minutes)
0098     elif minutes>0:
0099         return "%i minutes" % minutes
0100     return "%i seconds" % seconds
0101
0102 def write_default_T2H_config():
0103     """Writes configuration file for tree2hists"""
0104     defaultConfig = '''# Configuration file for tree2hists
0105 # Created %s.
0106 try:
0107     ## the normal way to import from rootplot
0108     from rootplot.tree2hists import RootTree, Plot
0109 except ImportError:
0110     ## special import for CMSSW installations of rootplot
0111     from PhysicsTools.PythonAnalysis.rootplot.tree2hists import RootTree, Plot
0112 from array import array      # to allow making Float_t arrays for ROOT hists
0113 from math import pi
0114 from ROOT import TH1F, TH2F  # import other kinds of hists as neeeded
0115
0116 list_of_files = [RootTree("Treename", fileName="photons.root", scale=1.0, cuts=""),
0117                  RootTree("Treename", fileName="photons2.root", scale=1.0, cuts="")]
0118
0119 output_filename = "Hists_photons.root"
0120
0121 cut_for_all_files = "(!TTBit[36] && !TTBit[37] && !TTBit[38] && !TTBit[39] && !vtxIsFake && vtxNdof>4 && abs(vtxZ)<=15)"
0122
0123 # All plots are made for each "cut set".
0124 # A "cut set" is 3 things: folder name to store hists in, string to add to hist titles, and cuts for these hists.
0125 # Let cut_sets = [] to make all plots.
0126 cut_sets = [
0127     ("barrel15to20", "(|#eta|<1.45, 15<E_{T}<20)", "et>15&&et<20&&abs(eta)<1.45"),
0128     ("barrel20to30", "(|#eta|<1.45, 20<E_{T}<30)", "et>20&&et<30&&abs(eta)<1.45"),
0129     ("endcap15to20", "(1.7<|#eta|<2.5, 15<E_{T}<20)", "et>15&&et<20&&abs(eta)>1.7&&abs(eta)<2.5"),
0130     ("endcap20to30", "(1.7<|#eta|<2.5, 20<E_{T}<30)", "et>20&&et<30&&abs(eta)>1.7&&abs(eta)<2.5")
0131     ]
0132
0133 # Define histograms to plot
0134 bins_et     = array("f", [15.0, 20.0, 30.0, 50.0, 80.0, 120.0]) # example custom bins
0135 list_of_plots = [
0136     Plot("et"           , TH1F("pho_et"           , "Lead #gamma: E_{T};E_{T} (GeV);entries/bin", 25, 0.0, 100.0)),
0137     Plot("eta"          , TH1F("pho_eta"          , "Lead #gamma: #eta;#eta;entries/bin"        , 25, -3.0, 3.0)),
0138     Plot("et"           , TH1F("pho_et_binned"    , "Lead #gamma: E_{T};E_{T} (GeV);entries/bin", len(bins_et)-1, bins_et)),
0139     Plot("sigmaIetaIeta", TH1F("pho_sigmaIEtaIEta", "Lead #gamma: #sigma_{i#etai#eta};#sigma_{i#etai#eta};entries/bin",20, 0, 0.06)),
0140     Plot("metEt/et"     , TH1F("metEt_over_phoEt" , "MET / E_{T}(#gamma);MET/E_{T}(sc);entries/bin"   , 20, 0.0, 3.0)),
0141     Plot("phi:eta"      , TH2F("phoPhi_vs_phoEta" , "Lead #gamma: #phi vs #eta;#eta;#phi"             , 50, -2.5, 2.5, 18, -pi, pi))
0142     ]
0143 ''' % datetime.now().strftime("%b %d, %Y")
0144     f = open('t2h_config.py', 'w')
0145     f.write(defaultConfig)
0146     f.close()
0147     print("Created default configuration file: t2h_config.py")
0148     print("Edit it, then run by typing:")
0149     print("  tree2hists t2h_config.py")
0150 ##############################################################################
0151
0152 def make_all_hists_all_files(list_of_RootTrees, list_of_Plots, cut_for_all_files, list_of_cutSets):
0153     '''Open root files one at a time, make plots, then close them.'''
0154     list_of_plots_to_write = []
0155
0156     # Create plots for each set of cuts
0157     for set_of_cuts in list_of_cutSets:
0158         histname_fix, title_fix, current_cut_set = set_of_cuts
0159         for plot in list_of_Plots:
0160             new_plot  = deepcopy(plot)
0161             new_title = ' '.join((plot.histogram.GetTitle(), title_fix))
0162             new_plot.histogram.SetTitle(new_title)
0163             list_of_plots_to_write.append((new_plot, set_of_cuts))
0164
0165     for j, rootTree in enumerate(list_of_RootTrees):
0166         rootTree.tfile = TFile(rootTree.fileName, "read")           # Open TFile
0167         if rootTree.tfile.IsZombie():
0168             print("Error opening %s, exiting..." % rootTree.fileName)
0169             sys.exit(0)
0170         try:                                      # Try to get TTree from file.
0171             rootTree.tfile.GetObject(rootTree.treeName, rootTree.ttree)
0172         except:
0173             print("Error: %s not found in %s, exiting..." % (rootTree.treeName,
0174                                                              rootTree.fileName))
0175             sys.exit(1)
0176
0177         print("\n%s: Opened %s  %i MB" % (datetime.now().strftime("%I:%M%p"),
0178                                           rootTree.fileName,
0179                                           path.getsize(rootTree.fileName)/1048576))
0180         print(" %s has %i entries, will plot with scale=%.2e, cuts='%s'" % (rootTree.treeName,
0181                                                                             rootTree.ttree.GetEntries(),
0182                                                                             rootTree.scale,
0183                                                                             rootTree.cuts))
0184
0185         # Loop over plots
0186         print("   # entries                  var >> histogram")
0187         for i, (plot, set_of_cuts) in enumerate(list_of_plots_to_write):
0188             histname_fix, title_fix, current_cut_set = set_of_cuts
0189             tmp_hist = plot.histogram.Clone("temp")    # Create temp hist
0190             all_cuts = join_cuts(cut_for_all_files, rootTree.cuts,
0191                                  current_cut_set, plot.cuts) # Set cuts
0192             rootTree.ttree.Draw( "%s >> temp" % plot.treeVariable, all_cuts,
0193                                  "goff")               # Draw with graphics off
0194             tmp_hist.Scale(rootTree.scale)             # Scale temp
0195             entries_before = plot.histogram.GetEntries()
0196             plot.histogram.Add(tmp_hist)               # Add temp hist to total
0197             entries_after = plot.histogram.GetEntries()
0198             print(" %3i %7i %20s >> %s/%s" % (i, entries_after-entries_before,
0199                                               plot.treeVariable, histname_fix,
0200                                               plot.histogram.GetName()), end=' ')
0201             if plot.cuts:
0202                 print("\textra cuts: %s" % plot.cuts, end=' ') # plot-specific cuts
0203             print()
0204
0205         rootTree.tfile.Close()                                    # Close TFile
0206         print("%s: Closed %s" % (datetime.now().strftime("%I:%M%p"),
0207                                  rootTree.fileName))
0208     return list_of_plots_to_write
0209
0210
0211 ######## Define the main program #############################################
0212 def tree2hists_main(config_file):
0213     try:
0214         # Import only certain variables
0215         sys.path.insert(0, '')
0216         _temp = __import__(config_file, globals(), locals(),
0217                            ['list_of_files','output_filename',
0218                             'cut_for_all_files','cut_sets','list_of_plots'], -1)
0219         list_of_files     = _temp.list_of_files
0220         output_filename   = _temp.output_filename
0221         cut_for_all_files = _temp.cut_for_all_files
0222         cut_sets          = _temp.cut_sets
0223         list_of_plots     = _temp.list_of_plots
0224         for rootTree in list_of_files:
0225             if not path.isfile(rootTree.fileName):
0226                 print("Error:\n  %s\nnot found for input." % rootTree.fileName)
0227                 sys.exit(1)
0228         hist_names = [plot.name for plot in list_of_plots]
0229         if len(hist_names)>len(set(hist_names)):
0230             print(hist_names)
0231             print("Error: Each plot needs a unique name, exiting...")
0232             sys.exit(1)
0233         if path.isfile(output_filename):
0234             print("Warning: %s exists" % output_filename)
0235     except Exception as e:
0236         print(e)
0237         print("Error with %s" % config_file)
0238         sys.exit(1)
0239
0240     if path.isfile('rootlogon.C'):
0241         print("Loading rootlogon.C")
0242         gROOT.Macro('rootlogon.C')    # Load functions from rootlogon script
0243
0244     if cut_sets:
0245         print("\n%i defined cut sets:" % len(cut_sets))
0246         for cut in cut_sets:
0247             name, title_fix, current_cut_set = cut
0248             print("  %s\t: '%s'" % (name, current_cut_set))
0249         cut_names = [name for name,num,cut in cut_sets]
0250         if len(cut_names)>len(set(cut_names)):
0251             print("Error: Each cut set needs a unique name, exiting...")
0252             sys.exit(1)
0253     else:
0254         cut_sets = [("","","")] # Make all plots, no extra cuts
0255
0256     print("\nCuts to apply to all files:\n\t'%s'" % cut_for_all_files)
0257
0258     start_time = datetime.now()
0259     list_of_plots_to_write = make_all_hists_all_files(list_of_files,
0260                                                       list_of_plots,
0261                                                       cut_for_all_files,
0262                                                       cut_sets)
0263     end_time = datetime.now()
0264     print("Done drawing all plots after %s." % duration_to_string(start_time, end_time))
0265
0266     #   Store and save/close files
0267     outputFile = TFile(output_filename, "recreate")
0268     if outputFile.IsZombie():
0269         print("Error opening %s for output exiting..." % output_filename)
0270         sys.exit(1)
0271     print("\nOpened output file. Saving histograms...")
0272     outputFile.cd()
0273     for set_of_cuts in cut_sets:
0274         outputFile.mkdir(set_of_cuts[0])
0275     print("   # entries  histogram")
0276     for i, (plot, cutset) in enumerate(list_of_plots_to_write):
0277         outputFile.cd(cutset[0])
0278         print(" %3i %7i  %s/%s" % (i, plot.histogram.GetEntries(),
0279                                    cutset[0],
0280                                    plot.histogram.GetName()))
0281         plot.histogram.Write()
0282     outputFile.Close()
0283     print("Done saving.")
0284     print("\nScaled & added histograms from %i TTrees saved in\n  %s" % (len(list_of_files), output_filename))
0285 ##############################################################################
0286
0287 def main():
0288     if len(sys.argv) > 1:
0289         if path.isfile(sys.argv[1]):
0290             config_file = sys.argv[1].split('.')[0]
0291             tree2hists_main(config_file)
0292         else:
0293             print("%s not found." % sys.argv[1])
0294             print("Create default config file by running tree2hists "
0295                   "with no argument.")
0296             sys.exit(1)
0297     else:
0298         if path.exists('t2h_config.py'):
0299             print("Run with specific config file, like:")
0300             print("  tree2hists t2h_config.py")
0301             sys.exit(1)
0302         write_default_T2H_config()
0303         sys.exit(0)
0304
0305 if __name__ == "__main__":
0306     main()