Caleb Fangmeier лет назад: 6
Родитель
Сommit
bbf73bd04b
5 измененных файлов с 442 добавлено и 410 удалено
  1. 164 86
      plotting/eff_plots.py
  2. 0 295
      plotting/plots.py
  3. 0 29
      plotting/publish.py
  4. 1 0
      plotting/requirements.txt
  5. 277 0
      plotting/sim_track_viz.py

+ 164 - 86
plotting/eff_plots.py

@@ -160,10 +160,10 @@ def hist_integral_ratio(num, den):
 def center_text(x, y, txt, **kwargs):
     plt.text(x, y, txt,
              horizontalalignment='center', verticalalignment='center',
-             transform=plt.gca().transAxes, **kwargs)
+             transform=plt.gca().transAxes, size=24, **kwargs)
 
 
-def hist_plot(h: Hist1D, *args, include_errors=False, **kwargs):
+def hist_plot(h: Hist1D, *args, include_errors=False, line_width=1, **kwargs):
     """ Plots a 1D ROOT histogram object using matplotlib """
 
     counts = h.get_counts()
@@ -172,7 +172,7 @@ def hist_plot(h: Hist1D, *args, include_errors=False, **kwargs):
     x = np.array([left, right]).T.flatten()
     y = np.array([counts, counts]).T.flatten()
 
-    plt.plot(x, y, *args, linewidth=2, **kwargs)
+    plt.plot(x, y, *args, linewidth=line_width, **kwargs)
     if include_errors:
         plt.errorbar(h.get_bin_centers(), h.counts, yerr=h.errors,
                      color='k', marker=None, linestyle='None',
@@ -301,7 +301,7 @@ def plot_roc_curve(pfx, ext=''):
 
 
 @mpb.decl_fig
-def plot_kinematic_eff_all(pref, ext=''):
+def plot_kinematic_eff(pref, ext='', xlim=(None, None), ylim=(None, None), norm=None, label_pfx='', incl_sel=True):
     load_samples()
     ax_pt = plt.subplot(221)
     ax_eta = plt.subplot(222)
@@ -311,38 +311,35 @@ def plot_kinematic_eff_all(pref, ext=''):
         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))
+        def do_plot(ax, name):
+            plt.sca(ax)
+            h = Hist1D(sample[name], no_overflow=True)
+            if norm:
+                h = h / (norm*h.get_integral())
+            hist_plot(h, include_errors=errors, label=l, color=c)
+
+        do_plot(ax_pt, f'{pref}_v_pt{ext}')
+        do_plot(ax_eta, f'{pref}_v_eta{ext}')
+        do_plot(ax_phi, f'{pref}_v_phi{ext}')
 
     plt.sca(ax_pt)
-    center_text(0.5, 0.3, r'$|\eta|<2.4$')
-    plt.xlabel(fr"{reference} $p_T$")
-    set_ylim()
+    if not incl_sel: center_text(0.5, 0.15, r'$|\eta|<2.4$')
+    plt.xlabel(fr"{label_pfx} $p_T$")
+    plt.ylim(ylim)
+    plt.xlim(xlim)
 
     plt.sca(ax_eta)
-    center_text(0.5, 0.3, r'$p_T>20$')
-    plt.xlabel(fr"{reference} $\eta$")
-    set_ylim()
+    if not incl_sel: center_text(0.5, 0.15, r'$p_T>20$')
+    plt.xlabel(fr"{label_pfx} $\eta$")
+    plt.ylim(ylim)
+    plt.xlim(xlim)
 
     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()
+    if not incl_sel: center_text(0.5, 0.15, r'$p_T>20$ and $|\eta|<2.4$')
+    plt.xlabel(fr"{label_pfx} $\phi$")
+    plt.ylim(ylim)
+    plt.xlim(xlim)
     plt.tight_layout()
     plt.legend(loc='upper left', bbox_to_anchor=(0.6, 0.45), bbox_transform=plt.gcf().transFigure)
 
@@ -385,24 +382,45 @@ def plot_res_contour(proc, hit_number, var, layers, ext='_TrackMatched'):
 
 
 @mpb.decl_fig
-def simple_dist(hist_name, xlabel, rebin=False, xmax=None):
+def simple_dist(hist_name, rebin=(), norm=1, xlabel="", ylabel="", xlim=None, ylim=None, line_width=1):
     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))
+            h.rebin(*rebin)
+
+        mean = np.sum(h.get_counts() * h.get_bin_centers()) / h.get_integral()
+        if norm is not None:
+            h = h * (norm / h.get_integral())
+        hist_plot(h, label=f'{sample_name} ($\\mu={mean:.2f}$)',
+                  color=color(proc, wp), line_width=line_width)
+    if xlim:
+        plt.xlim(xlim)
+    if ylim:
+        plt.ylim(ylim)
     plt.xlabel(xlabel)
+    plt.ylabel(ylabel)
     plt.legend()
 
 
+@mpb.decl_fig
+def simple_dist2d(hist_name, proc, wp, xlabel="", ylabel="", xlim=None, ylim=None, norm=None):
+    load_samples()
+    sample = samples[(proc, wp)]
+    # sample_name = f'{proc}-{wp}'
+    h = Hist2D(sample[hist_name])
+    if norm is not None:
+        h = h * (norm / h.get_integral())
+    plot_2d(h, colz_fmt='g')
+    if xlim:
+        plt.xlim(xlim)
+    if ylim:
+        plt.ylim(ylim)
+    plt.xlabel(xlabel)
+    plt.ylabel(ylabel)
+
+
 def all_cut_plots(refresh=True, publish=False):
 
     figures = {
@@ -410,50 +428,75 @@ def all_cut_plots(refresh=True, publish=False):
         '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$')),
+        'number_of_seeds': (simple_dist, ('n_seeds',),
+                            dict(xlabel='Number of Seeds', rebin=(50, -0.5, 200.5))),
+        'number_of_good_seeds': (simple_dist, ('n_good_seeds',),
+                                 dict(xlabel='Number of Good Seeds', rebin=(50, -0.5, 200.5))),
+        'number_of_scls': (simple_dist, ('n_scl',),
+                           dict(xlabel='Number of Super-Clusters', xlim=(-0.5, 25.5))),
+        'number_of_good_scls': (simple_dist, ('n_good_scl',),
+                                dict(xlabel='Number of Super-Clusters', xlim=(-0.5, 25.5))),
+
+        'number_of_sim_els': (simple_dist, ('n_good_sim',),
+                              dict(xlabel='Number of prompt(ish) electrons', xlim=(-0.5, 20.5))),
+        'number_of_gsf_tracks': (simple_dist, ('n_gsf_track',),
+                                 dict(xlabel='Number of reco electrons', xlim=(-0.5, 20.5))),
+
+        'number_of_matched': (simple_dist, ('n_matched',),
+                              dict(xlabel='Number of matched electrons', xlim=(-0.5, 10.5), line_width=4)),
+        'number_of_merged': (simple_dist, ('n_merged',),
+                             dict(xlabel='Number of merged electrons', xlim=(-0.5, 10.5), line_width=4)),
+        'number_of_lost': (simple_dist, ('n_lost',),
+                           dict(xlabel='Number of lost electrons', xlim=(-0.5, 10.5), line_width=4)),
+        'number_of_split': (simple_dist, ('n_split',),
+                            dict(xlabel='Number of split electrons', xlim=(-0.5, 10.5), line_width=4)),
+        'number_of_faked': (simple_dist, ('n_faked',),
+                            dict(xlabel='Number of faked electrons', xlim=(-0.5, 10.5), line_width=4)),
+        'number_of_flipped': (simple_dist, ('n_flipped',),
+                              dict(xlabel='Number of flipped electrons', xlim=(-0.5, 10.5), line_width=4)),
+        'matched_dR': (simple_dist, ('matched_dR',),
+                       dict(xlabel='dR between sim and reco')),
+        'matched_dpT': (simple_dist, ('matched_dpT',),
+                        dict(xlabel='dpT between sim and reco')),
+
+        'number_of_matched_dR': (simple_dist, ('n_matched_dR',),
+                                 dict(xlabel='Number of matched electrons - dR Matched', xlim=(-0.5, 10.5),
+                                      line_width=4)),
+        'number_of_merged_dR': (simple_dist, ('n_merged_dR',),
+                                dict(xlabel='Number of merged electrons - dR Matched', xlim=(-0.5, 10.5),
+                                     line_width=4)),
+        'number_of_lost_dR': (simple_dist, ('n_lost_dR',),
+                              dict(xlabel='Number of lost electrons - dR Matched', xlim=(-0.5, 10.5),
+                                   line_width=4)),
+        'number_of_split_dR': (simple_dist, ('n_split_dR',),
+                               dict(xlabel='Number of split electrons - dR Matched', xlim=(-0.5, 10.5),
+                                    line_width=4)),
+        'number_of_faked_dR': (simple_dist, ('n_faked_dR',),
+                               dict(xlabel='Number of faked electrons - dR Matched', xlim=(-0.5, 10.5),
+                                    line_width=4)),
+        'number_of_flipped_dR': (simple_dist, ('n_flipped_dR',),
+                                 dict(xlabel='Number of flipped electrons - dR Matched', xlim=(-0.5, 10.5),
+                                      line_width=4)),
+        'matched_dR_dR': (simple_dist, ('matched_dR_dR',),
+                          dict(xlabel='dR between sim and reco - dR Matched')),
+        'matched_dpT_dR': (simple_dist, ('matched_dpT_dR',),
+                           dict(xlabel='dpT between sim and reco - dR Matched')),
+
+        'sim_pt': (simple_dist, ('sim_pt',),
+                   dict(xlabel='Sim Track $p_T$', xlim=(0, None))),
+        'sim_eta': (simple_dist, ('sim_eta',),
+                    dict(xlabel='Sim Track $eta$', rebin=(20, -3, 3))),
+        'sim_phi': (simple_dist, ('sim_phi',),
+                    dict(xlabel='Sim Track $phi$', rebin=(20, -3.14, 3.14), ylim=(0, None))),
+        'reco_pt': (simple_dist, ('reco_pt',),
+                    dict(xlabel='Reco Track $p_T$', xlim=(0, None))),
+        'reco_eta': (simple_dist, ('reco_eta',),
+                     dict(xlabel='Reco Track $eta$', rebin=(20, -3, 3))),
+        'reco_phi': (simple_dist, ('reco_phi',),
+                     dict(xlabel='Reco Track $phi$', rebin=(20, -3.14, 3.14), ylim=(0, None))),
+
+        'tm_corr': (simple_dist2d, ('tm_corr', 'zee', 'old-default'),
+                    dict(xlabel='Seed Matched', ylabel='Track Matched', norm=1)),
 
         'ecal_rel_res': plot_ecal_rel_res,
         'hit_v_layer_BPIX_new-default_zee': (plot_hit_vs_layer, (('zee', 'new-default'), 'barrel')),
@@ -464,8 +507,43 @@ def all_cut_plots(refresh=True, publish=False):
         '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')),
+
+
+        'good_sim_kinem': (plot_kinematic_eff, ('good_sim',),
+                           dict(norm=1, ylim=(0, None))),
+        'gsf_track_kinem': (plot_kinematic_eff, ('gsf_track',),
+                            dict(norm=1, ylim=(0, None))),
+        'seed_kinem': (plot_kinematic_eff, ('seed',),
+                       dict(norm=1, ylim=(0, None))),
+        'scl_kinem': (plot_kinematic_eff, ('scl',),
+                      dict(norm=1, ylim=(0, None))),
     }
 
+    def add_num_den(key, func, args, kwargs):
+        figures[key] = (func, args, dict(**kwargs, ylim=(0, 1.1)))
+        base_ext = kwargs.get('ext', '')
+        kwargs_ = kwargs.copy()
+        kwargs_['ext'] = base_ext+'_num'
+        figures[key+'_num'] = (func, args, kwargs_)
+        kwargs_ = kwargs.copy()
+        kwargs_['ext'] = base_ext+'_den'
+        figures[key+'_den'] = (func, args, kwargs_)
+
+    add_num_den('tracking_eff', plot_kinematic_eff, ('tracking_eff',), dict(incl_sel=False))
+    add_num_den('tracking_pur', plot_kinematic_eff, ('tracking_pur',), dict(incl_sel=False))
+    add_num_den('tracking_eff_dR', plot_kinematic_eff, ('tracking_eff',), dict(ext='_dR', incl_sel=False))
+    add_num_den('tracking_pur_dR', plot_kinematic_eff, ('tracking_pur',), dict(ext='_dR', incl_sel=False))
+    add_num_den('seeding_eff', plot_kinematic_eff, ('seed_eff',), dict(incl_sel=False))
+    add_num_den('seeding_pur', plot_kinematic_eff, ('seed_pur',), dict(incl_sel=False))
+    add_num_den('fake_rate_incl', plot_kinematic_eff, ('fake_rate_incl',), {})
+    add_num_den('partial_fake_rate_incl', plot_kinematic_eff, ('partial_fake_rate_incl',), {})
+    add_num_den('full_fake_rate_incl', plot_kinematic_eff, ('full_fake_rate_incl',), {})
+    add_num_den('clean_fake_rate_incl', plot_kinematic_eff, ('clean_fake_rate_incl',), {})
+    add_num_den('fake_rate', plot_kinematic_eff, ('fake_rate',), dict(incl_sel=False))
+    add_num_den('partial_fake_rate', plot_kinematic_eff, ('partial_fake_rate',), dict(incl_sel=False))
+    add_num_den('full_fake_rate', plot_kinematic_eff, ('full_fake_rate',), dict(incl_sel=False))
+    add_num_den('clean_fake_rate', plot_kinematic_eff, ('clean_fake_rate',), dict(incl_sel=False))
+
     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,
@@ -490,9 +568,9 @@ def all_cut_plots(refresh=True, publish=False):
                         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')
+    mpb.generate_report(figures, 'Update',
+                        output='report.html',
+                        body='../docs/reports/report_2018_05_30.md')
     if publish:
         mpb.publish()
 
@@ -534,4 +612,4 @@ if __name__ == '__main__':
         'new-wide': 'Wide Settings',
         'old-default': 'Old Seeding',
     }
-    all_cut_plots(refresh=True, publish=False)
+    all_cut_plots(refresh=True, publish=True)

+ 0 - 295
plotting/plots.py

@@ -1,295 +0,0 @@
-#!/usr/bin/env python
-import sys
-
-import matplotlib.pyplot as plt
-
-from uproot import open as root_open
-from matplotboard import decl_fig, render, generate_report
-
-
-@decl_fig
-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_fig
-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_fig
-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
-    render(plots)
-    generate_report(plots)

+ 0 - 29
plotting/publish.py

@@ -1,29 +0,0 @@
-#!/usr/bin/env python
-from datetime import datetime as dt
-from subprocess import run
-# from glob import glob
-
-# dir_name = f'seeding_studies_{dt.strftime(dt.now(), "%Y_%m_%d_%H")}'
-# remote = 'caleb@fangmeier.tech:/var/www/eg/'
-# run(('rm', '-rf', dir_name))
-# run(('mkdir', '-p', dir_name+'/output'))
-# for gdir in glob('../hists*'):
-#     run(('cp', '-r', gdir, dir_name))
-# run(('cp', '-r', 'output', dir_name))
-#
-# run(('scp', '-r', dir_name, remote))
-# print(f'Data available at https://eg.fangmeier.tech/{dir_name}/')
-
-
-
-data_dir = "../hists/"
-par_dir = "dashboard"
-dir_name = f'{par_dir}_{dt.strftime(dt.now(), "%Y_%m_%d_%H")}'
-remote = f'caleb@fangmeier.tech:/var/www/eg/'
-run(('rm', '-rf', dir_name))
-# run(('mkdir', '-p', dir_name))
-run(('cp', '-r', par_dir, dir_name))
-run(('cp', '-r', data_dir, dir_name))
-
-run(('scp', '-r', dir_name, remote))
-print(f'Data available at https://eg.fangmeier.tech/{dir_name}/')

+ 1 - 0
plotting/requirements.txt

@@ -1,3 +1,4 @@
+numpy
 matplotlib
 ipython
 uproot

+ 277 - 0
plotting/sim_track_viz.py

@@ -0,0 +1,277 @@
+#!/usr/bin/env python
+import numpy as np
+from mpl_toolkits.mplot3d import Axes3D
+import matplotlib.pyplot as plt
+from uproot import open as root_open
+# plt.interactive(True)
+
+tree = None
+bsp = None
+
+def pdg_color(pdgId):
+    return {
+        11: 'b',
+        -11: 'b',
+        22: 'g',
+    }.get(pdgId, 'k')
+
+
+def get_roots(sourceSimIdxs):
+    roots = []
+    for idx, sourceSimIdx in enumerate(sourceSimIdxs):
+        if len(sourceSimIdx) == 0:
+            roots.append(idx)
+    return np.array(roots[:len(roots)//2])
+
+
+def p(px, py, pz):
+    return np.sqrt(px**2 + py**2 + pz**2)
+
+
+def plot_all_roots(sim_vtxs):
+    # ax = plt.gca(projection='3d')
+    ax = plt.gca()
+
+    xs = sim_vtxs[b'simvtx_x'] - bsp[b'bsp_x']
+    ys = sim_vtxs[b'simvtx_y'] - bsp[b'bsp_y']
+    zs = sim_vtxs[b'simvtx_z'] - bsp[b'bsp_z']
+    rs = np.sqrt(xs*xs + ys*ys)
+    roots = get_roots(sim_vtxs[b'simvtx_sourceSimIdx'])
+
+    bound_z = bsp[b'bsp_sigmaz']*5
+    # bound_z = 1
+    bound_r = np.sqrt(bsp[b'bsp_sigmax']**2 + bsp[b'bsp_sigmay']**2)
+    lumi_roots = []
+    for root in roots:
+        if abs(zs[root]) < bound_z and np.sqrt(xs[root]**2 + ys[root]**2) < bound_r:
+            lumi_roots.append(root)
+    # nvtx = len(xs)//2
+
+
+    # ax.plot(xs[roots], ys[roots], zs[roots], '.')
+    # ax.plot(zs[roots], rs[roots], '.')
+    ax.plot(zs[lumi_roots], rs[lumi_roots], '.')
+    # ax.plot([-bound_z, -bound_z, bound_z, bound_z], [0, bound_r, bound_r, 0], 'r')
+
+    ax.set_xlabel('z')
+    ax.set_ylabel('r')
+    ax.set_aspect('equal')
+    # ax.set_zlabel('z')
+
+
+def plot_all_vtx(sim_vtxs):
+    ax = plt.gca(projection='3d')
+
+    xs = sim_vtxs[b'simvtx_x']
+    ys = sim_vtxs[b'simvtx_y']
+    zs = sim_vtxs[b'simvtx_z']
+    # nvtx = len(xs)//2
+
+    ax.plot(xs, ys, zs, '.')
+
+    ax.set_xlabel('x')
+    ax.set_ylabel('y')
+    ax.set_zlabel('z')
+
+
+def plot_event_tree(sim_tracks, sim_vtxs, pv_idx):
+    print('='*80)
+    ax = plt.gca(projection='3d')
+    pdgIds = sim_tracks[b'sim_pdgId']
+    decayVtxIdxs = sim_tracks[b'sim_decayVtxIdx']
+    px = sim_tracks[b'sim_px']
+    py = sim_tracks[b'sim_py']
+    pz = sim_tracks[b'sim_pz']
+
+    xs = sim_vtxs[b'simvtx_x']
+    ys = sim_vtxs[b'simvtx_y']
+    zs = sim_vtxs[b'simvtx_z']
+    daughterSimIdxs = sim_vtxs[b'simvtx_daughterSimIdx']
+
+    vtx_idxs = [(0, pv_idx)]
+    while vtx_idxs:
+        depth, vtx_idx = vtx_idxs.pop()
+        print(' '*depth + str(vtx_idx), end=' -> ', flush=True)
+        ax.plot([xs[vtx_idx]], [ys[vtx_idx]], [zs[vtx_idx]], 'r.')
+        start_x = xs[vtx_idx]
+        start_y = ys[vtx_idx]
+        start_z = zs[vtx_idx]
+        for sim_idx in daughterSimIdxs[vtx_idx]:
+            pdgId = pdgIds[sim_idx]
+            # print(pdgId)
+            # if abs(pdgId) != 11:
+            #     continue
+            # if pdgId == 22:
+            #     continue
+            # if pdgId == 22 and p(px[sim_idx], py[sim_idx], pz[sim_idx]) < 1.0:
+            #     continue
+            for decay_vtx_idx in decayVtxIdxs[sim_idx]:
+                end_x = xs[decay_vtx_idx]
+                end_y = ys[decay_vtx_idx]
+                end_z = zs[decay_vtx_idx]
+                # if abs(pdgId) == 11:
+                if True:
+                    ax.plot([start_x, end_x], [start_y, end_y], [start_z, end_z],
+                            pdg_color(pdgId))
+                vtx_idxs.append((depth+1, decay_vtx_idx))
+                print(str(decay_vtx_idx) + ', ', end='')
+        print()
+
+    ax.plot([xs[pv_idx]], [ys[pv_idx]], [zs[pv_idx]], 'k.')
+    # ax.set_xlim((-5, 5))
+    # ax.set_ylim((-5, 5))
+    # ax.set_zlim((-5, 5))
+    ax.set_xlabel('x')
+    ax.set_ylabel('y')
+    ax.set_zlabel('z')
+
+
+def print_sim_vtxs(sim_vtxs, sim_tracks, sim_pvs, gens, gsf_tracks):
+    daughterSimIdxs = sim_vtxs[b'simvtx_daughterSimIdx']
+    sourceSimIdxs = sim_vtxs[b'simvtx_sourceSimIdx']
+    processTypes = sim_vtxs[b'simvtx_processType']
+    xs = sim_vtxs[b'simvtx_x'] - bsp[b'bsp_x']
+    ys = sim_vtxs[b'simvtx_y'] - bsp[b'bsp_y']
+    zs = sim_vtxs[b'simvtx_z'] - bsp[b'bsp_z']
+
+    parentVtxIdxs = sim_tracks[b'sim_parentVtxIdx']
+    decayVtxIdxs = sim_tracks[b'sim_decayVtxIdx']
+    sim_pdgIds = sim_tracks[b'sim_pdgId']
+    sim_pxs = sim_tracks[b'sim_px']
+    sim_pys = sim_tracks[b'sim_py']
+    sim_pzs = sim_tracks[b'sim_pz']
+
+    gen_pdgIds = gens[b'gen_pdgId']
+    vxs = gens[b'gen_vx'] - bsp[b'bsp_x']
+    vys = gens[b'gen_vy'] - bsp[b'bsp_y']
+    vzs = gens[b'gen_vz'] - bsp[b'bsp_z']
+    gen_pxs = gens[b'gen_px']
+    gen_pys = gens[b'gen_py']
+    gen_pzs = gens[b'gen_pz']
+
+    gsf_pdgIds = -11 * gsf_tracks[b'trk_q']
+    gsf_vxs = gsf_tracks[b'trk_vtxx'] - bsp[b'bsp_x']
+    gsf_vys = gsf_tracks[b'trk_vtxy'] - bsp[b'bsp_y']
+    gsf_vzs = gsf_tracks[b'trk_vtxz'] - bsp[b'bsp_z']
+    gsf_pxs = gsf_tracks[b'trk_px']
+    gsf_pys = gsf_tracks[b'trk_py']
+    gsf_pzs = gsf_tracks[b'trk_pz']
+
+
+    # print('VTX')
+    # for (idx, (sourceSimIdx, daughterSimIdx, processType)) in enumerate(zip(sourceSimIdxs, daughterSimIdxs, processTypes)):
+    #     if idx in sim_pvs:
+    #         print(f'*{idx}|{processType}', sourceSimIdx, daughterSimIdx, sep=" - ")
+    #     else:
+    #         print(f'{idx}|{processType}', sourceSimIdx, daughterSimIdx, sep=" - ")
+
+    print('GEN')
+    for (idx, (px, py, pz, vx, vy, vz, pdgId)) in enumerate(zip(gen_pxs, gen_pys, gen_pzs, vxs, vys, vzs, gen_pdgIds)):
+        if abs(pdgId) != 11: continue
+        p = np.sqrt(px**2 + py**2 + pz**2)
+        theta = np.arctan2(np.hypot(px, py), pz)
+        phi = np.arctan2(px, py)
+        print(f'{idx: 4d}|{pdgId: 3d} - ({vx:8.2f},{vy:8.2f},{vz:8.2f}) ({theta:5.2f},{phi:5.2f}) {p:.2f}GeV')
+
+    print('SIM')
+    for (idx, (px, py, pz, parentVtxIdx, decayVtxIdx, pdgId)) in enumerate(zip(gen_pxs, gen_pys, gen_pzs, parentVtxIdxs, decayVtxIdxs, sim_pdgIds)):
+        if abs(pdgId) != 11: continue
+        if len(sourceSimIdxs[parentVtxIdx]) > 0: continue
+        vx = xs[parentVtxIdx]
+        vy = ys[parentVtxIdx]
+        vz = zs[parentVtxIdx]
+        p = np.sqrt(px**2 + py**2 + pz**2)
+        theta = np.arctan2(np.hypot(px, py), pz)
+        phi = np.arctan2(px, py)
+        print(f'{idx: 4d}|{pdgId: 3d} - ({vx:8.2f},{vy:8.2f},{vz:8.2f}) ({theta:5.2f},{phi:5.2f}) {p:.2f}GeV')
+
+    print('RECO')
+    for (idx, (px, py, pz, vx, vy, vz, pdgId)) in enumerate(zip(gsf_pxs, gsf_pys, gsf_pzs, gsf_vxs, gsf_vys, gsf_vzs, gsf_pdgIds)):
+        p = np.sqrt(px**2 + py**2 + pz**2)
+        theta = np.arctan2(np.hypot(px, py), pz)
+        phi = np.arctan2(px, py)
+        print(f'{idx: 4d}|{pdgId: 3d} - ({vx:8.2f},{vy:8.2f},{vz:8.2f}) ({theta:5.2f},{phi:5.2f}) {p:.2f}GeV')
+
+    input()
+
+
+def plot_event(sim_tracks, sim_vtxs, sim_pvs, gens, gsf_tracks, event_idx):
+    print(f"Processing event {event_idx}")
+    sim_tracks = {k: v[event_idx] for k, v in sim_tracks.items()}
+    sim_vtxs = {k: v[event_idx] for k, v in sim_vtxs.items()}
+    gens = {k: v[event_idx] for k, v in gens.items()}
+    gsf_tracks = {k: v[event_idx] for k, v in gsf_tracks.items()}
+    sim_pvs = sim_pvs[event_idx]
+
+    print_sim_vtxs(sim_vtxs, sim_tracks, sim_pvs, gens, gsf_tracks)
+    # print(len(sim_pvs[event_idx]))
+
+    # plt.clf()
+    # plot_all_roots(sim_vtxs)
+    # plt.show()
+
+    # for pv_idx in get_roots(sim_vtxs[b'simvtx_sourceSimIdx']):
+    #     plt.clf()
+    #     plot_event_tree(sim_tracks, sim_vtxs, pv_idx)
+    #     # plot_all_vtx(sim_vtxs)
+    #     plt.show()
+
+
+def main():
+    global tree
+    global bsp
+    f = root_open('../data/zee.root')
+    tree = f['trackingNtuple/tree']
+    gsf_tracks = tree.arrays([
+        'trk_q',
+        'trk_vtxx',
+        'trk_vtxy',
+        'trk_vtxz',
+        'trk_px',
+        'trk_py',
+        'trk_pz',
+    ])
+    sim_tracks = tree.arrays([
+        'sim_pdgId',
+        'sim_parentVtxIdx',
+        'sim_decayVtxIdx',
+        'sim_pt',
+        'sim_eta',
+        'sim_px',
+        'sim_py',
+        'sim_pz',
+    ])
+    sim_vtxs = tree.arrays([
+        'simvtx_x',
+        'simvtx_y',
+        'simvtx_z',
+        'simvtx_sourceSimIdx',
+        'simvtx_daughterSimIdx',
+        'simvtx_processType',
+    ])
+    bsp = tree.arrays([
+        'bsp_x',
+        'bsp_y',
+        'bsp_z',
+        'bsp_sigmax',
+        'bsp_sigmay',
+        'bsp_sigmaz',
+    ])
+    gens = tree.arrays([
+        'gen_vx',
+        'gen_vy',
+        'gen_vz',
+        'gen_px',
+        'gen_py',
+        'gen_pz',
+        'gen_pdgId',
+    ])
+    bsp = {k: v[0] for k, v in bsp.items()}
+    sim_pvs = tree.array('simpv_idx')
+    for event_idx in range(tree.fEntries):
+        plot_event(sim_tracks, sim_vtxs, sim_pvs, gens, gsf_tracks, event_idx)
+        # break
+
+
+if __name__ == '__main__':
+    main()