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