rootanalyse

#!/usr/bin/env python
################################################################################
from RootAnalyser import lib_dir as RA_lib_dir
import time
import os
import sys
import argparse
import subprocess
import StringIO
import ROOT
from collections import OrderedDict
################################################################################
def process(args):
    
    for lib, path in lib_dir.items():
        
        libpath = '{}/{}'.format(path,lib)

        if not os.path.exists(libpath): 
            raise ParserError('Could not find "{}" in {}'.format(lib,path),parser)
        print 'Found {} in {}'.format(lib,path)
        print 'loading '+lib
        
        if lib == 'libExRootAnalysis.so':
            ROOT.gROOT.ProcessLine(".include {}".format(path))
        if lib == 'libDelphes.so':
            ROOT.gROOT.ProcessLine(".include {}/external".format(path))
        
        ROOT.gSystem.Load(libpath)


    # Get list of keys in file
    keys = []
    try:
        myfile = ROOT.TFile.Open(args.rootfile,'READ')
        list_of_keys = myfile.GetListOfKeys()
    except ReferenceError:
        args.rootfile = lhco_converter(args.rootfile)
        myfile = ROOT.TFile.Open(args.rootfile,'READ')
        list_of_keys = myfile.GetListOfKeys()
    
    for f in list_of_keys:
        keys.append(f.GetName())

    myfile.Close()

    try:# See if specified tree name is set of supported trees
        assert args.tree in keys
        treename = args.tree
   
    except AttributeError: # See if supported tree name is in list of keys
        trees = supported_trees.intersection(set(keys))
        if len(trees)==1: 
            treename = list(trees)[0]
        else:
            error_multitree='''
            More than one of the supported tree structures exists in {}. 
            Please specify which one to read using the "-t" option.
            '''.format(args.rootfile)
            raise ParserError(error_multitree, parser)
    except AssertionError:
        error_supported_trees = '''TREE option -t "{}" not supported.'''.format(args.tree)
        raise ParserError(error_supported_trees,parser)
    
    selector = args.selector.replace('.py','')
    selector_dir = os.path.dirname(args.selector)
    sys.path.append(selector_dir)
    print ''
    print 'Analysing {} with selector {}.py'.format(args.rootfile, selector)
    print 'Reading tree: {}'.format(treename)
    print ''
    # Process tree
    new_chain = ROOT.TChain(treename) 
    new_chain.Add(args.rootfile)
    start = time.time() # Timer
    print 'Begin analysis'
    new_chain.Process('TPySelector', selector) # selector file name without '.py'
    print 'Completed in {} seconds'.format((time.time() - start))
    print ''

def convert(args):
    lhco_converter(args.eventfile, outfile=args.rootfile)

    
def lhco_converter(infile, outfile=None):
    extension = infile.split('.')[-1].strip().lower()
    
    
    if extension.lower() == 'lhe':
        executable = 'ExRootLHEFConverter'
    elif extension.lower() == 'lhco':
        executable = 'ExRootLHCOlympicsConverter'
    else:
        print 'unrecognised file format: {}'.format(infile)
        sys.exit()

    if outfile is None:
        base = ''.join(os.path.basename(infile).split('.')[:-1])
        outfile = base + '.root'
        
    if os.path.exists(outfile):
        os.remove(outfile)
    
    # assumes the converters are in the same place as libExRootAnalysis.so
    exrootdir = RA_lib_dir['libExRootAnalysis.so']
    
    converter = subprocess.Popen(['{}/{}'.format(exrootdir,executable),
                                 infile, outfile], 
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE)

    print 'Converting {} to {} with {}'.format(infile, outfile, executable)
    out, err = converter.communicate()
    # print out
    if err: print err
    
    return outfile
        
################################################################################
class ParserError(Exception):
    def __init__(self, message, parser):
        # Call the base class constructor with the parameters it needs
        super(Exception, self).__init__(message)
        self.message = message
        self.parser = parser
    def __str__(self):
        rep = StringIO.StringIO()
        print >>  rep, self.message
        print >>  rep, ''
        print >>  rep, self.parser.format_help()
        return rep.getvalue()
        
################################################################################
supported_trees= {'LHCO','LHEF','Delphes'}
################################################################################
parser = argparse.ArgumentParser(description="Analyse an event file in LHCO or LHEF format using your TPySelector.")
parser.add_argument("--libexroot", metavar="EXROOT", type=str, default=argparse.SUPPRESS, help="Location of libExRootAnalysis.so")
parser.add_argument("--libdelphes", metavar="DELPHES", type=str, default=argparse.SUPPRESS, help="Location of libDelphes.so")
subparsers = parser.add_subparsers()
################################################################################
parser_proc = subparsers.add_parser('process', help='Analyse an event file in LHCO or LHEF format using your TPySelector.')
parser_proc.add_argument("rootfile", metavar="ROOTFILE", type=str,  help="Path to input root file.")
parser_proc.add_argument("selector", metavar="SELECTOR", type=str,  help="Path to user's TPySelector implementation without the '.py' extension")
parser_proc.add_argument("-t","--tree", metavar="TREE", type=str, default=argparse.SUPPRESS, help="Name of TTree to analyse (one of: {})".format(', '.join(supported_trees)))
parser_proc.add_argument("-o","--options", metavar="OPT", type=str, default="", help="Option string to feed into the analysis' Process() method")
parser_proc.set_defaults(run = process)
################################################################################
parser_conv = subparsers.add_parser('convert', help='Convert an .lhe or .lhco file into a ROOT tree using ExRootAnalysis.')
parser_conv.add_argument("eventfile", metavar="EVENTFILE", type=str,  help='Path to input file. Must have ".lhe" or ".lhco" extension')
parser_conv.add_argument("rootfile", metavar="ROOTFILE", type=str, default=None, nargs='?',help='Name of output root file.')
parser_conv.set_defaults(run = convert)
################################################################################
args = parser.parse_args()

################################################################################


################################################################################
error_lib_option='''
Location of {} unspecified. use "--{}" option, 
define an environment variable "{}" or modify 
lib_dir in __init__.py of the RootAnalyser package.
'''.format

lib_defaults = (
    ('libExRootAnalysis.so', {'parser_opt':'libexroot','env_var':'EXROOTANALYSIS'}),
    ('libDelphes.so', {'parser_opt':'libdelphes', 'env_var':'DELPHES_ROOT'})
)
    
try: # catch all ParserErrors raised
    # Load Libraries
    lib_dir = OrderedDict()
    for lib, defaults in lib_defaults:
        try:
            lib_dir[lib] = getattr(args, defaults['parser_opt'])
        except AttributeError:
            try:
                lib_dir[lib] = os.environ[defaults['env_var']]
                if not lib_dir[lib].strip(): 
                    raise KeyError
            except KeyError:
                lib_dir[lib] = RA_lib_dir[lib]
                print 'Trying default location for {}: {}'.format(lib, lib_dir[lib])
                if not os.path.exists(lib_dir[lib]):
                    errmsg = error_lib_option(lib, lib_dir['parser_opt'],
                                              lib_dir['env_var'])
                    raise ParserError(errmsg, parser)
    
    # call appropriate function
    args.run(args)
            
except ParserError as err:
    print ''
    print 'ParserError:'
    print ''
    print err
    sys.exit()

################################################################################