tracking_detectors/scripts/matscan_plot.py

0001 ## plot test_matscan.cxx results
0002 ## Shujie Li, 08, 2021
0003 ## usage: python matscan_plot.py zrange rrange
0004 import pandas as pd
0005 import matplotlib.pyplot as plt
0006 import numpy as np
0007 import math as m
0008 import sys
0009 from matplotlib import colors
0010 from scipy import stats
0011 from matplotlib.colors import LogNorm
0012
0013
0014 plt.rcParams['figure.figsize'] = [10.0, 6.0]
0015 # plt.rcParams['fure.dpi'] = 80
0016
0017 SMALL_SIZE = 14
0018 MEDIUM_SIZE = 16
0019 BIGGER_SIZE = 18
0020
0021 plt.rc('font', size=SMALL_SIZE)          # controls default text sizes
0022 plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title
0023 plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
0024 plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
0025 plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
0026 plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize
0027 plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title]
0028
0029
0030 def cart2sph(x,y,z):
0031     x=np.array(x)
0032     y=np.array(y)
0033     z=np.array(z)
0034     XsqPlusYsq = x**2 + y**2
0035     r = np.sqrt(XsqPlusYsq + z**2)               # r
0036     elev = np.arctan2(np.sqrt(XsqPlusYsq),z)     # theta (polar to z axis)
0037     az = np.arctan2(y,x)                           # phi (azimuthal)
0038     return r, elev, az
0039
0040 ## read a single scan table for cylinder volume
0041 def read_table_cylind(zmax=140,rmax=60,indir='./'):
0042     header  = np.array(['x', 'y', 'z', 'ind', 'material', 'Z', 'density', 'rad_length', 'thickness', 'path_length', 'sum_x0',"end_x","end_y","end_z"])
0043     df      = pd.read_csv(indir+"all_z%g_r%g.dat" %(zmax,rmax),names=header,delim_whitespace=True)
0044     df["x0"]=np.array(df["thickness"])/np.array(df["rad_length"])
0045     df=df.reset_index(drop=True).drop_duplicates()
0046
0047     # calcualte polar angle then eta
0048     ax = df["x"]
0049     ay = df["y"]
0050     az = df["z"]
0051     r,polar,azi=cart2sph(ax,ay,az)
0052     eta = -np.log(np.tan(polar/2.0))
0053     df["eta"] = eta
0054     df["azi_angle"] = azi
0055     df["pol_angle"] = polar
0056
0057     print(df.head())
0058     print(df.material.unique())
0059
0060 #     df=df[np.isfinite(df).all(1)] #drop in
0061     return df
0062
0063 def add_kin(df):
0064         # calcualte polar angle then eta
0065     ax = df["x"]
0066     ay = df["y"]
0067     az = df["z"]
0068     r,polar,azi=cart2sph(ax,ay,az)
0069     eta = -np.log(np.tan(polar/2.0))
0070     df["eta"] = eta
0071     df["azi_angle"] = azi
0072     df["pol_angle"] = polar
0073     return df
0074
0075 ## group by xyz coordinate to get sum x0 per track with selected materials
0076 def get_x0(df0,mat_name=""):
0077     if len(mat_name)>0:
0078         df=df0[df0["material"]==mat_name]
0079     else:
0080         df=df0
0081     # df=df[df["material"]!="CarbonFiber_25percent"]
0082     dfg=df.groupby(["x","y","z"])["x0"].sum()
0083     dfg=dfg.reset_index()
0084     dfg=dfg[np.isfinite(dfg).all(1)] #drop inf
0085     dfg=add_kin(dfg)
0086     return dfg
0087
0088
0089 ## plot mean,min,max X0 of each eta bin
0090 def plot_1d(dfgz,mat_name="",ax="eta"):
0091     xlow = 0
0092     values=np.array(dfgz["x0"])*100
0093     xx = np.array(dfgz[ax])
0094     if ax=="eta":
0095         xhi=3.5
0096         xlow = -xhi
0097     elif ax=="pol_angle":
0098         xhi=90
0099         xx=180/2-xx*180/np.pi
0100     elif ax=="z": # z projected at barrel radius
0101         xhi=140
0102         xx*=(43.23/60)
0103     ## with bin stats
0104     nbin = 50
0105     x0mean,binedge,binnumber=stats.binned_statistic(xx,values, 'mean', bins=nbin,range=[xlow,xhi])
0106     ## !!! possible bugs in this error bar calculation
0107     x0max ,binedge,binnumber=stats.binned_statistic(xx,values, 'max',  bins=nbin,range=[xlow,xhi])
0108     x0min ,binedge,binnumber=stats.binned_statistic(xx,values, 'min',  bins=nbin,range=[xlow,xhi])
0109
0110     bin_center = (binedge[0:-1]+binedge[1:])/2
0111     plt.plot(bin_center,x0mean,label=mat_name)
0112     plt.fill_between(bin_center,x0min,x0max,alpha=0.2)
0113     plt.xlim(xlow,xhi)
0114     plt.grid()
0115     plt.suptitle("total X0")
0116     plt.xlabel(ax)
0117
0118     return bin_center, x0mean, x0max, x0min
0119
0120
0121
0122 if __name__ == "__main__":
0123
0124     ## corresponding to the scan output filenmae from test_matscan.cxx, in cm
0125     zrange = int(sys.argv[1])
0126     rrange = int(sys.argv[2])
0127     outdir = './'
0128     indir  = './'
0129     cols = np.array(["x","y","z","material","thickness","path_length"])
0130     df   = read_table_cylind(zrange,rrange,indir)
0131
0132
0133     ## -----------------plot side view of geometry, z v.s. R-------------
0134     plt.figure()
0135     xe = df["end_x"]
0136     ye = df["end_y"]
0137     ze = df["end_z"]
0138     re = np.sqrt(xe**2+ye**2)
0139     # c5=(df["end_x"]**2+df["end_y"]**2)<11**2
0140     # c6=df["material"]=="Aluminum"
0141     plt.scatter(ze,re,s=0.001,marker=".")
0142
0143     # plt.xlabel("$\eta$")
0144     plt.xlabel("z [cm]")
0145     plt.ylabel("R [cm]")
0146     # plt.xlim(0,4)
0147     plt.grid()
0148     plt.savefig(outdir+"/matscan_geo_z_z%g_r%g.png" %(zrange,rrange))
0149
0150     ## -----------------plot side view of geometry, eta v.s. R-------------
0151     plt.figure()
0152     xe = df["end_x"]
0153     ye = df["end_y"]
0154     ze = df["end_z"]
0155     re = np.sqrt(xe**2+ye**2)
0156     plt.scatter(df["eta"],re,s=0.001,marker=".")
0157
0158     plt.xlabel("$\eta$")
0159     plt.ylabel("R [cm]")
0160     # plt.xlim(0,4)
0161     plt.grid()
0162     plt.savefig(outdir+"/matscan_geo_eta_z%g_r%g.png" %(zrange,rrange))
0163
0164
0165     ## -----------------plot 1 d matscan z v.s. X0-------------
0166     df_all=get_x0(df)
0167     plt.figure()
0168     ax="eta"
0169     print(df.columns,df_all.columns)
0170     print(df_all['eta'])
0171     _=plot_1d(df_all,"Total",ax)
0172
0173     mat_list=np.delete(df.material.unique(),0)
0174     # mat_list.drop("Vacuum")
0175     # mat_list=['NONE','Air','Beryllium','Aluminum','Silicon','CarbonFiber']
0176     # mat_list=["Aluminum","Silicon","CarbonFiber"]
0177     for mat_name in mat_list:
0178         df_mat = get_x0(df,mat_name)
0179         _=plot_1d(df_mat,mat_name,ax)
0180
0181     plt.legend(loc="upper right")
0182     plt.xlabel("$\eta$")
0183     plt.ylabel("X/X0 (%)")
0184     plt.savefig(outdir+"/matscan_1d_z%g_r%g.png" %(zrange,rrange))
0185
0186     ## -----------------plot 2 d matscan z v.s. phi-------------
0187
0188     dfgz=df_all
0189     values=dfgz["x0"]
0190     xx = dfgz["eta"]
0191     yy = dfgz["azi_angle"] # in rad
0192     ## with bin stats
0193     nbin = 50
0194     x0mean_2d,x_edge,y_edge,binnumber_2d=stats.binned_statistic_2d(xx,yy, values, 'mean', bins=[nbin,nbin],range=[[-5,5],[-5,5]])
0195     # x0max_2d ,_,_,_=stats.binned_statistic_2d(xx,yy,values, 'max',  bins=[nbin,nbin],range=[[-5,5],[-5,5]])
0196     # x0min_2d ,_,_,_=stats.binned_statistic_2d(xx,yy,values, 'min',  bins=[nbin,nbin],range=[[-5,5],[-5,5]])
0197
0198     x_center = (x_edge[0:-1]+x_edge[1:])/2
0199     y_center = (y_edge[0:-1]+y_edge[1:])/2
0200     # plt.contour(x0mean_2d.transpose(),extent=[x_edge[0],x_edge[-1],y_edge[0],y_edge[-1]],
0201     #     linewidths=3, cmap = plt.cm.rainbow)
0202     # plt.colorbar()
0203
0204
0205     fig=plt.figure(figsize=(7, 6))
0206     c  = plt.pcolormesh(x_edge,y_edge/np.pi*180,x0mean_2d.transpose(),norm=LogNorm(vmin=0.001, vmax=10))
0207     fig.colorbar(c)
0208     # plt.colorbar()
0209     # plt.ylim(-np.pi,np.pi)
0210     plt.ylim(-180,180)
0211     plt.xlim(-5,5)
0212     plt.xlabel("$\eta$")
0213     plt.ylabel("Azimuthal angle [degree]")
0214     plt.suptitle("total X0 [%]")
0215
0216     plt.savefig(outdir+"/matscan_2d_z%g_r%g.png" %(zrange,rrange))
0217
0218