insert_muon/analysis/muon_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 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 def Landau(x, normalization,location,stdev):
0020     #print(type(x))
0021     u=(x-location)*3.591/stdev/2.355
0022     renormalization = 1.64872*normalization
0023     return renormalization * np.exp(-u/2 - np.exp(-u)/2)
0024
0025 import uproot as ur
0026 arrays_sim={}
0027 momenta=50,
0028 for p in momenta:
0029     arrays_sim[p] = ur.concatenate({
0030         f'sim_output/insert_muon/{config}_sim_mu-_{p}GeV_{index}.edm4hep.root': 'events'
0031         for index in range(5)
0032     })
0033
0034 for array in arrays_sim.values():
0035     tilt=-0.025
0036     px=array['MCParticles.momentum.x'][:,2]
0037     py=array['MCParticles.momentum.y'][:,2]
0038     pz=array['MCParticles.momentum.z'][:,2]
0039     p=np.sqrt(px**2+py**2+pz**2)
0040
0041     pxp=px*np.cos(tilt)-pz*np.sin(tilt)
0042     pyp=py
0043     pzp=pz*np.cos(tilt)+px*np.sin(tilt)
0044
0045     array['eta_truth']=1/2*np.log((p+pzp)/(p-pzp))
0046     array['phi_truth']=np.arctan2(pyp,pxp)
0047
0048 for p in 50,:
0049     E=arrays_sim[p]["HcalEndcapPInsertHits.energy"]
0050     y, x,_=plt.hist(1e3*ak.flatten(E),bins=100, range=(0, 1.2), histtype='step')
0051     bc=(x[1:]+x[:-1])/2
0052     from scipy.optimize import curve_fit
0053     slc=abs(bc-.48)<.15
0054     fnc=Landau
0055     p0=[100, .5, .05]
0056     #print(list(y), list(x))
0057     MIP=p0[1]/1000
0058     try:
0059         coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0060                                      sigma=list(np.sqrt(y[slc])+(y[slc]==0)), maxfev=10000)
0061         print(coeff)
0062         xx=np.linspace(0,.7, 100)
0063         MIP=coeff[1]/1000
0064         plt.plot(xx, fnc(xx,*coeff), label=f'Landau fit:\nMIP={coeff[1]*1e3:.0f}$\\pm${1e3*np.sqrt(var_matrix[1][1]):.0f} keV')
0065         plt.xlabel("hit energy [MeV]")
0066         plt.ylabel("hits")
0067         plt.title(f"$E_{{\\mu^-}}=${p} GeV")
0068         plt.legend(fontsize=20)
0069         plt.savefig(outdir+"/MIP.pdf")
0070     except RuntimeError:
0071         print("fit failed")
0072
0073     plt.figure(figsize=(10,7))
0074     array=arrays_sim[p]
0075     bins=30; r=((-np.pi, np.pi),(2.8, 4.2))
0076     selection=np.sum(array["HcalEndcapPInsertHits.energy"]>0.5*MIP,axis=-1)>0
0077     h1, xedges, yedges = np.histogram2d(list(array[selection]['phi_truth']),list(array[selection]['eta_truth']), bins=bins, range=r)
0078     h2, xedges, yedges = np.histogram2d(list(array['phi_truth']),list(array['eta_truth']), bins=bins, range=r)
0079
0080     h = h1 / h2
0081     pc=plt.pcolor(xedges, yedges, h.T,linewidth=0)
0082     plt.xlabel("$\\phi^*$ [rad]")
0083     plt.ylabel("$\\eta^*$")
0084     cb = plt.colorbar(pc)
0085     cb.set_label("acceptance")
0086     plt.title(f"$E_{{\\mu^-}}=${p} GeV")
0087     plt.savefig(outdir+"/acceptance.pdf")