EGamma_ElectronTrackingValidation/analysis/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
from filval.plotter import (decl_plot, save_plots, hist_plot, hist2d_plot, Plot)


@decl_plot
def plot_residual(rs, var, layer, hit, projection=None):
    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(\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)]
    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)
    elif projection == 'x':  # projecting onto eta axis
        h = h.ProjectionX()
        hist_plot(hist(h), xlabel=r'$\eta$')
    else:
        hist2d_plot(hist2d(h, rescale_y=scale), ylabel=varlabel, xlabel=r"$\eta$")
    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**6,
             '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):
    ''' Raphael's plots '''
    errors = True

    # 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_either,
         ('phi', 'B2', 1, 'either'): rs.sc_first_hits_in_B2_dphi_either,
         ('z', 'B1', 1, 'either'): rs.sc_first_hits_in_B1_dz_either,
         ('z', 'B2', 1, 'either'): rs.sc_first_hits_in_B2_dz_either,
         ('phi', 'B2', 2, 'either'): rs.sc_second_hits_in_B2_dphi_either,
         ('phi', 'B3', 2, 'either'): rs.sc_second_hits_in_B3_dphi_either,
         ('z', 'B2', 2, 'either'): rs.sc_second_hits_in_B2_dz_either,
         ('z', 'B3', 2, 'either'): rs.sc_second_hits_in_B3_dz_either,

         ('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**6,
             '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)]

    h = hist(h.ProjectionY(), rescale_x=scale)
    if cut:
        h = hist_slice(h, (-cut, cut))
    hist_plot(h,
              include_errors=True,
              log=log,
              norm=norm,
              xlabel=varlabel)


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'])
    )

    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/",
                   quote=quote,
                   enumerate=enumerate)


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'), {})

    dphi1_sc_ex_B1 = (sc_extrapolation_first, (rs, 'phi', 'B1', 1, 'either'), {'log': True, 'norm': 1.0})
    dphi1_sc_ex_B2 = (sc_extrapolation_first, (rs, 'phi', 'B2', 1, 'either'), {'log': True, 'norm': 1.0})
    dz1_sc_ex_B1 = (sc_extrapolation_first, (rs, 'z', 'B1', 1, 'either'), {'norm': 1.0})
    dz1_sc_ex_B2 = (sc_extrapolation_first, (rs, 'z', 'B2', 1, 'either'), {'norm': 1.0})

    dphi1_sc_ex_B1_even = (sc_extrapolation_first, (rs, 'phi', 'B1', 1, 'even'), {'log': True, 'norm': 1.0})
    dphi1_sc_ex_B2_even = (sc_extrapolation_first, (rs, 'phi', 'B2', 1, 'even'), {'log': True, 'norm': 1.0})
    dz1_sc_ex_B1_even = (sc_extrapolation_first, (rs, 'z', 'B1', 1, 'even'), {'norm': 1.0})
    dz1_sc_ex_B2_even = (sc_extrapolation_first, (rs, 'z', 'B2', 1, 'even'), {'norm': 1.0})

    dphi1_sc_ex_B1_odd = (sc_extrapolation_first, (rs, 'phi', 'B1', 1, 'odd'), {'log': True, 'norm': 1.0})
    dphi1_sc_ex_B2_odd = (sc_extrapolation_first, (rs, 'phi', 'B2', 1, 'odd'), {'log': True, 'norm': 1.0})
    dz1_sc_ex_B1_odd = (sc_extrapolation_first, (rs, 'z', 'B1', 1, 'odd'), {'norm': 1.0})
    dz1_sc_ex_B2_odd = (sc_extrapolation_first, (rs, 'z', 'B2', 1, 'odd'), {'norm': 1.0})

    # 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'),
             ]

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