heppi-draw

#!/usr/bin/env python

from optparse   import OptionParser
from heppi      import heppi
from termcolor  import colored
import ROOT, logging, sys, logging, time

logging.basicConfig(format=colored('%(levelname)s:',attrs = ['bold'])
                    + colored('%(name)s:','blue') + ' %(message)s')
logger = logging.getLogger('heppi-draw')
heppi.ROOT.gROOT.ProcessLine(".x .root/rootlogon.C")

def get_options():
    parser = OptionParser()
    parser.add_option("-p", "--plotcard", dest="plotcard",default='plotcard.json',
                      help="""
                      Load the plot card in json format,
                      please use ./makeplotcard.py to create one
                      """,
                      metavar="FILE")

    parser.add_option("-s", "--sampledir", dest="sampledir",default='./data/',
                      help="""
                      Specify the detrectory where the trees are.
                      example: --filedir /data/trees
                      """)
    parser.add_option("-t", "--tree",
                      dest="treename", default='vbfTagDumper/trees/{sampleid}VBFDiJet',
                      help="Tree path in the root file that you want to use")
    parser.add_option("-a", "--all",
                      action="store_true", dest="draw_all", default=False,
                      help="draw all the variables specified in the plotcard")
    parser.add_option("-d", "--display",
                      action="store_true", dest="display", default=False,
                      help="draw all the variables specified in the plotcard")
    parser.add_option("-v", "--variable",
                      dest="variable",default="",
                      help="name of the variable you want to draw")
    parser.add_option("-n", "--allnorm",
                      action="store_true",dest="allnormhist",default=False,
                      help="all the histogram will be normilised")
    parser.add_option("-l", "--alllog",
                      action="store_true",dest="allloghist",default=False,
                      help="all the histogram will be in log scale")
    parser.add_option("--nocuts",
                      action="store_true",dest="nocuts",default=False,
                      help="makes histograms without applying any cut")
    parser.add_option("--scatter",
                      action="store_true",dest="scatter",default=False,
                      help="Draw the scatter plot booked in the plotcard")
    parser.add_option("--cutflow",
                      action="store_true",dest="print_cutflow",default=False,
                      help="Print cutflow: demograpgy and selection efficiencies")
    parser.add_option("--label", dest="label",default='VBF',
                      help="Label added in the plot file names")

    parser.add_option("--title", dest="title_on_plot",default=[],
                      help="""replace labels name from the plot card. Can specify each line seprated by ','
                      ex: --title 'selection','category 1' """)
    parser.add_option('--verbose', dest='verbose', default='INFO',
                     help="Increase verbosity (specify multiple times for more)")
    parser.add_option("--cut-card", dest="cut_card",default='',
                      help="""
                      Specify all the cut through a cut-card.
                      This might be included also in the plotcard.
                      Check the documentation, see the documentation for details.""")
    return parser.parse_args()

if __name__ == "__main__":
    (opt, args) = get_options()
    stack = None
    try:
        f = open(opt.plotcard)
        stack = heppi.instack(  plotcard  = opt.plotcard  ,
                                sampledir = opt.sampledir )
    except IOError as e:
        print("I/O error({0}): {1}".format(e.errno, e.strerror))

    stack.read_plotcard()
    stack.book_trees(make_sig_bkg_trees=opt.scatter)
    stack.print_cuts()
    stack.print_systematics()

    if opt.print_cutflow : stack.print_cutflow()
    if  opt.verbose == 'INFO':
        heppi.logger.setLevel(level=heppi.logging.INFO)
        logger.setLevel(level=logging.INFO)
    elif opt.verbose == 'DEBUG':
        heppi.logger.setLevel(level=heppi.logging.DEBUG)
        logger.setLevel(level=logging.DEBUG)


    if opt.display:
        ROOT.gROOT.SetBatch( ROOT.kFALSE )
    else:
        ROOT.gROOT.SetBatch( ROOT.kTRUE  )
        ROOT.gErrorIgnoreLevel = ROOT.kError

    # hadeling categories
    if opt.scatter:
        logger.info(colored(
            (' -- booked scatter plots :: (%i x %i)' % (len(stack.scatter_opt.xlist),len(stack.scatter_opt.ylist))),
            'red', attrs=['bold']
        ))
        for xvar in stack.scatter_opt.xlist :
            for yvar in stack.scatter_opt.ylist :
                if xvar != yvar:
                    logger.info(colored(' -- (%s x %s)' % (xvar, yvar),'green'))
                    if stack.scatter_opt.profile:stack.scatter(xvar,yvar,opt.label,make_profiles=True)
                    if stack.scatter_opt.scatter:stack.scatter(xvar,yvar,opt.label,make_profiles=False)


    # Part drawing 1D variables:
    if opt.draw_all and opt.variable == '':
        if len(stack.options.categories) > 0 :
            for cat in stack.options.categories:
                logger.info(colored(" --- Drawing category : %10s " % cat.get('name'), 'yellow', 'on_red'))
                logger.info(colored(
                    ('(%i) booked variables will be drawn :: ' % len(stack.variables)), 'red', attrs=['bold']
                ))

                for xvar in stack.variables :
                    if opt.allnormhist : stack.variables[xvar].norm = True
                    stack.options.label.append(cat.get('label'))
                    _cuts_ = cat.get('operator','&&').join(cat.get('cut')) if type(cat.get('cut')) == type([]) else cat.get('cut')
                    stack.draw( varkey = xvar,
                                label  = opt.label+'_'+cat.get('name'),
                                select = _cuts_)
                    stack.options.label.pop() # remove the last label after drawing
        else:
            logger.info(colored(
                ('(%i) booked variables will be drawn :: ' % len(stack.variables)), 'red', attrs=['bold']
            ))
            for xvar in stack.variables :
                if opt.allnormhist : stack.variables[xvar].norm = True
                stack.draw(
                    varkey=xvar,
                    label=opt.label
                )
    else:
        if opt.variable != '':
            if len(stack.options.categories) > 0 :
                for cat in stack.options.categories:
                    logger.info(colored(" --- Drawing category : %10s " % cat.get('name'), 'yellow', 'on_red'))
                    if opt.allnormhist : stack.variables[opt.variable].norm = True
                    stack.options.label.append(cat.get('label'))
                    _cuts_ = cat.get('operator','&&').join(cat.get('cut')) if type(cat.get('cut')) == type([]) else cat.get('cut')
                    stack.draw( varkey = opt.variable,
                                label  = opt.label+'_'+cat.get('name'),
                                select = _cuts_ )
                    stack.options.label.pop() # remove the last label after drawing
            else:
                if opt.allnormhist : stack.variables[opt.variable].norm = True
                stack.draw(opt.variable,opt.label)
            if opt.display:
                raw_input('... press any key to exit ...')
        else:
            logging.error('please specify the variable you what to draw ...')