merge-ntup

#!/usr/bin/env python

from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument('-s', '--student', required=True)
parser.add_argument('-o', '--output', default=None)
parser.add_argument('paths', nargs='+')
args = parser.parse_args()

import re
import os
from glob import glob
from rootpy.io import root_open as ropen
from rootpy.tree import Tree
import ROOT
import logging

log = logging.getLogger('merge-ntup')

if args.output is None:
    args.output = '%s.root' % args.student

ds_pattern = re.compile('%s\.(?P<name>.+)\.root$' % (args.student))

with ropen(args.output, 'UPDATE') as outfile:
    for path in args.paths:
        if os.path.isdir(path):
            filenames = glob(os.path.join(path, '*.root'))
        else:
            filenames = [path]
        for filename in filenames:
            if os.path.normpath(filename) == os.path.normpath(args.output):
                continue
            match = re.search(ds_pattern, filename)
            if not match:
                log.warning("%s is not a valid filename" % filename)
                continue
            # replace . and - with _ for natural naming in PyTables
            name = match.group('name').replace('.', '_').replace('-', '_')
            if name in outfile:
                log.warning("skipping %s" % filename)
                continue
            log.info("merging in %s ..." % filename)
            with ropen(filename, 'READ') as infile:
                intree = infile.tau
                intree.deactivate("jet_*_original")
                intree.deactivate("jet_antikt4truth*")
                intree.deactivate("mcevt_*")
                intree.deactivate("mc_*")
                intree.activate("mc_event_weight")
                intree.activate("mc_weight")
                outfile.cd()
                outtree = intree.CloneTree(-1, "fast SortBasketsByEntry")
                outtree.OptimizeBaskets()
                outtree.SetName(name)
                outtree.Write(outtree.GetName(), ROOT.TObject.kOverwrite)
                # will need to be upated for the next skim
                outcutflow = infile.cutflow.Clone(name=name + '_cutflow')
                #outcutflow_event = infile.cutflow_event.Clone(name=name + '_cutflow_event')
                outfile.cd()
                outcutflow.Write(outcutflow.GetName(), ROOT.TObject.kOverwrite)
                #outcutflow_event.Write(outcutflow_event.GetName(), ROOT.TObject.kOverwrite)