From 7fc670b850651d328c436ce67f0053d01924d65f Mon Sep 17 00:00:00 2001
From: "jihee.kim" <jihee.kim@anl.gov>
Date: Fri, 6 Aug 2021 15:13:09 -0500
Subject: [PATCH] Added a feature that draw true decaying particles onto
eta-phi plane
---
.../imaging_ecal/scripts/draw_cluster.py | 20 ++++++++++++-
benchmarks/imaging_ecal/scripts/utils.py | 29 +++++++++++++++++++
2 files changed, 48 insertions(+), 1 deletion(-)
diff --git a/benchmarks/imaging_ecal/scripts/draw_cluster.py b/benchmarks/imaging_ecal/scripts/draw_cluster.py
index 88f338d0..25974ee9 100644
--- a/benchmarks/imaging_ecal/scripts/draw_cluster.py
+++ b/benchmarks/imaging_ecal/scripts/draw_cluster.py
@@ -143,9 +143,25 @@ if __name__ == '__main__':
df['theta'] = np.arccos(df['z'].values/df['r'].values)*1000.
df['eta'] = -np.log(np.tan(df['theta'].values/1000./2.))
+ # Read all mc particles
+ ########################
+ dfallmcp = get_all_mcp(args.file, args.iev, 'mcparticles2')
+ # Select decaying particles
+ dfdecaymcp = dfallmcp[dfallmcp['g4Parent'] == 1.0]
+ pdgbase = ROOT.TDatabasePDG()
+ for iptl in [0, len(dfdecaymcp) - 1]:
+ infoptl = pdgbase.GetParticle(int(dfdecaymcp['pid'].iloc[iptl]))
+ print("{} Decaying particle = {}, pdgcode = {}, charge = {}, mass = {}"\
+ .format(iptl, infoptl.GetName(), infoptl.PdgCode(), infoptl.Charge(), infoptl.Mass()))
+ # Calculate geometric variables of decaying particles
+ dfdecaymcp['r'] = np.sqrt(dfdecaymcp['vex'].values**2 + dfdecaymcp['vey'].values**2 + dfdecaymcp['vez'].values**2)
+ dfdecaymcp['phi'] = np.arctan2(dfdecaymcp['vey'].values, dfdecaymcp['vex'].values)*1000.
+ dfdecaymcp['theta'] = np.arccos(dfdecaymcp['vez'].values/dfdecaymcp['r'].values)*1000.
+ dfdecaymcp['eta'] = -np.log(np.tan(dfdecaymcp['theta'].values/1000./2.))
+
# truth
dfmcp = get_mcp_simple(args.file, args.iev, 'mcparticles2').iloc[0]
- pdgbase = ROOT.TDatabasePDG()
+ #pdgbase = ROOT.TDatabasePDG()
inpart = pdgbase.GetParticle(int(dfmcp['pid']))
print("Incoming particle = {}, pdgcode = {}, charge = {}, mass = {}"\
.format(inpart.GetName(), inpart.PdgCode(), inpart.Charge(), inpart.Mass()))
@@ -216,6 +232,8 @@ if __name__ == '__main__':
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'))
+ # draw true decaying particle position
+ ax.scatter(dfdecaymcp['eta'].values, dfdecaymcp['phi'].values, marker='x', color='red', s=22**2, linewidth=5.0)
ax.set_ylabel(r'$\phi$ (mrad)', fontsize=32)
ax.set_xlabel(r'$\eta$', fontsize=32)
diff --git a/benchmarks/imaging_ecal/scripts/utils.py b/benchmarks/imaging_ecal/scripts/utils.py
index da9954ac..740e3ac1 100644
--- a/benchmarks/imaging_ecal/scripts/utils.py
+++ b/benchmarks/imaging_ecal/scripts/utils.py
@@ -3,6 +3,10 @@
Author: Chao Peng (ANL)
Date: 04/30/2021
+
+ Added all mc particle info to obtain decaying particles
+ Author: Jihee Kim (ANL)
+ Date: 08/06/2021
'''
import os
@@ -79,6 +83,31 @@ def get_mcp_simple(path, evnums=None, branch='mcparticles2'):
idb += 1
return pd.DataFrame(data=dbuf[:idb], columns=['event', 'px', 'py', 'pz', 'pid', 'status'])
+#######################################
+# read all mc particles from root file
+#######################################
+def get_all_mcp(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=(2000*len(evnums), 10))
+ 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 mc particle data
+ for ptl in getattr(events, branch):
+ dbuf[idb] = (iev, ptl.psx, ptl.psy, ptl.psz, ptl.pdgID, ptl.status, ptl.g4Parent, ptl.vex, ptl.vey, ptl.vez)
+ idb += 1
+
+ return pd.DataFrame(data=dbuf[:idb], columns=['event', 'px', 'py', 'pz', 'pid', 'status', 'g4Parent', 'vex', 'vey', 'vez'])
# read hits data from root file
def get_hits_data(path, evnums=None, branch='RecoEcalBarrelHits'):
--
GitLab