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     coeff, var_matrix = curve_fit(fnc, list(bc[slc]), list(y[slc]), p0=p0,
0058                                  sigma=list(np.sqrt(y[slc])+(y[slc]==0)), maxfev=10000)
0059     print(coeff)
0060     xx=np.linspace(0,.7, 100)
0061     MIP=coeff[1]/1000
0062     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')
0063     plt.xlabel("hit energy [MeV]")
0064     plt.ylabel("hits")
0065     plt.title(f"$E_{{\\mu^-}}=${p} GeV")
0066     plt.legend(fontsize=20)
0067     plt.savefig(outdir+"/MIP.pdf")
0068
0069     plt.figure(figsize=(10,7))
0070     array=arrays_sim[p]
0071     bins=30; r=((-np.pi, np.pi),(2.8, 4.2))
0072     selection=np.sum(array["HcalEndcapPInsertHits.energy"]>0.5*MIP,axis=-1)>0
0073     h1, xedges, yedges = np.histogram2d(list(array[selection]['phi_truth']),list(array[selection]['eta_truth']), bins=bins, range=r)
0074     h2, xedges, yedges = np.histogram2d(list(array['phi_truth']),list(array['eta_truth']), bins=bins, range=r)
0075
0076     h = h1 / h2
0077     pc=plt.pcolor(xedges, yedges, h.T,linewidth=0)
0078     plt.xlabel("$\\phi^*$ [rad]")
0079     plt.ylabel("$\\eta^*$")
0080     cb = plt.colorbar(pc)
0081     cb.set_label("acceptance")
0082     plt.title(f"$E_{{\\mu^-}}=${p} GeV")
0083     plt.savefig(outdir+"/acceptance.pdf")