adjust scripts for crystal-glass endcap ecal

benchmarks/clustering/scripts/cluster_plots.py

@@ -91,8 +91,8 @@ def general_clusters_figure(df, collection, save, min_nhits=3):
             data=np.vstack(list(data.values())).T)
     dfp.loc[:, 'evn'] = evns
     # select the max. energy cluster for each event
-    dfp = dfp.loc[dfp.groupby('evn')['edep'].idxmax()]
-    dfp = dfp.loc[dfp['nhits'] >= min_nhits]
+    dfc = dfp.loc[dfp.groupby('evn')['edep'].idxmax()]
+    dfc = dfc.loc[dfc['nhits'] >= min_nhits]
     # figure
     fig, axs = plt.subplots(2, 2, figsize=(16, 12), dpi=120)
     labels = [
@@ -101,7 +101,7 @@ def general_clusters_figure(df, collection, save, min_nhits=3):
         (r'$\theta$ (rad)', 'Counts'),
         (r'$\phi$ (rad)', 'Counts'),
     ]
-    for ax, label, vals in zip(axs.flat, labels, dfp[['nhits', 'edep', 'theta', 'phi']].values.T):
+    for ax, label, vals in zip(axs.flat, labels, dfc[['nhits', 'edep', 'theta', 'phi']].values.T):
         ax.hist(vals, bins=50, ec='k')
         ax.tick_params(labelsize=22, direction='in', which='both')
         ax.grid(linestyle=':', which='both')
@@ -112,6 +112,7 @@ def general_clusters_figure(df, collection, save, min_nhits=3):
     fig.text(0.5, 0.95, collection, ha='center', fontsize=24)
     fig.savefig(save)
     plt.close(fig)
+    return dfp
 
 
 if __name__ == '__main__':
@@ -143,12 +144,14 @@ if __name__ == '__main__':
     rdf_rec = ROOT.RDataFrame('events', args.rec_file)
 
     thrown_particles_figure(rdf_sim, save=os.path.join(args.outdir, 'thrown_particles.png'), mcbranch=args.mc)
+
     general_clusters_figure(rdf_rec, collection='EcalEndcapNClusters', min_nhits=10,
             save=os.path.join(args.outdir, 'ecal_electron_endcap_clusters.png'))
     general_clusters_figure(rdf_rec, collection='EcalEndcapPClusters', min_nhits=5,
             save=os.path.join(args.outdir, 'ecal_hadron_endcap_clusters.png'))
     general_clusters_figure(rdf_rec, collection='EcalBarrelClusters',
             save=os.path.join(args.outdir, 'ecal_barrel_clusters.png'))
+
     general_clusters_figure(rdf_rec, collection='HcalElectronEndcapClusters', min_nhits=5,
             save=os.path.join(args.outdir, 'hcal_electron_endcap_clusters.png'))
     general_clusters_figure(rdf_rec, collection='HcalHadronEndcapClusters', min_nhits=10,