python/rootplot/rootmath.py

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