EIC Software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
charm_jet_strange_helicity.py
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file charm_jet_strange_helicity.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 [default], eta, bjorken_x")
36 
37 args = parser.parse_args()
38 
39 
40 
41 # Units
42 global u_fb, u_mb, xsection
43 u_fb = 1
44 u_mb = 1e12*u_fb
45 u_pb = 1e3*u_fb
46 #xsection = 2.408e-08*u_mb
47 
48 
49 draw_config={}
50 if args.xvar == 'eta':
51  draw_config['xvar'] = 'Jet.Eta'
52  draw_config['xrange'] = [-4.0, 4.0]
53  draw_config['xbins'] = [-4.0, -3.5, -2.5, -1.0, 1.0, 2.5, 3.5, 4.0]
54  draw_config['ylimits'] = [1, 1e3]
55  draw_config['xlimits'] = [-5,5]
56  draw_config['yunits'] = ''
57  draw_config['xunits'] = ''
58 elif args.xvar == 'pt':
59  draw_config['xvar'] = 'Jet.PT'
60  draw_config['xrange'] = [10,50]
61  draw_config['xbins'] = [10,12.5,15,20,25,30,40,50]
62  draw_config['ylimits'] = [1, 1e3]
63  draw_config['xlimits'] = [0,60]
64  draw_config['yunits'] = '[$\\mathrm{GeV^{-1}}$]'
65  draw_config['xunits'] = '[GeV]'
66 elif args.xvar == 'bjorken_x':
67  draw_config['xvar'] = 'bjorken_x'
68  draw_config['xrange'] = [1e-3,1]
69  #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)])
70  draw_config['xbins'] = [5e-3, 1e-2, 1e-1, 1]
71  draw_config['ylimits'] = [1, 1e3]
72  draw_config['xlimits'] = [1e-3,1]
73  draw_config['yunits'] = ''
74  draw_config['xunits'] = ''
75 
76 
77 #DrawDiffTagYieldPlot(df, draw_config)
78 
79 
80 xvar=draw_config['xvar']
81 xrange=draw_config['xrange']
82 xbins=draw_config['xbins']
83 ylimits=draw_config['ylimits']
84 xlimits=draw_config['xlimits']
85 yunits=draw_config['yunits']
86 xunits=draw_config['xunits']
87 
88 branchlist=["Jet.PT", "Jet.Eta", "Jet.Flavor", "Jet.BTag", "Particle.Px", "Particle.Py", "Particle.Pz", "Particle.E"]
89 
90 print("Loading data...")
91 
92 df = eat.UprootLoad([f"../{args.input}/*/out.root"], "Delphes", branches=branchlist)
93 
94 #df = df[:50000]
95 
96 # Append Auxiliary Data
97 aux_data = {}
98 if xvar == 'bjorken_x':
99  aux_data['bjorken_x'] = eat.DISBjorkenX
100 
101 for aux_var in aux_data:
102  print(f"Appending auxiliary data {aux_var}")
103  df[aux_var] = df.apply(aux_data[aux_var], axis=1)
104 
105  var_array = np.concatenate(df[aux_var].to_numpy()).ravel()
106 
107  print(f"Min. and max. {aux_var}: ", np.min(var_array), " - ", np.max(var_array))
108 
109 
110 ucr_blue = "#2D6CC0"
111 ucr_gold = "#F1AB00"
112 ucr_gray = "#393841"
113 
114 
115 
116 # Flavour tagging study code
117 def DifferentialTaggingYield(df, x='Jet.PT', xrange=[10,50], xbins=[10,12.5,15,20,25,30,40,50], which='all', process='CC_DIS_e18_p275_B15'):
118  global u_fb, u_mb, n_gen
119  n_gen = len(df)
120  print(f"n_gen = {n_gen}")
121  jet_pt = np.concatenate(df['Jet.PT'].to_numpy()).ravel()
122  jet_eta = np.concatenate(df['Jet.Eta'].to_numpy()).ravel()
123  jet_flavor = np.concatenate(df['Jet.Flavor'].to_numpy()).ravel()
124  jet_tag = np.concatenate(df['Jet.BTag'].to_numpy()).ravel()
125 
126  jet_x = np.concatenate(df[x].to_numpy()).ravel()
127 
128 
129 
130  #jet_basics = (0.01 < jet_eta) & (jet_eta < 0.9)
131  jet_basics = (jet_tag == 1)
132 
133  all_flavor = ( jet_flavor > -999.0 ) & (jet_basics)
134  light_flavor = (( jet_flavor < 4 ) | ( jet_flavor == 21 )) & (jet_basics)
135  charm_flavor = ( jet_flavor == 4 ) & (jet_basics)
136 
137  selection = all_flavor
138  if which == 'charm':
139  selection = charm_flavor
140  elif which == 'light':
141  selection = light_flavor
142 
143  (counts, bins) = np.histogram(jet_x[ selection ], range=xrange,
144  bins=xbins)
145  #bins=40)
146 
147 
148  print("Raw Counts:")
149  print(counts)
150  print("Raw Efficiency:")
151  print(counts/n_gen)
152 
153  bin_widths = np.diff(bins)
154  bin_centers = bins[:-1] + bin_widths/2
155 
156  errors = np.sqrt(counts)
157  rel_errors = errors/counts
158 
159 
160  # convert counts to dsigma/dpT * 100/fb
161  print(f"Normalizing yields to x-section {eat.TotalXSection(process):.3f}fb")
162  dsigma_dx = counts * eat.TotalXSection(process) * 100*u_fb**(-1) / n_gen #/ bin_widths
163  dsigma_dx_errors = rel_errors * dsigma_dx
164 
165  return (bin_centers, bin_widths, dsigma_dx, dsigma_dx_errors)
166 
167 
168 
169 
170 
171 process='CC_DIS_e18_p275'
172 charm_ct18nnlo = DifferentialTaggingYield(df, x=xvar, xrange=xrange, xbins=xbins, which='charm', process=process)
173 
174 print(charm_ct18nnlo)
175 
176 
177 polarization_e=0.70
178 polarization_p=0.70
179 polarization=0.70 # single beam polarization
180 
181 N = charm_ct18nnlo[2]
182 errN = np.zeros(len(N))
183 
184 for index in range(len(errN)):
185  errN[index] = np.sqrt(N[index]*(1+polarization_e)/2.0)
186 
187 print("Yields")
188 print(N)
189 print("Errors")
190 print(errN)
191 print("Total Events:")
192 print(np.sum(N))
193 
194 
195 dA = 1.0/(errN*polarization_p)*100
196 
197 
198 fig, ax = plt.subplots(figsize=(12,8))
199 plt.axis('off')
200 
201 
202 gridspec = fig.add_gridspec(ncols=1, nrows=1, width_ratios=[1], height_ratios=[1])
203 
204 bins = charm_ct18nnlo[0]
205 bin_widths = charm_ct18nnlo[1]
206 
207 zero_line = np.zeros(len(bins))
208 
209 # Log plot of differential cross-sections
210 ax1 = fig.add_subplot(gridspec[0, 0])
211 
212 ax1.grid(which='both', axis='both')
213 
214 bandlabel=f"$\delta A$ ($p={polarization*100:.0f}\%$)"
215 
216 ax1.axhline(1.0, color='k', linestyle='-', linewidth=2)
217 errorboxes = [mpl.patches.Rectangle((x - xe, y - ye), 2*xe, 2*ye, facecolor=ucr_blue, alpha=0.8,
218  label=bandlabel)
219  for x, y, xe, ye in zip(bins, zero_line, (bin_widths/2), dA)]
220 
221 # Create patch collection with specified colour/alpha
222 pc = mpl.collections.PatchCollection(errorboxes, facecolor=ucr_blue, alpha=0.8, label=bandlabel)
223 
224 # Add collection to axes
225 ax1.add_collection(pc)
226 
227 
228 ax1.set_ylim([-100, 100])
229 ax1.set_xlim(xlimits)
230 plt.ylabel('Asymmetry Uncertainty [%]')
231 
232 xvar_symbol="p_T"
233 if xvar.find('Eta') != -1:
234  xvar_symbol='\\eta'
235 
236 if xvar == 'bjorken_x':
237  plt.xlabel(f'Bjorken x')
238 else:
239  plt.xlabel(f'Reconstructed Jet ${xvar_symbol}$ {xunits}')
240 
241 ax1.legend(handles=[errorboxes[0]])
242 
243 if xvar == 'bjorken_x':
244  plt.xscale('log')
245 
246 
247 
248 plt.tight_layout()
249 
250 plt.savefig(f"dA_100fb_{xvar}_{args.input}.png")
251 plt.savefig(f"dA_100fb_{xvar}_{args.input}.pdf")