EIC Software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
charm_jet_tagging_study.py
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file charm_jet_tagging_study.py
1 import matplotlib as mpl
2 import uproot
3 import matplotlib.pyplot as plt
4 import scipy
5 import numpy as np
6 import math
7 import pandas as pd
8 import seaborn as sns
9 import mplhep as hep
10 #import zfit
11 import inspect
12 import sys
13 import argparse
14 
15 from concurrent.futures import ThreadPoolExecutor
16 
17 plt.style.use(hep.style.ATLAS)
18 
19 plt.rcParams.update({'font.sans-serif': "Arial",
20  'font.family': "sans-serif",
21  'font.size': 30,
22  'mathtext.fontset': 'custom',
23  'mathtext.rm': 'Arial',
24  })
25 
26 import EICAnalysisTools as eat
27 
28 
29 # Parse arguments
30 parser = argparse.ArgumentParser()
31 
32 parser.add_argument("-n", "--input", type=str,
33  help="Directory containing input files")
34 parser.add_argument("-x", "--xvar", type=str, default='pt',
35  help="pt, eta, bjorken_x")
36 
37 args = parser.parse_args()
38 
39 
40 
41 branchlist=["Jet.PT", "Jet.Eta", "Jet.Flavor", "Jet.BTag", "Particle.Px", "Particle.Py", "Particle.Pz", "Particle.E"]
42 
43 print("Loading data...")
44 
45 df = eat.UprootLoad([f"../{args.input}/*/out.root"], "Delphes", branches=branchlist)
46 
47 #df = df[:1000000]
48 
49 n_gen = len(df)
50 print(f"n_gen = {n_gen}")
51 
52 # Append Auxiliary Data
53 aux_data = {}
54 if args.xvar == 'bjorken_x':
55  aux_data['bjorken_x'] = eat.DISBjorkenX
56 
57 for aux_var in aux_data:
58  print(f"Appending auxiliary data {aux_var}")
59  df[aux_var] = df.apply(aux_data[aux_var], axis=1)
60 
61  var_array = np.concatenate(df[aux_var].to_numpy()).ravel()
62 
63  print(f"Min. and max. {aux_var}: ", np.min(var_array), " - ", np.max(var_array))
64 
65 
66 def DrawDiffTagEfficiencyPlot(df, draw_config={}):
67  xvar=draw_config['xvar']
68  xrange=draw_config['xrange']
69  xbins=draw_config['xbins']
70  ylimits=draw_config['ylimits']
71  xlimits=draw_config['xlimits']
72  yunits=draw_config['yunits']
73  xunits=draw_config['xunits']
74 
75 
76  fig, ax = plt.subplots(figsize=(12,8))
77  plt.axis('off')
78 
79 
80  gridspec = fig.add_gridspec(ncols=1, nrows=1, width_ratios=[1], height_ratios=[1])
81 
82 
83  # Log plot of differential cross-sections
84  ax1 = fig.add_subplot(gridspec[0, 0])
85 
86  ax1.grid(which='both', axis='both')
87 
88  (bins, bin_widths, eff, eff_error) = eat.DifferentialTaggingEfficiency(df, x=xvar, xrange=xrange, xbins=xbins, which='light')
89  ax1.errorbar(bins, eff, xerr = bin_widths/2, yerr=eff_error, label='light jets (CT18NNLO)', marker='o', ms=10, ls='none', linewidth=2, color='red')
90 
91 
92  (bins, bin_widths, eff, eff_error) = eat.DifferentialTaggingEfficiency(df, x=xvar, xrange=xrange, xbins=xbins, which='charm')
93  ax1.errorbar(bins, eff, xerr = bin_widths/2, yerr=eff_error, label='charm jets (CT18NNLO)', marker='D', ms=10, ls='none', linewidth=2, color='blue')
94 
95 
96 
97  xvar_symbol="p_T"
98  if xvar.find('Eta') != -1:
99  xvar_symbol='\\eta'
100 
101  plt.ylabel(f'$\\varepsilon_{{tag}}$ {yunits}')
102  plt.xlabel(f'Reconstructed Jet ${xvar_symbol}$ {xunits}')
103 
104  beamconfig = "10x100"
105 
106  plt.title(f"CC-DIS, {beamconfig}GeV, $Q^2>100\\mathrm{{GeV^2}}$", fontsize=20)
107 
108  ax1.set_ylim(ylimits)
109  ax1.set_xlim(xlimits)
110 
111  if xvar == 'bjorken_x':
112  plt.xlabel(f'Bjorken x')
113  else:
114  plt.xlabel(f'Reconstructed Jet ${xvar_symbol}$ {xunits}')
115 
116  if xvar == 'bjorken_x':
117  plt.xscale('log')
118 
119  ax1.legend(fontsize=18)
120 
121  plt.yscale('log')
122 
123  plt.tight_layout()
124 
125  plt.savefig(f"jet_differential_tagefficiency_{xvar}_{args.input}.png")
126  plt.savefig(f"jet_differential_tagefficiency_{xvar}_{args.input}.pdf")
127 
128 
129 def DrawDiffTagYieldPlot(df, draw_config={}, process='CC_DIS_e10_p100_CT18NNLO'):
130  xvar=draw_config['xvar']
131  xrange=draw_config['xrange']
132  xbins=draw_config['xbins']
133  ylimits=draw_config['ylimits']
134  xlimits=draw_config['xlimits']
135  yunits=draw_config['yunits']
136  xunits=draw_config['xunits']
137 
138  fig, ax = plt.subplots(figsize=(12,8))
139  plt.axis('off')
140 
141 
142  gridspec = fig.add_gridspec(ncols=1, nrows=1, width_ratios=[1], height_ratios=[1])
143 
144 
145  # Log plot of differential cross-sections
146  ax1 = fig.add_subplot(gridspec[0, 0])
147 
148  ax1.grid(which='both', axis='both')
149 
150  (bins, bin_widths, xs, xs_error) = eat.DifferentialTaggingYield(df, x=xvar, xrange=xrange, xbins=xbins, which='light', process=process, target_lumi=100)
151  ax1.errorbar(bins, xs, xerr = bin_widths/2, yerr=xs_error, label='light jets (CT18NNLO)', marker='o', ms=10, ls='none', linewidth=2, color='red')
152 
153 
154  charm_ct18nnlo = eat.DifferentialTaggingYield(df, x=xvar, xrange=xrange, xbins=xbins, which='charm', process=process, target_lumi=100)
155  (bins, bin_widths, xs, xs_error) = charm_ct18nnlo
156  ax1.errorbar(bins, xs, xerr = bin_widths/2, yerr=xs_error, label='charm jets (CT18NNLO)', marker='D', ms=10, ls='none', linewidth=2, fillstyle='none', color='blue')
157 
158 
159 
160  xvar_symbol="p_T"
161  if xvar.find('Eta') != -1:
162  xvar_symbol='\\eta'
163 
164  plt.ylabel(f'$\\varepsilon_{{tag}} \\times \\sigma_{{\\mathrm{{CC-DIS}}}} \\times 100\\mathrm{{fb^{{-1}}}}$')
165  plt.xlabel(f'Reconstructed Jet ${xvar_symbol}$ {xunits}')
166  plt.yscale('log')
167 
168  beamconfig = "10x100"
169  process_name = 'CC-DIS'
170 
171  if process.find('NC') != -1:
172  process_name = 'NC-DIS'
173  plt.title(f"{process_name}, {beamconfig}GeV, $Q^2>100\\mathrm{{GeV^2}}$", fontsize=20)
174 
175  ax1.set_ylim(ylimits)
176  ax1.set_xlim(xlimits)
177 
178  if xvar == 'bjorken_x':
179  plt.xlabel(f'Bjorken x')
180  else:
181  plt.xlabel(f'Reconstructed Jet ${xvar_symbol}$ {xunits}')
182 
183  if xvar == 'bjorken_x':
184  plt.xscale('log')
185 
186  ax1.legend(fontsize=18)
187 
188 
189  plt.tight_layout()
190 
191  plt.savefig(f"jet_differential_tagyield_{xvar}_{args.input}.png")
192  plt.savefig(f"jet_differential_tagyield_{xvar}_{args.input}.pdf")
193 
194 
195 # Inputs
196 draw_config={}
197 
198 if args.xvar == 'eta':
199  draw_config['xvar'] = 'Jet.Eta'
200  draw_config['xrange'] = [-4.0, 4.0]
201  draw_config['xbins'] = [-4.0, -3.5, -2.5, -1.0, 1.0, 2.5, 3.5, 4.0]
202  draw_config['ylimits'] = [1e-5, 1]
203  draw_config['xlimits'] = [-5,5]
204  draw_config['yunits'] = ''
205  draw_config['xunits'] = ''
206 elif args.xvar == 'pt':
207  draw_config['xvar'] = 'Jet.PT'
208  draw_config['xrange'] = [10,50]
209  draw_config['xbins'] = [10,12.5,15,20,25,30,40,50]
210  draw_config['ylimits'] = [1e-5, 1]
211  draw_config['xlimits'] = [0,60]
212  draw_config['yunits'] = '[$\\mathrm{GeV^{-1}}$]'
213  draw_config['xunits'] = '[GeV]'
214 elif args.xvar == 'bjorken_x':
215  draw_config['xvar'] = 'bjorken_x'
216  draw_config['xrange'] = [5e-3,1]
217  draw_config['xbins'] = np.concatenate([np.arange(5e-3,1e-2,1e-3),np.arange(1e-2,1e-1,1e-2),np.arange(1e-1,1.1,1e-1)])
218  #draw_config['xbins'] = [5e-3, 1e-2, 1e-1, 1]
219  draw_config['ylimits'] = [1e-5, 1]
220  draw_config['xlimits'] = [5e-3,1]
221  draw_config['yunits'] = ''
222  draw_config['xunits'] = ''
223 
224 
225 
226 DrawDiffTagEfficiencyPlot(df,draw_config)
227 
228 draw_config={}
229 if args.xvar == 'eta':
230  draw_config['xvar'] = 'Jet.Eta'
231  draw_config['xrange'] = [-4.0, 4.0]
232  draw_config['xbins'] = [-4.0, -3.5, -2.5, -1.0, 1.0, 2.5, 3.5, 4.0]
233  draw_config['ylimits'] = [1, 1e3]
234  draw_config['xlimits'] = [-5,5]
235  draw_config['yunits'] = ''
236  draw_config['xunits'] = ''
237 elif args.xvar == 'pt':
238  draw_config['xvar'] = 'Jet.PT'
239  draw_config['xrange'] = [10,50]
240  draw_config['xbins'] = [10,12.5,15,20,25,30,40,50]
241  draw_config['ylimits'] = [1, 1e3]
242  draw_config['xlimits'] = [0,60]
243  draw_config['yunits'] = '[$\\mathrm{GeV^{-1}}$]'
244  draw_config['xunits'] = '[GeV]'
245 elif args.xvar == 'bjorken_x':
246  draw_config['xvar'] = 'bjorken_x'
247  draw_config['xrange'] = [5e-3,1]
248  draw_config['xbins'] = np.concatenate([np.arange(5e-3,1e-2,1e-3),np.arange(1e-2,1e-1,1e-2),np.arange(1e-1,1.1,1e-1)])
249  #draw_config['xbins'] = [5e-3, 1e-2, 1e-1, 1]
250  draw_config['ylimits'] = [1, 1e3]
251  draw_config['xlimits'] = [5e-3,1]
252  draw_config['yunits'] = ''
253  draw_config['xunits'] = ''
254 
255 DrawDiffTagYieldPlot(df, draw_config)
256 
257 
258 if args.xvar == "bjorken_x":
259  sys.exit()
260 
261 # Finally, project the data statistical uncertainties to 100/fb of EIC data
262 # Overlay the PDF range variations with these predicted stat. errors
263 
264 df_20rs2 = eat.UprootLoad([f"../CC_DIS_e10_p100_B15_dR5_maxIP3mm_trkpt10_22sigmin_lha_20Rs2/*/out.root"], "Delphes", branches=branchlist)
265 df_21rs2 = eat.UprootLoad([f"../CC_DIS_e10_p100_B15_dR5_maxIP3mm_trkpt10_22sigmin_lha_21Rs2/*/out.root"], "Delphes", branches=branchlist)
266 
267 #df_20rs2 = df_20rs2[:5000]
268 #df_21rs2 = df_21rs2[:5000]
269 
270 xvar=draw_config['xvar']
271 xrange=draw_config['xrange']
272 xbins=draw_config['xbins']
273 ylimits=draw_config['ylimits']
274 xlimits=draw_config['xlimits']
275 yunits=draw_config['yunits']
276 xunits=draw_config['xunits']
277 
278 charm_ct18nnlo = eat.DifferentialTaggingYield(df, x=xvar, xrange=xrange, xbins=xbins, which='charm', process='CC_DIS_e10_p100_CT18NNLO')
279 charm_ct18nnlo_20rs2 = eat.DifferentialTaggingYield(df_20rs2, x=xvar, xrange=xrange, xbins=xbins, which='charm', process='CC_DIS_e10_p100_CT1820Rs2')
280 charm_ct18nnlo_21rs2 = eat.DifferentialTaggingYield(df_21rs2, x=xvar, xrange=xrange, xbins=xbins, which='charm', process='CC_DIS_e10_p100_CT1821Rs2')
281 
282 
283 # Calculate data statistical errors for 100/fb
284 #(bins, bin_widths, xs, xs_error)
285 
286 print(charm_ct18nnlo[2])
287 print(charm_ct18nnlo[3])
288 
289 
290 
291 N_20 = charm_ct18nnlo_20rs2[2]
292 errN_20 = np.zeros(len(N_20))
293 
294 for index in range(len(errN_20)):
295  errN_20[index] = np.sqrt(N_20[index])
296 
297 R_N_20 = np.ones(len(N_20)) + errN_20/N_20
298 
299 N_21 = charm_ct18nnlo_21rs2[2]
300 
301 diff_20_21 = N_21 - N_20
302 
303 R_N_diff = np.ones(len(N_20)) + diff_20_21/N_20
304 
305 
306 fig, ax = plt.subplots(figsize=(12,8))
307 plt.axis('off')
308 
309 
310 gridspec = fig.add_gridspec(ncols=1, nrows=1, width_ratios=[1], height_ratios=[1])
311 
312 bins = charm_ct18nnlo[0]
313 bin_widths = charm_ct18nnlo[1]
314 
315 one_line = np.ones(len(bins))
316 
317 # Log plot of differential cross-sections
318 ax1 = fig.add_subplot(gridspec[0, 0])
319 
320 ax1.grid(which='both', axis='both')
321 
322 R_N_20_label = "Stat. Uncertainty [CT18NNLO, $R_s=2s/(\overline{u}+\overline{d})=0.325$ (suppressed)]"
323 
324 ax1.axhline(1.0, color='k', linestyle='-', linewidth=2)
325 errorboxes = [mpl.patches.Rectangle((x - xe, y - ye), 2*xe, 2*ye, facecolor='#7A6E67', alpha=0.35,
326  label=R_N_20_label)
327  for x, y, xe, ye in zip(bins, one_line, (bin_widths/2), errN_20/N_20)]
328 
329 # Create patch collection with specified colour/alpha
330 pc = mpl.collections.PatchCollection(errorboxes, facecolor='#7A6E67', alpha=0.35, label=R_N_20_label)
331 
332 # Add collection to axes
333 ax1.add_collection(pc)
334 
335 enhanced = ax1.errorbar(bins, R_N_diff, xerr = bin_widths/2, marker='s', ms=10, ls='none', linewidth=2, fillstyle='none', color='#003066', label='CT18ZNNLO with enhanced strangeness, $R_s=2s/(\overline{u}+\overline{d})=0.863$')
336 
337 #ax1.fill_between(bins, pct_diff_ct18nnlo_20rs2, pct_diff_ct18nnlo_21rs2, color='#2D6CC0', alpha=0.35)
338 
339 ax1.set_ylim([0.0, 2.00])
340 ax1.set_xlim(xlimits)
341 plt.ylabel('')
342 
343 xvar_symbol="p_T"
344 if xvar.find('Eta') != -1:
345  xvar_symbol='\\eta'
346 
347 plt.xlabel(f'Reconstructed Jet ${xvar_symbol}$ {xunits}')
348 
349 if xvar == 'bjorken_x':
350  plt.xlabel(f'Bjorken x')
351 else:
352  plt.xlabel(f'Reconstructed Jet ${xvar_symbol}$ {xunits}')
353 
354 if xvar == 'bjorken_x':
355  plt.xscale('log')
356 
357 ax1.legend(handles=[errorboxes[0],enhanced], fontsize=18)
358 
359 
360 plt.tight_layout()
361 
362 plt.savefig(f"jet_differential_tagyield_100fb_{xvar}_{args.input}.png")
363 plt.savefig(f"jet_differential_tagyield_100fb_{xvar}_{args.input}.pdf")