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     try:
0136         coeff, var_matrix = curve_fit(fnc, list(bcs[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
0137     except RuntimeError as e:
0138         print(e)
0139         print("x:", list(bcs[slc]), "y:", list(y[slc]))
0140     #res=np.abs(coeff[2]/coeff[1])
0141     if p==50:
0142         xx=np.linspace(15*p/20,22*p/20, 100)
0143
0144         plt.plot(xx, fnc(xx,*coeff), label=f"$\\sigma_E/E={abs(coeff[2])/coeff[1]*100:.1f}\%$")
0145         plt.axvline(p, color='g', ls='--', alpha=0.7)
0146         plt.legend()
0147         #plt.xlim(0,60)
0148     #plt.show()
0149     pvals.append(p)
0150     res.append(abs(coeff[2])/coeff[1])
0151     dres.append(np.sqrt(var_matrix[2][2])/coeff[1])
0152     scale.append(abs(coeff[1])/p)
0153     dscale.append(np.sqrt(var_matrix[1][1])/p)
0154 plt.sca(axs[1])
0155 plt.errorbar(pvals, 100*np.array(res), 100*np.array(dres), ls='', marker='o')
0156 fnc = lambda E, a, b: np.hypot(a,b/np.sqrt(E))
0157 p0=(.05, .12)
0158 coeff, var_matrix = curve_fit(fnc, pvals, res, p0=p0, sigma=dres, maxfev=10000)
0159 xx=np.linspace(7, 85, 100)
0160 plt.plot(xx, 100*fnc(xx,*coeff), label=f'fit:{100*coeff[0]:.1f}%$\\oplus\\frac{{{100*coeff[1]:.0f}\\%}}{{\\sqrt{{E}}}}$')
0161 plt.legend()
0162 plt.ylim(0)
0163 plt.ylabel("E resolution [%]")
0164 plt.xlabel("E truth [GeV]")
0165 plt.sca(axs[2])
0166
0167 plt.errorbar(pvals, 100*np.array(scale), 100*np.array(dscale), ls='', marker='o')
0168 plt.ylabel("energy scale [%]")
0169 plt.xlabel("E truth [GeV]")
0170 plt.axhline(100, color='0.5', alpha=0.5, ls='--')
0171 plt.ylim(0, 110)
0172 plt.tight_layout()
0173 plt.savefig(outdir+"/energy_res.pdf")
0174
0175 #energy res in eta bins
0176 fig, axs=plt.subplots(1,2, figsize=(16,8))
0177
0178 partitions=[1.5, 2.0, 2.5, 3.0]
0179 for eta_min, eta_max in zip(partitions[:-1], partitions[1:]):
0180     pvals=[]
0181     res=[]
0182     dres=[]
0183     scale=[]
0184     dscale=[]
0185     for p in arrays_sim:
0186         bins=np.linspace(15*p/20,22*p/20, 50)
0187         eta_truth=arrays_sim[p]['eta_truth']
0188         y,x=np.histogram(ak.flatten(arrays_sim[p][(eta_truth<eta_max)&(eta_truth>eta_min)]['EcalEndcapPClusters.energy']), bins=bins)
0189         bcs=(x[1:]+x[:-1])/2
0190
0191         fnc=gauss
0192         slc=abs(bcs-p)<3
0193         sigma=np.sqrt(y[slc])+0.5*(y[slc]==0)
0194         p0=(100, p, 3)
0195
0196         try:
0197             coeff, var_matrix = curve_fit(fnc, list(bcs[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
0198             if abs(coeff[1])>100 or np.sqrt(var_matrix[1][1])>100:
0199                 continue
0200             pvals.append(p)
0201             res.append(abs(coeff[2])/coeff[1])
0202             dres.append(np.sqrt(var_matrix[2][2])/coeff[1])
0203             scale.append(abs(coeff[1])/p)
0204             dscale.append(np.sqrt(var_matrix[1][1])/p)
0205         except:
0206             pass
0207     plt.sca(axs[0])
0208     plt.errorbar(pvals, 100*np.array(res), 100*np.array(dres), ls='', marker='o', label=f'${eta_min:.1f}<\\eta<{eta_max:.1f}$')
0209
0210     plt.sca(axs[1])
0211     plt.errorbar(pvals, 100*np.array(scale), 100*np.array(dscale), ls='', marker='o', label=f'${eta_min:.1f}<\\eta<{eta_max:.1f}$')
0212
0213 plt.sca(axs[0])
0214 plt.legend()
0215 plt.ylim(0)
0216 plt.ylabel("E resolution [%]")
0217 plt.xlabel("E truth [GeV]")
0218
0219 plt.sca(axs[1])
0220 plt.ylabel("energy scale [%]")
0221 plt.xlabel("E truth [GeV]")
0222 plt.axhline(100, color='0.5', alpha=0.5, ls='--')
0223 plt.ylim(0, 110)
0224
0225 plt.tight_layout()
0226 plt.savefig(outdir+"/energy_res_eta_partitions.pdf")