diff --git a/benchmarks/imaging_ecal/scripts/draw_cluster.py b/benchmarks/imaging_ecal/scripts/draw_cluster.py index 88f338d0befa96e45b6a16e4754e69d1cbfee2f3..25974ee9c4f9185a3b272423828411f441334826 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 da9954ac3e2dc1b0b36009c16fe97c25195394de..740e3ac19f0343a1fc87589d365bc07ac9f24bd5 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'):