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