Skip to content
Snippets Groups Projects
Commit cbfe2274 authored by Chao Peng's avatar Chao Peng
Browse files

update draw_cluster center finding

parent 8bf1ce09
No related branches found
No related tags found
1 merge request!81remove compact files from barrel emcal benchmarks
Hide whitespace changes
Inline Side-by-side
benchmarks/sampling_ecal/scripts/draw_cluster.py
+ 23
− 19
View file @ cbfe2274
......@@ -107,6 +107,8 @@ if __name__ == '__main__':
help='bin size for projection plot (mrad)')
parser.add_argument('--topo-range', type=float, default=50.0, dest='topo_range',
help='half range for projection plot (mrad)')
parser.add_argument('--zoom-factor', type=float, default=1.0, dest='zoom_factor',
help='factor for zoom-in')
args = parser.parse_args()
......@@ -115,7 +117,7 @@ if __name__ == '__main__':
'EcalBarrelAstroPix_RMin',
'EcalBarrelAstroPix_ReadoutLayerThickness',
'EcalBarrelAstroPix_ReadoutLayerNumber',
'EcalBarrelAstroPix_Length'
'EcalBarrelAstroPix_Length',
])
if not len(desc):
# or define Ecal shapes
......@@ -130,10 +132,10 @@ if __name__ == '__main__':
# read data
load_root_macros(args.macros)
df = get_hits_data(args.file, args.iev, branch=args.branch)
dfmcp = get_mcp_data(args.file, args.iev, 'mcparticles2')
vec = dfmcp.loc[dfmcp['status'] == 24578, ['px', 'py', 'pz']].iloc[0].values
dfmcp = get_mcp_simple(args.file, args.iev, 'mcparticles2')
vec = dfmcp[['px', 'py', 'pz']].values[0]
vec = vec/np.linalg.norm(vec)
df = df[df['cluster'] == args.icl]
# df = df[df['cluster'] == args.icl]
if not len(df):
print("Error: do not find any hits for cluster {:d} in event {:d}".format(args.icl, args.iev))
exit(-1)
......@@ -145,7 +147,7 @@ if __name__ == '__main__':
os.makedirs(args.outdir, exist_ok=True)
# cluster plot
cmap = truncate_colormap(plt.get_cmap('jet'), 0.1, 0.9)
fig = plt.figure(figsize=(20, 16), dpi=160)
fig = plt.figure(figsize=(15, 12), dpi=160)
ax = fig.add_subplot(111, projection='3d')
# draw particle line
ax.plot(*pline[[2, 1]], '--', zs=pline[0], color='green')
......@@ -161,17 +163,17 @@ if __name__ == '__main__':
# zoomed-in cluster plot
fig = plt.figure(figsize=(20, 16), dpi=160)
fig = plt.figure(figsize=(15, 12), dpi=160)
ax = fig.add_subplot(111, projection='3d')
# draw particle line
ax.plot(*pline[[2, 1]], '--', zs=pline[0], color='green')
# draw hits
draw_hits3d(ax, df[['x', 'y', 'z', 'edep']].values, cmap, s=20.0)
draw_track_fit(ax, df, stop_layer=args.stop,
scat_kw=dict(color='k', s=50.0), line_kw=dict(linestyle=':', color='k', lw=3))
# draw_track_fit(ax, df, stop_layer=args.stop,
# scat_kw=dict(color='k', s=50.0), line_kw=dict(linestyle=':', color='k', lw=3))
# view range
center = (length + thickness/2.)*vec
ranges = np.vstack([center - thickness/4., center + thickness/4.]).T
ranges = np.vstack([center - thickness/args.zoom_factor, center + thickness/args.zoom_factor]).T
ax.set_zlim(*ranges[0])
ax.set_ylim(*ranges[1])
ax.set_xlim(*ranges[2])
......@@ -185,27 +187,29 @@ if __name__ == '__main__':
df['r'] = np.sqrt(df['x'].values**2 + df['y'].values**2 + df['z'].values**2)
df['phi'] = np.arctan2(df['y'].values, df['x'].values)*1000.
df['theta'] = np.arccos(df['z'].values/df['r'].values)*1000.
df['eta'] = -np.log(np.tan(df['theta'].values/1000./2.))
# convert to mrad
vecp = np.asarray([np.arccos(vec[2]), np.arctan2(vec[1], vec[0])])*1000.
phi_rg = np.asarray([vecp[1] - args.topo_range, vecp[1] + args.topo_range])
th_rg = np.asarray([vecp[0] - args.topo_range, vecp[0] + args.topo_range])
eta_rg = np.resize(-np.log(np.tan(vecp[0]/1000./2.)), 2) + np.asarray([-args.topo_range, args.topo_range])/1000.
fig, axs = plt.subplots(1, 2, figsize=(17, 16), dpi=160, gridspec_kw={'wspace':0., 'width_ratios': [16, 1]})
ax, sm = draw_heatmap(axs[0], df['theta'].values, df['phi'].values, weights=df['edep'].values,
bins=(np.arange(*th_rg, step=args.topo_size), np.arange(*phi_rg, step=args.topo_size)),
fig, axs = plt.subplots(1, 2, figsize=(13, 12), dpi=160, gridspec_kw={'wspace':0., 'width_ratios': [12, 1]})
ax, sm = draw_heatmap(axs[0], df['eta'].values, df['phi'].values, weights=df['edep'].values,
bins=(np.arange(*eta_rg, step=args.topo_size/1000.), np.arange(*phi_rg, step=args.topo_size)),
cmap=cmap, cmin=0., pc_kw=dict(alpha=0.8, edgecolor='k'))
ax.set_xlabel(r'$\theta$ (mrad)', fontsize=28)
ax.set_ylabel(r'$\phi$ (mrad)', fontsize=28)
ax.tick_params(labelsize=24)
ax.set_ylabel(r'$\phi$ (mrad)', fontsize=32)
ax.set_xlabel(r'$\eta$', fontsize=32)
ax.tick_params(labelsize=28)
ax.xaxis.set_minor_locator(MultipleLocator(5))
ax.yaxis.set_minor_locator(MultipleLocator(5))
ax.grid(linestyle=':', which='both')
ax.set_axisbelow(True)
cb = plt.colorbar(sm, cax=axs[1], shrink=0.85)
cb.ax.tick_params(labelsize=24)
cb.ax.get_yaxis().labelpad = 24
cb.ax.set_ylabel('Energy Deposit (MeV)', rotation=90, fontsize=28)
cb = plt.colorbar(sm, cax=axs[1], shrink=0.85, aspect=1.2*20)
cb.ax.tick_params(labelsize=28)
cb.ax.get_yaxis().labelpad = 10
cb.ax.set_ylabel('Energy Deposit (MeV)', rotation=90, fontsize=32)
fig.savefig(os.path.join(args.outdir, 'e{}_topo.png'.format(args.iev)))
benchmarks/sampling_ecal/scripts/draw_cluster_layers.py
+ 22
− 11
View file @ cbfe2274
......@@ -71,9 +71,8 @@ if __name__ == '__main__':
'EcalBarrelAstroPix_RMin',
'EcalBarrelAstroPix_ReadoutLayerThickness',
'EcalBarrelAstroPix_ReadoutLayerNumber',
'EcalBarrelAstroPix_Length'
'EcalBarrelAstroPix_Length',
])
if not len(desc):
# or define Ecal shapes
rmin, thickness, length = 890, 20*(10. + 1.65), 860*2+500
......@@ -87,14 +86,25 @@ if __name__ == '__main__':
# read data
load_root_macros(args.macros)
df = get_hits_data(args.file, args.iev, args.branch)
dfmcp = get_mcp_data(args.file, args.iev, 'mcparticles2')
vec = dfmcp.loc[dfmcp['status'] == 24578, ['px', 'py', 'pz']].iloc[0].values
vec = vec/np.linalg.norm(vec)
df = df[df['cluster'] == args.icl]
# convert to polar coordinates (mrad), and stack all r values
df['r'] = np.sqrt(df['x'].values**2 + df['y'].values**2 + df['z'].values**2)
df['phi'] = np.arctan2(df['y'].values, df['x'].values)*1000.
df['theta'] = np.arccos(df['z'].values/df['r'].values)*1000.
df['eta'] = -np.log(np.tan(df['theta'].values/1000./2.))
dfmcp = get_mcp_simple(args.file, args.iev, 'mcparticles2').iloc[0]
pdgbase = ROOT.TDatabasePDG()
inpart = pdgbase.GetParticle(int(dfmcp['pid']))
print("Incoming particle = {}, pdgcode = {}, charge = {}, mass = {}".format(inpart.GetName(), inpart.PdgCode(), inpart.Charge(), inpart.Mass()))
# neutral particle, no need to consider magnetic field
if np.isclose(inpart.Charge(), 0., rtol=1e-5):
vec = dfmcp[['px', 'py', 'pz']].values
# charge particle, use the cluster center
else:
flayer = df[df['layer'] == df['layer'].min()]
vec = flayer[['x', 'y', 'z']].mean().values
vec = vec/np.linalg.norm(vec)
if not len(df):
print("Error: do not find any hits for cluster {:d} in event {:d}".format(args.icl, args.iev))
exit(-1)
......@@ -108,6 +118,7 @@ if __name__ == '__main__':
vecp = np.asarray([np.arccos(vec[2]), np.arctan2(vec[1], vec[0])])*1000.
phi_rg = np.asarray([vecp[1] - args.topo_range, vecp[1] + args.topo_range])
th_rg = np.asarray([vecp[0] - args.topo_range, vecp[0] + args.topo_range])
eta_rg = np.resize(-np.log(np.tan(vecp[0]/1000./2.)), 2) + np.asarray([-args.topo_range, args.topo_range])/1000.
os.makedirs(args.outdir, exist_ok=True)
# cluster plot by layers (rebinned)
......@@ -117,11 +128,11 @@ if __name__ == '__main__':
for i in np.arange(1, df['layer'].max() + 1, dtype=int):
data = df[df['layer'] == i]
fig, axs = plt.subplots(1, 2, figsize=(17, 16), dpi=160, gridspec_kw={'wspace':0., 'width_ratios': [16, 1]})
ax, sm = draw_heatmap(axs[0], data['theta'].values, data['phi'].values, weights=data['edep'].values,
bins=(np.arange(*th_rg, step=args.topo_size), np.arange(*phi_rg, step=args.topo_size)),
ax, sm = draw_heatmap(axs[0], data['eta'].values, data['phi'].values, weights=data['edep'].values,
bins=(np.arange(*eta_rg, step=args.topo_size/1000.), np.arange(*phi_rg, step=args.topo_size)),
cmap=cmap, vmin=0.0, vmax=1.0, pc_kw=dict(alpha=0.8, edgecolor='k'))
ax.set_ylabel(r'$\phi$ (mrad)', fontsize=28)
ax.set_xlabel(r'$\theta$ (mrad)', fontsize=28)
ax.set_xlabel(r'$\eta$', fontsize=28)
ax.tick_params(labelsize=24)
ax.xaxis.set_minor_locator(MultipleLocator(5))
ax.yaxis.set_minor_locator(MultipleLocator(5))
......@@ -155,11 +166,11 @@ if __name__ == '__main__':
fig, axs = plt.subplots(1, 2, figsize=(17, 16), dpi=160, gridspec_kw={'wspace':0., 'width_ratios': [16, 1]})
ax = axs[0]
colors = cmap(data['edep']/1.0)
ax.scatter(data['theta'].values, data['phi'].values, c=colors, marker='s', s=15.0)
ax.set_xlim(*th_rg)
ax.scatter(data['eta'].values, data['phi'].values, c=colors, marker='s', s=15.0)
ax.set_xlim(*eta_rg)
ax.set_ylim(*phi_rg)
ax.set_ylabel(r'$\phi$ (mrad)', fontsize=28)
ax.set_xlabel(r'$\theta$ (mrad)', fontsize=28)
ax.set_xlabel(r'$\eta$', fontsize=28)
ax.tick_params(labelsize=24)
ax.xaxis.set_minor_locator(MultipleLocator(5))
ax.yaxis.set_minor_locator(MultipleLocator(5))
......
benchmarks/sampling_ecal/scripts/utils.py
+ 51
− 1
View file @ cbfe2274
......@@ -41,7 +41,7 @@ def get_mcp_data(path, evnums=None, branch='mcparticles2'):
elif isinstance(evnums, int):
evnums = [evnums]
dbuf = np.zeros(shape=(2000*len(evnums), 6))
dbuf = np.zeros(shape=(5000*len(evnums), 6))
idb = 0
for iev in evnums:
if iev >= events.GetEntries():
......@@ -122,3 +122,53 @@ def compact_constants(path, names):
kernel.terminate()
return vals
# read clusters data from root file
def get_clusters_data(path, evnums=None, branch='EcalBarrelClustersReco'):
f = ROOT.TFile(path)
events = f.events
if evnums is None:
evnums = np.arange(events.GetEntries())
elif isinstance(evnums, int):
evnums = [evnums]
dbuf = np.zeros(shape=(20*len(evnums), 6))
idb = 0
for iev in evnums:
if iev >= events.GetEntries():
print('Error: event {:d} is out of range (0 - {:d})'.format(iev, events.GetEntries() - 1))
continue
events.GetEntry(iev)
for k, cl in enumerate(getattr(events, branch)):
dbuf[idb] = (iev, k, cl.nhits, cl.edep, cl.cl_theta, cl.cl_phi)
idb += 1
return pd.DataFrame(data=dbuf[:idb], columns=['event', 'cluster', 'nhits', 'edep', 'cl_theta', 'cl_phi'])
# read mc particles from root file
def get_mcp_simple(path, evnums=None, branch='mcparticles2'):
f = ROOT.TFile(path)
events = f.events
if evnums is None:
evnums = np.arange(events.GetEntries())
elif isinstance(evnums, int):
evnums = [evnums]
dbuf = np.zeros(shape=(len(evnums), 6))
idb = 0
for iev in evnums:
if iev >= events.GetEntries():
print('Error: event {:d} is out of range (0 - {:d})'.format(iev, events.GetEntries() - 1))
continue
events.GetEntry(iev)
# extract full mc particle data
part = getattr(events, branch)[2]
dbuf[idb] = (iev, part.psx, part.psy, part.psz, part.pdgID, part.status)
idb += 1
return pd.DataFrame(data=dbuf[:idb], columns=['event', 'px', 'py', 'pz', 'pid', 'status'])
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment