delphes_EIC/notebooks/charm_jet_coverage.py

0001 import matplotlib as mpl
0002 import uproot
0003 import matplotlib.pyplot as plt
0004 import scipy
0005 import numpy as np
0006 import math
0007 import pandas as pd
0008 import seaborn as sns
0009 import mplhep as hep
0010 #import zfit
0011 import inspect
0012 import sys
0013 import argparse
0014
0015 from concurrent.futures import ThreadPoolExecutor
0016
0017 plt.style.use(hep.style.ATLAS)
0018
0019 plt.rcParams.update({'font.sans-serif': "Arial",
0020                      'font.family': "sans-serif",
0021                      'font.size': 30,
0022                      'mathtext.fontset': 'custom',
0023                      'mathtext.rm': 'Arial',
0024                      })
0025
0026 import EICAnalysisTools as eat
0027
0028
0029 # Parse arguments
0030 parser = argparse.ArgumentParser()
0031
0032 parser.add_argument("-n", "--input", type=str,
0033                     help="Directory containing input files")
0034 parser.add_argument("-x", "--xvar", type=str, default='jet_p',
0035                     help="jet_pt, jet_p, etc.")
0036
0037 args = parser.parse_args()
0038
0039 pt_name = "Jet.PT"
0040 eta_name = "Jet.Eta"
0041 flavor_name = "Jet.Flavor"
0042
0043 if args.xvar == "genjet_p":
0044     pt_name = "GenJet.PT"
0045     eta_name = "GenJet.Eta"
0046     flavor_name = "GenJet.Flavor"
0047
0048 branchlist=[pt_name, eta_name, flavor_name]
0049
0050 print("Loading data...")
0051
0052 df = eat.UprootLoad([f"../{args.input}/*/out.root"], "Delphes", branches=branchlist)
0053
0054 #df = df[:10000]
0055
0056 n_gen = len(df)
0057 print(f"n_gen = {n_gen}")
0058
0059
0060 # Code by D. Higinbotham
0061 #========================
0062 # Example Contour Plot
0063
0064 # Bins in angle and momentum
0065
0066 jet_pt = np.concatenate(df[pt_name].to_numpy()).ravel()
0067 jet_eta = np.concatenate(df[eta_name].to_numpy()).ravel()
0068 #jet_tag = np.concatenate(df['GenJet.BTag'].to_numpy()).ravel()
0069 jet_flavor = np.concatenate(df[flavor_name].to_numpy()).ravel()
0070 jet_theta = 2*np.arctan(np.exp(-jet_eta))
0071 jet_p = jet_pt*np.cosh(jet_eta)
0072
0073 #jet_tagged = (jet_tag == 1)
0074 charm_flavor = ( jet_flavor == 4 )
0075
0076
0077
0078 angles = np.radians(np.linspace(0, 180, 90))
0079
0080 mom = np.linspace(0,100,10)
0081 xlabel = "Jet $p_T$ [GeV]"
0082 xvals = jet_pt[charm_flavor]
0083 thetavals = jet_theta[charm_flavor]
0084
0085 if args.xvar == "jet_pt":
0086     mom=np.linspace(0,50,10)
0087     xlabel = "Charm Jet $p_T$ [GeV]"
0088     xvals = jet_pt[charm_flavor]
0089 elif args.xvar == "jet_p":
0090     mom=np.linspace(0,80,16)
0091     xlabel = "Charm Jet Momentum [GeV]"
0092     xvals = jet_p[charm_flavor]
0093 elif args.xvar == "genjet_p":
0094     mom=np.linspace(0,80,16)
0095     xlabel = "Generator-Level Charm Jet Momentum [GeV]"
0096     xvals = jet_p[charm_flavor]
0097 else:
0098     print("Unknown x variable")
0099     sys.exit()
0100
0101
0102 values, thetaedges, redges = np.histogram2d(thetavals, xvals, bins=[angles, mom])
0103 r, theta = np.meshgrid( redges[:-1], thetaedges[:-1])
0104
0105
0106 # Make the plot
0107 fig,ax = plt.subplots(subplot_kw=dict(projection='polar'),dpi=300)
0108 #ax.contourf(theta, r, values,levels=18)
0109 ax.contourf(theta, r, values)
0110 list=[0,np.pi/6,np.pi/3,np.pi/2,4*np.pi/6,5*np.pi/6,np.pi]
0111 ax.set_xticks(list)
0112 ax.set_thetamin(0)
0113 ax.set_thetamax(180)
0114 plt.xlabel(xlabel,labelpad=-40,fontsize=18)
0115 plt.title(f"CC-DIS, 10GeVx275GeV, $Q^2>100\\mathrm{{GeV^2}}$", fontsize=24)
0116 plt.text(3.45/2,np.max(mom)+12.5,'Degrees',fontsize=18,multialignment='right')
0117 ax.tick_params(axis='x', labelsize=12 )
0118 ax.tick_params(axis='y', labelsize=12 )
0119 plt.tight_layout()
0120
0121 plt.savefig(f"charm_jet_coverage_{args.xvar}_{args.input}.png")
0122 plt.savefig(f"charm_jet_coverage_{args.xvar}_{args.input}.pdf")