modify clustering to use the latest components

benchmarks/clustering/scripts/cluster_plots.py

@@ -21,7 +21,7 @@ def load_root_macros(arg_macros):
             print('Loading root macros: \'{}\' does not exist, skip loading it.'.format(path))
 
 
-def thrown_particles_figure(df):
+def thrown_particles_figure(df, save):
     # define truth particle info
     df = df.Define('isThrown', 'mcparticles2.genStatus == 1')\
            .Define('thrownParticles', 'mcparticles2[isThrown]')\
@@ -70,7 +70,37 @@ def thrown_particles_figure(df):
         ax.set_ylabel(label[1], fontsize=24)
 
     fig.text(0.5, 0.95, 'Thrown Particles', ha='center', fontsize=24)
-    return fig
+    fig.savefig(save)
+    plt.close(fig)
+
+
+def general_clusters_figure(df, collection, save):
+    data = df.AsNumpy([
+        '{}.nhits'.format(collection),
+        '{}.energy'.format(collection),
+        '{}.polar.theta'.format(collection),
+        '{}.polar.phi'.format(collection),
+        ])
+    # figure
+    fig, axs = plt.subplots(2, 2, figsize=(16, 12), dpi=120)
+    labels = [
+        ('Number of Hits', 'Counts'),
+        ('Energy (MeV)', 'Counts'),
+        ('Theta (rad)', 'Counts'),
+        ('Phi (rad)', 'Counts'),
+    ]
+    for ax, label, (name, vals) in zip(axs.flat, labels, data.items()):
+        hvals = [v for vec in vals for v in vec]
+        ax.hist(hvals, bins=50, ec='k')
+        ax.tick_params(labelsize=22, direction='in', which='both')
+        ax.grid(linestyle=':', which='both')
+        ax.set_axisbelow(True)
+        ax.set_xlabel(label[0], fontsize=24)
+        ax.set_ylabel(label[1], fontsize=24)
+
+    fig.text(0.5, 0.95, collection, ha='center', fontsize=24)
+    fig.savefig(save)
+    plt.close(fig)
 
 
 if __name__ == '__main__':
@@ -97,6 +127,8 @@ if __name__ == '__main__':
 
     rdf = ROOT.RDataFrame('events', args.file)
 
-    fig = thrown_particles_figure(rdf)
-    fig.savefig(os.path.join(args.outdir, 'thrown_particles.png'))
+    thrown_particles_figure(rdf, save=os.path.join(args.outdir, 'thrown_particles.png'))
+    general_clusters_figure(rdf, collection='CrystalEcalClusters', save=os.path.join(args.outdir, 'crystal_ecal_clusters.png'))
+    general_clusters_figure(rdf, collection='EcalBarrelClusters', save=os.path.join(args.outdir, 'ecal_barrel_clusters.png'))
+    general_clusters_figure(rdf, collection='HcalBarrelClusters', save=os.path.join(args.outdir, 'hcal_barrel_clusters.png'))