TTTT2/plots.py

#!/usr/bin/env python
import sys

import matplotlib.pyplot as plt

from filval.result_set import ResultSet
from filval.histogram_utils import hist, hist2d, hist_slice, hist_add
from filval.plotter import (decl_plot, render_plots, hist_plot, hist2d_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(l, n=2):
        l = list(l)
        while l:
            yield l[:n]
            l = l[n:]

    render_to_file('dashboard.htm', plots=get_by_n(plots, 3),
                   outdir="figures/")


@decl_plot
def plot_eff(rs, old):
    r''' '''
    if old:
        accept = hist(rs.accept_eff)
        reco = hist(rs.reco_eff)
        id_ = hist(rs.id_eff_old)
        ult = hist(rs.ult_eff_old)
    else:
        accept = hist(rs.accept_eff)
        reco = hist(rs.reco_eff)
        id_ = hist(rs.id_eff)
        ult = hist(rs.ult_eff)
    hist_plot(accept, include_errors=True, stats=False,
              color='k', label='acceptance eff wrt gen')
    hist_plot(reco, include_errors=True, stats=False,
              color='g', label='reco eff wrt accept')
    hist_plot(id_, include_errors=True, stats=False,
              color='b', label='id eff wrt reco')
    hist_plot(ult, include_errors=True, stats=False,
              color='r', label='id eff wrt gen')
    plt.xlabel('$p_T$(GeV)')
    plt.legend()


@decl_plot
def plot_eta(rs):
    r''' '''
    gen_eta = hist(rs.gen_eta)
    pf_eta = hist(rs.pf_eta)
    hist_plot(gen_eta, include_errors=True, stats=False,
              color='r', label='gen taus')
    hist_plot(pf_eta, include_errors=True, stats=False,
              color='g', label='pf taus(no id req)')
    plt.xlabel(r'$\eta$')
    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")
    rs_ft = ResultSet("ft", 'tau_ft.root')
    rs_dy = ResultSet("dy", 'tau_dy.root')

    # Next, declare all of the (sub)plots that will be assembled into full
    # figures later
    eff_ft = (plot_eff, (rs_ft, False), {})
    eff_dy = (plot_eff, (rs_dy, False), {})

    eff_ft_old = (plot_eff, (rs_ft, True), {})
    eff_dy_old = (plot_eff, (rs_dy, True), {})

    eta_ft = (plot_eta, (rs_ft,), {})
    eta_dy = (plot_eta, (rs_dy,), {})

    # Now assemble the plots into figures.
    plots = [
        Plot([[eff_ft],
              [eff_dy]],
             'eff_byTightIsolationMVArun2v1DBdR03oldDMwLT'),
        Plot([[eff_ft_old],
              [eff_dy_old]],
             'eff_byTightIsolationMVArun2v1DBoldDMwLT'),
        Plot([[eta_ft],
              [eta_dy]],
             'eta'),
    ]

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