zdc_lambda/analysis/lambda_plots.py

0001 import numpy as np, pandas as pd, matplotlib.pyplot as plt, matplotlib as mpl, awkward as ak, sys
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 outdir=sys.argv[1]+"/"
0012 config=outdir.split("/")[1]
0013 try:
0014     import os
0015     os.mkdir(outdir[:-1])
0016 except:
0017     pass
0018 import uproot as ur
0019 arrays_sim={}
0020 momenta=100, 125, 150, 175,200,225,250,275
0021 for p in momenta:
0022     arrays_sim[p] = ur.concatenate({
0023         f'sim_output/zdc_lambda/{config}_rec_lambda_dec_{p}GeV_{index}.edm4eic.root': 'events'
0024         for index in range(5)
0025     })
0026
0027 def gauss(x, A,mu, sigma):
0028     return A * np.exp(-(x-mu)**2/(2*sigma**2))
0029
0030 #keep track of the number of clusters per event
0031 nclusters={}
0032
0033 for p in momenta:
0034     plt.figure()
0035     nclusters[p]=[]
0036     for i in range(len(arrays_sim[p])):
0037         nclusters[p].append(len(arrays_sim[p]["HcalFarForwardZDCClusters.position.x"][i]))
0038     nclusters[p]=np.array(nclusters[p])
0039     plt.hist(nclusters[p],bins=20, range=(0,20))
0040     plt.xlabel("number of clusters")
0041     plt.yscale('log')
0042     plt.title(f"$p_\Lambda={p}$ GeV")
0043     plt.ylim(1)
0044     plt.savefig(outdir+f"nclust_{p}GeV_recon.pdf")
0045     print("saved file ", outdir+f"nclust_{p}GeV_recon.pdf")
0046
0047
0048
0049 pt_truth={}
0050 theta_truth={}
0051
0052 for p in momenta:
0053     #get the truth value of theta* and pt*
0054     px=arrays_sim[p]["MCParticles.momentum.x"][:,2]
0055     py=arrays_sim[p]["MCParticles.momentum.y"][:,2]
0056     pz=arrays_sim[p]["MCParticles.momentum.z"][:,2]
0057     tilt=-0.025
0058     pt_truth[p]=np.hypot(px*np.cos(tilt)-pz*np.sin(tilt), py)
0059     theta_truth[p]=np.arctan2(pt_truth[p],pz*np.cos(tilt)+px*np.sin(tilt))
0060
0061
0062 #create an array with the same shape as the cluster-level arrays
0063 is_neutron_cand={}
0064 for p in momenta:
0065     is_neutron_cand[p]=(0*arrays_sim[p][f"HcalFarForwardZDCClusters.energy"]).to_list()
0066
0067     #largest_eigenvalues
0068     for i in range(len(arrays_sim[p])):
0069         pars=arrays_sim[p]["_HcalFarForwardZDCClusters_shapeParameters"][i]
0070         index_of_max=-1
0071         max_val=0
0072         eigs=[]
0073         shape_params_per_cluster=7
0074         for j in range(len(pars)//shape_params_per_cluster):
0075             largest_eigenvalue=max(pars[shape_params_per_cluster*j+4:shape_params_per_cluster*j+7])
0076             eigs.append(largest_eigenvalue)
0077             if(largest_eigenvalue>max_val):
0078                 max_val=largest_eigenvalue
0079                 index_of_max=j
0080         if index_of_max >=0:
0081             is_neutron_cand[p][i][index_of_max]=1
0082         eigs.sort()
0083
0084     is_neutron_cand[p]=ak.Array(is_neutron_cand[p])
0085
0086
0087 #with the position of the vertex determined by assuming the mass of the pi0
0088 #corrected pt* and theta* recon
0089 pt_recon_corr={}
0090 theta_recon_corr={}
0091 mass_recon_corr={}
0092 mass_pi0_recon_corr={}
0093 pi0_converged={}
0094 zvtx_recon={}
0095
0096 maxZ=35800
0097 for p in momenta:
0098     xvtx=0
0099     yvtx=0
0100     zvtx=0
0101
0102     for iteration in range(20):
0103
0104         #compute the value of theta* using the clusters in the ZDC
0105         xc=arrays_sim[p][f"HcalFarForwardZDCClusters.position.x"]
0106         yc=arrays_sim[p][f"HcalFarForwardZDCClusters.position.y"]
0107         zc=arrays_sim[p][f"HcalFarForwardZDCClusters.position.z"]
0108         E=arrays_sim[p][f"HcalFarForwardZDCClusters.energy"]
0109         #apply correction to the neutron candidates only
0110         A,B,C=-0.0756, -1.91,  2.30
0111         neutron_corr=(1-is_neutron_cand[p])+is_neutron_cand[p]/(1+A+B/np.sqrt(E)+C/E)
0112         E=E*neutron_corr
0113
0114         E_recon=np.sum(E, axis=-1)
0115         pabs=np.sqrt(E**2-is_neutron_cand[p]*.9406**2)
0116         tilt=-0.025
0117         xcp=xc*np.cos(tilt)-zc*np.sin(tilt)
0118         ycp=yc
0119         zcp=zc*np.cos(tilt)+xc*np.sin(tilt)
0120         rcp=np.sqrt(xcp**2+ycp**2+zcp**2)
0121
0122         ux=(xcp-xvtx)
0123         uy=(ycp-yvtx)
0124         uz=(zcp-zvtx)
0125
0126         norm=np.sqrt(ux**2+uy**2+uz**2)
0127         ux=ux/norm
0128         uy=uy/norm
0129         uz=uz/norm
0130
0131         px_recon,py_recon,pz_recon=np.sum(pabs*ux, axis=-1),np.sum(pabs*uy, axis=-1),np.sum(pabs*uz, axis=-1)
0132
0133         pt_recon_corr[p]=np.hypot(px_recon,py_recon)
0134         theta_recon_corr[p]=np.arctan2(pt_recon_corr[p], pz_recon)
0135
0136         mass_recon_corr[p]=np.sqrt((E_recon)**2\
0137                                 -(px_recon)**2\
0138                                 -(py_recon)**2\
0139                                 -(pz_recon)**2)
0140         mass_pi0_recon_corr[p]=np.sqrt(np.sum(pabs*(1-is_neutron_cand[p]), axis=-1)**2\
0141                                     -np.sum(pabs*ux*(1-is_neutron_cand[p]), axis=-1)**2\
0142                                     -np.sum(pabs*uy*(1-is_neutron_cand[p]), axis=-1)**2\
0143                                     -np.sum(pabs*uz*(1-is_neutron_cand[p]), axis=-1)**2)
0144         alpha=1
0145         if iteration==0:
0146             u=np.sqrt(px_recon**2+py_recon**2+pz_recon**2)
0147             ux=px_recon/u
0148             uy=py_recon/u
0149             uz=pz_recon/u
0150             zeta=1/2
0151             zvtx=maxZ*np.power(zeta,alpha)
0152             xvtx=ux/uz*zvtx
0153             yvtx=uy/uz*zvtx
0154         else :
0155             u=np.sqrt(px_recon**2+py_recon**2+pz_recon**2)
0156             ux=px_recon/u
0157             uy=py_recon/u
0158             uz=pz_recon/u
0159             s=2*(mass_pi0_recon_corr[p]<0.135)-1
0160             zeta=np.power(zvtx/maxZ, 1/alpha)
0161             zeta=zeta+s*1/2**(1+iteration)
0162             zvtx=maxZ*np.power(zeta,alpha)
0163             xvtx=ux/uz*zvtx
0164             yvtx=uy/uz*zvtx
0165         #print(zvtx)
0166     pi0_converged[p]=np.abs(mass_pi0_recon_corr[p]-0.135)<0.01
0167     zvtx_recon[p]=zvtx
0168
0169 fig,axs=plt.subplots(1,3, figsize=(24, 8))
0170 plt.sca(axs[0])
0171 plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
0172 x=[]
0173 y=[]
0174 for p in momenta:
0175     accept=(nclusters[p]==3) &(pi0_converged[p])
0176     x+=list(theta_truth[p][accept]*1000)
0177     y+=list(theta_recon_corr[p][accept]*1000)
0178 plt.scatter(x,y)
0179 plt.xlabel("$\\theta^{*\\rm truth}_{\\Lambda}$ [mrad]")
0180 plt.ylabel("$\\theta^{*\\rm recon}_{\\Lambda}$ [mrad]")
0181 plt.xlim(0,3.2)
0182 plt.ylim(0,3.2)
0183
0184 plt.sca(axs[1])
0185 plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
0186 y,x,_=plt.hist(y-np.array(x), bins=50, range=(-1,1))
0187 bc=(x[1:]+x[:-1])/2
0188
0189 from scipy.optimize import curve_fit
0190 slc=abs(bc)<0.3
0191 fnc=gauss
0192 p0=[100, 0, 0.05]
0193 coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0194                                  sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
0195 x=np.linspace(-1, 1)
0196 plt.plot(x, gauss(x, *coeff), color='tab:orange')
0197 plt.xlabel("$\\theta^{*\\rm recon}_{\\Lambda}-\\theta^{*\\rm truth}_{\\Lambda}$ [mrad]")
0198 plt.ylabel("events")
0199
0200 plt.sca(axs[2])
0201 sigmas=[]
0202 dsigmas=[]
0203 xvals=[]
0204 for p in momenta:
0205
0206     accept=(nclusters[p]==3) &(pi0_converged[p])
0207     y,x=np.histogram((theta_recon_corr[p]-theta_truth[p])[accept]*1000, bins=100, range=(-1,1))
0208     bc=(x[1:]+x[:-1])/2
0209
0210     from scipy.optimize import curve_fit
0211     slc=abs(bc)<0.3
0212     fnc=gauss
0213     p0=(100, 0, 0.06)
0214     #print(bc[slc],y[slc])
0215     sigma=np.sqrt(y[slc])+(y[slc]==0)
0216     try:
0217         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
0218         sigmas.append(coeff[2])
0219         dsigmas.append(np.sqrt(var_matrix[2][2]))
0220         xvals.append(p)
0221     except:
0222         print("fit failed")
0223 plt.ylim(0, 0.3)
0224
0225 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0226 x=np.linspace(100, 275, 100)
0227 plt.plot(x, 3/np.sqrt(x), color='tab:orange')
0228 plt.text(170, .23, "YR requirement:\n   3 mrad/$\\sqrt{E}$")
0229 plt.xlabel("$E_{\\Lambda}$ [GeV]")
0230 plt.ylabel("$\\sigma[\\theta^*_{\\Lambda}]$ [mrad]")
0231 plt.tight_layout()
0232 plt.savefig(outdir+"thetastar_recon.pdf")
0233 #plt.show()
0234
0235
0236 fig,axs=plt.subplots(1,3, figsize=(24, 8))
0237 plt.sca(axs[0])
0238 plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
0239 x=[]
0240 y=[]
0241 for p in momenta:
0242     accept=(nclusters[p]==3) &(pi0_converged[p])
0243     x+=list(arrays_sim[p]['MCParticles.vertex.z'][:,3][accept]/1000)
0244     y+=list(zvtx_recon[p][accept]/1000)
0245 plt.scatter(x,y)
0246 #print(x,y)
0247 plt.xlabel("$z^{\\rm truth}_{\\rm vtx}$ [m]")
0248 plt.ylabel("$z^{\\rm recon}_{\\rm vtx}$  [m]")
0249 plt.xlim(0,40)
0250 plt.ylim(0,40)
0251
0252 plt.sca(axs[1])
0253 plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
0254 y,x,_=plt.hist(y-np.array(x), bins=50, range=(-10,10))
0255 bc=(x[1:]+x[:-1])/2
0256
0257 from scipy.optimize import curve_fit
0258 slc=abs(bc)<5
0259 fnc=gauss
0260 p0=[100, 0, 1]
0261 coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0262                                  sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
0263 x=np.linspace(-5, 5)
0264 plt.plot(x, gauss(x, *coeff), color='tab:orange')
0265 print(coeff[2], np.sqrt(var_matrix[2][2]))
0266 plt.xlabel("$z^{*\\rm recon}_{\\rm vtx}-z^{*\\rm truth}_{\\rm vtx}$ [m]")
0267 plt.ylabel("events")
0268
0269 plt.sca(axs[2])
0270 sigmas=[]
0271 dsigmas=[]
0272 xvals=[]
0273 for p in momenta:
0274
0275     accept=(nclusters[p]==3) &(pi0_converged[p])
0276     a=list((zvtx_recon[p]-arrays_sim[p]['MCParticles.vertex.z'][:,3])[accept]/1000)
0277     y,x=np.histogram(a, bins=100, range=(-10,10))
0278     bc=(x[1:]+x[:-1])/2
0279
0280     from scipy.optimize import curve_fit
0281     slc=abs(bc)<5
0282     fnc=gauss
0283     p0=(100, 0, 1)
0284     #print(bc[slc],y[slc])
0285     sigma=np.sqrt(y[slc])+(y[slc]==0)
0286     try:
0287         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
0288         sigmas.append(coeff[2])
0289         dsigmas.append(np.sqrt(var_matrix[2][2]))
0290         xvals.append(p)
0291     except:
0292         print("fit failed")
0293 plt.ylim(0, 2)
0294
0295 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0296 x=np.linspace(100, 275, 100)
0297
0298 avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
0299 plt.axhline(avg, color='tab:orange')
0300 plt.text(150, 1.25,f"$\\sigma\\approx${avg:.1f} m")
0301
0302 plt.xlabel("$E_{\\Lambda}$ [GeV]")
0303 plt.ylabel("$\\sigma[z_{\\rm vtx}]$ [m]")
0304 plt.tight_layout()
0305 plt.savefig(outdir+"zvtx_recon.pdf")
0306 #plt.show()
0307
0308 p=100
0309 fig,axs=plt.subplots(1,2, figsize=(16, 8))
0310 plt.sca(axs[0])
0311 lambda_mass=1.115683
0312 vals=[]
0313 for p in momenta:
0314     accept=(nclusters[p]==3) &(pi0_converged[p])
0315     vals+=list(mass_recon_corr[p][accept])
0316
0317 y,x,_= plt.hist(vals, bins=100, range=(1.0, 1.25))
0318 bc=(x[1:]+x[:-1])/2
0319 plt.axvline(lambda_mass, ls='--', color='tab:green', lw=3)
0320 plt.text(lambda_mass+.01, np.max(y)*1.05, "PDG mass", color='tab:green')
0321 plt.xlabel("$m_{\\Lambda}^{\\rm recon}$ [GeV]")
0322 plt.ylim(0, np.max(y)*1.2)
0323 plt.xlim(1.0, 1.25)
0324
0325 from scipy.optimize import curve_fit
0326 slc=abs(bc-lambda_mass)<0.07
0327 fnc=gauss
0328 p0=[100, lambda_mass, 0.04]
0329 coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0330                                  sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
0331 x=np.linspace(0.8, 1.3, 200)
0332 plt.plot(x, gauss(x, *coeff), color='tab:orange')
0333 print(coeff[2], np.sqrt(var_matrix[2][2]))
0334 plt.xlabel("$m^{\\rm recon}_{\\Lambda}$ [GeV]")
0335 plt.ylabel("events")
0336 plt.title(f"$E_{{\\Lambda}}=100-275$ GeV")
0337
0338 plt.sca(axs[1])
0339 xvals=[]
0340 sigmas=[]
0341 dsigmas=[]
0342 for p in momenta:
0343     accept=(nclusters[p]==3) &(pi0_converged[p])
0344     y,x= np.histogram(mass_recon_corr[p][accept], bins=100, range=(0.6,1.4))
0345     bc=(x[1:]+x[:-1])/2
0346
0347     from scipy.optimize import curve_fit
0348     slc=abs(bc-lambda_mass)<0.07
0349     fnc=gauss
0350     p0=[100, lambda_mass, 0.05]
0351     try:
0352         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0353                                        sigma=list(np.sqrt(y[slc])+(y[slc]==0)), maxfev=10000)
0354         x=np.linspace(0.8, 1.3, 200)
0355         sigmas.append(coeff[2])
0356         dsigmas.append(np.sqrt(var_matrix[2][2]))
0357         xvals.append(p)
0358     except:
0359         print("fit failed")
0360
0361 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0362 avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
0363 plt.axhline(avg, color='tab:orange')
0364 plt.text(150, 0.01,f"$\\sigma\\approx${avg*1000:.0f} MeV")
0365 plt.xlabel("$E_{\\Lambda}$ [GeV]")
0366 plt.ylabel("$\\sigma[m_{\\Lambda}]$ [GeV]")
0367 plt.ylim(0, 0.02)
0368 plt.tight_layout()
0369 plt.savefig(outdir+"lambda_mass_rec.pdf")