diff --git a/benchmarks/tracking/scripts/tracking_performance.py b/benchmarks/tracking/scripts/tracking_performance.py index 3166f92108aca57e65f6e02cfb5548c05f4d7d9d..69e2affccd95f0bc788711895449e8f5d80878f0 100644 --- a/benchmarks/tracking/scripts/tracking_performance.py +++ b/benchmarks/tracking/scripts/tracking_performance.py @@ -47,7 +47,6 @@ def flatten_collection(rdf, collection, cols=None): dfp.loc[:, 'event'] = evns return dfp - def thrown_particles_figure(rdf, save, mcbranch="mcparticles2"): # define truth particle info dft = flatten_collection(rdf, mcbranch, ['genStatus', 'pdgID', 'ps.x', 'ps.y', 'ps.z', 'mass']) @@ -137,6 +136,8 @@ if __name__ == '__main__': dfm = thrown_particles_figure(rdf_rec, save=os.path.join(args.outdir, '{}_thrown_particles.png'.format(args.nametag)), mcbranch=args.mc) df = flatten_collection(rdf_rec, args.coll) df.rename(columns={c: c.replace(args.coll + '.', '') for c in df.columns}, inplace=True) + df['theta'] = df['direction.theta'] + df['phi'] = df['direction.phi'] fig, axs = plt.subplots(2, 2, figsize=(16, 12), dpi=120) for ax in axs.flat: