EDM4hep: mcparticles -> MCParticles

benchmarks/clustering/scripts/cluster_plots.py

@@ -48,12 +48,12 @@ def flatten_collection(rdf, collection, cols=None):
     return dfp
 
 
-def thrown_particles_figure(rdf, save, mcbranch="mcparticles"):
+def thrown_particles_figure(rdf, save, mcbranch="MCParticles"):
     # define truth particle info
-    dft = flatten_collection(rdf, mcbranch, ['genStatus', 'pdgID', 'ps.x', 'ps.y', 'ps.z', 'mass'])
+    dft = flatten_collection(rdf, mcbranch, ['generatorStatus', 'PDG', 'momentum.x', 'momentum.y', 'momentum.z', 'mass'])
     dft.rename(columns={c: c.replace(mcbranch + '.', '') for c in dft.columns}, inplace=True)
     # select thrown particles
-    dft = dft[dft['genStatus'] == 1]
+    dft = dft[dft['generatorStatus'] == 1]
 
     # figure
     fig, axs = plt.subplots(2, 2, figsize=(16, 12), dpi=120)
@@ -62,7 +62,7 @@ def thrown_particles_figure(rdf, save, mcbranch="mcparticles"):
     get_pname = np.vectorize(lambda pid: pdgbase.GetParticle(int(pid)).GetName())
 
     # enumerate particle names
-    dft.loc[:, 'pname'] = get_pname(dft['pdgID'].values)
+    dft.loc[:, 'pname'] = get_pname(dft['PDG'].values)
     penum = {pname: i for i, pname in enumerate(dft['pname'].unique())}
     dft.loc[:, 'pname_id'] = dft['pname'].map(penum)
 
@@ -77,7 +77,7 @@ def thrown_particles_figure(rdf, save, mcbranch="mcparticles"):
 
     # calculate kinematics
     get_4vecs = np.vectorize(lambda x, y, z, m: ROOT.Math.PxPyPzMVector(x, y, z, m))
-    fourvecs = get_4vecs(*dft[['ps.x', 'ps.y', 'ps.z', 'mass']].values.T)
+    fourvecs = get_4vecs(*dft[['momentum.x', 'momentum.y', 'momentum.z', 'mass']].values.T)
 
     dft.loc[:, 'p'] = [v.P() for v in fourvecs]
     dft.loc[:, 'eta'] = [v.Eta() for v in fourvecs]
@@ -109,7 +109,7 @@ if __name__ == '__main__':
     parser.add_argument('-o', dest='outdir', default='.', help='Output directory.')
     parser.add_argument('-m', '--macros', dest='macros', default='rootlogon.C',
             help='Macro files to be loaded by root, separated by \",\".')
-    parser.add_argument('--mc-branch', dest='mc', default='mcparticles',
+    parser.add_argument('--mc-branch', dest='mc', default='MCParticles',
             help='Branch name of MC particles truth info.')
     parser.add_argument('--collections', dest='coll', default='',
             help='Collection name of clusters to plot')