Skip to content
Snippets Groups Projects

EDM4hep: mcparticles -> MCParticles

Merged Wouter Deconinck requested to merge edm4hep-mcparticles into master
2 files
+ 7
5
Compare changes
  • Side-by-side
  • Inline
Files
2
@@ -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')
Loading