File indexing completed on 2024-11-25 02:29:52
0001
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
0009
0010
0011
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
0036
0037 import sys
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
0055 from math import pi
0056 from array import array
0057 from datetime import datetime
0058 from os import path
0059
0060
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
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")
0166 if rootTree.tfile.IsZombie():
0167 print("Error opening %s, exiting..." % rootTree.fileName)
0168 sys.exit(0)
0169 try:
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
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")
0189 all_cuts = join_cuts(cut_for_all_files, rootTree.cuts,
0190 current_cut_set, plot.cuts)
0191 rootTree.ttree.Draw( "%s >> temp" % plot.treeVariable, all_cuts,
0192 "goff")
0193 tmp_hist.Scale(rootTree.scale)
0194 entries_before = plot.histogram.GetEntries()
0195 plot.histogram.Add(tmp_hist)
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=' ')
0202 print()
0203
0204 rootTree.tfile.Close()
0205 print("%s: Closed %s" % (datetime.now().strftime("%I:%M%p"),
0206 rootTree.fileName))
0207 return list_of_plots_to_write
0208
0209
0210
0211 def tree2hists_main(config_file):
0212 try:
0213
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')
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 = [("","","")]
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
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()