TTTT2/yields.py

#!/usr/bin/env python
import matplotlib.pyplot as plt

from filval.result_set import ResultSet
from filval.histogram_utils import hist
from filval.plotter import (decl_plot, render_plots, hist_plot, Plot)


def generate_dashboard(plots, output='dashboard.htm'):
    from jinja2 import Environment, PackageLoader, select_autoescape
    from os.path import join
    from urllib.parse import quote

    env = Environment(
        loader=PackageLoader('plots', 'templates'),
        autoescape=select_autoescape(['htm', 'html', 'xml']),
    )
    env.globals.update({'quote': quote,
                        'enumerate': enumerate,
                        'zip': zip,
                        })

    def render_to_file(template_name, **kwargs):
        with open(join('output', template_name), 'w') as tempout:
            template = env.get_template(template_name)
            tempout.write(template.render(**kwargs))

    def get_by_n(objs, n=2):
        objs = list(objs)
        while objs:
            yield objs[:n]
            objs = objs[n:]

    render_to_file(output, plots=get_by_n(plots, 3),
                   outdir="figures/")


@decl_plot
def plot_yield(rss):
    r''' '''
    from filval.plotter import StackHist
    ft, ttw, ttz, tth = map(lambda rs: hist(rs.SRs), rss)
    # ft10 = ft[0]*10, ft[1]*10, ft[2]
    sh = StackHist()
    sh.add_mc_background(rss[1].SRs, 'TTW')
    sh.set_mc_signal(rss[0].SRs, 'TTTT')
    sh.draw(plt.gca())
    # hist_plot(ft10, include_errors=False, stats=False,
    #           color='k', label='TTTT (x10)')
    
    # bg = 
    # plt.hist(
    # hist_plot(ttw, include_errors=False, stats=False,
    #           color='g', label='TTW')
    # hist_plot(ttz, include_errors=False, stats=False,
    #           color='r', label='TTZ')
    # hist_plot(tth, include_errors=False, stats=False,
    #           color='b', label='TTH')
    # plt.xlabel('Signal Region')
    # plt.legend()


if __name__ == '__main__':
    # First create a ResultSet object which loads all of the objects from output.root
    # into memory and makes them available as attributes
    # if len(sys.argv) != 2:
    #     raise ValueError("please supply root file")
    rss = (ResultSet("ft", 'yield_ft.root'),
           ResultSet("ttw", 'yield_ttw.root'),
           ResultSet("ttz", 'yield_ttz.root'),
           ResultSet("tth", 'yield_tth.root'))
    rss_notau = (ResultSet("ft_notau", 'yield_ft_notau.root'),
                 ResultSet("ttw_notau", 'yield_ttw_notau.root'),
                 ResultSet("ttz_notau", 'yield_ttz_notau.root'),
                 ResultSet("tth_notau", 'yield_tth_notau.root'))

    # Next, declare all of the (sub)plots that will be assembled into full
    # figures later
    yield_tau = (plot_yield, (rss,), {})
    yield_notau = (plot_yield, (rss_notau,), {})

    # Now assemble the plots into figures.
    plots = [
        Plot([[yield_tau]],
             'Yield With Tau'),
        Plot([[yield_notau]],
             'Yield Without Tau'),
    ]

    # Finally, render and save the plots and generate the html+bootstrap
    # dashboard to view them
    render_plots(plots, to_disk=False)
    generate_dashboard(plots)