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