femc_electron/analysis/femc_electron_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_electron/{config}_rec_e-_{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
0064 #number of hits per cluster
0065 fig, axs=plt.subplots(1,2, figsize=(16,8))
0066 avgs=[]
0067 stds=[]
0068 pvals=[]
0069
0070 for p in arrays_sim:
0071
0072     a=arrays_sim[p]
0073     n=[]
0074     nn=-a['EcalEndcapPClusters.hits_begin']+a['EcalEndcapPClusters.hits_end']
0075     E=a['EcalEndcapPClusters.energy']
0076     for evt in range(len(array)):
0077         if len(E[evt])==0:
0078             continue
0079         maxE=np.max(E[evt])
0080         found=False
0081         for i in range(len(E[evt])):
0082             if E[evt][i]==maxE:
0083                 n.append(nn[evt][i])
0084                 found=True
0085                 break
0086         #if not found:
0087         #    n.append(0)
0088
0089     if p ==50:
0090         plt.sca(axs[0])
0091         y,x,_=plt.hist(n, range=(0,100), bins=100, histtype='step', label=f"E={p} GeV")
0092         plt.ylabel("events")
0093         plt.xlabel("# hits in cluster")
0094         plt.title(f"e-, E={p} GeV")
0095     pvals.append(p)
0096     avgs.append(np.mean(n))
0097     stds.append(np.std(n))
0098
0099 plt.sca(axs[1])
0100 plt.errorbar(pvals, avgs, stds, marker='o',ls='')
0101 plt.xlabel("E [GeV]")
0102 plt.ylabel("# hits in cluster [mean$\\pm$std]")
0103 plt.ylim(0)
0104 plt.savefig(outdir+"/nhits_per_cluster.pdf")
0105
0106
0107 #energy resolution
0108 def gauss(x, A,mu, sigma):
0109     return A * np.exp(-(x-mu)**2/(2*sigma**2))
0110 from scipy.optimize import curve_fit
0111
0112 fig, axs=plt.subplots(1,3, figsize=(24,8))
0113 pvals=[]
0114 res=[]
0115 dres=[]
0116 scale=[]
0117 dscale=[]
0118 for p in arrays_sim:
0119     bins=np.linspace(15*p/20,22*p/20, 50)
0120     if p==50:
0121         plt.sca(axs[0])
0122         plt.title(f"E={p} GeV")
0123         y,x,_=plt.hist(ak.flatten(arrays_sim[p]['EcalEndcapPClusters.energy']), bins=bins, histtype='step')
0124         plt.ylabel("events")
0125         plt.xlabel("$E^{rec}_e$ [GeV]")
0126     else:
0127         y,x=np.histogram(ak.flatten(arrays_sim[p]['EcalEndcapPClusters.energy']), bins=bins)
0128     bcs=(x[1:]+x[:-1])/2
0129
0130     fnc=gauss
0131     slc=abs(bcs-p)<3
0132     sigma=np.sqrt(y[slc])+0.5*(y[slc]==0)
0133     p0=(100, p, 3)
0134
0135     coeff, var_matrix = curve_fit(fnc, list(bcs[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
0136     #res=np.abs(coeff[2]/coeff[1])
0137     if p==50:
0138         xx=np.linspace(15*p/20,22*p/20, 100)
0139
0140         plt.plot(xx, fnc(xx,*coeff), label=f"$\\sigma_E/E={abs(coeff[2])/coeff[1]*100:.1f}\%$")
0141         plt.axvline(p, color='g', ls='--', alpha=0.7)
0142         plt.legend()
0143         #plt.xlim(0,60)
0144     #plt.show()
0145     pvals.append(p)
0146     res.append(abs(coeff[2])/coeff[1])
0147     dres.append(np.sqrt(var_matrix[2][2])/coeff[1])
0148     scale.append(abs(coeff[1])/p)
0149     dscale.append(np.sqrt(var_matrix[1][1])/p)
0150 plt.sca(axs[1])
0151 plt.errorbar(pvals, 100*np.array(res), 100*np.array(dres), ls='', marker='o')
0152 fnc = lambda E, a, b: np.hypot(a,b/np.sqrt(E))
0153 p0=(.05, .12)
0154 coeff, var_matrix = curve_fit(fnc, pvals, res, p0=p0, sigma=dres, maxfev=10000)
0155 xx=np.linspace(7, 85, 100)
0156 plt.plot(xx, 100*fnc(xx,*coeff), label=f'fit:{100*coeff[0]:.1f}%$\\oplus\\frac{{{100*coeff[1]:.0f}\\%}}{{\\sqrt{{E}}}}$')
0157 plt.legend()
0158 plt.ylim(0)
0159 plt.ylabel("E resolution [%]")
0160 plt.xlabel("E truth [GeV]")
0161 plt.sca(axs[2])
0162
0163 plt.errorbar(pvals, 100*np.array(scale), 100*np.array(dscale), ls='', marker='o')
0164 plt.ylabel("energy scale [%]")
0165 plt.xlabel("E truth [GeV]")
0166 plt.axhline(100, color='0.5', alpha=0.5, ls='--')
0167 plt.ylim(0, 110)
0168 plt.tight_layout()
0169 plt.savefig(outdir+"/energy_res.pdf")
0170
0171 #energy res in eta bins
0172 fig, axs=plt.subplots(1,2, figsize=(16,8))
0173
0174 partitions=[1.5, 2.0, 2.5, 3.0]
0175 for eta_min, eta_max in zip(partitions[:-1], partitions[1:]):
0176     pvals=[]
0177     res=[]
0178     dres=[]
0179     scale=[]
0180     dscale=[]
0181     for p in arrays_sim:
0182         bins=np.linspace(15*p/20,22*p/20, 50)
0183         eta_truth=arrays_sim[p]['eta_truth']
0184         y,x=np.histogram(ak.flatten(arrays_sim[p][(eta_truth<eta_max)&(eta_truth>eta_min)]['EcalEndcapPClusters.energy']), bins=bins)
0185         bcs=(x[1:]+x[:-1])/2
0186
0187         fnc=gauss
0188         slc=abs(bcs-p)<3
0189         sigma=np.sqrt(y[slc])+0.5*(y[slc]==0)
0190         p0=(100, p, 3)
0191
0192         try:
0193             coeff, var_matrix = curve_fit(fnc, list(bcs[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
0194             if abs(coeff[1])>100 or np.sqrt(var_matrix[1][1])>100:
0195                 continue
0196             pvals.append(p)
0197             res.append(abs(coeff[2])/coeff[1])
0198             dres.append(np.sqrt(var_matrix[2][2])/coeff[1])
0199             scale.append(abs(coeff[1])/p)
0200             dscale.append(np.sqrt(var_matrix[1][1])/p)
0201         except:
0202             pass
0203     plt.sca(axs[0])
0204     plt.errorbar(pvals, 100*np.array(res), 100*np.array(dres), ls='', marker='o', label=f'${eta_min:.1f}<\\eta<{eta_max:.1f}$')
0205
0206     plt.sca(axs[1])
0207     plt.errorbar(pvals, 100*np.array(scale), 100*np.array(dscale), ls='', marker='o', label=f'${eta_min:.1f}<\\eta<{eta_max:.1f}$')
0208
0209 plt.sca(axs[0])
0210 plt.legend()
0211 plt.ylim(0)
0212 plt.ylabel("E resolution [%]")
0213 plt.xlabel("E truth [GeV]")
0214
0215 plt.sca(axs[1])
0216 plt.ylabel("energy scale [%]")
0217 plt.xlabel("E truth [GeV]")
0218 plt.axhline(100, color='0.5', alpha=0.5, ls='--')
0219 plt.ylim(0, 110)
0220
0221 plt.tight_layout()
0222 plt.savefig(outdir+"/energy_res_eta_partitions.pdf")