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
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
0048
0049 from .version import __version__
0050
0051
0052
0053
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)
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 = {}
0088 dict_tfiles = {}
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
0093 dict_tfiles[target.filename] = ROOT.TFile(target.filename, 'read')
0094
0095
0096
0097 f = ROOT.TFile(targets[0].filename, 'read')
0098 if f.GetDirectory(targets[0].path):
0099
0100 destdir = '/'
0101 histnames = [os.path.basename(x) for x in
0102 rootglob(f, targets[0].path + '/*')]
0103 f.Close()
0104
0105
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
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
0127
0128 def walk_rootfile(rootfile, path=''):
0129
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
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
0272 for i in range(1,len(hists)):
0273 sumhist.Add(scale_with_error(hists[i], scales[i], scale_errors[i]))
0274
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()