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     filename=f'sim_output/zdc_sigma/{config}_rec_sigma_dec_{p}GeV.edm4eic.root'
0023     print("opening file", filename)
0024     events = ur.open(filename+':events')
0025     arrays_sim[p] = events.arrays()[:-1] #remove last event, which for some reason is blank
0026     import gc
0027     gc.collect()
0028     print("read", filename)
0029
0030 def gauss(x, A,mu, sigma):
0031     return A * np.exp(-(x-mu)**2/(2*sigma**2))
0032
0033 #keep track of the number of clusters per event
0034 nclusters={}
0035
0036 for p in momenta:
0037     plt.figure()
0038     nclusters[p]=[]
0039     for i in range(len(arrays_sim[p])):
0040         nclusters[p].append(len(arrays_sim[p]["HcalFarForwardZDCClusters.position.x"][i]))
0041     nclusters[p]=np.array(nclusters[p])
0042     plt.hist(nclusters[p],bins=20, range=(0,20))
0043     plt.xlabel("number of clusters")
0044     plt.yscale('log')
0045     plt.title(f"$p_\Sigma={p}$ GeV")
0046     plt.ylim(1)
0047     plt.savefig(outdir+f"nclust_{p}GeV_recon.pdf")
0048     print("saved file ", outdir+f"nclust_{p}GeV_recon.pdf")
0049
0050
0051
0052 pt_truth={}
0053 theta_truth={}
0054
0055 for p in momenta:
0056     #get the truth value of theta* and pt*
0057     px=arrays_sim[p]["MCParticles.momentum.x"][:,2]
0058     py=arrays_sim[p]["MCParticles.momentum.y"][:,2]
0059     pz=arrays_sim[p]["MCParticles.momentum.z"][:,2]
0060     tilt=-0.025
0061     pt_truth[p]=np.hypot(px*np.cos(tilt)-pz*np.sin(tilt), py)
0062     theta_truth[p]=np.arctan2(pt_truth[p],pz*np.cos(tilt)+px*np.sin(tilt))
0063
0064 #create an array with the same shape as the cluster-level arrays
0065 is_neutron_cand={}
0066 for p in momenta:
0067     is_neutron_cand[p]=(0*arrays_sim[p][f"HcalFarForwardZDCClusters.energy"]).to_list()
0068
0069     #largest_eigenvalues
0070     for i in range(len(arrays_sim[p])):
0071         pars=arrays_sim[p]["_HcalFarForwardZDCClusters_shapeParameters"][i]
0072         index_of_max=-1
0073         max_val=0
0074         eigs=[]
0075         #Must make sure this doesn't get messed up if someone changes the number of shape parameters in EICrecon.
0076         nClust=nclusters[p][i]
0077         nShapePars=len(pars)//nClust
0078         for j in range(nClust):
0079             largest_eigenvalue=max(pars[nShapePars*j+4:nShapePars*j+7])
0080             eigs.append(largest_eigenvalue)
0081             if(largest_eigenvalue>max_val):
0082                 max_val=largest_eigenvalue
0083                 index_of_max=j
0084         if index_of_max >=0:
0085             is_neutron_cand[p][i][index_of_max]=1
0086         eigs.sort()
0087
0088     is_neutron_cand[p]=ak.Array(is_neutron_cand[p])
0089
0090 import ROOT
0091
0092 lambda_mass=1.115683
0093 pi0_mass=0.1349768
0094 pt_recon_corr={}
0095 theta_recon_corr={}
0096 mass_recon_corr={}
0097 mass_lambda_recon_corr={}
0098 mass_pi0_recon_corr={}
0099 pi0_converged={}
0100 zvtx_recon={}
0101
0102 #run this event-by-event:
0103 maxZ=35800
0104 for p in momenta:
0105     pt_recon_corr[p]=[]
0106     theta_recon_corr[p]=[]
0107     mass_recon_corr[p]=[]
0108     mass_lambda_recon_corr[p]=[]
0109     mass_pi0_recon_corr[p]=[]
0110     zvtx_recon[p]=[]
0111     for evt in range(len(arrays_sim[p])):
0112         if nclusters[p][evt]!=4:
0113             nan=-1
0114             pt_recon_corr[p].append(nan)
0115             theta_recon_corr[p].append(nan)
0116             mass_recon_corr[p].append(nan)
0117             mass_lambda_recon_corr[p].append(nan)
0118             mass_pi0_recon_corr[p].append(nan)
0119             zvtx_recon[p].append(nan)
0120             continue
0121         xc=arrays_sim[p][f"HcalFarForwardZDCClusters.position.x"][evt]
0122         yc=arrays_sim[p][f"HcalFarForwardZDCClusters.position.y"][evt]
0123         zc=arrays_sim[p][f"HcalFarForwardZDCClusters.position.z"][evt]
0124         E=arrays_sim[p][f"HcalFarForwardZDCClusters.energy"][evt]
0125
0126         #apply correction to the neutron candidates only
0127         A,B,C=-0.0756, -1.91,  2.30
0128         neutron_corr=(1-is_neutron_cand[p][evt])+is_neutron_cand[p][evt]/(1+A+B/np.sqrt(E)+C/E)
0129         E=E*neutron_corr
0130
0131         pabs=np.sqrt(E**2-is_neutron_cand[p][evt]*.9406**2)
0132         tilt=-0.025
0133         xcp=xc*np.cos(tilt)-zc*np.sin(tilt)
0134         ycp=yc
0135         zcp=zc*np.cos(tilt)+xc*np.sin(tilt)
0136
0137         #search for the combination of photons that would give the best lambda mass
0138         pt_best=-999
0139         theta_best=-999
0140         mass_lambda_best=-999
0141         mass_sigma_best=-999
0142         mass_pi0_best=-999
0143         zvtx_best=-999
0144         for hypothesis in range(4):
0145             if is_neutron_cand[p][evt][hypothesis]:
0146                 continue
0147
0148             xvtx=0
0149             yvtx=0
0150             zvtx=0
0151             #find the vertex position that reconstructs the pi0 mass
0152             for iteration in range(20):
0153                 tot=ROOT.TLorentzVector(0,0,0,0)
0154                 Lambda=ROOT.TLorentzVector(0,0,0,0)
0155                 pi0=ROOT.TLorentzVector(0,0,0,0)
0156
0157                 for i in range(4):
0158
0159                     if i!=hypothesis:
0160                         ux=xcp[i]-xvtx
0161                         uy=ycp[i]-yvtx
0162                         uz=zcp[i]-zvtx
0163                     else:
0164                         ux=xcp[i]
0165                         uy=ycp[i]
0166                         uz=zcp[i]
0167                     u=np.sqrt(ux**2+uy**2+uz**2)
0168                     ux/=u
0169                     uy/=u
0170                     uz/=u
0171
0172                     P=ROOT.TLorentzVector(pabs[i]*ux, pabs[i]*uy, pabs[i]*uz, E[i])
0173                     tot+=P
0174                     if not is_neutron_cand[p][evt][i] and i!=hypothesis:
0175                         pi0+=P
0176                     if i!=hypothesis:
0177                         Lambda+=P
0178                 alpha=1
0179                 if iteration==0:
0180                     zeta=1/2
0181                     zvtx=maxZ*np.power(zeta,alpha)
0182                     xvtx=Lambda.X()/Lambda.Z()*zvtx
0183                     yvtx=Lambda.Y()/Lambda.Z()*zvtx
0184                 else :
0185                     s=2*(pi0.M()<pi0_mass)-1
0186                     zeta=np.power(zvtx/maxZ, 1/alpha)
0187                     zeta=zeta+s*1/2**(1+iteration)
0188                     zvtx=maxZ*np.power(zeta,alpha)
0189                     xvtx=Lambda.X()/Lambda.Z()*zvtx
0190                     yvtx=Lambda.Y()/Lambda.Z()*zvtx
0191
0192             if abs(Lambda.M()-lambda_mass)< abs(mass_lambda_best-lambda_mass):
0193                 pt_best=tot.Pt()
0194                 theta_best=tot.Theta()
0195                 mass_lambda_best=Lambda.M()
0196                 mass_sigma_best=tot.M()
0197                 mass_pi0_best=pi0.M()
0198                 zvtx_best=zvtx
0199
0200         pt_recon_corr[p].append(pt_best)
0201         theta_recon_corr[p].append(theta_best)
0202         mass_recon_corr[p].append(mass_sigma_best)
0203         mass_lambda_recon_corr[p].append(mass_lambda_best)
0204         mass_pi0_recon_corr[p].append(mass_pi0_best)
0205         zvtx_recon[p].append(zvtx_best)
0206     pt_recon_corr[p]=ak.Array(pt_recon_corr[p])
0207     theta_recon_corr[p]=ak.Array(theta_recon_corr[p])
0208     mass_recon_corr[p]=ak.Array(mass_recon_corr[p])
0209     mass_lambda_recon_corr[p]=ak.Array(mass_lambda_recon_corr[p])
0210     mass_pi0_recon_corr[p]=ak.Array(mass_pi0_recon_corr[p])
0211     zvtx_recon[p]=ak.Array(zvtx_recon[p])
0212
0213 #now make plots
0214
0215 #reconstructed vertex position plot
0216 fig,axs=plt.subplots(1,3, figsize=(24, 8))
0217 plt.sca(axs[0])
0218 plt.title(f"$E_{{\\Sigma}}=100-275$ GeV")
0219 x=[]
0220 y=[]
0221 for p in momenta:
0222     accept=(nclusters[p]==4)# &(pi0_converged[p])
0223     x+=list(theta_truth[p][accept]*1000)
0224     y+=list(theta_recon_corr[p][accept]*1000)
0225 #print(x)
0226 plt.scatter(x,y)
0227 plt.xlabel("$\\theta^{*\\rm truth}_{\\Sigma}$ [mrad]")
0228 plt.ylabel("$\\theta^{*\\rm recon}_{\\Sigma}$ [mrad]")
0229 plt.xlim(0,3.2)
0230 plt.ylim(0,3.2)
0231
0232 plt.sca(axs[1])
0233 plt.title(f"$E_{{\\Sigma}}=100-275$ GeV")
0234 y,x,_=plt.hist(y-np.array(x), bins=50, range=(-1,1))
0235 bc=(x[1:]+x[:-1])/2
0236
0237 from scipy.optimize import curve_fit
0238 slc=abs(bc)<0.6
0239 fnc=gauss
0240 p0=[100, 0, 0.5]
0241 coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0242                                  sigma=np.sqrt(y[slc])+(y[slc]==0))
0243 x=np.linspace(-1, 1)
0244 plt.plot(x, gauss(x, *coeff), color='tab:orange')
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))
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 coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0313                                  sigma=np.sqrt(y[slc])+(y[slc]==0))
0314 x=np.linspace(-5, 5)
0315 plt.plot(x, gauss(x, *coeff), color='tab:orange')
0316 plt.xlabel("$z^{*\\rm recon}_{\\rm vtx}-z^{*\\rm truth}_{\\rm vtx}$ [m]")
0317 plt.ylabel("events")
0318
0319 plt.sca(axs[2])
0320 sigmas=[]
0321 dsigmas=[]
0322 xvals=[]
0323 for p in momenta:
0324
0325     accept=(nclusters[p]==4)&(abs(mass_pi0_recon_corr[p]-pi0_mass)<.01)
0326     a=list((zvtx_recon[p]-arrays_sim[p]['MCParticles.vertex.z'][:,5])[accept]/1000)
0327     y,x=np.histogram(a, bins=100, range=(-10,10))
0328     bc=(x[1:]+x[:-1])/2
0329
0330     from scipy.optimize import curve_fit
0331     slc=abs(bc)<5
0332     fnc=gauss
0333     p0=(100, 0, 1)
0334     #print(bc[slc],y[slc])
0335     sigma=np.sqrt(y[slc])+(y[slc]==0)
0336     try:
0337         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,sigma=list(sigma))
0338         sigmas.append(abs(coeff[2]))
0339         dsigmas.append(np.sqrt(var_matrix[2][2]))
0340         xvals.append(p)
0341     except:
0342         print(f"fit failed for p={p}")
0343 plt.ylim(0, 3)
0344
0345 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0346 x=np.linspace(100, 275, 100)
0347
0348 avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
0349 plt.axhline(avg, color='tab:orange')
0350 plt.text(150, 1.25,f"$\\sigma\\approx${avg:.1f} m")
0351
0352 plt.xlabel("$E_{\\Sigma}$ [GeV]")
0353 plt.ylabel("$\\sigma[z_{\\rm vtx}]$ [m]")
0354 plt.tight_layout()
0355 plt.savefig(outdir+"zvtx_recon.pdf")
0356
0357 #lambda mass reconstruction
0358 fig,axs=plt.subplots(1,2, figsize=(16, 8))
0359 plt.sca(axs[0])
0360 lambda_mass=1.115683
0361 vals=[]
0362 for p in momenta:
0363     accept=(nclusters[p]==4)&(abs(mass_pi0_recon_corr[p]-pi0_mass)<.01)
0364     vals+=list(mass_lambda_recon_corr[p][accept])
0365
0366 y,x,_= plt.hist(vals, bins=100, range=(0.9, 1.3))
0367 bc=(x[1:]+x[:-1])/2
0368 plt.axvline(lambda_mass, ls='--', color='tab:green', lw=3)
0369 plt.text(lambda_mass+.01, np.max(y)*1.05, "PDG mass", color='tab:green')
0370 plt.xlabel("$m_{\\Lambda}^{\\rm recon}$ [GeV]")
0371 plt.ylim(0, np.max(y)*1.2)
0372 plt.xlim(0.9, 1.3)
0373
0374 from scipy.optimize import curve_fit
0375 slc=abs(bc-lambda_mass)<0.05
0376 fnc=gauss
0377 p0=[100, lambda_mass, 0.03]
0378 coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0379                                  sigma=np.sqrt(y[slc])+(y[slc]==0))
0380 x=np.linspace(0.8, 1.3, 200)
0381 plt.plot(x, gauss(x, *coeff), color='tab:orange')
0382 print(coeff[2], np.sqrt(var_matrix[2][2]))
0383 plt.xlabel("$m^{\\rm recon}_{\\Lambda}$ [GeV]")
0384 plt.ylabel("events")
0385 plt.title(f"$E_{{\\Sigma}}=100-275$ GeV")
0386
0387 plt.sca(axs[1])
0388 xvals=[]
0389 sigmas=[]
0390 dsigmas=[]
0391 for p in momenta:
0392     accept=(nclusters[p]==4)&(abs(mass_pi0_recon_corr[p]-pi0_mass)<.01)
0393     y,x= np.histogram(mass_lambda_recon_corr[p][accept], bins=100, range=(0.6,1.4))
0394     bc=(x[1:]+x[:-1])/2
0395
0396     from scipy.optimize import curve_fit
0397     slc=abs(bc-lambda_mass)<0.05
0398     fnc=gauss
0399     p0=[100, lambda_mass, 0.03]
0400     try:
0401         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0402                                        sigma=list(np.sqrt(y[slc])+(y[slc]==0)))
0403         x=np.linspace(0.8, 1.3, 200)
0404         sigmas.append(coeff[2])
0405         dsigmas.append(np.sqrt(var_matrix[2][2]))
0406         xvals.append(p)
0407     except:
0408         print("fit failed")
0409
0410 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0411 avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
0412 plt.axhline(avg, color='tab:orange')
0413 plt.text(150, 0.01,f"$\\sigma\\approx${avg*1000:.0f} MeV")
0414 plt.xlabel("$E_{\\Sigma}$ [GeV]")
0415 plt.ylabel("$\\sigma[m_{\\Lambda}]$ [GeV]")
0416 plt.ylim(0, 0.02)
0417 plt.tight_layout()
0418 plt.savefig(outdir+"lambda_mass_rec_from_sigma_decay.pdf")
0419
0420 #sigma mass reconstruction
0421 p=100
0422 fig,axs=plt.subplots(1,2, figsize=(16, 8))
0423 plt.sca(axs[0])
0424 sigma_mass=1.192
0425 vals=[]
0426 for p in momenta:
0427     accept=(nclusters[p]==4)&(abs(mass_pi0_recon_corr[p]-pi0_mass)<.01)
0428     vals+=list(mass_recon_corr[p][accept])
0429
0430 y,x,_= plt.hist(vals, bins=100, range=(1.0, 1.4))
0431 bc=(x[1:]+x[:-1])/2
0432 plt.axvline(sigma_mass, ls='--', color='tab:green', lw=3)
0433 plt.text(sigma_mass+.01, np.max(y)*1.05, "PDG mass", color='tab:green')
0434 plt.xlabel("$m_{\\Sigma}^{\\rm recon}$ [GeV]")
0435 plt.ylim(0, np.max(y)*1.2)
0436 plt.xlim(1.0, 1.45)
0437
0438 from scipy.optimize import curve_fit
0439 slc=abs(bc-sigma_mass)<0.02
0440 fnc=gauss
0441 p0=[100, sigma_mass, 0.03]
0442 coeff, var_matrix = curve_fit(fnc, bc[slc], y[slc], p0=p0,
0443                                  sigma=np.sqrt(y[slc])+(y[slc]==0))
0444 x=np.linspace(0.8, 1.3, 200)
0445 plt.plot(x, gauss(x, *coeff), color='tab:orange')
0446 print(coeff[2], np.sqrt(var_matrix[2][2]))
0447 plt.xlabel("$m^{\\rm recon}_{\\Sigma}$ [GeV]")
0448 plt.ylabel("events")
0449 plt.title(f"$E_{{\\Sigma}}=100-275$ GeV")
0450
0451 plt.sca(axs[1])
0452 xvals=[]
0453 sigmas=[]
0454 dsigmas=[]
0455 for p in momenta:
0456     accept=(nclusters[p]==4)&(abs(mass_pi0_recon_corr[p]-pi0_mass)<.01)
0457     y,x= np.histogram(mass_recon_corr[p][accept], bins=100, range=(1.0,1.4))
0458     bc=(x[1:]+x[:-1])/2
0459
0460     from scipy.optimize import curve_fit
0461     slc=abs(bc-sigma_mass)<0.02
0462     fnc=gauss
0463     p0=[100, sigma_mass, 0.03]
0464     try:
0465         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0466                                        sigma=list(np.sqrt(y[slc])+(y[slc]==0)))
0467         sigmas.append(abs(coeff[2]))
0468         dsigmas.append(np.sqrt(var_matrix[2][2]))
0469         xvals.append(p)
0470     except:
0471         print("fit failed")
0472
0473 plt.errorbar(xvals, sigmas, dsigmas, ls='', marker='o', color='k')
0474 avg=np.sum(sigmas/np.array(dsigmas)**2)/np.sum(1/np.array(dsigmas)**2)
0475 plt.axhline(avg, color='tab:orange')
0476 plt.text(150, 0.01,f"$\\sigma\\approx${avg*1000:.0f} MeV")
0477 plt.xlabel("$E_{\\Sigma}$ [GeV]")
0478 plt.ylabel("$\\sigma[m_{\\Sigma}]$ [GeV]")
0479 plt.ylim(0, 0.1)
0480 plt.tight_layout()
0481 plt.savefig(outdir+"sigma_mass_rec.pdf")