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