EGamma_ElectronTrackingValidation/plotting/eff_plots.py

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

from filval.result_set import ResultSet
from filval.histogram import hist, hist_add, hist_norm, hist_scale, hist2d
from filval.plotting import (decl_plot, render_plots, hist_plot, hist2d_plot,
                             Plot, generate_dashboard, simple_plot)


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


@decl_plot
def plot_residuals(rs, layer, hit, variable, plot_cuts=True):
    matching_cuts = [
        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,
        )
    ]

    h = hist2d(getattr(rs, f'{variable}_BPIX_L{layer}_H{hit}_v_Et'))

    def calc_window(et):
        idx = min(hit-1, 2)
        high_et = matching_cuts[idx][f'{variable}MaxHighEt']
        high_et_thres = matching_cuts[idx][f'{variable}MaxHighEtThres']
        low_et_grad = matching_cuts[idx][f'{variable}MaxLowEtGrad']
        return high_et + min(0, et-high_et_thres)*low_et_grad

    hist2d_plot(h, colorbar=True)

    if plot_cuts:
        ets = h[3][:, 0]
        cuts = [calc_window(et) for et in ets]
        plt.plot(cuts, ets, color='red', label='Cut Value')
        plt.legend()

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



@decl_plot
def plot_seed_eff(rs):
    r"""## ECAL-Driven Seeding Efficiency

    The proportion of gen-level electrons originating in the luminous region that have
    an associated Seed, matched via rechit-simhit associations in the pixel detector. Cuts are on simtrack quantities.
    """
    ax_pt = plt.subplot(221)
    ax_eta = plt.subplot(222)
    ax_phi = plt.subplot(223)

    errors = True
    plt.sca(ax_pt)
    hist_plot(hist(rs.seed_eff_v_pt), include_errors=errors)
    center_text(0.5, 0.3, r'$|\eta|<2.4$')
    plt.xlabel(r"Sim-Track $p_T$")
    plt.ylim((0, 1.1))

    plt.sca(ax_eta)
    hist_plot(hist(rs.seed_eff_v_eta), include_errors=errors)
    center_text(0.5, 0.3, r'$p_T>20$')
    plt.xlabel(r"Sim-Track $\eta$")
    plt.ylim((0, 1.1))

    plt.sca(ax_phi)
    hist_plot(hist(rs.seed_eff_v_phi), include_errors=errors)
    center_text(0.5, 0.3, r'$p_T>20$ and $|\eta|<2.4$')
    plt.xlabel(r"Sim-Track $\phi$")
    plt.ylim((0, 1.1))


@decl_plot
def plot_tracking_eff(rs):
    r"""## GSF Tracking Efficiency

    The proportion of electrons origination in the luminous region from the that have
    an associated GSF track. Cuts are on simtrack quantities.
    """
    ax_pt = plt.subplot(221)
    ax_eta = plt.subplot(222)
    ax_phi = plt.subplot(223)

    errors = True
    plt.sca(ax_pt)
    hist_plot(hist(rs.tracking_eff_v_pt), include_errors=errors)
    center_text(0.5, 0.3, r'$|\eta|<2.4$')
    plt.xlabel(r"Sim-Track $p_T$")
    # plt.ylim((0, 1.1))

    plt.sca(ax_eta)
    hist_plot(hist(rs.tracking_eff_v_eta), include_errors=errors)
    center_text(0.5, 0.3, r'$p_T>20$')
    plt.xlabel(r"Sim-Track $\eta$")
    # plt.ylim((0, 1.1))

    plt.sca(ax_phi)
    hist_plot(hist(rs.tracking_eff_v_phi), include_errors=errors)
    center_text(0.5, 0.3, r'$p_T>20$ and $|\eta|<2.4$')
    plt.xlabel(r"Sim-Track $\phi$")
    # plt.ylim((0, 1.1))


@decl_plot
def plot_seed_purity(rs, ext=""):
    r"""## ECAL-Driven Seed Purity

    The proportion of ECAL-driven seeds that have a matched gen-level electron originating in
    the luminous region. Cuts are on seed quantities.
    """
    ax_pt = plt.subplot(221)
    ax_eta = plt.subplot(222)
    ax_phi = plt.subplot(223)

    def get_hist(base_name):
        return hist(getattr(rs, base_name+ext))

    errors = True
    plt.sca(ax_pt)
    hist_plot(get_hist("seed_pur_v_pt"), include_errors=errors)
    center_text(0.5, 0.3, r'$|\eta|<2.4$')
    plt.xlabel(r"Seed $p_T$")
    if not ext:
        plt.ylim((0, 1.1))

    plt.sca(ax_eta)
    hist_plot(get_hist("seed_pur_v_eta"), include_errors=errors)
    center_text(0.5, 0.3, r'$p_T>20$')
    plt.xlabel(r"Seed $\eta$")
    if not ext:
        plt.ylim((0, 1.1))

    plt.sca(ax_phi)
    hist_plot(get_hist("seed_pur_v_phi"), include_errors=errors)
    center_text(0.5, 0.3, r'$p_T>20$ and $|\eta|<2.4$')
    plt.xlabel(r"Seed $\phi$")
    if not ext:
        plt.ylim((0, 1.1))


@decl_plot
def plot_track_purity(rs, ext=""):
    r"""## GSF Track Purity

    The proportion of GSF-tracks w\ ECAL-driven seeds that have a matched gen-level electron originating in
    the luminous region. Cuts are on GSF track quantities.
    """
    ax_pt = plt.subplot(221)
    ax_eta = plt.subplot(222)
    ax_phi = plt.subplot(223)

    def get_hist( base_name):
        return hist(getattr(rs, base_name+ext))

    errors = True
    plt.sca(ax_pt)
    hist_plot(get_hist("tracking_pur_v_pt"), include_errors=errors)
    center_text(0.5, 0.3, r'$|\eta|<2.4$')
    plt.xlabel(r"GSF-Track $p_T$")
    if not ext:
        plt.ylim((0, 1.1))

    plt.sca(ax_eta)
    hist_plot(get_hist("tracking_pur_v_eta"), include_errors=errors)
    center_text(0.5, 0.3, r'$p_T>20$')
    plt.xlabel(r"GSF-Track $\eta$")
    if not ext:
        plt.ylim((0, 1.1))

    plt.sca(ax_phi)
    hist_plot(get_hist("tracking_pur_v_phi"), include_errors=errors)
    center_text(0.5, 0.3, r'$p_T>20$ and $|\eta|<2.4$')
    plt.xlabel(r"GSF-Track $\phi$")
    if not ext:
        plt.ylim((0, 1.1))


@decl_plot
def plot_hit_vs_layer(rs, region):

    h = hist2d(getattr(rs, f'hit_vs_layer_{region}'))

    hist2d_plot(h, txt_format='{:2.0f}')
    plt.xlabel('Layer \#')
    plt.ylabel('Hit \#')


if __name__ == '__main__':
    rs = ResultSet('seeds', '../hists/new_seeding.root')

    # Next, declare all of the (sub)plots that will be assembled into full
    # figures later
    seed_eff = plot_seed_eff, (rs, )
    tracking_eff = plot_tracking_eff, (rs,)

    seed_pur = plot_seed_purity, rs
    seed_pur_num = plot_seed_purity, (rs,), {'ext': '_num'}
    seed_pur_den = plot_seed_purity, (rs,), {'ext': '_den'}

    track_pur = plot_track_purity, (rs,)
    track_pur_num = plot_track_purity, (rs,), {'ext': '_num'}
    track_pur_den = plot_track_purity, (rs,), {'ext': '_den'}

    residuals_L1_H1_dPhi = plot_residuals, (rs, 1, 1, 'dPhi')
    residuals_L2_H2_dPhi = plot_residuals, (rs, 2, 2, 'dPhi')
    residuals_L3_H3_dPhi = plot_residuals, (rs, 3, 3, 'dPhi')

    residuals_L1_H1_dRz = plot_residuals, (rs, 1, 1, 'dRz')
    residuals_L1_H1_dRz_no_cuts = plot_residuals, (rs, 1, 1, 'dRz'), {'plot_cuts': False}
    residuals_L2_H2_dRz = plot_residuals, (rs, 2, 2, 'dRz')
    residuals_L3_H3_dRz = plot_residuals, (rs, 3, 3, 'dRz')

    hit_vs_layer_barrel = plot_hit_vs_layer, (rs, 'barrel')
    hit_vs_layer_forward = plot_hit_vs_layer, (rs, 'forward')

    # Now assemble the plots into figures.
    plots = [
        Plot(residuals_L1_H1_dPhi, 'Phi Residuals Layer 1 Hit 1'),
        Plot(residuals_L2_H2_dPhi, 'Phi Residuals Layer 2 Hit 2'),
        Plot(residuals_L3_H3_dPhi, 'Phi Residuals Layer 3 Hit 3'),
        Plot(residuals_L1_H1_dRz, 'dZ Residuals Layer 1 Hit 1'),
        Plot(residuals_L1_H1_dRz_no_cuts, 'dZ Residuals Layer 1 Hit 1 w/o cuts'),
        Plot(residuals_L2_H2_dRz, 'dZ Residuals Layer 2 Hit 2'),
        Plot(residuals_L3_H3_dRz, 'dZ Residuals Layer 3 Hit 3'),
        Plot(seed_eff, 'ECAL-Driven Seeding Efficiency'),
        Plot(tracking_eff, 'GSF Tracking Efficiency'),
        Plot(hit_vs_layer_barrel, 'Hit vs Layer - Barrel'),
        Plot(hit_vs_layer_forward, 'Hit vs Layer - Forward'),
        Plot(track_pur, 'GSF Track Purity'),
        Plot(track_pur_num, 'GSF Track Purity Numerator'),
        Plot(track_pur_den, 'GSF Track Purity Denominator'),
    ]

    # Finally, render and save the plots and generate the html+bootstrap
    # dashboard to view them
    render_plots(plots, to_disk=False)
    generate_dashboard(plots, 'Seeding Efficiency',
                       output='eff_plots.html',
                       source=__file__,
                       config=rs.config)