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