SinglePion/7etabins/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, -3.0, -2.5, -1.0, 1.0, 2.5, 3.0, 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=r"\s+",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_slices.txt", dir_path="./", setting1="", setting2=""):
0122     df= pd.read_csv(f'{dir_path}{resol_file}',header=None,sep=" ")
0123     df.columns = ["tag","deg_low", "deg_hi", "sig_dp", "err_dp", "sig_th", "err_th", "sig_ph", "err_ph", "sig_dca", "err_dca"]
0124
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'] = df['tag'].str.split('_').str[1]
0130
0131     print(df)
0132
0133     eta_lo_pwg = [-3.5,   -3.0,  -2.5,   -1.0,   1.0,    2.5  ,  3.0]
0134     eta_hi_pwg = [-3.0,   -2.5,  -1.0,    1.0,   2.5,    3.0  ,  3.5]
0135
0136     if varname  == "th":
0137         y_hi     = 0.01
0138         yname    = r"$\theta$ [rad]"
0139         xname    = "momentum [GeV]"
0140         x_hi      = [20,20,20,20,20,20,20]
0141
0142     elif varname  == "ph":
0143         y_hi     = 0.01
0144         yname    = r"$\phi$ [rad]"
0145         xname    = "momentum [GeV]"
0146         x_hi      = [20,20,20,20,20,20,20]
0147
0148     elif varname  == "dp":
0149         y_hi     = 12
0150         yname    = r"$\delta p/p$ [%]"
0151         xname    = "momentum [GeV/c]"
0152         x_hi      = [20,20,20,20,20,20,20]
0153
0154     elif varname  ==  "dca":
0155         y_hi     = 1
0156         yname    = "DCA$_r$ [mm]"
0157         xname    = "pT [GeV]"
0158         x_hi     = [1.5,2.5,5,10,5,2.5,1.5]
0159     else:
0160         print("ERROR(plot_resol):please use a valid varname: th, ph, dp, dca")
0161         return -1
0162
0163     fig,axs=plt.subplots(2,4,figsize=(16,8))
0164     axs = axs.flat
0165     ii=0
0166     for ii, e_lo in enumerate(eta_lo_pwg):
0167         e_hi = eta_hi_pwg[ii]
0168         ax   = axs[ii]
0169
0170         # select rows within eta range,
0171         c1 = df.eta_lo>=e_lo - 0.01
0172         c2 = df.eta_hi<=e_hi + 0.01
0173         dft = df[c1&c2]
0174
0175         print("---")
0176         print(dft)
0177
0178         if len(dft)==0:
0179             continue
0180         # 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
0181         dft = dft[['name','mom',"sig_"+varname,"err_"+varname]]
0182         dft = dft.groupby(['name','mom']).mean().reset_index()
0183         cond = (dft.name==setting1) & (dft["sig_"+varname]>0)
0184
0185         if varname=="dca": ## p to pT
0186             xdata = dft.mom[cond]*np.sin(theta2eta((e_lo+e_hi)/2,1))
0187         else:
0188             xdata = dft.mom[cond]
0189         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]}")
0190
0191         if len(setting2)>0:
0192             cond = (dft.name==setting2) & (dft["sig_"+varname]>0)
0193             if varname=="dca": ## p to pT
0194                 xdata = dft.mom[cond]*np.sin(theta2eta((e_lo+e_hi)/2,1))
0195             else:
0196                 xdata = dft.mom[cond]
0197             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]}")
0198
0199
0200         ## PWG curve
0201         xline      = np.arange(0.001,20,0.001)
0202         y_pwg=pwg_value(varname, e_lo, xline)
0203         if varname=="dca" or varname=="dp":
0204             ax.plot(xline, y_pwg, 'k--',zorder=10)
0205         ax.set_ylim(-y_hi*0.05,y_hi)
0206         ax.set_xlim(0,x_hi[ii]*1.05)
0207         ax.text(x_hi[ii]*0.1,y_hi*0.9, f"{e_lo}<$\\eta$<{e_hi}",fontsize=14)
0208
0209     ## create legend on an empty sub-panel
0210     ax = axs[7]
0211     ax.axis('off')
0212     ax.plot(xline, y_pwg-100000, "k--",label="PWG Requirements")
0213     ax.errorbar(xline,y_pwg-100000,ls="none",marker="o",color="blue",label="Simulation")
0214     #ax.errorbar(xline,y_pwg-100000,ls="none",marker="o",color="blue",label=setting1)#,label=f"{e_lo} to {eta_hi_pwg[ii]}")
0215     if len(setting2)>0:
0216         ax.errorbar(xline,y_pwg-100000,ls="none",marker="x",color="r",label=setting2)#,label=f"{e_lo} to {eta_hi_pwg[ii]}")
0217     ax.set_ylim(0,1)
0218     ax.legend(frameon=0,loc="upper left",fontsize=16)
0219
0220     axs[4].set_xlabel(xname)
0221     axs[5].set_xlabel(xname)
0222     axs[6].set_xlabel(xname)
0223
0224     axs[0].set_ylabel(yname)
0225     axs[4].set_ylabel(yname)
0226
0227     #plt.subplots_adjust(left=0.1, right=0.9, top=0.99, bottom=0.3)
0228     plt.tight_layout()
0229
0230     plt.savefig(f"{dir_path}tracking_single_resol_{varname}_eta.pdf")
0231
0232
0233 if __name__ == "__main__":
0234     plot_resol(varname="dca",resol_file="resol_out_slices.txt", dir_path="./", setting1="default", setting2="")
0235     plot_eff(eff_file="eff_out.txt", dir_path="./", setting="default")