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 make_plot, plot_registry, hist_plot, hist2d_plot


@make_plot(scale=0.85)
def first_hits_v_eta(rs):
    r'''
    Plots of $\Delta \phi$, $\Delta z$, and $\Delta r$ vs $\eta$ between RecHit and SimHit for
    the first matched hit in the seed.
    '''

    plt.subplot(321)
    hist2d_plot(hist2d(rs.dphi_v_eta_first_hits_in_B1),
                ylabel=r"$\Delta \phi_1$(rad)",
                title="BPIX - Layer 1")

    plt.subplot(322)
    hist2d_plot(hist2d(rs.dphi_v_eta_first_hits_in_B2),
                title="BPIX - Layer 2")

    plt.subplot(323)
    hist2d_plot(hist2d(rs.dz_v_eta_first_hits_in_B1),
                ylabel=r"$\Delta \textrm{z}_1$(cm)",
                )

    plt.subplot(324)
    hist2d_plot(hist2d(rs.dz_v_eta_first_hits_in_B2))

    plt.subplot(325)
    hist2d_plot(hist2d(rs.dr_v_eta_first_hits_in_B1),
                ylabel=r"$\Delta r_1$(cm)",
                xlabel=r"$\eta$"
                )

    plt.subplot(326)
    hist2d_plot(hist2d(rs.dr_v_eta_first_hits_in_B2),
                xlabel=r"$\eta$"
                )


@make_plot(scale=0.85)
def second_hits_v_eta(rs):
    plt.subplot(321)
    hist2d_plot(hist2d(rs.dphi_v_eta_second_hits_in_B2),
                ylabel=r"$\Delta \phi_2$(rad)",
                title="BPIX - Layer 2")

    plt.subplot(322)
    hist2d_plot(hist2d(rs.dphi_v_eta_second_hits_in_B3),
                title="BPIX - Layer 3")

    plt.subplot(323)
    hist2d_plot(hist2d(rs.dz_v_eta_second_hits_in_B2),
                ylabel=r"$\Delta \textrm{z}_2$(cm)")

    plt.subplot(324)
    hist2d_plot(hist2d(rs.dz_v_eta_second_hits_in_B3))

    plt.subplot(325)
    hist2d_plot(hist2d(rs.dr_v_eta_second_hits_in_B2),
                ylabel=r"$\Delta r_2$(cm)",
                xlabel=r"$\eta$")

    plt.subplot(326)
    hist2d_plot(hist2d(rs.dr_v_eta_second_hits_in_B3),
                xlabel=r"$\eta$")


@make_plot(scale=0.85)
def first_hits(rs):
    plt.subplot(321)
    hist_plot(hist(rs.dphi_v_eta_first_hits_in_B1.ProjectionY()),
              include_errors=True, xlabel=r"$\Delta \phi_1$(rad)",
              title="BPIX - Layer 1")

    plt.subplot(322)
    hist_plot(hist(rs.dphi_v_eta_first_hits_in_B2.ProjectionY()),
              include_errors=True,
              xlabel=r"$\Delta \phi_1$(rad)",
              title="BPIX - Layer 2")

    plt.subplot(323)
    hist_plot(hist(rs.dz_v_eta_first_hits_in_B1.ProjectionY()),
              include_errors=True,
              xlabel=r"$\Delta z_1$(cm)")

    plt.subplot(324)
    hist_plot(hist(rs.dz_v_eta_first_hits_in_B2.ProjectionY()),
              include_errors=True,
              xlabel=r"$\Delta z_1$(cm)")

    plt.subplot(325)
    hist_plot(hist(rs.dr_v_eta_first_hits_in_B1.ProjectionY()),
              include_errors=True,
              xlabel=r"$\Delta r_1$(cm)")

    plt.subplot(326)
    hist_plot(hist(rs.dr_v_eta_first_hits_in_B2.ProjectionY()),
              include_errors=True,
              xlabel="r$\Delta r_1$(cm)")


@make_plot(scale=0.85)
def second_hits(rs):
    plt.subplot(321)
    hist_plot(hist(rs.dphi_v_eta_second_hits_in_B2.ProjectionY()),
              include_errors=True,
              xlabel=r"$\Delta \phi_2$(rad)",
              title="BPIX - Layer 2")

    plt.subplot(322)
    hist_plot(hist(rs.dphi_v_eta_second_hits_in_B3.ProjectionY()),
              include_errors=True,
              xlabel=r"$\Delta \phi_2$(rad)",
              title="BPIX - Layer 3")

    plt.subplot(323)
    hist_plot(hist(rs.dz_v_eta_second_hits_in_B2.ProjectionY()),
              include_errors=True,
              xlabel=r"$\Delta z_2$(cm)")

    plt.subplot(324)
    hist_plot(hist(rs.dz_v_eta_second_hits_in_B3.ProjectionY()),
              include_errors=True,
              xlabel=r"$\Delta z_2$(cm)")

    plt.subplot(325)
    hist_plot(hist(rs.dr_v_eta_second_hits_in_B2.ProjectionY()),
              include_errors=True,
              xlabel=r"$\Delta r_2$(cm)")

    plt.subplot(326)
    hist_plot(hist(rs.dr_v_eta_second_hits_in_B3.ProjectionY()),
              include_errors=True,
              xlabel=r"$\Delta r_2$(cm)")


@make_plot(scale=0.85)
def delta_phi_z_v_ladder(rs):
    def even_odd_plot(even, odd, var, scale, unit, **kwargs):
        even_dist = even.ProjectionY()
        odd_dist = odd.ProjectionY()
        hist_plot(hist(even_dist), include_errors=True, label="even ladders")
        hist_plot(hist(odd_dist), include_errors=True, color='r', label="odd ladders",
                  **kwargs)

        even_mean = even_dist.GetMean()*scale
        odd_mean = odd_dist.GetMean()*scale
        axes = plt.gca()
        txt = r"$ \hat{{\Delta {0} }}_{{1,\textrm{{ {1} }} }}={2:4.2g}${3}"
        axes.text(0.05, .7, txt.format(var, "odd", odd_mean, unit), transform=axes.transAxes)
        axes.text(0.05, .6, txt.format(var, "even", even_mean, unit), transform=axes.transAxes)
    plt.subplot(221)
    even_odd_plot(rs.dphi_v_eta_first_hits_in_B1_even_ladder, rs.dphi_v_eta_first_hits_in_B1_odd_ladder,
                  r"\phi", 10**6, "urad", xlabel=r"$\Delta \phi_1$(rad)", title="BPIX - Layer 1")

    plt.subplot(222)
    even_odd_plot(rs.dphi_v_eta_first_hits_in_B2_even_ladder, rs.dphi_v_eta_first_hits_in_B2_odd_ladder,
                  r"\phi", 10**6, "urad", xlabel=r"$\Delta \phi_1$(rad)", title="BPIX - Layer 2")
    plt.legend()

    plt.subplot(223)
    even_odd_plot(rs.dz_v_eta_first_hits_in_B1_even_ladder, rs.dz_v_eta_first_hits_in_B1_odd_ladder,
                  "z", 10**4, "um", xlabel=r"$\Delta z_1$(cm)")

    plt.subplot(224)
    even_odd_plot(rs.dz_v_eta_first_hits_in_B2_even_ladder, rs.dz_v_eta_first_hits_in_B2_odd_ladder,
                  "z", 10**4, "um", xlabel=r"$\Delta z_1$(cm)")


@make_plot(scale=0.85)
def sc_extrapolation_first(rs):
    ''' Raphael's plots '''
    norm = 1
    errors = True

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

    plt.subplot(222)
    hist_plot(hist(rs.sc_first_hits_in_B2_dz.ProjectionY()),
              include_errors=errors,
              norm=norm,
              log=False,
              xlabel=r"$\Delta z_1$(cm)",
              title="BPIX - Layer 2")

    plt.subplot(223)
    hist_plot(preproc(rs.sc_first_hits_in_B1_dphi),
              include_errors=errors,
              norm=norm,
              log=False,
              # ylim=(0.5E-3, 5),
              xlabel=r"$\Delta \phi_1$(rad)",
              # xlim=(-0.06, 0.06)
              )

    plt.subplot(224)
    hist_plot(preproc(rs.sc_first_hits_in_B2_dphi),
              include_errors=errors,
              norm=norm,
              log=False,
              # ylim=(0.5E-3, 5),
              xlabel=r"$\Delta \phi_1$(rad)",
              # xlim=(-0.06, 0.06)
              )


@make_plot(scale=0.85)
def sc_extrapolation_second(rs):
    from scipy.stats import norm

    def gauss(x, A, mu, sigma):
        return A*norm.pdf(x, mu, sigma)

    def gauss2(x, A1, mu1, sigma1, A2, mu2, sigma2):
        return (A1*norm.pdf(x, mu1, sigma1) +
                A2*norm.pdf(x, mu2, sigma2))

    integral = 1
    errors = True

    def preproc(h, slice_=None):
        h = hist(h.ProjectionY())
        if slice_:
            h = hist_slice(h, slice_)
        return h

    plt.subplot(221)
    hist_plot(preproc(rs.sc_second_hits_in_B2_dz),
              include_errors=errors,
              norm=integral,
              log=False,
              fit=(gauss2, (1, 0, 0.025, 1, 0, 0.005)),
              xlabel=r"$\Delta z_2$(cm)",
              title="BPIX - Layer 2")

    plt.subplot(222)
    hist_plot(preproc(rs.sc_second_hits_in_B3_dz),
              include_errors=errors,
              norm=integral,
              log=False,
              fit=(gauss2, (1, 0, 0.025, 1, 0, 0.005)),
              xlabel=r"$\Delta z_2$(cm)",
              title="BPIX - Layer 3")

    plt.subplot(223)
    hist_plot(preproc(rs.sc_second_hits_in_B2_dphi, (-0.09, 0.09)),
              include_errors=errors,
              norm=integral,
              log=False,
              fit=(gauss2, (1, 0, 0.015, 1, 0, 0.005)),
              xlabel=r"$\Delta \phi_2$(rad)")

    plt.subplot(224)
    hist_plot(preproc(rs.sc_second_hits_in_B3_dphi, (-0.09, 0.09)),
              include_errors=errors,
              norm=integral,
              log=False,
              fit=(gauss2, (1, 0, 0.025, 1, 0, 0.005)),
              xlabel=r"$\Delta \phi_2$(rad)")



def generate_dashboard():
    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(plot_registry.values(), 3),
                   quote=quote)


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

    first_hits_v_eta(rs)
    second_hits_v_eta(rs)
    first_hits(rs)
    second_hits(rs)
    delta_phi_z_v_ladder(rs)
    sc_extrapolation_first(rs)
    sc_extrapolation_second(rs)

    generate_dashboard()