ONNXRuntime/test/plot_particle_net.py

0001 #!/usr/bin/env python
0002 import ROOT as rt
0003 from DataFormats.FWLite import Events,Handle
0004 import itertools as it
0005 from ROOT import btagbtvdeep
0006 import pandas as pd
0007 import numpy as np
0008 import matplotlib.pyplot as plt
0009
0010 # file evaluated with DeepJet framework
0011 # jets redone from AOD using CMSSW TF modules
0012 #cmssw_miniaod = "test_particle_net_MINIAODSIM.root"
0013 cmssw_miniaod = "test_particle_net_MINIAODSIM_noragged.root"
0014
0015 jetsLabel = "selectedUpdatedPatJets"
0016
0017 from RecoBTag.ONNXRuntime.pfParticleNet_cff import _pfParticleNetJetTagsProbs
0018 from RecoBTag.ONNXRuntime.pfParticleNet_cff import _pfParticleNetSonicJetTagsProbs
0019 disc_names = _pfParticleNetJetTagsProbs+_pfParticleNetSonicJetTagsProbs
0020
0021 jet_pt = "fj_pt"
0022 jet_eta = "fj_eta"
0023
0024 c_numbers = ['event_n']
0025
0026 c_cmssw = { d_name : []  for d_name in disc_names + [jet_pt, jet_eta] + c_numbers }
0027 jetsHandle = Handle("std::vector<pat::Jet>")
0028 cmssw_evs = Events(cmssw_miniaod)
0029
0030 max_n_jets = 1000000
0031 max_n_events = 500000
0032 n_jets = 0
0033
0034 for i, ev in enumerate(cmssw_evs):
0035     event_number = ev.object().id().event()
0036     if (n_jets >= max_n_jets): break
0037     ev.getByLabel(jetsLabel, jetsHandle)
0038     jets = jetsHandle.product()
0039     for i_j,j in enumerate(jets):
0040         uncorr = j.jecFactor("Uncorrected")
0041         ptRaw = j.pt()*uncorr
0042         if ptRaw < 200.0 or abs(j.eta()) > 2.4: continue
0043         if (n_jets >= max_n_jets): break
0044         c_cmssw["event_n"].append(event_number)
0045         c_cmssw[jet_pt].append(ptRaw)
0046         c_cmssw[jet_eta].append(j.eta())
0047         discs = j.getPairDiscri()
0048         for d in discs:
0049             if d.first in disc_names:
0050                 c_cmssw[d.first].append(d.second)
0051         n_jets +=1
0052
0053 df_cmssw = pd.DataFrame(c_cmssw)
0054 df_cmssw.sort_values(['event_n', jet_pt], ascending=[True, False], inplace=True)
0055 df_cmssw.reset_index(drop=True)
0056 print(df_cmssw[['event_n','fj_eta','fj_pt',
0057                 'pfParticleNetJetTags:probTbq','pfParticleNetSonicJetTags:probTbq',
0058             ]])
0059
0060 n_bins = 50
0061
0062 print('number of tags', len(disc_names))
0063 fig, axs = plt.subplots(5,4,figsize=(50,40))
0064 for i,ax in enumerate(axs.flatten()):
0065     cmssw_col = disc_names[i]
0066     ax.hist(df_cmssw[cmssw_col], bins=np.linspace(np.amin(df_cmssw[cmssw_col]), np.amax(df_cmssw[cmssw_col]), n_bins))
0067     ax.set_yscale('log')
0068     ax.set_ylim(0.5, 1000)
0069     ax.set_xlim(0, 1)
0070     ax.set_xlabel(cmssw_col)
0071     ax.set_ylabel('Jets')
0072 #fig.savefig('particle_net_hist.png')
0073 fig.savefig('particle_net_hist_noragged.png')