EGamma_ElectronTrackingValidation/plotting/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, save_plots, hist_plot, hist2d_plot, Plot)


@decl_plot
def plot_residual(rs, var, layer, hit, projection=None, display_layer=True):
    r'''
    Plots of $\Delta \phi$, $\Delta z$, or $\Delta r$ (vs $\eta$) between RecHit
    and SimHit for the nth hit in the matched seed.
    '''
    h = {('z', 'B1', 1): rs.dz_v_eta_first_hits_in_B1,
         ('z', 'B2', 1): rs.dz_v_eta_first_hits_in_B2,
         ('phi', 'B1', 1): rs.dphi_v_eta_first_hits_in_B1,
         ('phi', 'B2', 1): rs.dphi_v_eta_first_hits_in_B2,
         ('z', 'B2', 2): rs.dz_v_eta_second_hits_in_B2,
         ('z', 'B3', 2): rs.dz_v_eta_second_hits_in_B3,
         ('phi', 'B2', 2): rs.dphi_v_eta_second_hits_in_B2,
         ('phi', 'B3', 2): rs.dphi_v_eta_second_hits_in_B3,
         }[(var, layer, hit)]
    varlabel = {('z', 1): r"$\Delta z_1^{\textrm{sim}}(\mu \mathrm{m})$",
                ('z', 2): r"$\Delta z_2^{\textrm{sim}}(\mu \mathrm{m})$",
                ('phi', 1): r"$\Delta \phi_1^{\textrm{sim}}(\mathrm{millirad})$",
                ('phi', 2): r"$\Delta \phi_2^{\textrm{sim}}(\mathrm{millirad})$"
                }[(var, hit)]
    scale = {'z': 10**4,  # cm -> um
             'phi': 10**3  # rad -> millirad
             }[var]
    if projection == 'y':  # projecting onto var axis
        h = h.ProjectionY()
        hist_plot(hist(h, rescale_x=scale), xlabel=varlabel, title='Layer - ' + layer)
    elif projection == 'x':  # projecting onto eta axis
        h = h.ProjectionX()
        hist_plot(hist(h), xlabel=r'$\eta$', title=layer)
    else:
        hist2d_plot(hist2d(h, rescale_y=scale), ylabel=varlabel, xlabel=r"$\eta$")
    # if display_layer:
    #     plt.text(0.9, 0.9, layer,
    #              bbox={'facecolor': 'white', 'alpha': 0.7},
    #              transform=plt.gca().transAxes)


@decl_plot
def plot_residuals_v_ladder(rs, var, layer):
    even, odd = {('phi', 'B1'): (rs.dphi_v_eta_first_hits_in_B1_even_ladder, rs.dphi_v_eta_first_hits_in_B1_odd_ladder),
                 ('phi', 'B2'): (rs.dphi_v_eta_first_hits_in_B2_even_ladder, rs.dphi_v_eta_first_hits_in_B2_odd_ladder),
                 ('z', 'B1'): (rs.dz_v_eta_first_hits_in_B1_even_ladder, rs.dz_v_eta_first_hits_in_B1_odd_ladder),
                 ('z', 'B2'): (rs.dz_v_eta_first_hits_in_B2_even_ladder, rs.dz_v_eta_first_hits_in_B2_odd_ladder),
                 }[(var, layer)]
    unit = {'phi': 'urad',
            'z': 'um',
            }[var]
    fmt_var = {'phi': r'\phi',
               'z': 'z',
               }[var]
    scale = {'phi': 10**3,
             'z': 10**4,
             }[var]
    xlabel = {'phi': r'$\Delta \phi_1$(rad)',
              'z': r"$\Delta z_1$(cm)",
              }[var]
    even = even.ProjectionY()
    odd = odd.ProjectionY()
    hist_plot(hist(even), include_errors=True, label='even ladders')
    hist_plot(hist(odd), include_errors=True, color='r', label='odd ladders')
    plt.xlabel(xlabel)

    even_mean = even.GetMean()*scale
    odd_mean = odd.GetMean()*scale
    txt = r'$ \hat{{\Delta {0} }}_{{1,\textrm{{ {1} }} }}={2:4.2g}${3}'
    plt.text(0.05, .8, txt.format(fmt_var, 'odd', odd_mean, unit), transform=plt.gca().transAxes)
    plt.text(0.05, .7, txt.format(fmt_var, 'even', even_mean, unit), transform=plt.gca().transAxes)
    plt.legend()


@decl_plot
def sc_extrapolation_first(rs, var, layer, hit, even_odd, log=False, norm=None, cut=None, display_layer=True):
    ''' Raphael's plots '''

    # def preproc(h):
    #     return hist_slice(hist(h.ProjectionY()), (-0.07, 0.07))
    # hist_plot(hist(rs.sc_first_hits_in_B1_dz.ProjectionY()),
    #           include_errors=errors,
    #           norm=norm,
    #           log=log,
    #           xlabel=r"$\Delta z_1$(cm)",
    #           title="BPIX - Layer 1")

    h = {('phi', 'B1', 1, 'either'): (rs.sc_first_hits_in_B1_dphi_even, rs.sc_first_hits_in_B1_dphi_odd),
         ('phi', 'B2', 1, 'either'): (rs.sc_first_hits_in_B2_dphi_even, rs.sc_first_hits_in_B1_dphi_odd),
         ('z', 'B1', 1, 'either'): (rs.sc_first_hits_in_B1_dz_even, rs.sc_first_hits_in_B1_dz_odd),
         ('z', 'B2', 1, 'either'): (rs.sc_first_hits_in_B2_dz_even, rs.sc_first_hits_in_B2_dz_odd),
         ('phi', 'B2', 2, 'either'): (rs.sc_second_hits_in_B2_dphi_even, rs.sc_second_hits_in_B2_dphi_odd),
         ('phi', 'B3', 2, 'either'): (rs.sc_second_hits_in_B3_dphi_even, rs.sc_second_hits_in_B3_dphi_odd),
         ('z', 'B2', 2, 'either'): (rs.sc_second_hits_in_B2_dz_even, rs.sc_second_hits_in_B2_dz_odd),
         ('z', 'B3', 2, 'either'): (rs.sc_second_hits_in_B3_dz_even, rs.sc_second_hits_in_B3_dz_odd),

         ('phi', 'B1', 1, 'even'): rs.sc_first_hits_in_B1_dphi_even,
         ('phi', 'B2', 1, 'even'): rs.sc_first_hits_in_B2_dphi_even,
         ('z', 'B1', 1, 'even'): rs.sc_first_hits_in_B1_dz_even,
         ('z', 'B2', 1, 'even'): rs.sc_first_hits_in_B2_dz_even,
         ('phi', 'B2', 2, 'even'): rs.sc_second_hits_in_B2_dphi_even,
         ('phi', 'B3', 2, 'even'): rs.sc_second_hits_in_B3_dphi_even,
         ('z', 'B2', 2, 'even'): rs.sc_second_hits_in_B2_dz_even,
         ('z', 'B3', 2, 'even'): rs.sc_second_hits_in_B3_dz_even,

         ('phi', 'B1', 1, 'odd'): rs.sc_first_hits_in_B1_dphi_odd,
         ('phi', 'B2', 1, 'odd'): rs.sc_first_hits_in_B2_dphi_odd,
         ('z', 'B1', 1, 'odd'): rs.sc_first_hits_in_B1_dz_odd,
         ('z', 'B2', 1, 'odd'): rs.sc_first_hits_in_B2_dz_odd,
         ('phi', 'B2', 2, 'odd'): rs.sc_second_hits_in_B2_dphi_odd,
         ('phi', 'B3', 2, 'odd'): rs.sc_second_hits_in_B3_dphi_odd,
         ('z', 'B2', 2, 'odd'): rs.sc_second_hits_in_B2_dz_odd,
         ('z', 'B3', 2, 'odd'): rs.sc_second_hits_in_B3_dz_odd,
         }[(var, layer, hit, even_odd)]
    scale = {'phi': 10**3,
             'z': 10**4,
             }[var]
    varlabel = {('z', 1): r"$\Delta z_1(\mu \mathrm{m})$",
                ('z', 2): r"$\Delta z_2(\mu \mathrm{m})$",
                ('phi', 1): r"$\Delta \phi_1(\mathrm{millirad})$",
                ('phi', 2): r"$\Delta \phi_2(\mathrm{millirad})$"
                }[(var, hit)]

    try:
        h = hist(h.ProjectionY(), rescale_x=scale)
    except AttributeError:
        h = hist_add(hist(h[0].ProjectionY(), rescale_x=scale), hist(h[1].ProjectionY(), rescale_x=scale))
    if cut:
        h = hist_slice(h, (-cut, cut))
    hist_plot(h,
              include_errors=True,
              log=log,
              norm=norm,
              xlabel=varlabel,
              title='Layer - ' + layer)
    # if display_layer:
    #     plt.text(0.9, 0.9, layer,
    #              bbox={'facecolor': 'white', 'alpha': 0.7},
    #              transform=plt.gca().transAxes)


def generate_dashboard(plots):
    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/")


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 = ResultSet("DY2LL", sys.argv[1])

    # Next, declare all of the (sub)plots that will be assembled into full
    # figures later
    dphi1_v_eta_B1 = (plot_residual, (rs, 'phi', 'B1', 1), {})
    dphi1_v_eta_B2 = (plot_residual, (rs, 'phi', 'B2', 1), {})
    dz1_v_eta_B1 = (plot_residual, (rs, 'z', 'B1', 1), {})
    dz1_v_eta_B2 = (plot_residual, (rs, 'z', 'B2', 1), {})

    dphi2_v_eta_B2 = (plot_residual, (rs, 'phi', 'B2', 2), {})
    dphi2_v_eta_B3 = (plot_residual, (rs, 'phi', 'B3', 2), {})
    dz2_v_eta_B2 = (plot_residual, (rs, 'z', 'B2', 2), {})
    dz2_v_eta_B3 = (plot_residual, (rs, 'z', 'B3', 2), {})

    projectY = {'projection': 'y'}
    dphi1_B1 = (plot_residual, (rs, 'phi', 'B1', 1), projectY)
    dphi1_B2 = (plot_residual, (rs, 'phi', 'B2', 1), projectY)
    dz1_B1 = (plot_residual, (rs, 'z', 'B1', 1), projectY)
    dz1_B2 = (plot_residual, (rs, 'z', 'B2', 1), projectY)

    dphi2_B2 = (plot_residual, (rs, 'phi', 'B2', 2), projectY)
    dphi2_B3 = (plot_residual, (rs, 'phi', 'B3', 2), projectY)
    dz2_B2 = (plot_residual, (rs, 'z', 'B2', 2), projectY)
    dz2_B3 = (plot_residual, (rs, 'z', 'B3', 2), projectY)

    projectX = {'projection': 'x'}
    eta1_B1 = (plot_residual, (rs, 'phi', 'B1', 1), projectX)
    eta1_B2 = (plot_residual, (rs, 'phi', 'B2', 1), projectX)
    eta2_B2 = (plot_residual, (rs, 'phi', 'B2', 2), projectX)
    eta2_B3 = (plot_residual, (rs, 'phi', 'B3', 2), projectX)

    dphi1_v_ladder_B1 = (plot_residuals_v_ladder, (rs, 'phi', 'B1'), {})
    dphi1_v_ladder_B2 = (plot_residuals_v_ladder, (rs, 'phi', 'B2'), {})
    dz1_v_ladder_B1 = (plot_residuals_v_ladder, (rs, 'z', 'B1'), {})
    dz1_v_ladder_B2 = (plot_residuals_v_ladder, (rs, 'z', 'B2'), {})

    opts = {'log': False, 'norm': None, 'cut': 150}
    dphi1_sc_ex_B1 = (sc_extrapolation_first, (rs, 'phi', 'B1', 1, 'either'), opts)
    dphi1_sc_ex_B2 = (sc_extrapolation_first, (rs, 'phi', 'B2', 1, 'either'), opts)
    dz1_sc_ex_B1 = (sc_extrapolation_first, (rs, 'z', 'B1', 1, 'either'), {})
    dz1_sc_ex_B2 = (sc_extrapolation_first, (rs, 'z', 'B2', 1, 'either'), {})

    dphi2_sc_ex_B1 = (sc_extrapolation_first, (rs, 'phi', 'B1', 2, 'either'), opts)
    dphi2_sc_ex_B2 = (sc_extrapolation_first, (rs, 'phi', 'B2', 2, 'either'), opts)
    dz2_sc_ex_B1 = (sc_extrapolation_first, (rs, 'z', 'B1', 2, 'either'), {})
    dz2_sc_ex_B2 = (sc_extrapolation_first, (rs, 'z', 'B2', 2, 'either'), {})

    dphi1_sc_ex_B1_even = (sc_extrapolation_first, (rs, 'phi', 'B1', 1, 'even'), opts)
    dphi1_sc_ex_B2_even = (sc_extrapolation_first, (rs, 'phi', 'B2', 1, 'even'), opts)
    dz1_sc_ex_B1_even = (sc_extrapolation_first, (rs, 'z', 'B1', 1, 'even'), {})
    dz1_sc_ex_B2_even = (sc_extrapolation_first, (rs, 'z', 'B2', 1, 'even'), {})

    dphi1_sc_ex_B1_odd = (sc_extrapolation_first, (rs, 'phi', 'B1', 1, 'odd'), opts)
    dphi1_sc_ex_B2_odd = (sc_extrapolation_first, (rs, 'phi', 'B2', 1, 'odd'), opts)
    dz1_sc_ex_B1_odd = (sc_extrapolation_first, (rs, 'z', 'B1', 1, 'odd'), {})
    dz1_sc_ex_B2_odd = (sc_extrapolation_first, (rs, 'z', 'B2', 1, 'odd'), {})

    # Now assemble the plots into figures.
    plots = [
             # Plot([[dphi1_v_eta_B1, dphi1_v_eta_B2],
             #       [dz1_v_eta_B1,   dz1_v_eta_B2  ]],
             #      'res1_v_eta'),
             # Plot([[dphi2_v_eta_B2, dphi2_v_eta_B3],
             #       [dz2_v_eta_B2,   dz2_v_eta_B3  ]],
             #      'res2_v_eta'),
             # Plot([[dphi1_B1, dphi1_B2],
             #       [dz1_B1,   dz1_B2  ]],
             #      'res1'),
             # Plot([[eta1_B1, eta1_B2],
             #       [eta2_B2, eta2_B3  ]],
             #      'eta_dist'),
             # Plot([[dphi2_B2, dphi2_B3],
             #       [dz2_B2,   dz2_B3  ]],
             #      'res2'),
             # Plot([[dphi1_v_ladder_B1, dphi1_v_ladder_B2],
             #       [dz1_v_ladder_B1,   dz1_v_ladder_B2]],
             #      'res_v_ladder'),
             # Plot([[dphi1_sc_ex_B1, dphi1_sc_ex_B2],
             #       [dz1_sc_ex_B1, dz1_sc_ex_B2]],
             #      'sc_ex_1'),
             # Plot([[dphi1_sc_ex_B1_even, dphi1_sc_ex_B2_even],
             #       [dz1_sc_ex_B1_even, dz1_sc_ex_B2_even]],
             #      'sc_ex_1_even'),
             # Plot([[dphi1_sc_ex_B1_odd, dphi1_sc_ex_B2_odd],
             #       [dz1_sc_ex_B1_odd, dz1_sc_ex_B2_odd]],
             #      'sc_ex_1_odd'),
             # Plot([[(dphi1_sc_ex_B1_odd, dphi1_sc_ex_B1_even), "FL"],
             #       [dz1_sc_ex_B1_odd, dz1_sc_ex_B1_even]],
             #      'sc_ex_1_odd_v_even'),

             Plot([[dphi1_sc_ex_B1],
                   [dphi1_B1]],
                  'sc_ex_v_sim_phi1_B1'),
             Plot([[dphi1_sc_ex_B2],
                   [dphi1_B2]],
                  'sc_ex_v_sim_phi1_B2'),
             Plot([[dphi2_sc_ex_B2],
                   [dphi2_B2]],
                  'sc_ex_v_sim_phi2_B2'),
             Plot([[dz1_sc_ex_B1],
                   [dz1_B1]],
                  'sc_ex_v_sim_z1_B1'),
             Plot([[dz2_sc_ex_B2],
                   [dz2_B2]],
                  'sc_ex_v_sim_z2_B2'),
             Plot([[dz1_sc_ex_B1],
                   [dz1_sc_ex_B1_even],
                   [dz1_sc_ex_B1_odd]],
                  'even_odd'),
             ]

    # Finally, render and save the plots and generate the html+bootstrap
    # dashboard to view them
    save_plots(plots)
    generate_dashboard(plots)