Back to home page

Project CMSSW displayed by LXR

 
 

    


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

0001 """
0002 rootmath description
0003 """
0004 
0005 from builtins import range
0006 __license__ = '''\
0007 Copyright (c) 2009-2010 Jeff Klukas <klukas@wisc.edu>
0008 
0009 Permission is hereby granted, free of charge, to any person obtaining a copy
0010 of this software and associated documentation files (the "Software"), to deal
0011 in the Software without restriction, including without limitation the rights
0012 to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
0013 copies of the Software, and to permit persons to whom the Software is
0014 furnished to do so, subject to the following conditions:
0015 
0016 The above copyright notice and this permission notice shall be included in
0017 all copies or substantial portions of the Software.
0018 
0019 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
0020 IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
0021 FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
0022 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
0023 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
0024 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
0025 THE SOFTWARE.
0026 '''
0027 
0028 
0029 ##############################################################################
0030 ######## Import python libraries #############################################
0031 
0032 import sys
0033 import shutil
0034 import math
0035 import os
0036 import re
0037 import tempfile
0038 import copy
0039 import fnmatch
0040 from . import argparse
0041 from os.path import join as joined
0042 from .utilities import rootglob, loadROOT
0043 
0044 ROOT = loadROOT()
0045 
0046 ##############################################################################
0047 ######## Define globals ######################################################
0048 
0049 from .version import __version__          # version number
0050 
0051 
0052 ##############################################################################
0053 ######## Classes #############################################################
0054 
0055 class Target(object):
0056     """Description."""
0057     def __init__(self, filename, path='', scale=1., scale_error=None):
0058         self.filename = filename
0059         self.path = path
0060         self.scale = scale
0061         self.scale_error = scale_error
0062     def __repr__(self):
0063         return "%s:%s:%f" % (self.filename, self.path, self.scale)
0064 
0065 def newadd(outfile, targets, dest_path=""):
0066     """Description."""
0067     if allsame([x.filename for x in targets]):
0068         f = ROOT.TFile(targets[0].filename, 'read')
0069         paths = [x.path for x in targets]
0070         scales = [x.scale for x in targets]
0071         scale_errors = [x.scale_error for x in targets]
0072         if f.GetDirectory(paths[0]):
0073             destdir = pathdiff2(paths)    # What does this do?
0074             for h in [os.path.basename(x) for x in
0075                       rootglob(f, paths[0] + '/*')]:
0076                 hists = [f.GetDirectory(x).Get(h) for x in paths]
0077                 if not alltrue([x and x.InheritsFrom('TH1') for x in hists]):
0078                     continue
0079                 dest = joined(destdir, h)
0080                 add(outfile, dest, hists, scales, dest_path, scale_errors=scale_errors)
0081         else:
0082             hists = [f.Get(x) for x in paths]
0083             if alltrue([x and x.InheritsFrom('TH1') for x in hists]):
0084                 dest = pathdiff2(paths)
0085                 add(outfile, dest, hists, scales, scale_errors=scale_errors)
0086     else:
0087         dict_targets = {}  # Stores paths and scales, key = filename
0088         dict_tfiles = {}   # Stores map from filenames to Root.TFile() objects
0089         for target in targets:
0090             dict_targets.setdefault(target.filename, []).append((target.path, target.scale))
0091             if (target.filename not in dict_tfiles):
0092                 # Only open root files once
0093                 dict_tfiles[target.filename] = ROOT.TFile(target.filename, 'read')
0094         # dict_targets now a dictionary, with keys the filenames, example:
0095         # {'fileA.root': [('path0',scale0), ('path1', scale1)],
0096         #  'fileB.root': [('path3', scale3)]}
0097         f = ROOT.TFile(targets[0].filename, 'read')
0098         if f.GetDirectory(targets[0].path):
0099             # Create list of histograms to get
0100             destdir = '/'               # should probably use pathdiff2 somehow
0101             histnames = [os.path.basename(x) for x in
0102                          rootglob(f, targets[0].path + '/*')]
0103             f.Close()
0104             # For each histogram name found, grab it from
0105             # every file & path
0106             for histname in histnames:
0107                 hists = []
0108                 scales = []
0109                 for filename in dict_targets:
0110                     tfile_cur = dict_tfiles[filename]
0111                     for path, scale in dict_targets[filename]:
0112                         hists.append(tfile_cur.GetDirectory(path).Get(histname))
0113                         scales.append(scale)
0114                         #print "%s:%s:%s:%f" % (filename, path, histname, scale)
0115                 if not alltrue([x and x.InheritsFrom('TH1') for x in hists]):
0116                     continue
0117                 dest = joined(destdir, histname)
0118                 add(outfile, dest, hists, scales, dest_path)
0119         else:
0120             print("Code not written yet to add histograms from multiple files")
0121             return
0122         return
0123 
0124 
0125 ##############################################################################
0126 ######## Implementation ######################################################
0127 
0128 def walk_rootfile(rootfile, path=''):
0129     #### Yield (path, folders, objects) for each directory under path.
0130     keys = rootfile.GetDirectory(path).GetListOfKeys()
0131     folders, objects = [], []
0132     for key in keys:
0133         name = key.GetName()
0134         classname = key.GetClassName()
0135         newpath = joined(path, name)
0136         dimension = 0
0137         if 'TDirectory' in classname:
0138             folders.append(name)
0139         else:
0140             objects.append(name)
0141     yield path, folders, objects
0142     for folder in folders:
0143         for x in walk_rootfile(rootfile, joined(path, folder)):
0144             yield x
0145 
0146 def allsame(iterable):
0147     for element in iterable:
0148         if element != iterable[0]:
0149             return False
0150     return True
0151 
0152 def alltrue(iterable):
0153     for element in iterable:
0154         if element != True:
0155             return False
0156     return True
0157 
0158 def pathdiff(paths, joiner):
0159     """
0160     Return the appropriate destination for an object.
0161     
0162     In all cases, the result will be placed in the deepest directory shared by
0163     all paths.  If the histogram names are the same, the result will be named
0164     based on the first directories that they do not share.  Otherwise, the 
0165     result will be named based on the names of the other histograms.
0166 
0167     >>> pathdiff(['/dirA/dirB/dirX/hist', '/dirA/dirB/dirY/hist'], '_div_')
0168     '/dirA/dirB/dirX_div_dirY'
0169     >>> pathdiff(['/dirA/hist1', '/dirA/hist2', '/dirA/hist3'], '_plus_')
0170     '/dirA/hist1_plus_hist2_plus_hist3'
0171     >>> pathdiff(['/hist1', '/dirA/hist2'], '_minus_')
0172     '/hist1_minus_hist2'
0173     """
0174     paths = [x.split('/') for x in paths]
0175     dest = '/'
0176     for i in range(len(paths[0])):
0177         if allsame([p[i] for p in paths]):
0178             dest = joined(dest, paths[0][i])
0179         else:
0180             break
0181     name = joiner.join([p[-1] for p in paths])
0182     if allsame([p[-1] for p in paths]):
0183         for i in range(len(paths[0])):
0184             if not allsame([p[i] for p in paths]):
0185                 name = joiner.join([p[i] for p in paths])
0186     return joined(dest, name)
0187 
0188 def pathdiff2(paths, joiner='__', truncate=False):
0189     """
0190     Placeholder.
0191     """
0192     paths = [x.split('/') for x in paths]
0193     commonbeginning = ''
0194     for i in range(len(paths[0])):
0195         if allsame([p[i] for p in paths]):
0196             commonbeginning = joined(commonbeginning, paths[0][i])
0197         else:
0198             break
0199     commonending = ''
0200     for i in range(-1, -1 * len(paths[0]), -1):
0201         if allsame([p[i] for p in paths]):
0202             commonending = joined(paths[0][i], commonending)
0203         else:
0204             break
0205     #return commonbeginning, commonending
0206     if truncate:
0207         return commonending
0208     else:
0209         return joined(commonbeginning, commonending)
0210 
0211 def pathdiff3(paths, joiner='__'):
0212     """
0213     Return the appropriate destination for an object.
0214     
0215     If the final objects in each path match, then the return value will be the
0216     matching part of the paths.  Otherwise, the output path will simply be those
0217     names joined together with *joiner*.  See the examples below.
0218     
0219     >>> pathdiff3(['/dirA/dirX/hist', '/dirA/dirY/hist'])
0220     '/hist'
0221     >>> pathdiff3(['/dirA/dirX/dirB/hist', '/dirA/dirY/dirB/hist'])
0222     '/dirB/hist'
0223     >>> pathdiff3(['/dirA/hist1', '/dirA/hist2', '/dirA/hist3'], '_plus_')
0224     '/hist1_plus_hist2_plus_hist3'
0225     >>> pathdiff3(['/hist1', '/dirA/hist2'], '_div_')
0226     '/hist1_div_hist2'
0227     """
0228     paths = [x.split('/') for x in paths]
0229     if allsame([x[-1] for x in paths]):
0230         dest = paths[0][-1]
0231         for i in range(-2, min([len(x) for x in paths]) * -1, -1):
0232             if allsame([p[i] for p in paths]):
0233                 dest = joined(paths[0][i], dest)
0234             else:
0235                 break
0236         return '/' + dest
0237     else:
0238         return '/' + joiner.join([x[-1] for x in paths])
0239 
0240 def operator_func(fn):
0241     def newfunc(outfile, dest, hists, scales=None, dest_path="", scale_errors=None):
0242         outfile.cd()
0243         for d in os.path.dirname(dest).split('/'):
0244             if not ROOT.gDirectory.GetDirectory(d):
0245                 ROOT.gDirectory.mkdir(d)
0246             ROOT.gDirectory.cd(d)
0247         fn(outfile, dest, hists, scales, dest_path, scale_errors)
0248     return newfunc
0249 
0250 def scale_with_error(hist, scale, scale_error=None):
0251     '''Scale a histogram by a scale factor that has an error.
0252     This takes into account the scale error to set new error bars.'''
0253     hist_new = hist.Clone()
0254     if scale_error:
0255         for i in range(hist_new.GetNbinsX()+2):
0256             hist_new.SetBinContent(i, scale)
0257             hist_new.SetBinError(i, scale_error)
0258         hist_new.Multiply(hist)
0259     else:
0260         hist_new.Scale(scale)
0261     return hist_new
0262 
0263 @operator_func
0264 def add(outfile, dest, hists, scales=None, dest_path="", scale_errors=None):
0265     if not scales:
0266         scales = [1. for i in range(len(hists))]
0267     if not scale_errors:
0268         scale_errors = [None for i in range(len(hists))]
0269     sumhist = hists[0].Clone(os.path.basename(dest))
0270     sumhist = scale_with_error(sumhist, scales[0], scale_errors[0])
0271     #sumhist.Scale(scales[0])
0272     for i in range(1,len(hists)):
0273         sumhist.Add(scale_with_error(hists[i], scales[i], scale_errors[i]))
0274         #sumhist.Add(hists[i], scales[i])
0275     if dest_path:
0276         outfile.cd()
0277         if not ROOT.gDirectory.GetDirectory(dest_path):
0278             ROOT.gDirectory.mkdir(dest_path)
0279         ROOT.gDirectory.cd(dest_path)
0280     sumhist.Write()
0281     ROOT.gDirectory.cd("/")
0282 
0283 @operator_func
0284 def subtract(outfile, dest, hists):
0285     diffhist = hists[0].Clone(os.path.basename(dest))
0286     for hist in hists[1:]:
0287         diffhist.Add(hist, -1)
0288     diffhist.Write()
0289 
0290 @operator_func
0291 def divide(outfile, dest, numer, denom):
0292     quotient = numer.Clone(os.path.basename(dest))
0293     quotient.Divide(numer, denom)
0294     quotient.Write()
0295 
0296 @operator_func
0297 def bayes_divide(outfile, dest, numer, denom):
0298     quotient = ROOT.TGraphAsymmErrors()
0299     quotient.SetName(os.path.basename(dest))
0300     quotient.BayesDivide(numer, denom)
0301     quotient.Write()
0302 
0303 def main():
0304     parser = argparse.ArgumentParser()
0305     parser.add_argument('filenames', type=str, nargs='+',
0306                        help='root files to process')
0307     parser.add_argument('--dirs', type=str, nargs='+', default=['/'],
0308                         help='target directories in the root files; paths to '
0309                         'histograms will be relative to these')
0310     parser.add_argument('--add', default=[], action='append', nargs='*',
0311                         help='a list of directories or histograms to add')
0312     parser.add_argument('--subtract', default=[], action='append', nargs='*',
0313                         help='a list of directories or histograms to subtract')
0314     parser.add_argument('--divide', default=[], action='append', nargs='*',
0315                         help='2 directories or histograms to divide')
0316     parser.add_argument('--bayes-divide', default=[], action='append', nargs='*',
0317                         help='2 directories or histograms from which to make '
0318                         'an efficiency plot')
0319     args = parser.parse_args()
0320     separators = {'add' : '_plus_',
0321                   'subtract' : '_minus_',
0322                   'divide' : '_div_',
0323                   'bayes_divide' : '_eff_'}
0324 
0325     files = [ROOT.TFile(x, 'read') for x in args.filenames]
0326     outfile = ROOT.TFile('out.root', 'recreate')
0327     dirs = []
0328     for d in args.dirs:
0329         dirs += rootglob(files[0], d)
0330 
0331     if len(files) == 1:
0332         f = files[0]
0333         for thisdir in dirs:
0334             for operation_type, separator in separators.items():
0335                 for arg_set in getattr(args, operation_type):
0336                     paths = [joined(thisdir, x) for x in arg_set]
0337                     if f.GetDirectory(paths[0]):
0338                         destdir = pathdiff(paths, separator)
0339                         for target in [os.path.basename(x) for x in
0340                                        rootglob(f, paths[0] + '/*')]:
0341                             hists = [f.GetDirectory(x).Get(target)
0342                                      for x in paths]
0343                             if not alltrue([x and x.InheritsFrom('TH1')
0344                                             for x in hists]):
0345                                 continue
0346                             dest = joined(destdir, target)
0347                             math_func = globals()[operation_type]
0348                             math_func(outfile, dest, hists)
0349                     else:
0350                         hists = [f.GetDirectory(thisdir).Get(x) for x in paths]
0351                         if not alltrue([x and x.InheritsFrom('TH1') 
0352                                         for x in hists]):
0353                             continue
0354                         dest = pathdiff(paths, separator)
0355                         math_func = globals()[operation_type]
0356                         math_func(outfile, dest, hists)
0357     else:
0358         for operation_type, separator in separators.items():
0359             arg_sets = getattr(args, operation_type)
0360             if arg_sets and arg_sets != [[]]:
0361                 raise ValueError("No arguments to --%s allowed when multiple "
0362                                  "files are specified" % operation_type)
0363             elif arg_sets:
0364                 if 'divide' in operation_type and len(files) != 2:
0365                     raise ValueError("Exactly 2 files are expected with --%s; "
0366                                      "%i given" % (operation_type, len(files)))
0367                 for path, folders, objects in walk_rootfile(files[0]):
0368                     for obj in objects:
0369                         hists = [x.GetDirectory(path).Get(obj) for x in files]
0370                         if not alltrue([x and x.InheritsFrom('TH1') 
0371                                         for x in hists]):
0372                             continue
0373                         math_func = globals()[operation_type]
0374                         math_func(outfile, joined(path, obj), hists)
0375 
0376     outfile.Close()
0377 
0378 if __name__ == '__main__':
0379     import doctest
0380     doctest.testmod()