Skip to content
Find file
Fetching contributors…
Cannot retrieve contributors at this time
executable file 199 lines (159 sloc) 6.15 KB
#!/usr/bin/env python
import argparse
import rootpy
import ROOT
import fnmatch
import os, sys
class formatter_class(argparse.ArgumentDefaultsHelpFormatter,
argparse.RawTextHelpFormatter):
pass
parser = argparse.ArgumentParser(
formatter_class=formatter_class)
parser.add_argument('-v', '--verbose', action='store_true', default=False)
parser.add_argument('--version', action='version', version=rootpy.__version__,
help="show the version number and exit")
parser.add_argument('-p', '--pattern', default='*.root*',
help="files must match this pattern when searching in directories")
subparsers = parser.add_subparsers()
def entries(args):
chain = make_chain(args)
if args.selection is None:
print "%i total entries" % chain.GetEntries()
else:
from rootpy.tree import Cut
selection = str(Cut(args.selection))
print "%i entries after selection %s" % (
chain.GetEntries(selection), selection)
parser_entries = subparsers.add_parser('entries')
parser_entries.add_argument('-s', '--selection', default=None,
help="only entries satisfying this cut will be included in total")
parser_entries.add_argument('tree',
help="name of tree (including path) in each file")
parser_entries.add_argument('files', nargs='+')
parser_entries.set_defaults(op=entries)
def draw(args):
from rootpy.interactive import wait
wait(True)
from rootpy.plotting import Canvas
canvas = Canvas()
def update(*args, **kwargs):
canvas.Modified()
canvas.Update()
chain = make_chain(args,
onfilechange=[(update, ())])
if args.selection is None:
chain.Draw(args.expression, '', args.draw)
else:
from rootpy.tree import Cut
selection = str(Cut(args.selection))
chain.Draw(args.expression, selection, args.draw)
parser_draw = subparsers.add_parser('draw')
parser_draw.add_argument('-e', '--expression', required=True,
help="expression to be drawn")
parser_draw.add_argument('-s', '--selection', default=None,
help="only entries satisfying this cut will be drawn")
parser_draw.add_argument('-d', '--draw', default='',
help="draw options")
parser_draw.add_argument('--staged', action='store_true', default=False,
help="update the canvas after each file is drawn")
parser_draw.add_argument('tree',
help="name of tree (including path) in each file")
parser_draw.add_argument('files', nargs='+')
parser_draw.set_defaults(op=draw)
def sum(args):
from rootpy.interactive import wait
wait(True)
from rootpy.plotting import Canvas
canvas = Canvas()
from rootpy.io import open as ropen
total = None
for filename in args.files:
f = ropen(filename)
h = f.get(args.hist)
if total is None:
total = h.Clone()
total.SetDirectory(0)
else:
total += h
f.close()
if total is not None:
total.Draw(args.draw)
# wait not working...
raw_input()
parser_sum = subparsers.add_parser('sum')
parser_sum.add_argument('-d', '--draw', default='',
help="draw options")
parser_sum.add_argument('hist',
help="name of histogram (including path) in each file")
parser_sum.add_argument('files', nargs='+')
parser_sum.set_defaults(op=sum)
def merge(args):
chain = make_chain(args)
if os.path.exists(args.output):
sys.exit("Output destination already exists.")
print "Merging tree %s in %d files into %s ..." % (
args.tree, len(args.files), args.output)
chain.Merge(args.output,
'fast SortBasketsBy%s' % args.sort_by.capitalize())
parser_merge = subparsers.add_parser('merge')
parser_merge.add_argument('-o', '--output', required=True,
help="output file name")
parser_merge.add_argument('tree',
help="name of tree (including path) in each file")
parser_merge.add_argument('--sort-by', choices=('offset', 'branch', 'entry'),
default='offset',
help="""\
When using 'offset' the baskets are written in
the output file in the same order as in the original file
(i.e. the basket are sorted on their offset in the original
file; Usually this also means that the baskets are sorted
on the index/number of the *last* entry they contain)
When using 'branch' all the baskets of each
individual branches are stored contiguously. This tends to
optimize reading speed when reading a small number (1->5) of
branches, since all their baskets will be clustered together
instead of being spread across the file. However it might
decrease the performance when reading more branches (or the full
entry).
When using 'entry' the baskets with the lowest
starting entry are written first. (i.e. the baskets are
sorted on the index/number of the first entry they contain).
This means that on the file the baskets will be in the order
in which they will be needed when reading the whole tree
sequentially.""")
parser_merge.add_argument('files', nargs='+')
parser_merge.set_defaults(op=merge)
args = parser.parse_args()
def find_files(dirs, pattern=None):
for arg in dirs:
if os.path.isfile(arg):
yield arg
continue
elif os.path.isdir(arg):
for root, dirnames, filenames in os.walk(arg):
if pattern is not None:
for filename in fnmatch.filter(filenames, pattern):
yield os.path.join(root, filename)
else:
for filename in filenames:
yield filename
def make_chain(args, **kwargs):
if hasattr(args, 'staged') and args.staged:
from rootpy.tree import TreeChain
files = list(find_files(args.files, args.pattern))
chain = TreeChain(
args.tree,
files,
verbose=args.verbose,
**kwargs)
nfiles = len(files)
else:
nfiles = 0
chain = ROOT.TChain(args.tree, '')
for filename in find_files(args.files, args.pattern):
nfiles += 1
if args.verbose:
print filename
chain.Add(filename)
return chain
args.op(args)
Something went wrong with that request. Please try again.