TTTT2/yields.py

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

from filval.result_set import ResultSet
from filval.histogram_utils import hist, hist_add, hist_normalize, hist_scale
from filval.plotter import (decl_plot, render_plots, hist_plot, hist_plot_stack, Plot, generate_dashboard)


@decl_plot
def plot_yield_grid(rss):
    r"""## Event Yield

    The event yield for the eight signal regions defined in AN-17-115. Data is normalized
    to the Moriond 2018 integrated luminosity ($35.9\textrm{fb}^{-1}$). Code for the histogram generation is
    here: <https://github.com/cfangmeier/FTAnalysis/blob/master/studies/tau/Yield.C>
    """
    _, ((ax_tttt, ax_ttw), (ax_ttz, ax_tth)) = plt.subplots(2, 2)
    ft, ttw, ttz, tth = map(lambda rs: hist(rs.SRs), rss)
    # ft, ttw = map(lambda rs: hist(rs.SRs), rss[:2])
    plt.sca(ax_tttt)
    an = ((0.47, 0.33, 0.18, 0.78, 0.49, 0.52, 0.33, 0.49),
          (0, 0, 0,  0, 0, 0, 0, 0), [0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5])
    hist_plot(ft, title='TTTT', stats=False, label='Mock')
    hist_plot(an, title='TTTT', stats=False, label='AN')

    plt.sca(ax_ttw)
    an = ([2.29663, 0.508494, 0.161166, 1.03811, 0.256401, 0.127582, 0.181522, 0.141659],
          [0, 0, 0,  0, 0, 0, 0, 0], [0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5])
    hist_plot(ttw, title='TTW', stats=False, label='Mock')
    hist_plot(an, title='TTW', stats=False, label='AN')
    plt.legend()

    plt.sca(ax_ttz)
    an = ([0.974751, 0.269195, 1e-06, 0.395831, 0.0264703, 0.06816, 0.8804, 0.274265],
          [0, 0, 0,  0, 0, 0, 0, 0], [0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5])
    hist_plot(ttz, title='TTZ', stats=False, label='Mock')
    hist_plot(an, title='TTZ', stats=False, label='AN')
    plt.xlabel('Signal Region')

    plt.sca(ax_tth)
    an = ([1.13826, 0.361824, 0.162123, 0.683917, 0.137608, 0.0632719, 0.554491, 0.197864],
          [0, 0, 0,  0, 0, 0, 0, 0], [0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5])
    hist_plot(tth, title='TTH', stats=False, label='Mock')
    hist_plot(an, title='TTH', stats=False, label='AN')
    plt.xlabel('Signal Region')



@decl_plot
def plot_yield_stack(rss):
    r"""## Event Yield - Stacked

    The event yield for the eight signal regions defined in AN-17-115. Data is normalized
    to the Moriond 2018 integrated luminosity ($35.9\textrm{fb}^{-1}$). Code for the histogram generation is
    here: <https://github.com/cfangmeier/FTAnalysis/blob/master/studies/tau/Yield.C>
    """
    ft, ttw, ttz, tth = map(lambda rs: hist(rs.SRs), rss)
    hist_plot_stack([ttw, ttz, tth], labels=['TTW', 'TTZ', 'TTH'])
    ft = ft[0]*10, ft[1], ft[2]
    hist_plot(ft, label='TTTT (x10)', stats=False, color='k')

    plt.ylim((0, 60))
    plt.xlabel('Signal Region')
    plt.legend()


@decl_plot
def plot_lep_multi(rss, dataset):
    _, (ax_els, ax_mus, ax_taus) = plt.subplots(3, 1)
    els = list(map(lambda rs: hist_normalize(hist(rs.nEls)), rss))
    mus = list(map(lambda rs: hist_normalize(hist(rs.nMus)), rss))
    taus = list(map(lambda rs: hist_normalize(hist(rs.nTaus)), rss))

    def _plot(ax, procs):
        plt.sca(ax)
        ft, ttw, ttz, tth = procs
        h = {'TTTT': ft,
             'TTW': ttw,
             'TTZ': ttz,
             'TTH': tth}[dataset]
        hist_plot(h, stats=False, label=dataset)

    _plot(ax_els, els)
    plt.xlabel('\\# Good Electrons')
    plt.legend()
    _plot(ax_mus, mus)
    plt.xlabel('\\# Good Muons')
    _plot(ax_taus, taus)
    plt.xlabel('\\# Good Taus')

@decl_plot
def plot_sig_strength(rss):
    r""" The signal strength of the TTTT signal defined as

    $\frac{S}{\sqrt{S+B}}$

    """
    ft, ttw, ttz, tth = map(lambda rs: hist(rs.SRs), rss)
    bg = hist_add(ttw, ttz, tth)
    strength = ft[0] / np.sqrt(ft[0] + bg[0])
    hist_plot((strength, ft[1], ft[2]), stats=False)

@decl_plot
def plot_event_obs(rss, dataset):
    r"""

    """
    _, ((ax_njet, ax_nbjet), (ax_ht, ax_met)) = plt.subplots(2, 2)
    # ft, ttw, ttz, tth = map(lambda rs: hist(rs.SRs), rss)
    ft, ttw, ttz, tth = rss
    rs = {'TTTT': ft,
          'TTW': ttw,
          'TTZ': ttz,
          'TTH': tth}[dataset]

    def _plot(ax, obs):
        plt.sca(ax)
        h = {'MET': rs.met_in_SR,
             'HT': rs.ht_in_SR,
             'NJET': rs.njet_in_SR,
             'NBJET': rs.nbjet_in_SR}[obs]
        hist_plot(hist(h), stats=False, label=dataset, xlabel=obs)

    _plot(ax_njet, 'NJET')
    _plot(ax_nbjet, 'NBJET')
    _plot(ax_ht, 'HT')
    _plot(ax_met, 'MET')

@decl_plot
def plot_event_obs_stack(rss):
    r"""

    """
    _, ((ax_njet, ax_nbjet), (ax_ht, ax_met)) = plt.subplots(2, 2)

    def _plot(ax, obs):
        plt.sca(ax)
        attr = {'MET': 'met_in_SR',
                'HT': 'ht_in_SR',
                'NJET': 'njet_in_SR',
                'NBJET': 'nbjet_in_SR'}[obs]
        ft, ttw, ttz, tth = map(lambda rs: hist(getattr(rs, attr)), rss)
        hist_plot_stack([ttw, ttz, tth], labels=["TTW", "TTZ", "TTH"])
        hist_plot(hist_scale(ft, 5), label="TTTT (x5)", color='k')
        plt.xlabel(obs)

    _plot(ax_njet, 'NJET')
    _plot(ax_nbjet, 'NBJET')
    plt.legend()
    _plot(ax_ht, 'HT')
    _plot(ax_met, 'MET')

@decl_plot
def plot_tau_purity(rss):
    _, ((ax_ft, ax_ttw), (ax_ttz, ax_tth)) = plt.subplots(2, 2)
    ft, ttw, ttz, tth = list(map(lambda rs: hist(rs.tau_purity_v_pt), rss))

    def _plot(ax, dataset):
        plt.sca(ax)
        h = {'TTTT': ft,
             'TTW': ttw,
             'TTZ': ttz,
             'TTH': tth}[dataset]
        hist_plot(h, stats=False, label=dataset)
        plt.text(200, 0.05, dataset)
        plt.xlabel(r"$P_T$(GeV)")

    _plot(ax_ft, 'TTTT')
    _plot(ax_ttw, 'TTW')
    _plot(ax_ttz, 'TTZ')
    _plot(ax_tth, 'TTH')


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
    rss = (ResultSet("ft", 'data/yield_ft.root'),
           ResultSet("ttw", 'data/yield_ttw.root'),
           ResultSet("ttz", 'data/yield_ttz.root'),
           ResultSet("tth", 'data/yield_tth.root'))
    rss_notau = (ResultSet("ft_notau", 'data/yield_ft_notau.root'),
                 ResultSet("ttw_notau", 'data/yield_ttw_notau.root'),
                 ResultSet("ttz_notau", 'data/yield_ttz_notau.root'),
                 ResultSet("tth_notau", 'data/yield_tth_notau.root'))

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

    yield_tau_stack = (plot_yield_stack, (rss,), {})
    yield_notau_stack = (plot_yield_stack, (rss_notau,), {})

    sig_strength_tau = (plot_sig_strength, (rss,), {})
    sig_strength_notau = (plot_sig_strength, (rss_notau,), {})

    ft_lep_multi = (plot_lep_multi, (rss, 'TTTT'), {})
    ttw_lep_multi = (plot_lep_multi, (rss, 'TTW'), {})
    ttz_lep_multi = (plot_lep_multi, (rss, 'TTZ'), {})
    tth_lep_multi = (plot_lep_multi, (rss, 'TTH'), {})

    ft_event_obs = (plot_event_obs, (rss_notau, 'TTTT'), {})
    ttw_event_obs = (plot_event_obs, (rss_notau, 'TTW'), {})
    ttz_event_obs = (plot_event_obs, (rss_notau, 'TTZ'), {})
    tth_event_obs = (plot_event_obs, (rss_notau, 'TTH'), {})

    event_obs_stack = (plot_event_obs_stack, (rss_notau,), {})

    tau_purity = (plot_tau_purity, (rss,), {})

    # Now assemble the plots into figures.
    plots = [
        Plot([[yield_tau]],
             'Yield With Tau'),
        Plot([[yield_notau]],
             'Yield Without Tau'),
        Plot([[yield_tau_stack]],
             'Yield With Tau Stacked'),
        Plot([[yield_notau_stack]],
             'Yield Without Tau Stacked'),
        Plot([[yield_tau_stack],
              [yield_notau_stack]],
             'Event Yield, top: with tau, bottom: no tau'),
        Plot([[sig_strength_tau],
              [sig_strength_notau]],
             'Signal Strength'),
        Plot([[ft_lep_multi]],
             'Lepton Multiplicity - TTTT'),
        Plot([[ttw_lep_multi]],
             'Lepton Multiplicity - TTW'),
        Plot([[ttz_lep_multi]],
             'Lepton Multiplicity - TTZ'),
        Plot([[tth_lep_multi]],
             'Lepton Multiplicity - TTH'),
        Plot([[ft_event_obs]],
             'TTTT - Event Observables'),
        Plot([[ttw_event_obs]],
             'TTW - Event Observables'),
        Plot([[ttz_event_obs]],
             'TTZ - Event Observables'),
        Plot([[tth_event_obs]],
             'TTH - Event Observables'),
        Plot([[event_obs_stack]],
             'Event Observables'),
        Plot([[tau_purity]],
             'Tau Purity'),
    ]

    # Finally, render and save the plots and generate the html+bootstrap
    # dashboard to view them
    render_plots(plots, to_disk=False)
    generate_dashboard(plots, 'TTTT Yields', output='yield2.html', source_file=__file__)