plot-postfit

#!/usr/bin/env python

# python imports
import os
import pickle
import logging

# ROOT/rootpy imports
from ROOT import TLatex, TGraphAsymmErrors, TH1 
from rootpy import asrootpy
from rootpy.io import root_open
from rootpy.plotting import Canvas, HistStack, Legend, set_style
from rootpy.plotting.utils import draw
from rootpy.plotting.style.atlas import ATLAS_style, ATLAS_label
from rootpy.utils.path import mkdir_p
from rootpy.extern.tabulartext import PrettyTable

# local imports
from mva import set_hsg4_style
from mva.categories import CATEGORIES
from mva.analysis import Analysis
from statstools.fitresult import Prefit_RooFitResult
from statstools.postfit import FitModel, ModelCalculator
from statstools.jobs import run_pool
from statstools.plotting import (get_uncertainty_graph, parse_name,
                                 get_category, get_binning, get_blinding,
                                 get_rebinned_graph, get_rebinned_hist)

log = logging.getLogger(os.path.basename(__file__))


def fit_workspace(file_name, ws_name, output_basename='frames', n_jobs=-1):
    """
    Fit the WS and compute the histograms and TGraphAssymErrors
    for the final plotting drawing

    Parameters
    ----------
    workspace : RooWorkspace
        HSG4 like workspace

    """
    with root_open(file_name) as file0:
        workspace = file0[ws_name]
        mc = workspace.obj("ModelConfig")
        obsData = workspace.data('obsData')
        workspace.saveSnapshot('StartingPoint', mc.GetPdf().getParameters(obsData))
        roo_min = workspace.fit()
        fit_res = roo_min.save()
        cats =  [cat for cat in mc.GetPdf()]
    
        output_root = output_basename + '.root'
        output_pickle = output_basename + '.pickle'
        log.info('Output: {0}'.format(output_root))
        log.info('Output: {0}'.format(output_pickle))
        root_open(output_root, 'recreate').Close()
        with open(output_pickle, 'w') as pickle_file:
            pickle.dump({}, pickle_file)
        workers = []
        # workspace.loadSnapshot('StartingPoint')
        # fit_res = Prefit_RooFitResult(fit_res, True)
        for cat in cats:
            log.info('retrieve plotting objects of {0} ...'.format(cat.name))
            workers.append(ModelCalculator(file0, workspace, cat, fit_res,
                                           output_root, output_pickle))
        run_pool(workers, n_jobs=n_jobs)
        

def plot_from_frame(file, frame, fit_var, category, year, binning=None, unblind=True):
    """
    """
    Ana = Analysis(2012)

    hbkg = file.Get('h_sum_bkg_{0}'.format(frame.GetName()))
    curve_uncert_bkg = frame.getCurve('FitError_AfterFit_sum_bkg_{0}'.format(frame.GetName()))
    graph_bkg_ws = get_uncertainty_graph(hbkg, curve_uncert_bkg)
    graph_bkg = get_rebinned_graph(graph_bkg_ws, binning, unblind=True) 
    graph_bkg.fillstyle='//'
    graph_bkg.color='black'
    graph_bkg.name = 'Background'
    graph_bkg.title = 'Uncert.'
    graph_bkg.legendstyle = 'F'

    data_ws = frame.getHist('Data')
    # HACK HACK HACK
    data_ws.__class__ = TGraphAsymmErrors
    data = get_rebinned_graph(asrootpy(data_ws), binning, unblind=unblind)
    data.name = 'Data'
    data.title = 'Data'

    hist_fake_ws = file.Get('h_Fakes_{0}'.format(frame.GetName()))
    hist_fake = get_rebinned_hist(hist_fake_ws, binning)
    hist_fake.color = Ana.qcd.hist_decor['color']
    hist_fake.fillstyle = 'solid'
    hist_fake.name = 'Fakes'
    hist_fake.title = Ana.qcd.label
    hist_fake.legendstyle = 'F'
    
    hist_others_ws = file.Get('h_Others_{0}'.format(frame.GetName()))
    hist_others = get_rebinned_hist(hist_others_ws, binning)
    hist_others.color = Ana.others.hist_decor['color']
    hist_others.fillstyle = 'solid'
    hist_others.name = 'Others'
    hist_others.title = Ana.others.label
    hist_others.legendstyle = 'F'

    hist_ztautau_ws = file.Get('h_Ztautau_{0}'.format(frame.GetName()))
    hist_ztautau = get_rebinned_hist(hist_ztautau_ws, binning)
    hist_ztautau.color = Ana.ztautau.hist_decor['color']
    hist_ztautau.fillstyle = 'solid'
    hist_ztautau.name = 'Ztautau'
    hist_ztautau.title = Ana.ztautau.label
    hist_ztautau.legendstyle = 'F'
    
    hist_signal_ws = file.Get('h_sum_sig_{0}'.format(frame.GetName()))
    hist_signal = get_rebinned_hist(hist_signal_ws, binning)
    hist_signal.color = Ana.higgs_125.hist_decor['linecolor']
    hist_signal.linestyle = 'solid'
    hist_signal.linewidth = 4
    hist_signal.name = 'Signal'
    hist_signal.title = Ana.higgs_125.label+ ' (best fit #mu)'
    hist_signal.legendstyle = 'F'
    


    hists = [hist_fake, hist_others, hist_ztautau, hist_signal]
    if (unblind is not True) and isinstance(unblind, (list, tuple)):
        hists = [hist_fake, hist_others, hist_ztautau]
    Model = HistStack(hists=hists, name='Model_{0}'.format(frame.GetName()))
    plotables = [Model]

    log.info(category.name)
    postfit_table = PrettyTable(['Sample']+['bin_{0}'.format(bin.idx) for bin in hist_ztautau.bins()]) 
    for h in [data, graph_bkg]+hists:
        row = [h.name]
        #if isinstance(h, Graph):
        for val in list(h.y()):
            row.append('{0:1.3f}'.format(val))
        postfit_table.add_row(row)
    log.info(postfit_table.get_string())


    # TODO: make plot identical to standard variable plots
    if fit_var == 'bdt':
        xtitle = 'BDT Score'
        categories = CATEGORIES['mva_all']
        logy = True
    else:
        xtitle = 'MMC MASS [GeV]'
        categories=CATEGORIES['cuts']+CATEGORIES['cuts_2011']
        logy = False

    c = Canvas()
    draw(plotables, pad=c, ypadding=(0.3, 0.3),
         xtitle=xtitle, ytitle= 'Events', logy=logy)
    if year==2011:
        text='Internal 2011'
        sqrts=7
    else:
        text='Internal 2012'
        sqrts=8
        
    ATLAS_label(0.2, 0.89, pad=c, sep=0.132, text=text, sqrts=sqrts)
    graph_bkg.Draw('sameE2')
    data.Draw('samePE')
    leg = Legend([data]+Model.GetHists()+[graph_bkg],
                 rightmargin=0.2,
                 margin=0.35,
                 topmargin=0.01,
                 textsize=20,
                 entrysep=0.02,
                 entryheight=0.04)
    leg.Draw()
    latex = TLatex(0.2, 0.85, category.label)
    latex.SetNDC()
    latex.SetTextSize(20)
    latex.Draw()
    c.RedrawAxis()
    if fit_var == 'bdt':
        c.SetLogy()
    return c


if __name__ == '__main__':
    from rootpy.extern.argparse import ArgumentParser
    parser = ArgumentParser()
    parser.add_argument('file')
    parser.add_argument('--name', default='combined')
    parser.add_argument('--fit-var', default='bdt', choices=['bdt', 'mmc'])
    parser.add_argument('--force-fit', action='store_true', default=False)
    parser.add_argument('--blind', action='store_true', default=False)
    parser.add_argument('--jobs', default=-1)
    args = parser.parse_args()

    if args.fit_var=='bdt':
        categories = CATEGORIES['mva_all']
    else:
        categories = CATEGORIES['cuts'] + CATEGORIES['cuts_2011']

    input = os.path.splitext(args.file)[0]
    output = input + '_postfit'
    plots = input + '_plots'

    if not os.path.exists(output+'.root') or args.force_fit:
        fit_workspace(args.file, args.name, output, n_jobs=args.jobs)
    
    if not os.path.exists(plots):
        mkdir_p(plots)

    set_hsg4_style(shape='rect')
    file = root_open(output + '.root')
    for _, _, names in file.walk(class_pattern='*RooPlot*'):
        for name in names:
            log.info('Channel: {0}'.format(name))
            year, category_name, mass = parse_name(name)
            log.info('Year: {0}; Mass: {1}; Category: {2}'.format(
                year, mass, category_name))
            category = get_category(category_name, categories)
            if category is None:
                # not found
                log.warning("skipping category: {0}".format(category_name))
                continue
            binning = get_binning(category, year, fit_var=args.fit_var)
            if args.blind:
                unblind = get_blinding(category, year, fit_var=args.fit_var)
                log.info(unblind)
            else:
                unblind=True
            canvas = plot_from_frame(file, file[name], args.fit_var,
                                     category, year, binning, unblind=unblind)
            canvas.SaveAs(os.path.join(plots, 'postfit_{0}.png'.format(name)))
    file.Close()