Resolve "Add true decaying particle info onto event display"

Closes #43 (closed)

benchmarks/imaging_ecal/scripts/draw_cluster.py

@@ -5,6 +5,10 @@
 
     Author: Chao Peng (ANL)
     Date: 04/30/2021
+
+    Added true decaying particles on eta-phi plane projection plot
+    Author: Jihee Kim (ANL)
+    Data: 08/06/2021
 '''
 
 import os
@@ -143,9 +147,26 @@ 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')
+    pdgbase = ROOT.TDatabasePDG()
+    # Select decaying particles
+    dftemp = dfallmcp[dfallmcp['g4Parent'] == 1.0]
+    if len(dftemp) > 0:
+        dfdecaymcp = dftemp.copy()
+        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 +237,9 @@ if __name__ == '__main__':
     ax, sm = draw_heatmap(axs[0], df['eta'].values, df['phi'].values, weights=df['energy'].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 
+    if len(dftemp) > 0:
+        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)