zdc_sigma/analysis/sigma_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_sigma/{config}_rec_sigma_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(rf"$p_\Sigma={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 #create an array with the same shape as the cluster-level arrays
0062 is_neutron_cand={}
0063 for p in momenta:
0064     is_neutron_cand[p]=(0*arrays_sim[p][f"HcalFarForwardZDCClusters.energy"]).to_list()
0065
0066     #largest_eigenvalues
0067     for i in range(len(arrays_sim[p])):
0068         pars=arrays_sim[p]["_HcalFarForwardZDCClusters_shapeParameters"][i]
0069         index_of_max=-1
0070         max_val=0
0071         eigs=[]
0072         #Must make sure this doesn't get messed up if someone changes the number of shape parameters in EICrecon.
0073         nClust=nclusters[p][i]
0074         nShapePars=len(pars)//nClust
0075         for j in range(nClust):
0076             largest_eigenvalue=max(pars[nShapePars*j+4:nShapePars*j+7])
0077             eigs.append(largest_eigenvalue)
0078             if(largest_eigenvalue>max_val):
0079                 max_val=largest_eigenvalue
0080                 index_of_max=j
0081         if index_of_max >=0:
0082             is_neutron_cand[p][i][index_of_max]=1
0083         eigs.sort()
0084
0085     is_neutron_cand[p]=ak.Array(is_neutron_cand[p])
0086
0087 import ROOT
0088
0089 lambda_mass=1.115683
0090 pi0_mass=0.1349768
0091 pt_recon_corr={}
0092 theta_recon_corr={}
0093 mass_recon_corr={}
0094 mass_lambda_recon_corr={}
0095 mass_pi0_recon_corr={}
0096 pi0_converged={}
0097 zvtx_recon={}
0098
0099 #run this event-by-event:
0100 maxZ=35800
0101 for p in momenta:
0102     pt_recon_corr[p]=[]
0103     theta_recon_corr[p]=[]
0104     mass_recon_corr[p]=[]
0105     mass_lambda_recon_corr[p]=[]
0106     mass_pi0_recon_corr[p]=[]
0107     zvtx_recon[p]=[]
0108     for evt in range(len(arrays_sim[p])):
0109         if nclusters[p][evt]!=4:
0110             nan=-1
0111             pt_recon_corr[p].append(nan)
0112             theta_recon_corr[p].append(nan)
0113             mass_recon_corr[p].append(nan)
0114             mass_lambda_recon_corr[p].append(nan)
0115             mass_pi0_recon_corr[p].append(nan)
0116             zvtx_recon[p].append(nan)
0117             continue
0118         xc=arrays_sim[p][f"HcalFarForwardZDCClusters.position.x"][evt]
0119         yc=arrays_sim[p][f"HcalFarForwardZDCClusters.position.y"][evt]
0120         zc=arrays_sim[p][f"HcalFarForwardZDCClusters.position.z"][evt]
0121         E=arrays_sim[p][f"HcalFarForwardZDCClusters.energy"][evt]
0122
0123         #apply correction to the neutron candidates only
0124         A,B,C=-0.0756, -1.91,  2.30
0125         neutron_corr=(1-is_neutron_cand[p][evt])+is_neutron_cand[p][evt]/(1+A+B/np.sqrt(E)+C/E)
0126         E=E*neutron_corr
0127
0128         pabs=np.sqrt(E**2-is_neutron_cand[p][evt]*.9406**2)
0129         tilt=-0.025
0130         xcp=xc*np.cos(tilt)-zc*np.sin(tilt)
0131         ycp=yc
0132         zcp=zc*np.cos(tilt)+xc*np.sin(tilt)
0133
0134         #search for the combination of photons that would give the best lambda mass
0135         pt_best=-999
0136         theta_best=-999
0137         mass_lambda_best=-999
0138         mass_sigma_best=-999
0139         mass_pi0_best=-999
0140         zvtx_best=-999
0141         for hypothesis in range(4):
0142             if is_neutron_cand[p][evt][hypothesis]:
0143                 continue
0144
0145             xvtx=0
0146             yvtx=0
0147             zvtx=0
0148             #find the vertex position that reconstructs the pi0 mass
0149             for iteration in range(20):
0150                 tot=ROOT.TLorentzVector(0,0,0,0)
0151                 Lambda=ROOT.TLorentzVector(0,0,0,0)
0152                 pi0=ROOT.TLorentzVector(0,0,0,0)
0153
0154                 for i in range(4):
0155
0156                     if i!=hypothesis:
0157                         ux=xcp[i]-xvtx
0158                         uy=ycp[i]-yvtx
0159                         uz=zcp[i]-zvtx
0160                     else:
0161                         ux=xcp[i]
0162                         uy=ycp[i]
0163                         uz=zcp[i]
0164                     u=np.sqrt(ux**2+uy**2+uz**2)
0165                     ux/=u
0166                     uy/=u
0167                     uz/=u
0168
0169                     P=ROOT.TLorentzVector(pabs[i]*ux, pabs[i]*uy, pabs[i]*uz, E[i])
0170                     tot+=P
0171                     if not is_neutron_cand[p][evt][i] and i!=hypothesis:
0172                         pi0+=P
0173                     if i!=hypothesis:
0174                         Lambda+=P
0175                 alpha=1
0176                 if iteration==0:
0177                     zeta=1/2
0178                     zvtx=maxZ*np.power(zeta,alpha)
0179                     xvtx=Lambda.X()/Lambda.Z()*zvtx
0180                     yvtx=Lambda.Y()/Lambda.Z()*zvtx
0181                 else :
0182                     s=2*(pi0.M()<pi0_mass)-1
0183                     zeta=np.power(zvtx/maxZ, 1/alpha)
0184                     zeta=zeta+s*1/2**(1+iteration)
0185                     zvtx=maxZ*np.power(zeta,alpha)
0186                     xvtx=Lambda.X()/Lambda.Z()*zvtx
0187                     yvtx=Lambda.Y()/Lambda.Z()*zvtx
0188
0189             if abs(Lambda.M()-lambda_mass)< abs(mass_lambda_best-lambda_mass):
0190                 pt_best=tot.Pt()
0191                 theta_best=tot.Theta()
0192                 mass_lambda_best=Lambda.M()
0193                 mass_sigma_best=tot.M()
0194                 mass_pi0_best=pi0.M()
0195                 zvtx_best=zvtx
0196
0197         pt_recon_corr[p].append(pt_best)
0198         theta_recon_corr[p].append(theta_best)
0199         mass_recon_corr[p].append(mass_sigma_best)
0200         mass_lambda_recon_corr[p].append(mass_lambda_best)
0201         mass_pi0_recon_corr[p].append(mass_pi0_best)
0202         zvtx_recon[p].append(zvtx_best)
0203     pt_recon_corr[p]=ak.Array(pt_recon_corr[p])
0204     theta_recon_corr[p]=ak.Array(theta_recon_corr[p])
0205     mass_recon_corr[p]=ak.Array(mass_recon_corr[p])
0206     mass_lambda_recon_corr[p]=ak.Array(mass_lambda_recon_corr[p])
0207     mass_pi0_recon_corr[p]=ak.Array(mass_pi0_recon_corr[p])
0208     zvtx_recon[p]=ak.Array(zvtx_recon[p])
0209
0210 #now make plots
0211
0212 #reconstructed vertex position plot
0213 fig,axs=plt.subplots(1,3, figsize=(24, 8))
0214 plt.sca(axs[0])
0215 plt.title(f"$E_{{\\Sigma}}=100-275$ GeV")
0216 x=[]
0217 y=[]
0218 for p in momenta:
0219     accept=(nclusters[p]==4)# &(pi0_converged[p])
0220     x+=list(theta_truth[p][accept]*1000)
0221     y+=list(theta_recon_corr[p][accept]*1000)
0222 #print(x)
0223 plt.scatter(x,y)
0224 plt.xlabel("$\\theta^{*\\rm truth}_{\\Sigma}$ [mrad]")
0225 plt.ylabel("$\\theta^{*\\rm recon}_{\\Sigma}$ [mrad]")
0226 plt.xlim(0,3.2)
0227 plt.ylim(0,3.2)
0228
0229 plt.sca(axs[1])
0230 plt.title(f"$E_{{\\Sigma}}=100-275$ GeV")
0231 y,x,_=plt.hist(y-np.array(x), bins=50, range=(-1,1))
0232 bc=(x[1:]+x[:-1])/2
0233
0234 from scipy.optimize import curve_fit
0235 slc=abs(bc)<0.6
0236 fnc=gauss
0237 p0=[100, 0, 0.5]
0238 try:
0239     coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0240                                      sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
0241     x=np.linspace(-1, 1)
0242     plt.plot(x, gauss(x, *coeff), color='tab:orange')
0243 except RuntimeError:
0244     print("fit failed")
0245 plt.xlabel("$\\theta^{*\\rm recon}_{\\Sigma}-\\theta^{*\\rm truth}_{\\Sigma}$ [mrad]")
0246 plt.ylabel("events")
0247
0248 plt.sca(axs[2])
0249 sigmas=[]
0250 dsigmas=[]
0251 xvals=[]
0252 for p in momenta:
0253
0254     accept=(nclusters[p]==4)
0255     y,x=np.histogram((theta_recon_corr[p]-theta_truth[p])[accept]*1000, bins=100, range=(-0.5,0.5))
0256     bc=(x[1:]+x[:-1])/2
0257
0258     from scipy.optimize import curve_fit
0259     slc=abs(bc)<0.3
0260     fnc=gauss
0261     p0=(100, 0, 0.15)
0262     #print(bc[slc],y[slc])
0263     sigma=np.sqrt(y[slc])+(y[slc]==0)
0264     try:
0265         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
0266         sigmas.append(np.abs(coeff[2]))
0267         dsigmas.append(np.sqrt(var_matrix[2][2]))
0268         xvals.append(p)
0269     except:
0270         print(f"fit failed for p={p}")
0271 print(xvals)
0272 print(sigmas)
0273 plt.ylim(0, 0.3)
0274
0275 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0276 x=np.linspace(100, 275, 100)
0277 plt.plot(x, 3/np.sqrt(x), color='tab:orange')
0278 plt.text(170, .23, "YR requirement:\n   3 mrad/$\\sqrt{E}$")
0279 plt.xlabel("$E_{\\Sigma}$ [GeV]")
0280 plt.ylabel("$\\sigma[\\theta^*_{\\Sigma}]$ [mrad]")
0281 plt.tight_layout()
0282 plt.savefig(outdir+"thetastar_recon.pdf")
0283
0284 #reconstructed vertex position plot
0285 fig,axs=plt.subplots(1,3, figsize=(24, 8))
0286 plt.sca(axs[0])
0287 plt.title(f"$E_{{\\Sigma}}=100-275$ GeV")
0288 x=[]
0289 y=[]
0290 for p in momenta:
0291     accept=(nclusters[p]==4)&(abs(mass_pi0_recon_corr[p]-pi0_mass)<.01)
0292     tilt=-0.025
0293     x+=list(arrays_sim[p]['MCParticles.vertex.z'][:,5][accept]*np.cos(tilt)/1000
0294             +np.sin(tilt)*arrays_sim[p]['MCParticles.vertex.z'][:,5][accept]/1000)
0295     y+=list(zvtx_recon[p][accept]/1000)
0296 plt.scatter(x,y)
0297 #print(x,y)
0298 plt.xlabel("$z^{\\rm truth}_{\\rm vtx}$ [m]")
0299 plt.ylabel("$z^{\\rm recon}_{\\rm vtx}$  [m]")
0300 plt.xlim(0,40)
0301 plt.ylim(0,40)
0302
0303 plt.sca(axs[1])
0304 plt.title(f"$E_{{\\Sigma}}=100-275$ GeV")
0305 y,x,_=plt.hist(y-np.array(x), bins=50, range=(-10,10))
0306 bc=(x[1:]+x[:-1])/2
0307
0308 from scipy.optimize import curve_fit
0309 slc=abs(bc)<5
0310 fnc=gauss
0311 p0=[100, 0, 1]
0312 try:
0313     coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0314                                      sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
0315     x=np.linspace(-5, 5)
0316     plt.plot(x, gauss(x, *coeff), color='tab:orange')
0317 except RuntimeError:
0318     print("fit failed")
0319 plt.xlabel("$z^{*\\rm recon}_{\\rm vtx}-z^{*\\rm truth}_{\\rm vtx}$ [m]")
0320 plt.ylabel("events")
0321
0322 plt.sca(axs[2])
0323 sigmas=[]
0324 dsigmas=[]
0325 xvals=[]
0326 for p in momenta:
0327
0328     accept=(nclusters[p]==4)&(abs(mass_pi0_recon_corr[p]-pi0_mass)<.01)
0329     a=list((zvtx_recon[p]-arrays_sim[p]['MCParticles.vertex.z'][:,5])[accept]/1000)
0330     y,x=np.histogram(a, bins=100, range=(-10,10))
0331     bc=(x[1:]+x[:-1])/2
0332
0333     from scipy.optimize import curve_fit
0334     slc=abs(bc)<5
0335     fnc=gauss
0336     p0=(100, 0, 1)
0337     #print(bc[slc],y[slc])
0338     sigma=np.sqrt(y[slc])+(y[slc]==0)
0339     try:
0340         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0, sigma=list(sigma), maxfev=10000)
0341         sigmas.append(abs(coeff[2]))
0342         dsigmas.append(np.sqrt(var_matrix[2][2]))
0343         xvals.append(p)
0344     except:
0345         print(f"fit failed for p={p}")
0346 plt.ylim(0, 3)
0347
0348 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0349 x=np.linspace(100, 275, 100)
0350
0351 avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
0352 plt.axhline(avg, color='tab:orange')
0353 plt.text(150, 1.25,f"$\\sigma\\approx${avg:.1f} m")
0354
0355 plt.xlabel("$E_{\\Sigma}$ [GeV]")
0356 plt.ylabel("$\\sigma[z_{\\rm vtx}]$ [m]")
0357 plt.tight_layout()
0358 plt.savefig(outdir+"zvtx_recon.pdf")
0359
0360 #lambda mass reconstruction
0361 fig,axs=plt.subplots(1,2, figsize=(16, 8))
0362 plt.sca(axs[0])
0363 lambda_mass=1.115683
0364 vals=[]
0365 for p in momenta:
0366     accept=(nclusters[p]==4)&(abs(mass_pi0_recon_corr[p]-pi0_mass)<.01)
0367     vals+=list(mass_lambda_recon_corr[p][accept])
0368
0369 y,x,_= plt.hist(vals, bins=100, range=(0.9, 1.3))
0370 bc=(x[1:]+x[:-1])/2
0371 plt.axvline(lambda_mass, ls='--', color='tab:green', lw=3)
0372 plt.text(lambda_mass+.01, np.max(y)*1.05, "PDG mass", color='tab:green')
0373 plt.xlabel("$m_{\\Lambda}^{\\rm recon}$ [GeV]")
0374 plt.ylim(0, np.max(y)*1.2)
0375 plt.xlim(0.9, 1.3)
0376
0377 from scipy.optimize import curve_fit
0378 slc=abs(bc-lambda_mass)<0.05
0379 fnc=gauss
0380 p0=[100, lambda_mass, 0.03]
0381 coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0382                                  sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
0383 x=np.linspace(0.8, 1.3, 200)
0384 plt.plot(x, gauss(x, *coeff), color='tab:orange')
0385 print(coeff[2], np.sqrt(var_matrix[2][2]))
0386 plt.xlabel("$m^{\\rm recon}_{\\Lambda}$ [GeV]")
0387 plt.ylabel("events")
0388 plt.title(f"$E_{{\\Sigma}}=100-275$ GeV")
0389
0390 plt.sca(axs[1])
0391 xvals=[]
0392 sigmas=[]
0393 dsigmas=[]
0394 for p in momenta:
0395     accept=(nclusters[p]==4)&(abs(mass_pi0_recon_corr[p]-pi0_mass)<.01)
0396     y,x= np.histogram(mass_lambda_recon_corr[p][accept], bins=100, range=(0.6,1.4))
0397     bc=(x[1:]+x[:-1])/2
0398
0399     from scipy.optimize import curve_fit
0400     slc=abs(bc-lambda_mass)<0.05
0401     fnc=gauss
0402     p0=[100, lambda_mass, 0.03]
0403     try:
0404         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0405                                        sigma=list(np.sqrt(y[slc])+(y[slc]==0)), maxfev=10000)
0406         x=np.linspace(0.8, 1.3, 200)
0407         sigmas.append(coeff[2])
0408         dsigmas.append(np.sqrt(var_matrix[2][2]))
0409         xvals.append(p)
0410     except:
0411         print("fit failed")
0412
0413 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0414 avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
0415 plt.axhline(avg, color='tab:orange')
0416 plt.text(150, 0.01,f"$\\sigma\\approx${avg*1000:.0f} MeV")
0417 plt.xlabel("$E_{\\Sigma}$ [GeV]")
0418 plt.ylabel("$\\sigma[m_{\\Lambda}]$ [GeV]")
0419 plt.ylim(0, 0.02)
0420 plt.tight_layout()
0421 plt.savefig(outdir+"lambda_mass_rec_from_sigma_decay.pdf")
0422
0423 #sigma mass reconstruction
0424 p=100
0425 fig,axs=plt.subplots(1,2, figsize=(16, 8))
0426 plt.sca(axs[0])
0427 sigma_mass=1.192
0428 vals=[]
0429 for p in momenta:
0430     accept=(nclusters[p]==4)&(abs(mass_pi0_recon_corr[p]-pi0_mass)<.01)
0431     vals+=list(mass_recon_corr[p][accept])
0432
0433 y,x,_= plt.hist(vals, bins=100, range=(1.0, 1.4))
0434 bc=(x[1:]+x[:-1])/2
0435 plt.axvline(sigma_mass, ls='--', color='tab:green', lw=3)
0436 plt.text(sigma_mass+.01, np.max(y)*1.05, "PDG mass", color='tab:green')
0437 plt.xlabel("$m_{\\Sigma}^{\\rm recon}$ [GeV]")
0438 plt.ylim(0, np.max(y)*1.2)
0439 plt.xlim(1.0, 1.45)
0440
0441 from scipy.optimize import curve_fit
0442 slc=abs(bc-sigma_mass)<0.02
0443 fnc=gauss
0444 p0=[100, sigma_mass, 0.03]
0445 coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0446                                  sigma=np.sqrt(y[slc])+(y[slc]==0), maxfev=10000)
0447 x=np.linspace(0.8, 1.3, 200)
0448 plt.plot(x, gauss(x, *coeff), color='tab:orange')
0449 print(coeff[2], np.sqrt(var_matrix[2][2]))
0450 plt.xlabel("$m^{\\rm recon}_{\\Sigma}$ [GeV]")
0451 plt.ylabel("events")
0452 plt.title(f"$E_{{\\Sigma}}=100-275$ GeV")
0453
0454 plt.sca(axs[1])
0455 xvals=[]
0456 sigmas=[]
0457 dsigmas=[]
0458 for p in momenta:
0459     accept=(nclusters[p]==4)&(abs(mass_pi0_recon_corr[p]-pi0_mass)<.01)
0460     y,x= np.histogram(mass_recon_corr[p][accept], bins=100, range=(1.0,1.4))
0461     bc=(x[1:]+x[:-1])/2
0462
0463     from scipy.optimize import curve_fit
0464     slc=abs(bc-sigma_mass)<0.02
0465     fnc=gauss
0466     p0=[100, sigma_mass, 0.03]
0467     try:
0468         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0469                                        sigma=list(np.sqrt(y[slc])+(y[slc]==0)), maxfev=10000)
0470         sigmas.append(abs(coeff[2]))
0471         dsigmas.append(np.sqrt(var_matrix[2][2]))
0472         xvals.append(p)
0473     except:
0474         print("fit failed")
0475
0476 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0477 avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
0478 plt.axhline(avg, color='tab:orange')
0479 plt.text(150, 0.01,f"$\\sigma\\approx${avg*1000:.0f} MeV")
0480 plt.xlabel("$E_{\\Sigma}$ [GeV]")
0481 plt.ylabel("$\\sigma[m_{\\Sigma}]$ [GeV]")
0482 plt.ylim(0, 0.1)
0483 plt.tight_layout()
0484 plt.savefig(outdir+"sigma_mass_rec.pdf")