EGamma_ElectronTrackingValidation/plotting/eff_plots.py

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

from uproot import open as root_open
from matplottery.utils import Hist1D, Hist2D, to_html_table
from matplottery.plotter import set_defaults, plot_2d

import matplotboard as mpb

matching_cuts = {
    'new-extra-narrow': [
        dict(
            dPhiMaxHighEt=0.025,
            dPhiMaxHighEtThres=20.0,
            dPhiMaxLowEtGrad=-0.002,
            dRzMaxHighEt=9999.0,
            dRzMaxHighEtThres=0.0,
            dRzMaxLowEtGrad=0.0,
        ),
        dict(
            dPhiMaxHighEt=0.0015,
            dPhiMaxHighEtThres=0.0,
            dPhiMaxLowEtGrad=0.0,
            dRzMaxHighEt=0.025,
            dRzMaxHighEtThres=30.0,
            dRzMaxLowEtGrad=-0.002,
        ),
        dict(
            dPhiMaxHighEt=0.0015,
            dPhiMaxHighEtThres=0.0,
            dPhiMaxLowEtGrad=0.0,
            dRzMaxHighEt=0.025,
            dRzMaxHighEtThres=30.0,
            dRzMaxLowEtGrad=-0.002,
        )
    ],
    'new-default': [
        dict(
            dPhiMaxHighEt=0.05,
            dPhiMaxHighEtThres=20.0,
            dPhiMaxLowEtGrad=-0.002,
            dRzMaxHighEt=9999.0,
            dRzMaxHighEtThres=0.0,
            dRzMaxLowEtGrad=0.0,
        ),
        dict(
            dPhiMaxHighEt=0.003,
            dPhiMaxHighEtThres=0.0,
            dPhiMaxLowEtGrad=0.0,
            dRzMaxHighEt=0.05,
            dRzMaxHighEtThres=30.0,
            dRzMaxLowEtGrad=-0.002,
        ),
        dict(
            dPhiMaxHighEt=0.003,
            dPhiMaxHighEtThres=0.0,
            dPhiMaxLowEtGrad=0.0,
            dRzMaxHighEt=0.05,
            dRzMaxHighEtThres=30.0,
            dRzMaxLowEtGrad=-0.002,
        )
    ],
    'new-wide': [
        dict(
            dPhiMaxHighEt=0.10,
            dPhiMaxHighEtThres=20.0,
            dPhiMaxLowEtGrad=-0.002,
            dRzMaxHighEt=9999.0,
            dRzMaxHighEtThres=0.0,
            dRzMaxLowEtGrad=0.0,
        ),
        dict(
            dPhiMaxHighEt=0.006,
            dPhiMaxHighEtThres=0.0,
            dPhiMaxLowEtGrad=0.0,
            dRzMaxHighEt=0.10,
            dRzMaxHighEtThres=30.0,
            dRzMaxLowEtGrad=-0.002,
        ),
        dict(
            dPhiMaxHighEt=0.006,
            dPhiMaxHighEtThres=0.0,
            dPhiMaxLowEtGrad=0.0,
            dRzMaxHighEt=0.10,
            dRzMaxHighEtThres=30.0,
            dRzMaxLowEtGrad=-0.002,
        )
    ],
    'new-extra-wide': [
        dict(
            dPhiMaxHighEt=0.15,
            dPhiMaxHighEtThres=20.0,
            dPhiMaxLowEtGrad=-0.002,
            dRzMaxHighEt=9999.0,
            dRzMaxHighEtThres=0.0,
            dRzMaxLowEtGrad=0.0,
        ),
        dict(
            dPhiMaxHighEt=0.009,
            dPhiMaxHighEtThres=0.0,
            dPhiMaxLowEtGrad=0.0,
            dRzMaxHighEt=0.15,
            dRzMaxHighEtThres=30.0,
            dRzMaxLowEtGrad=-0.002,
        ),
        dict(
            dPhiMaxHighEt=0.009,
            dPhiMaxHighEtThres=0.0,
            dPhiMaxLowEtGrad=0.0,
            dRzMaxHighEt=0.15,
            dRzMaxHighEtThres=30.0,
            dRzMaxLowEtGrad=-0.002,
        )
    ],
}

samples = None


def load_samples():
    global samples
    if samples is None:
        print("loading samples")
        samples = {(proc, wp): root_open(f'../hists/{proc}-{wp}.root') for proc, wp in product(procs, wps)}
    return samples


def calc_window(et, eta, hit, variable, cut_sel):
    idx = min(hit-1, 2)
    cuts = matching_cuts[cut_sel][idx]
    if 'etaBins' in cuts:
        for eta_idx, bin_high in enumerate(cuts['etaBins']):
            if eta < bin_high:
                high_et = cuts[f'{variable}MaxHighEt'][eta_idx]
                high_et_thres = cuts[f'{variable}MaxHighEtThres'][eta_idx]
                low_et_grad = cuts[f'{variable}MaxLowEtGrad'][eta_idx]
                break
        else:  # highest bin
            high_et = cuts[f'{variable}MaxHighEt'][-1]
            high_et_thres = cuts[f'{variable}MaxHighEtThres'][-1]
            low_et_grad = cuts[f'{variable}MaxLowEtGrad'][-1]
    else:
        high_et = cuts[f'{variable}MaxHighEt']
        high_et_thres = cuts[f'{variable}MaxHighEtThres']
        low_et_grad = cuts[f'{variable}MaxLowEtGrad']
    return high_et + min(0, et-high_et_thres)*low_et_grad


def hist_integral_ratio(num, den):
    num_int = num.get_integral()
    den_int = den.get_integral()

    ratio = num_int / den_int
    error = np.sqrt(den_int) / den_int  # TODO: Check this definition of error
    return ratio, error


def center_text(x, y, txt, **kwargs):
    plt.text(x, y, txt,
             horizontalalignment='center', verticalalignment='center',
             transform=plt.gca().transAxes, **kwargs)


def hist_plot(h: Hist1D, *args, include_errors=False, **kwargs):
    """ Plots a 1D ROOT histogram object using matplotlib """

    counts = h.get_counts()
    edges = h.get_edges()
    left, right = edges[:-1], edges[1:]
    x = np.array([left, right]).T.flatten()
    y = np.array([counts, counts]).T.flatten()

    plt.plot(x, y, *args, linewidth=2, **kwargs)
    if include_errors:
        plt.errorbar(h.get_bin_centers(), h.counts, yerr=h.errors,
                     color='k', marker=None, linestyle='None',
                     barsabove=True, elinewidth=.7, capsize=1)


def hist2d_percent_contour(h: Hist1D, percent: float, axis: str):
    values = h.counts
    try:
        axis = axis.lower()
        axis_idx = {'x': 1, 'y': 0}[axis]
    except KeyError:
        raise ValueError('axis must be \'x\' or \'y\'')
    if percent < 0 or percent > 1:
        raise ValueError('percent must be in [0,1]')

    with np.warnings.catch_warnings():
        np.warnings.filterwarnings('ignore', 'invalid value encountered in true_divide')
        values = values / np.sum(values, axis=axis_idx, keepdims=True)
        np.nan_to_num(values, copy=False)
    values = np.cumsum(values, axis=axis_idx)
    idxs = np.argmax(values > percent, axis=axis_idx)

    x_centers, y_centers = h.get_bin_centers()

    if axis == 'x':
        return x_centers[idxs], y_centers
    else:
        return x_centers, y_centers[idxs]


@mpb.decl_fig
def plot_residuals(sample, layer, hit, variable, subdet):
    load_samples()
    proc, wp = sample

    track_matched = samples[sample][f'{variable}_{subdet}_L{layer}_H{hit}_v_Et_TrackMatched']
    # no_match = samples[sample][f'{variable}_{subdet}_L{layer}_H{hit}_v_Et_NoMatch']

    h_real = Hist2D(track_matched)
    # h_fake = Hist2D(no_match)

    def do_plot(h):
        plot_2d(h, cmap='viridis')

        xs, ys = hist2d_percent_contour(h, .90, 'x')
        plt.plot(xs, ys, color='green', label='90\% contour')
        xs, ys = hist2d_percent_contour(h, .995, 'x')
        plt.plot(xs, ys, color='darkgreen', label='99.5\% contour')

        ets = h.edges[1]
        cuts = [calc_window(et, 0, hit, variable, wp) for et in ets]
        plt.plot(cuts, ets, color='red', label='Cut Value')
        plt.xlabel({'dPhi': r'$\delta \phi$ (rads)',
                    'dRz': r'$\delta R/z$ (cm)'}[variable])

    # plt.sca(plt.subplot(1, 2, 1))
    do_plot(h_real)
    plt.title('Truth-Matched Seeds')
    plt.ylabel('$E_T$ (GeV)')

    # plt.sca(plt.subplot(1, 2, 2))
    # do_plot(h_fake)
    # plt.title('Not Truth-Matched Seeds')
    plt.legend(loc='upper right')


@mpb.decl_fig
def plot_residuals_eta(sample, hit, variable):
    load_samples()

    h = Hist2D(samples[sample][f'{variable}_residuals_v_eta_H{hit}'])

    plot_2d(h)

    xs, ys = hist2d_percent_contour(h, .90, 'x')
    plt.plot(xs, ys, color='green', label='90\% contour')
    xs, ys = hist2d_percent_contour(h, .995, 'x')
    plt.plot(xs, ys, color='darkgreen', label='99.5\% contour')

    plt.xlabel({'dPhi': r'$\delta \phi$ (rads)',
                'dRz': r'$\delta R/z$ (cm)'}[variable])
    plt.ylabel(r'$|\eta|$')


@mpb.decl_fig
def plot_hit_vs_layer(sample, region):
    load_samples()
    h = Hist2D(samples[sample][f'hit_vs_layer_{region}'])
    plot_2d(h, colz_fmt='2.0f')
    plt.xlabel('Layer #')
    plt.ylabel('Hit #')


@mpb.decl_fig
def plot_roc_curve(pfx, ext=''):
    load_samples()

    def get_num_den(sample, basename):
        num = Hist1D(sample[f'{basename}_num'])
        den = Hist1D(sample[f'{basename}_den'])
        return hist_integral_ratio(num, den)
    rows = []
    for (proc, wp), sample in samples.items():
        sample_name = f'{proc}-{wp}'
        eff, eff_err = get_num_den(sample, f'{pfx}_eff_v_phi{ext}')
        pur, pur_err = get_num_den(sample, f'{pfx}_pur_v_phi{ext}')
        rows.append([wp,
                     rf'${eff*100:0.2f}\pm{eff_err*100:0.2f}\%$',
                     rf'${pur*100:0.2f}\pm{pur_err*100:0.2f}\%$'])

        plt.errorbar([pur], [eff], xerr=[pur_err], yerr=[eff_err],
                     label=sample_name, marker='o', color=color(proc, wp))
    center_text(0.3, 0.3, r'$p_T>20$ and $|\eta|<2.4$')
    plt.axis('equal')
    plt.xlim((0.5, 1.02))
    plt.ylim((0.5, 1.02))
    plt.xlabel('Purity')
    plt.ylabel('Efficiency')
    plt.grid()
    plt.legend(loc='lower right')

    col_labels = ['Sample', 'Working Point', 'Efficiency', 'Purity']
    row_labels = [r'$Z \rightarrow ee$', '', '', r'$t\bar{t}$', '', '']
    return to_html_table(rows, col_labels, row_labels, 'table-condensed')


@mpb.decl_fig
def plot_kinematic_eff_all(pref, ext=''):
    load_samples()
    ax_pt = plt.subplot(221)
    ax_eta = plt.subplot(222)
    ax_phi = plt.subplot(223)
    errors = True
    for (proc, wp), sample in samples.items():
        sample_name = f'{proc}-{wp}'
        l = sample_name
        c = color(proc, wp)
        plt.sca(ax_pt)
        hist_plot(Hist1D(sample[f'{pref}_v_pt{ext}']), include_errors=errors, label=l, color=c)
        plt.sca(ax_eta)
        hist_plot(Hist1D(sample[f'{pref}_v_eta{ext}']), include_errors=errors, label=l, color=c)
        plt.sca(ax_phi)
        hist_plot(Hist1D(sample[f'{pref}_v_phi{ext}']), include_errors=errors, label=l, color=c)

    if 'eff' in pref:
        reference = 'Sim-Track'
    elif 'fake' in pref:
        reference = 'Supercluster'
    else:
        reference = 'GSF-Track'

    def set_ylim():
        if 'num' not in ext and 'den' not in ext:
            plt.ylim((0, 1.1))

    plt.sca(ax_pt)
    center_text(0.5, 0.3, r'$|\eta|<2.4$')
    plt.xlabel(fr"{reference} $p_T$")
    set_ylim()

    plt.sca(ax_eta)
    center_text(0.5, 0.3, r'$p_T>20$')
    plt.xlabel(fr"{reference} $\eta$")
    set_ylim()

    plt.sca(ax_phi)
    center_text(0.5, 0.3, r'$p_T>20$ and $|\eta|<2.4$')
    plt.xlabel(fr"{reference} $\phi$")
    set_ylim()
    plt.tight_layout()
    plt.legend(loc='upper left', bbox_to_anchor=(0.6, 0.45), bbox_transform=plt.gcf().transFigure)


@mpb.decl_fig
def plot_ecal_rel_res():
    load_samples()
    for sample_name, sample in samples.items():
        h = Hist1D(sample['ecal_energy_resolution'])
        h = h / h.get_integral()
        hist_plot(h, label=sample_name)
    plt.xlabel(r"ECAL $E_T$ relative error")
    plt.legend()


@mpb.decl_fig
def plot_res_contour(proc, hit_number, var, layers, ext='_TrackMatched'):
    load_samples()
    _, axs = plt.subplots(1, 2, sharey=True)

    def do_plot(ax, sample):
        plt.sca(ax)
        wp = sample[1]
        plt.title(wp)
        h = None
        for subdet, layer in layers:
            h = Hist2D(samples[sample][f'{var}_{subdet}_L{layer}_H{hit_number}_v_Et{ext}'])
            xs, ys = hist2d_percent_contour(h, .99, 'x')
            plt.plot(xs, ys, label=f'{subdet} - L{layer}')

        ets = h.edges[1]
        cuts = [calc_window(et, 0, hit_number, var, wp) for et in ets]
        plt.plot(cuts, ets, color='k', label='Cut Value')

    for ax, wp in zip(axs, ['new-default', 'new-wide']):
        do_plot(ax, (proc, wp))

    plt.sca(axs[-1])
    plt.legend(loc='upper right')


@mpb.decl_fig
def simple_dist(hist_name, xlabel, rebin=False, xmax=None):
    load_samples()
    for (proc, wp), sample in samples.items():
        sample_name = f'{proc}-{wp}'
        h = Hist1D(sample[hist_name])
        if rebin:
            h.rebin(50, -0.5, 200.5)

        h = h / h.get_integral()
        mean = int(np.sum(h.get_counts() * h.get_bin_centers()))
        hist_plot(h, label=f'{sample_name} ($\\mu={mean:g}$)',
                  color=color(proc, wp))
    if xmax:
        plt.xlim((0, xmax))
    plt.xlabel(xlabel)
    plt.legend()


def all_cut_plots(refresh=True, publish=False):

    figures = {
        'tracking_roc_curve': (plot_roc_curve, ('tracking',)),
        'tracking_roc_curve_dR': (plot_roc_curve, ('tracking',), {'ext': '_dR'}),
        'seeding_roc_curve': (plot_roc_curve, ('seed',)),

        'tracking_eff_all': (plot_kinematic_eff_all, ('tracking_eff',)),
        'tracking_pur_all': (plot_kinematic_eff_all, ('tracking_pur',)),
        'tracking_eff_all_dR': (plot_kinematic_eff_all, ('tracking_eff',), {'ext': '_dR'}),
        'tracking_pur_all_dR': (plot_kinematic_eff_all, ('tracking_pur',), {'ext': '_dR'}),
        'seeding_eff_all': (plot_kinematic_eff_all, ('seed_eff',)),
        'seeding_pur_all': (plot_kinematic_eff_all, ('seed_pur',)),
        'partial_fake_rate_all': (plot_kinematic_eff_all, ('partial_fake_rate',)),
        'partial_fake_rate_all_num': (plot_kinematic_eff_all, ('partial_fake_rate',), {'ext': '_num'}),
        'partial_fake_rate_all_den': (plot_kinematic_eff_all, ('partial_fake_rate',), {'ext': '_den'}),
        'full_fake_rate_all': (plot_kinematic_eff_all, ('full_fake_rate',)),
        'full_fake_rate_all_num': (plot_kinematic_eff_all, ('full_fake_rate',), {'ext': '_num'}),
        'full_fake_rate_all_den': (plot_kinematic_eff_all, ('full_fake_rate',), {'ext': '_den'}),
        'clean_fake_rate_all': (plot_kinematic_eff_all, ('clean_fake_rate',)),
        'clean_fake_rate_all_num': (plot_kinematic_eff_all, ('clean_fake_rate',), {'ext': '_num'}),
        'clean_fake_rate_all_den': (plot_kinematic_eff_all, ('clean_fake_rate',), {'ext': '_den'}),

        'number_of_seeds': (simple_dist, ('n_seeds', 'Number of Seeds'), {'rebin': True}),
        'number_of_sim_els': (simple_dist, ('n_good_sim', 'Number of prompt(ish) electrons'), {'xmax': 10.5}),
        'number_of_gsf_tracks': (simple_dist, ('n_gsf_track', 'Number of reco electrons'), {'xmax': 10.5}),

        'number_of_matched': (simple_dist, ('n_matched', 'Number of matched electrons'), {'xmax': 10.5}),
        'number_of_merged': (simple_dist, ('n_merged', 'Number of merged electrons'), {'xmax': 10.5}),
        'number_of_lost': (simple_dist, ('n_lost', 'Number of lost electrons'), {'xmax': 10.5}),
        'number_of_split': (simple_dist, ('n_split', 'Number of split electrons'), {'xmax': 10.5}),
        'number_of_faked': (simple_dist, ('n_faked', 'Number of faked electrons'), {'xmax': 10.5}),
        'number_of_flipped': (simple_dist, ('n_flipped', 'Number of flipped electrons'), {'xmax': 10.5}),
        'matched_dR': (simple_dist, ('matched_dR', 'dR between sim and reco')),
        'matched_dpT': (simple_dist, ('matched_dpT', 'dpT between sim and reco')),

        'number_of_matched_dR': (simple_dist, ('n_matched_dR', 'Number of matched electrons - dR Matched'), {'xmax': 10.5}),
        'number_of_merged_dR': (simple_dist, ('n_merged_dR', 'Number of merged electrons - dR Matched'), {'xmax': 10.5}),
        'number_of_lost_dR': (simple_dist, ('n_lost_dR', 'Number of lost electrons - dR Matched'), {'xmax': 10.5}),
        'number_of_split_dR': (simple_dist, ('n_split_dR', 'Number of split electrons - dR Matched'), {'xmax': 10.5}),
        'number_of_faked_dR': (simple_dist, ('n_faked_dR', 'Number of faked electrons - dR Matched'), {'xmax': 10.5}),
        'number_of_flipped_dR': (simple_dist, ('n_flipped_dR', 'Number of flipped electrons - dR Matched'), {'xmax': 10.5}),
        'matched_dR_dR': (simple_dist, ('matched_dR_dR', 'dR between sim and reco - dR Matched')),
        'matched_dpT_dR': (simple_dist, ('matched_dpT_dR', 'dpT between sim and reco - dR Matched')),

        'sim_pt': (simple_dist, ('sim_pt', 'Sim Track $p_T$')),
        'sim_eta': (simple_dist, ('sim_eta', 'Sim Track $eta$')),
        'sim_phi': (simple_dist, ('sim_phi', 'Sim Track $phi$')),
        'reco_pt': (simple_dist, ('reco_pt', 'Reco Track $p_T$')),
        'reco_eta': (simple_dist, ('reco_eta', 'Reco Track $eta$')),
        'reco_phi': (simple_dist, ('reco_phi', 'Reco Track $phi$')),

        'ecal_rel_res': plot_ecal_rel_res,
        'hit_v_layer_BPIX_new-default_zee': (plot_hit_vs_layer, (('zee', 'new-default'), 'barrel')),
        'hit_v_layer_FPIX_new-default_zee': (plot_hit_vs_layer, (('zee', 'new-default'), 'forward')),
        'hit_v_layer_BPIX_new-default_tt': (plot_hit_vs_layer, (('tt', 'new-default'), 'barrel')),
        'hit_v_layer_FPIX_new-default_tt': (plot_hit_vs_layer, (('tt', 'new-default'), 'forward')),
        'hit_v_layer_BPIX_new-wide_zee': (plot_hit_vs_layer, (('zee', 'new-wide'), 'barrel')),
        'hit_v_layer_FPIX_new-wide_zee': (plot_hit_vs_layer, (('zee', 'new-wide'), 'forward')),
        'hit_v_layer_BPIX_new-wide_tt': (plot_hit_vs_layer, (('tt', 'new-wide'), 'barrel')),
        'hit_v_layer_FPIX_new-wide_tt': (plot_hit_vs_layer, (('tt', 'new-wide'), 'forward')),
    }

    hit_layers = [(1, 1), (1, 2), (2, 2), (2, 3), (3, 3), (3, 4)]
    for proc, wp, (hit, layer), var, subdet in product(['zee', 'tt'], ['new-default', 'new-wide'],
                                                       hit_layers,
                                                       ['dPhi', 'dRz'], ['BPIX', 'FPIX']):
        figures.update({
            f'res_{subdet}_L{layer}_H{hit}_{var}_{proc}_{wp}':
                (plot_residuals, ((proc, wp), layer, hit, var, subdet))})

    rel_layers = {1: [('BPIX', 1), ('BPIX', 2), ('FPIX', 1), ('FPIX', 2)],
                  2: [('BPIX', 2), ('BPIX', 3), ('FPIX', 2), ('FPIX', 3)],
                  3: [('BPIX', 3), ('BPIX', 4), ('FPIX', 3)], }
    for proc, hit, var in product(['zee', 'tt'], [1, 2, 3], ['dPhi', 'dRz']):
        figures.update({
            f'resall_H{hit}_{var}_{proc}': (plot_res_contour, (proc, hit, var, rel_layers[hit]))})

    for proc, wp, hit, var in product(['zee', 'tt'], ['new-default', 'new-wide'], [1, 2, 3], ['dPhi', 'dRz']):
        figures.update({
            f'res_v_eta_H{hit}_{var}_{proc}_{wp}': (plot_residuals_eta, ((proc, wp), hit, var))})

    mpb.render(figures, refresh=refresh)
    mpb.generate_report(figures, 'Electron Seeding Studies',
                        output=f'hists.html',
                        source=__file__)

    # mpb.generate_report(figures, 'Fake Rate Investigation',
    #                     output=f'report.html',
    #                     body='../docs/reports/report_2018_05_18.md')
    if publish:
        mpb.publish()


def color(proc, wp):
    from matplotlib.colors import XKCD_COLORS

    def f(name):
        return XKCD_COLORS['xkcd:'+name]

    return {
        ('zee', 'new-extra-narrow'): f('kelly green'),
        ('tt',  'new-extra-narrow'): f('grey green'),
        ('zee', 'new-default'):      f('red'),
        ('tt',  'new-default'):      f('pink'),
        ('zee', 'new-wide'):         f('strong blue'),
        ('tt',  'new-wide'):         f('bright blue'),
        ('zee', 'new-extra-wide'):   f('indigo'),
        ('tt',  'new-extra-wide'):   f('bright purple'),
        ('zee', 'old-default'):      f('black'),
        ('tt',  'old-default'):      f('blue grey'),
    }[(proc, wp)]


if __name__ == '__main__':
    set_defaults()
    mpb.configure(output_dir='seeding_studies',
                  multiprocess=True,
                  publish_remote="caleb@fangmeier.tech",
                  publish_dir="/var/www/eg",
                  publish_url="eg.fangmeier.tech/",
                  )
    procs = {
        'zee': r'$Z\rightarrow e^+e_-}$',
        'tt': r'$t\bar{t}$'
    }
    wps = {
        'new-default': 'HLT Settings',
        'new-wide': 'Wide Settings',
        'old-default': 'Old Seeding',
    }
    all_cut_plots(refresh=True, publish=False)