From cbfe2274c876be1e84d30adc621bfc1af2d3b13b Mon Sep 17 00:00:00 2001
From: Chao Peng <cpeng@anl.gov>
Date: Mon, 24 May 2021 19:08:02 -0500
Subject: [PATCH] update draw_cluster center finding
---
.../sampling_ecal/scripts/draw_cluster.py | 42 ++++++++-------
.../scripts/draw_cluster_layers.py | 33 ++++++++----
benchmarks/sampling_ecal/scripts/utils.py | 52 ++++++++++++++++++-
3 files changed, 96 insertions(+), 31 deletions(-)
diff --git a/benchmarks/sampling_ecal/scripts/draw_cluster.py b/benchmarks/sampling_ecal/scripts/draw_cluster.py
index 31aafb78..ff599f06 100644
--- a/benchmarks/sampling_ecal/scripts/draw_cluster.py
+++ b/benchmarks/sampling_ecal/scripts/draw_cluster.py
@@ -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)))
diff --git a/benchmarks/sampling_ecal/scripts/draw_cluster_layers.py b/benchmarks/sampling_ecal/scripts/draw_cluster_layers.py
index e3b9ed3d..0e6e1fda 100644
--- a/benchmarks/sampling_ecal/scripts/draw_cluster_layers.py
+++ b/benchmarks/sampling_ecal/scripts/draw_cluster_layers.py
@@ -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))
diff --git a/benchmarks/sampling_ecal/scripts/utils.py b/benchmarks/sampling_ecal/scripts/utils.py
index 2e81c508..a1e31a80 100644
--- a/benchmarks/sampling_ecal/scripts/utils.py
+++ b/benchmarks/sampling_ecal/scripts/utils.py
@@ -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'])
+
+
+
--
GitLab