heppi-make-plotcard

#!/usr/bin/env python


import sys, os, glob, json, re
print("version python : ",  sys.version)
import ROOT
from optparse import OptionParser
from heppi import heppi


# --------------------
# read arguments
# --------------------

def create_json(rootfile='file.root', treename=None, jout=''):
    """
    Create a json template file for
    plotng all the branches in a given
    tree
    """
    root = ROOT.TFile.Open(rootfile)
    treenm = []

    def getall(d, basepath=""):
        try:
            for key in d.GetListOfKeys():
                kname = key.GetName()
                if key.IsFolder():
                    for i in getall(d.Get(kname), os.path.join(basepath,kname)):
                        yield i
                else:
                    yield basepath+kname, d.Get(kname)
        except :
             yield basepath, d

    trees = []
    for k, o in getall(root):
        if type(o) == type(ROOT.TTree()) and o.GetEntries() > 0:
            if o.GetName() not in treename:
                continue
            trees.append(k.split(o.GetName())[0] + o.GetName())
    if len(trees) > 0:
        heppi.logger.info(' The following trees have been found on this root file :')
        heppi.logger.info( trees )
        heppi.logger.info(' Only "' + trees[0] + '" will be used to build the plotcard')
    else:
        heppi.logger.error( 'No TTree found on this root file or no empty TTree' )

    tree = ROOT.gDirectory.Get(trees[0])
    varlist={}
    for var in tree.GetListOfLeaves():
        L    = tree.GetBranch(str(var))
        if '[' in var.GetTitle() and ']' in var.GetTitle():
            dim = [int(s) for s in re.findall('[-+]?\d*\.\d+|\d+', var.GetTitle().split('[')[1])]
            for idim in range(0, dim[0]):
                _var_ = var.GetTitle().replace('[%i]'%dim[0], '[%i]'%idim)
                tree.Project("_h_"+_var_.split('[')[0],_var_)
                _h_  = ROOT.gDirectory.Get('_h_' + _var_.split('[')[0])
                nbin = _h_.GetNbinsX()
                xmin = _h_.GetXaxis().GetXmin()
                xmax = _h_.GetXaxis().GetXmax()
                hist = "(%i,%1.3f,%1.3f)" % (nbin,xmin,xmax)
                varlist[_var_]={
                    'cut'  :'',
                    "hist" : hist,
                    "cut"  : "",
                    "blind": "",
                    "log"  : True ,
                    "norm" : False,
                    "title": var.GetTitle().replace('[%i]'%dim[0], '_%i'%idim)
                }
                print ('variable: %25s' % var.GetTitle().replace('[%i]'%dim[0], '[%i]'%idim))
        else:

            tree.Project("_h_"+var.GetTitle(),var.GetTitle())
            _h_ = ROOT.gDirectory.Get('_h_'+var.GetTitle())
            nbin = _h_.GetNbinsX()
            xmin = _h_.GetXaxis().GetXmin()
            xmax = _h_.GetXaxis().GetXmax()
            hist = "(%i,%1.1f,%1.1f)" % (nbin,xmin,xmax)
            varlist[var.GetTitle()]={
                'cut'  :'',
                "hist" : hist,
                "cut"  : "",
                "blind": "",
                "log"  : True ,
                "norm" : False,
                "title": var.GetTitle()
            }
            print ('variable: %25s' % var.GetTitle())
    samples_new = {}
    sample_tree = {'name':treename}
    selection   = {'cutflow':''}
    plot_options = heppi.options()
    with open(jout, "w") as file:
        json.dump({'variables':varlist,
                   'processes':{
                        "SMprocess": {
                            "color": 138,
                            "order": 0,
                            "files" : rootfile,
                            "title" : 'SMprocess',
                            "tree"  : trees[0],
                            "cut"   : "",
                            "label" : "background"
                        },
                   },
                   'option':plot_options.__dict__ } ,
                  file, indent=2)
        file.close()

# --------------------
# read arguments
# --------------------
if __name__ == "__main__":
    parser = OptionParser()
    parser.add_option("-r", "--load", dest="rootfile",default='rootfile.root',
                      help="create a plot card from tree", metavar="FILE")
    parser.add_option("-t", "--tree", dest="treename",default=None,
                      help="Dumper tree that you wnat to plot", metavar="FILE")
    parser.add_option("-o", "--out", dest="jsonfile",default="plotcard.json",
                      help="specify output json file")

    (options, args) = parser.parse_args()

    if os.path.exists(options.rootfile) :
        create_json( options.rootfile,
                     options.treename,
                     options.jsonfile)
    else:
        parser.print_help()