EIC Software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
charm_jet_coverage.py
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file charm_jet_coverage.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='jet_p',
35  help="jet_pt, jet_p, etc.")
36 
37 args = parser.parse_args()
38 
39 pt_name = "Jet.PT"
40 eta_name = "Jet.Eta"
41 flavor_name = "Jet.Flavor"
42 
43 if args.xvar == "genjet_p":
44  pt_name = "GenJet.PT"
45  eta_name = "GenJet.Eta"
46  flavor_name = "GenJet.Flavor"
47 
48 branchlist=[pt_name, eta_name, flavor_name]
49 
50 print("Loading data...")
51 
52 df = eat.UprootLoad([f"../{args.input}/*/out.root"], "Delphes", branches=branchlist)
53 
54 #df = df[:10000]
55 
56 n_gen = len(df)
57 print(f"n_gen = {n_gen}")
58 
59 
60 # Code by D. Higinbotham
61 #========================
62 # Example Contour Plot
63 
64 # Bins in angle and momentum
65 
66 jet_pt = np.concatenate(df[pt_name].to_numpy()).ravel()
67 jet_eta = np.concatenate(df[eta_name].to_numpy()).ravel()
68 #jet_tag = np.concatenate(df['GenJet.BTag'].to_numpy()).ravel()
69 jet_flavor = np.concatenate(df[flavor_name].to_numpy()).ravel()
70 jet_theta = 2*np.arctan(np.exp(-jet_eta))
71 jet_p = jet_pt*np.cosh(jet_eta)
72 
73 #jet_tagged = (jet_tag == 1)
74 charm_flavor = ( jet_flavor == 4 )
75 
76 
77 
78 angles = np.radians(np.linspace(0, 180, 90))
79 
80 mom = np.linspace(0,100,10)
81 xlabel = "Jet $p_T$ [GeV]"
82 xvals = jet_pt[charm_flavor]
83 thetavals = jet_theta[charm_flavor]
84 
85 if args.xvar == "jet_pt":
86  mom=np.linspace(0,50,10)
87  xlabel = "Charm Jet $p_T$ [GeV]"
88  xvals = jet_pt[charm_flavor]
89 elif args.xvar == "jet_p":
90  mom=np.linspace(0,80,16)
91  xlabel = "Charm Jet Momentum [GeV]"
92  xvals = jet_p[charm_flavor]
93 elif args.xvar == "genjet_p":
94  mom=np.linspace(0,80,16)
95  xlabel = "Generator-Level Charm Jet Momentum [GeV]"
96  xvals = jet_p[charm_flavor]
97 else:
98  print("Unknown x variable")
99  sys.exit()
100 
101 
102 values, thetaedges, redges = np.histogram2d(thetavals, xvals, bins=[angles, mom])
103 r, theta = np.meshgrid( redges[:-1], thetaedges[:-1])
104 
105 
106 # Make the plot
107 fig,ax = plt.subplots(subplot_kw=dict(projection='polar'),dpi=300)
108 #ax.contourf(theta, r, values,levels=18)
109 ax.contourf(theta, r, values)
110 list=[0,np.pi/6,np.pi/3,np.pi/2,4*np.pi/6,5*np.pi/6,np.pi]
111 ax.set_xticks(list)
112 ax.set_thetamin(0)
113 ax.set_thetamax(180)
114 plt.xlabel(xlabel,labelpad=-40,fontsize=18)
115 plt.title(f"CC-DIS, 10GeVx275GeV, $Q^2>100\\mathrm{{GeV^2}}$", fontsize=24)
116 plt.text(3.45/2,np.max(mom)+12.5,'Degrees',fontsize=18,multialignment='right')
117 ax.tick_params(axis='x', labelsize=12 )
118 ax.tick_params(axis='y', labelsize=12 )
119 plt.tight_layout()
120 
121 plt.savefig(f"charm_jet_coverage_{args.xvar}_{args.input}.png")
122 plt.savefig(f"charm_jet_coverage_{args.xvar}_{args.input}.pdf")