PerformanceStudy/SinglePion/run_performance_plot.py

0001
0002 '''
0003 Generate preTDR style single particle performance plot for tracking study.
0004 Run after run_performance_study.py
0005 Shujie Li, 10.2025
0006
0007 Note:
0008 * both functions extract kinematic info from your rec file name, e.g. rec_defaul_0.5GeV_eta_-1_0_10000
0009 * the uncertainties pacage is required for error propagation.  Install:
0010 pip install uncertainties
0011
0012 plot_eff(): plot eff v.s. mom for one setting generated by run_performance_study.py
0013
0014 plot_resol(): plot dp/p, theta, phi, DCAr resol v.s. mom (or pT) in five eta ranges [-3.5, -2.5, -1, 1, 2.5, 3.5]. dp/p and DCAr plots come with pwg requirement curves.
0015 * varname=th, ph, dp, dca
0016 * can plot results from one or two settings
0017 * IF each eta range are further binned in finner eta in resolution calculation, we need to combine for plotting
0018
0019
0020 '''
0021
0022 import ast
0023 import pandas as pd
0024 import scipy
0025 from scipy.signal import find_peaks
0026 from uncertainties import unumpy as unp
0027 import numpy as np
0028 from matplotlib import pyplot as plt
0029
0030 plt.rcParams['figure.figsize'] = [8.0, 6.0]
0031 plt.rcParams['ytick.direction'] = 'in'
0032 plt.rcParams['xtick.direction'] = 'in'
0033 plt.rcParams['xaxis.labellocation'] = 'right'
0034 plt.rcParams['yaxis.labellocation'] = 'top'
0035 SMALL_SIZE = 10
0036 MEDIUM_SIZE = 12
0037 BIGGER_SIZE = 16
0038 plt.rc('font', size=SMALL_SIZE)          # controls default text sizes
0039 plt.rc('axes', titlesize=MEDIUM_SIZE)     # fontsize of the axes title
0040 plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
0041 plt.rc('xtick', labelsize=MEDIUM_SIZE)    # fontsize of the tick labels
0042 plt.rc('ytick', labelsize=MEDIUM_SIZE)    # fontsize of the tick labels
0043 plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize
0044 plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title]
0045
0046 deg2rad = np.pi/180.0
0047 ## convert theta to eta
0048 def theta2eta(xx, inverse=0):
0049     xx = np.array(xx)
0050     if inverse==1:
0051         return np.arctan((np.e)**(-xx))*2
0052     else:
0053         return -np.log(np.tan(xx/2.))
0054
0055 ## pwg requirement parameters are available for dp/p and DCAr
0056 pwg = pd.read_csv('./pwg_requirements.txt',sep="\t",skiprows=1)
0057 def pwg_value(varname, eta, mom):
0058     if varname=="dca" or varname=="dp":
0059         cond = (pwg.eta_lo <=eta)&(pwg.eta_hi>eta)
0060         a = pwg[varname+"_par1"].values[cond][0]
0061         b = pwg[varname+"_par2"].values[cond][0]
0062         x = mom
0063         if varname=="dca":
0064             x*=np.sin(theta2eta(eta,1))
0065             return np.sqrt((a/1000.0 / x)**2 + (b/1000.0)**2)
0066         elif varname=="dp":
0067             return np.sqrt((a * x)**2 + b**2)
0068     else:
0069         return -1*np.ones(len(mom))
0070
0071
0072 def plot_eff(eff_file="eff_out.txt", dir_path="./", setting="default"):
0073     plt.figure()
0074     df= pd.read_csv(dir_path+eff_file,header=None,sep=";")
0075     mom_list=[0.5, 1, 2, 5, 10, 15]
0076
0077     df.columns= ['tag', 'eff', 'err','eta_bins']
0078     ## select results from corresponding setting
0079     df = df[df['tag'].str.contains(setting, na=False)].reset_index()
0080
0081     ## this depends on your actual customized file format
0082     ## here assumed rec_official_0.5GeV_eta_-1_0_10000
0083     df["mom_range"]=[ float(dd.split('GeV')[0].split('_')[-1])  for dd in df.tag]
0084     df["eta_lo"]=[ float(dd.split('eta_')[1].split('_')[0])  for dd in df.tag]
0085     df["eta_hi"]=[ float(dd.split('eta_')[1].split('_')[1])  for dd in df.tag]
0086     # df["deg_range"]=[ dd.split('eta_')[1].split('_')[0:1]  for dd in df.tag]
0087     # df["deg_range"]=[ dd.split('eta')[1].split(']_[')[0][1:]  for dd in df.tag]
0088
0089     eta_bins=np.array(ast.literal_eval(df.eta_bins[0]))
0090
0091     line_styles = [ (0, (3, 3, 1, 2)),  '--', '-.', ':','-',(0, (3, 1, 1, 1))]
0092     for ii,mom in enumerate(mom_list):
0093     # for ii,mom in enumerate(['0.5GeV', '1GeV', '2GeV', '5GeV', '10GeV', '15GeV']):
0094         cond = df.mom_range==mom
0095         dft  = df[cond]
0096         ys   = 0
0097         cnt  = 0
0098         for i in dft.index:
0099             row   = df.iloc[i]
0100             yy    = np.array(ast.literal_eval(row.eff))
0101             ee    = np.array(ast.literal_eval(row.err))
0102             cnt   += (yy != 0).astype(int)
0103             ys    += unp.uarray(yy,ee)
0104         cnt[cnt==0]=1 # set cnt to 1 when there's no non-zero values in that bin to avoid division error 0/0
0105         ys   = np.divide(ys,cnt)
0106         cond = eta_bins<10
0107         plt.plot(eta_bins[cond], unp.nominal_values(ys[cond]),ls=line_styles[ii],label=f"{mom} GeV")
0108
0109     plt.legend(frameon=0, loc="upper left",ncol=2, fontsize=13)
0110     plt.ylim(0.,1.4)
0111     plt.xlim(-4,4)
0112     # plt.axvline(-1.5,c='grey')
0113     plt.xlabel("$\\eta$")
0114     plt.ylabel("efficiency")
0115     # plt.title(setting)
0116     plt.grid()
0117     plt.savefig(f"{dir_path}tracking_eff_{setting}.pdf")
0118
0119 # varname=th, ph, dp, dca
0120 # setting2 is optional if you want to compare performance from two settings
0121 def plot_resol(varname="ph",resol_file="resol_out_whole.txt", dir_path="./", setting1="", setting2=""):
0122     df= pd.read_csv(f'{dir_path}resol_out_whole.txt',header=None,sep=" ")
0123
0124     df.columns = ["tag","sig_dp", "err_dp", "sig_th", "err_th", "sig_ph", "err_ph", "sig_dca", "err_dca"]
0125     # df = df.replace(-1,0)
0126     df["eta_lo"]=[ float(dd.split('eta_')[1].split('_')[0])  for dd in df.tag]
0127     df["eta_hi"]=[ float(dd.split('eta_')[1].split('_')[1])  for dd in df.tag]
0128     df["mom"]=[ float(dd.split('GeV')[0].split('_')[-1])  for dd in df.tag]
0129     df["name"]=["_".join(dd.split('GeV')[0].split('_')[1:-1]) for dd in df.tag]
0130
0131     eta_lo_pwg = [-3.5,   -2.5,   -1,   1,    2.5  ]
0132     eta_hi_pwg = [-2.5,   -1,   1,    2.5,    3.5  ]
0133     xline      = np.arange(0,20,0.001)
0134
0135     if varname  == "th":
0136         y_hi     = 0.01
0137         yname    = r"$\theta$ [rad]"
0138         xname    = "momentum [GeV]"
0139         x_hi      = [20,20,20,20,20]
0140
0141     elif varname  == "ph":
0142         y_hi     = 0.01
0143         yname    = r"$\phi$ [rad]"
0144         xname    = "momentum [GeV]"
0145         x_hi      = [20,20,20,20,20]
0146
0147     elif varname  == "dp":
0148         y_hi     = 3
0149         yname    = r"$\delta p/p$ [%]"
0150         xname    = "momentum [GeV]"
0151         x_hi      = [20,20,20,20,20]
0152
0153     elif varname  ==  "dca":
0154         y_hi     = 1
0155         yname    = "DCA$_r$ [mm]"
0156         xname    = "pT [GeV]"
0157         x_hi     = [2,5,10,5,2]
0158     else:
0159         print("ERROR(plot_resol):please use a valid varname: th, ph, dp, dca")
0160         return -1
0161
0162     fig,axs=plt.subplots(2,3,figsize=(12,9), dpi=100,sharex=0)
0163     axs = axs.flat
0164     ii=0
0165     for ii, e_lo in enumerate(eta_lo_pwg):
0166         e_hi = eta_hi_pwg[ii]
0167         ax   = axs[ii]
0168         if ii>1:
0169             ax   = axs[ii+1]
0170
0171         # select rows within eta range,
0172         c1 = df.eta_lo>=e_lo
0173         c2 = df.eta_hi<=e_hi
0174         dft = df[c1&c2]
0175         if len(dft)==0:
0176             continue
0177         # take a quick average of all eta bins for plot: e.g. eta -1 to 0 are the average of -1,-0.5 and -0.5, 0
0178         dft = dft[['name','mom',"sig_"+varname,"err_"+varname]]
0179         dft = dft.groupby(['name','mom']).mean().reset_index()
0180         cond = (dft.name==setting1) & (dft["sig_"+varname]>0)
0181
0182         if varname=="dca": ## p to pT
0183             xdata = dft.mom[cond]*np.sin(theta2eta((e_lo+e_hi)/2,1))
0184         else:
0185             xdata = dft.mom[cond]
0186         ax.errorbar(xdata,dft["sig_"+varname][cond],yerr=dft["err_"+varname][cond],color="b",ls="none",marker="o")#,label=f"{e_lo} to {eta_hi_pwg[ii]}")
0187
0188         if len(setting2)>0:
0189             cond = (dft.name==setting2) & (dft["sig_"+varname]>0)
0190             if varname=="dca": ## p to pT
0191                 xdata = dft.mom[cond]*np.sin(theta2eta((e_lo+e_hi)/2,1))
0192             else:
0193                 xdata = dft.mom[cond]
0194             ax.errorbar(xdata,dft["sig_"+varname][cond],yerr=dft["err_"+varname][cond],color="r",ls="none",marker="x")#,label=f"{e_lo} to {eta_hi_pwg[ii]}")
0195
0196
0197         ## PWG curve
0198         y_pwg=pwg_value(varname, e_lo, xline)
0199         if varname=="dca" or varname=="dp":
0200             ax.plot(xline, y_pwg, 'k--',zorder=10)
0201         ax.set_ylim(-y_hi*0.05,y_hi)
0202         ax.set_xlim(0,x_hi[ii]*1.05)
0203         ax.text(x_hi[ii]*0.1,y_hi*0.9, f"{e_lo}<$\\eta$<{e_hi}",fontsize=14)
0204
0205     ## create legend on an empty sub-panel
0206     ax = axs[2]
0207     ax.axis('off')
0208     ax.plot(xline, y_pwg-100000, "k--",label="PWG Requirements")
0209     ax.errorbar(xline,y_pwg-100000,ls="none",marker="o",color="blue",label=setting1)#,label=f"{e_lo} to {eta_hi_pwg[ii]}")
0210     if len(setting2)>0:
0211         ax.errorbar(xline,y_pwg-100000,ls="none",marker="x",color="r",label=setting2)#,label=f"{e_lo} to {eta_hi_pwg[ii]}")
0212     ax.set_ylim(0,1)
0213     ax.legend(frameon=0,loc="upper left",fontsize=16)
0214
0215     axs[3].set_xlabel(xname)
0216     axs[3].set_ylabel(yname)
0217
0218     plt.subplots_adjust(left=0.1, right=0.9, top=0.99, bottom=0.3)
0219
0220     plt.savefig(f"{dir_path}tracking_single_resol_{varname}_eta.pdf")
0221
0222
0223 if __name__ == "__main__":
0224     plot_resol(varname="ph",resol_file="resol_out_whole.txt", dir_path="./", setting1="official", setting2="")
0225     plot_eff(eff_file="eff_out.txt", dir_path="./", setting="default")