femc_pi0/analysis/femc_pi0_plots.py

0001 import numpy as np, pandas as pd, matplotlib.pyplot as plt, matplotlib as mpl, awkward as ak, sys, uproot as ur
0002 import mplhep as hep
0003 hep.style.use("CMS")
0004
0005 plt.rcParams['figure.facecolor']='white'
0006 plt.rcParams['savefig.facecolor']='white'
0007 plt.rcParams['savefig.bbox']='tight'
0008
0009 plt.rcParams["figure.figsize"] = (7, 7)
0010
0011 config=sys.argv[1].split("/")[1]  #results/{config}/{benchmark_name}
0012 outdir=sys.argv[1]+"/"
0013 try:
0014     import os
0015     os.mkdir(outdir[:-1])
0016 except:
0017     pass
0018
0019 import uproot as ur
0020 arrays_sim={p:ur.open(f'sim_output/femc_pi0/{config}_rec_pi0_{p}GeV.edm4eic.root:events').arrays() for p in (10, 20, 30, 40, 50, 60,70,80)}
0021
0022 for p in arrays_sim:
0023     array=arrays_sim[p]
0024     tilt=-.025
0025     px=array['MCParticles.momentum.x'][:,2]
0026     py=array['MCParticles.momentum.y'][:,2]
0027     pz=array['MCParticles.momentum.z'][:,2]
0028     p=np.sqrt(px**2+py**2+pz**2)
0029
0030     pxp=px*np.cos(tilt)-pz*np.sin(tilt)
0031     pyp=py
0032     pzp=pz*np.cos(tilt)+px*np.sin(tilt)
0033
0034     array['eta_truth']=1/2*np.log((p+pzp)/(p-pzp))
0035     array['nclust_endcap']=[len(array['EcalEndcapPClusters.energy'][i]) for i in range(len(array))]
0036
0037 for array in arrays_sim.values():
0038     tilt=-0.025
0039     px=array['MCParticles.momentum.x'][:,2]
0040     py=array['MCParticles.momentum.y'][:,2]
0041     pz=array['MCParticles.momentum.z'][:,2]
0042     p=np.sqrt(px**2+py**2+pz**2)
0043
0044     pxp=px*np.cos(tilt)-pz*np.sin(tilt)
0045     pyp=py
0046     pzp=pz*np.cos(tilt)+px*np.sin(tilt)
0047
0048     array['eta_truth']=1/2*np.log((p+pzp)/(p-pzp))
0049     array['phi_truth']=np.arctan2(pyp,pxp)
0050
0051 #number of clusters
0052 plt.figure()
0053 for eta_min, eta_max, field in (1.5, 2.8, 'nclust_endcap'),:
0054     for p in arrays_sim:
0055         array=arrays_sim[p]
0056         plt.hist(array[field][(array['eta_truth']>eta_min)&(array['eta_truth']<eta_max)],
0057                  bins=np.linspace(0,10,11), histtype='step', label=f'{p} GeV', density=True)
0058     plt.ylabel("events")
0059     plt.xlabel("# of Ecal clusters")
0060     plt.legend()
0061     plt.savefig(outdir+f"/{field}.pdf")
0062
0063 #number of clusters
0064 plt.figure()
0065 for eta_min, eta_max, field in (1.5, 2.8, 'nclust_endcap'),:
0066     pvals=[]
0067     f0=[]
0068     f1=[]
0069     f2=[]
0070     f3=[]
0071     f4=[]
0072     f5p=[]
0073     for p in arrays_sim:
0074         array=arrays_sim[p]
0075         n=array[field][(array['eta_truth']>eta_min)&(array['eta_truth']<eta_max)]
0076         pvals.append(p)
0077         tot=len(n)
0078         f0.append(np.sum(1*(n==0))/tot)
0079         f1.append(np.sum(1*(n==1))/tot)
0080         f2.append(np.sum(1*(n==2))/tot)
0081         f3.append(np.sum(1*(n==3))/tot)
0082         f4.append(np.sum(1*(n==4))/tot)
0083         f5p.append(np.sum(1*(n>=5))/tot)
0084     plt.errorbar(pvals, f0, label="$n_{cl}$=0")
0085     plt.errorbar(pvals, f1, label="$n_{cl}$=1")
0086     plt.errorbar(pvals, f2, label="$n_{cl}$=2")
0087     plt.errorbar(pvals, f3, label="$n_{cl}$=3")
0088     plt.errorbar(pvals, f4, label="$n_{cl}$=4")
0089     plt.errorbar(pvals, f5p, label="$n_{cl}\\geq$5")
0090     plt.legend()
0091     plt.ylabel("fraction of events")
0092     plt.xlabel("$E_{\\pi^0}$ [GeV]")
0093     plt.ylim(0, 1)
0094     plt.legend(fontsize=20, ncol=2)
0095     plt.savefig(outdir+f"/nclust.pdf")
0096
0097
0098
0099 #number of hits per cluster
0100 fig, axs=plt.subplots(1,2, figsize=(16,8))
0101 avgs=[]
0102 stds=[]
0103 pvals=[]
0104
0105 for p in arrays_sim:
0106
0107     a=arrays_sim[p]
0108     n=[]
0109     nn=-a['EcalEndcapPClusters.hits_begin']+a['EcalEndcapPClusters.hits_end']
0110     E=a['EcalEndcapPClusters.energy']
0111     for evt in range(len(array)):
0112         if len(E[evt])==0:
0113             continue
0114         maxE=np.max(E[evt])
0115         found=False
0116         for i in range(len(E[evt])):
0117             if E[evt][i]==maxE:
0118                 n.append(nn[evt][i])
0119                 found=True
0120                 break
0121         #if not found:
0122         #    n.append(0)
0123
0124     if p ==50:
0125         plt.sca(axs[0])
0126         y,x,_=plt.hist(n, range=(0,100), bins=100, histtype='step', label=f"E={p} GeV")
0127         plt.ylabel("events")
0128         plt.xlabel("# hits in cluster")
0129         plt.title(f"$\\pi^0$, E={p} GeV")
0130     pvals.append(p)
0131     avgs.append(np.mean(n))
0132     stds.append(np.std(n))
0133
0134 plt.sca(axs[1])
0135 plt.errorbar(pvals, avgs, stds, marker='o',ls='')
0136 plt.xlabel("E [GeV]")
0137 plt.ylabel("# hits in cluster [mean$\\pm$std]")
0138 plt.ylim(0)
0139 plt.savefig(outdir+"/nhits_per_cluster.pdf")
0140
0141
0142 #energy resolution
0143 def gauss(x, A,mu, sigma):
0144     return A * np.exp(-(x-mu)**2/(2*sigma**2))
0145 from scipy.optimize import curve_fit
0146
0147 fig, axs=plt.subplots(1,3, figsize=(24,8))
0148 pvals=[]
0149 res=[]
0150 dres=[]
0151 scale=[]
0152 dscale=[]
0153 for p in arrays_sim:
0154     bins=np.linspace(15*p/20,22*p/20, 50)
0155     if p==50:
0156         plt.sca(axs[0])
0157         plt.title(f"E={p} GeV")
0158         y,x,_=plt.hist(np.sum(arrays_sim[p]['EcalEndcapPClusters.energy'], axis=-1), bins=bins, histtype='step')
0159         plt.ylabel("events")
0160         plt.xlabel("$E^{rec}_{\\pi^0}$ [GeV]")
0161     else:
0162         y,x=np.histogram(np.sum(arrays_sim[p]['EcalEndcapPClusters.energy'], axis=-1), bins=bins)
0163     bcs=(x[1:]+x[:-1])/2
0164
0165     fnc=gauss
0166     slc=abs(bcs-p)<3
0167     sigma=np.sqrt(y[slc])+0.5*(y[slc]==0)
0168     p0=(100, p, 3)
0169
0170     try:
0171         coeff, var_matrix = curve_fit(fnc, list(bcs[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
0172     except RuntimeError as e:
0173         print(e)
0174         print("x:", list(bcs[slc]), "y:", list(y[slc]))
0175     #res=np.abs(coeff[2]/coeff[1])
0176     if p==50:
0177         xx=np.linspace(15*p/20,22*p/20, 100)
0178
0179         plt.plot(xx, fnc(xx,*coeff), label=f"$\\sigma_E/E={abs(coeff[2])/coeff[1]*100:.1f}\%$")
0180         plt.axvline(p, color='g', ls='--', alpha=0.7)
0181         plt.legend()
0182         #plt.xlim(0,60)
0183     #plt.show()
0184     pvals.append(p)
0185     res.append(abs(coeff[2])/coeff[1])
0186     dres.append(np.sqrt(var_matrix[2][2])/coeff[1])
0187     scale.append(abs(coeff[1])/p)
0188     dscale.append(np.sqrt(var_matrix[1][1])/p)
0189 plt.sca(axs[1])
0190 plt.errorbar(pvals, 100*np.array(res), 100*np.array(dres), ls='', marker='o')
0191 fnc = lambda E, a, b: np.hypot(a,b/np.sqrt(E))
0192 p0=(.05, .12)
0193 coeff, var_matrix = curve_fit(fnc, pvals, res, p0=p0, sigma=dres, maxfev=10000)
0194 xx=np.linspace(15, 85, 100)
0195 plt.plot(xx, 100*fnc(xx,*coeff), label=f'fit:{100*coeff[0]:.1f}%$\\oplus\\frac{{{100*coeff[1]:.0f}\\%}}{{\\sqrt{{E}}}}$')
0196 plt.legend()
0197 plt.ylim(0)
0198 plt.ylabel("E resolution [%]")
0199 plt.xlabel("E truth [GeV]")
0200 plt.sca(axs[2])
0201
0202 plt.errorbar(pvals, 100*np.array(scale), 100*np.array(dscale), ls='', marker='o')
0203 plt.ylabel("energy scale [%]")
0204 plt.xlabel("E truth [GeV]")
0205 plt.axhline(100, color='0.5', alpha=0.5, ls='--')
0206 plt.ylim(0, 110)
0207 plt.tight_layout()
0208 plt.savefig(outdir+"/energy_res.pdf")
0209
0210 #energy res in eta bins
0211 fig, axs=plt.subplots(1,2, figsize=(16,8))
0212
0213 partitions=[1.5, 2.0, 2.5, 3.0]
0214 for eta_min, eta_max in zip(partitions[:-1], partitions[1:]):
0215     pvals=[]
0216     res=[]
0217     dres=[]
0218     scale=[]
0219     dscale=[]
0220     for p in arrays_sim:
0221         bins=np.linspace(15*p/20,22*p/20, 50)
0222         eta_truth=arrays_sim[p]['eta_truth']
0223         y,x=np.histogram(np.sum(arrays_sim[p][(eta_truth<eta_max)&(eta_truth>eta_min)]['EcalEndcapPClusters.energy'], axis=-1), bins=bins)
0224         bcs=(x[1:]+x[:-1])/2
0225
0226         fnc=gauss
0227         slc=abs(bcs-p)<3
0228         sigma=np.sqrt(y[slc])+0.5*(y[slc]==0)
0229         p0=(100, p, 3)
0230
0231         try:
0232             coeff, var_matrix = curve_fit(fnc, list(bcs[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
0233             if abs(coeff[1])>100 or np.sqrt(var_matrix[1][1])>100:
0234                 continue
0235             pvals.append(p)
0236             res.append(abs(coeff[2])/coeff[1])
0237             dres.append(np.sqrt(var_matrix[2][2])/coeff[1])
0238             scale.append(abs(coeff[1])/p)
0239             dscale.append(np.sqrt(var_matrix[1][1])/p)
0240         except:
0241             pass
0242     plt.sca(axs[0])
0243     plt.errorbar(pvals, 100*np.array(res), 100*np.array(dres), ls='', marker='o', label=f'${eta_min:.1f}<\\eta<{eta_max:.1f}$')
0244
0245     plt.sca(axs[1])
0246     plt.errorbar(pvals, 100*np.array(scale), 100*np.array(dscale), ls='', marker='o', label=f'${eta_min:.1f}<\\eta<{eta_max:.1f}$')
0247
0248 plt.sca(axs[0])
0249 plt.legend()
0250 plt.ylim(0)
0251 plt.ylabel("E resolution [%]")
0252 plt.xlabel("E truth [GeV]")
0253
0254 plt.sca(axs[1])
0255 plt.ylabel("energy scale [%]")
0256 plt.xlabel("E truth [GeV]")
0257 plt.axhline(100, color='0.5', alpha=0.5, ls='--')
0258 plt.ylim(0, 110)
0259
0260 plt.tight_layout()
0261 plt.savefig(outdir+"/energy_res_eta_partitions.pdf")