Skip to content
Snippets Groups Projects
Commit b16f46e0 authored by Chao Peng's avatar Chao Peng
Browse files

add dataframe return for the script

parent 1a5a9334
No related branches found
No related tags found
1 merge request!114adjust scripts for crystal-glass endcap ecal
This commit is part of merge request !114. Comments created here will be created in the context of that merge request.
...@@ -91,8 +91,8 @@ def general_clusters_figure(df, collection, save, min_nhits=3): ...@@ -91,8 +91,8 @@ def general_clusters_figure(df, collection, save, min_nhits=3):
data=np.vstack(list(data.values())).T) data=np.vstack(list(data.values())).T)
dfp.loc[:, 'evn'] = evns dfp.loc[:, 'evn'] = evns
# select the max. energy cluster for each event # select the max. energy cluster for each event
dfp = dfp.loc[dfp.groupby('evn')['edep'].idxmax()] dfc = dfp.loc[dfp.groupby('evn')['edep'].idxmax()]
dfp = dfp.loc[dfp['nhits'] >= min_nhits] dfc = dfc.loc[dfc['nhits'] >= min_nhits]
# figure # figure
fig, axs = plt.subplots(2, 2, figsize=(16, 12), dpi=120) fig, axs = plt.subplots(2, 2, figsize=(16, 12), dpi=120)
labels = [ labels = [
...@@ -101,7 +101,7 @@ def general_clusters_figure(df, collection, save, min_nhits=3): ...@@ -101,7 +101,7 @@ def general_clusters_figure(df, collection, save, min_nhits=3):
(r'$\theta$ (rad)', 'Counts'), (r'$\theta$ (rad)', 'Counts'),
(r'$\phi$ (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.hist(vals, bins=50, ec='k')
ax.tick_params(labelsize=22, direction='in', which='both') ax.tick_params(labelsize=22, direction='in', which='both')
ax.grid(linestyle=':', which='both') ax.grid(linestyle=':', which='both')
...@@ -112,6 +112,7 @@ def general_clusters_figure(df, collection, save, min_nhits=3): ...@@ -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.text(0.5, 0.95, collection, ha='center', fontsize=24)
fig.savefig(save) fig.savefig(save)
plt.close(fig) plt.close(fig)
return dfp
if __name__ == '__main__': if __name__ == '__main__':
...@@ -143,12 +144,14 @@ if __name__ == '__main__': ...@@ -143,12 +144,14 @@ if __name__ == '__main__':
rdf_rec = ROOT.RDataFrame('events', args.rec_file) 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) 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, general_clusters_figure(rdf_rec, collection='EcalEndcapNClusters', min_nhits=10,
save=os.path.join(args.outdir, 'ecal_electron_endcap_clusters.png')) save=os.path.join(args.outdir, 'ecal_electron_endcap_clusters.png'))
general_clusters_figure(rdf_rec, collection='EcalEndcapPClusters', min_nhits=5, general_clusters_figure(rdf_rec, collection='EcalEndcapPClusters', min_nhits=5,
save=os.path.join(args.outdir, 'ecal_hadron_endcap_clusters.png')) save=os.path.join(args.outdir, 'ecal_hadron_endcap_clusters.png'))
general_clusters_figure(rdf_rec, collection='EcalBarrelClusters', general_clusters_figure(rdf_rec, collection='EcalBarrelClusters',
save=os.path.join(args.outdir, 'ecal_barrel_clusters.png')) save=os.path.join(args.outdir, 'ecal_barrel_clusters.png'))
general_clusters_figure(rdf_rec, collection='HcalElectronEndcapClusters', min_nhits=5, general_clusters_figure(rdf_rec, collection='HcalElectronEndcapClusters', min_nhits=5,
save=os.path.join(args.outdir, 'hcal_electron_endcap_clusters.png')) save=os.path.join(args.outdir, 'hcal_electron_endcap_clusters.png'))
general_clusters_figure(rdf_rec, collection='HcalHadronEndcapClusters', min_nhits=10, general_clusters_figure(rdf_rec, collection='HcalHadronEndcapClusters', min_nhits=10,
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment