Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-11-25 02:29:52

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