scripts/imaging_calorimeter/energy_profile.py

0001 '''
0002     A script to generate the energy profile (layer-wise)
0003     It reads the output from the Juggler component ImagingClusterReco, which is supposed to be clusters of hits after
0004     digitization, reconstruction, and clustering
0005
0006     Author: Chao Peng (ANL)
0007     Date: 04/30/2021
0008 '''
0009
0010 import os
0011 import numpy as np
0012 import pandas as pd
0013 import ROOT
0014 from ROOT import gROOT, gInterpreter
0015 import argparse
0016 import matplotlib
0017 from matplotlib import pyplot as plt
0018 from matplotlib.ticker import MultipleLocator, MaxNLocator
0019 import sys
0020 from utils import *
0021
0022
0023 def find_start_layer(grp, min_edep=0.5):
0024     ids, edeps = grp.sort_values('layer')[['layer', 'edep']].values.T
0025     edeps = np.cumsum(edeps)
0026     return min(ids[edeps > min_edep]) if sum(edeps > min_edep) > 0 else 20
0027
0028
0029 # execute this script
0030 if __name__ == '__main__':
0031     parser = argparse.ArgumentParser(description='Generate energy profiles')
0032     parser.add_argument('file', type=str, help='path to root file')
0033     parser.add_argument('-o', '--plot-dir', type=str, default='./plots', dest='outdir', help='output directory')
0034     parser.add_argument('-t', '--type', type=str, default='unknown', dest='type', help='profile type (used in save)')
0035     parser.add_argument('-e', '--energy', type=float, default=5., dest='energy', help='incident particle energy (GeV)')
0036     parser.add_argument('-s', '--save', type=str, default='', dest='save', help='path to save profile')
0037     parser.add_argument('-c', '--color', type=str, default='royalblue', dest='color', help='colors for bar plots')
0038     parser.add_argument('-b', '--branch-name', type=str, default='EcalBarrelClustersLayers', dest='branch',
0039                         help='branch name in the root file (outputLayerCollection from ImagingClusterReco)')
0040     parser.add_argument('-m', '--macros', type=str, default='rootlogon.C', dest='macros',
0041                         help='root macros to load (accept multiple paths separated by \",\")')
0042     args = parser.parse_args()
0043
0044     load_root_macros(args.macros)
0045     # prepare data
0046     dfe = get_layers_data(args.file, branch=args.branch)
0047     dfe.loc[:, 'total_edep'] = dfe.groupby(['event', 'cluster'])['edep'].transform('sum')
0048     # dfe.loc[:, 'efrac'] = dfe['edep']/dfe['total_edep']
0049     dfe.loc[:, 'efrac'] = dfe['edep']/(args.energy*1000.)
0050     dfe = dfe[dfe['cluster'] == 0]
0051
0052     os.makedirs(args.outdir, exist_ok=True)
0053
0054     slayer = dfe.groupby('event').apply(lambda x: find_start_layer(x, 1.0)).astype(int)
0055     dfe = dfe.merge(slayer.to_frame(name='slayer'), on='event')
0056     # dfe.loc[:, 'eff_layer'] = dfe['layer'] - dfe['slayer']
0057     # prof = dfe[dfe['eff_layer'] > 0].groupby('eff_layer')['edep'].describe()
0058     prof = dfe.groupby('layer')['efrac'].describe()
0059
0060     # print(prof['mean'])
0061     # plot profiles
0062     bpos, bprops = (0.5, 0.95), dict(boxstyle='round', facecolor='wheat', alpha=0.3)
0063     fig, ax = plt.subplots(figsize=(16, 9), dpi=160)
0064     nev = len(dfe['event'].unique())
0065     ax.hist(dfe.groupby('event')['slayer'].min(), weights=[1/float(nev)]*nev,
0066             ec='black', bins=np.arange(0.5, 10.5, step=1.0))
0067     ax.xaxis.set_major_locator(MaxNLocator(integer=True))
0068     ax.xaxis.set_minor_locator(MultipleLocator(1))
0069     ax.yaxis.set_minor_locator(MultipleLocator(1))
0070     ax.grid(linestyle=':', which='both')
0071     ax.tick_params(labelsize=24)
0072     ax.set_xlabel('Start Layer', fontsize=26)
0073     ax.set_ylabel('Normalized Counts', fontsize=26)
0074     ax.text(*bpos, 'Mininum Edep\n' + '{:.1f} MeV'.format(1.0),
0075             transform=ax.transAxes, fontsize=26, verticalalignment='top', bbox=bprops)
0076     fig.savefig(os.path.join(args.outdir, 'edep_start_{}_{}.png'.format(args.type, int(args.energy))))
0077
0078
0079     fig, ax = plt.subplots(figsize=(16, 9), dpi=160)
0080     ax.plot(prof.index, prof['mean'].values*100., color=args.color, lw=2)
0081     # ax.fill_between(prof.index, prof['25%'], prof['75%'], color=args.color, alpha=0.3)
0082     ax.fill_between(prof.index, (prof['mean'] - prof['std'])*100., (prof['mean'] + prof['std'])*100.,
0083                     color=args.color, alpha=0.3)
0084     ax.xaxis.set_major_locator(MaxNLocator(integer=True))
0085     ax.xaxis.set_minor_locator(MultipleLocator(1))
0086     ax.yaxis.set_minor_locator(MultipleLocator(1))
0087     ax.grid(linestyle=':', which='both')
0088     ax.tick_params(labelsize=24)
0089     ax.set_xlabel('Layer', fontsize=26)
0090     ax.set_ylabel('Energy Deposit Percentage', fontsize=26)
0091     fig.savefig(os.path.join(args.outdir, 'efrac_{}_{}.png'.format(args.type, int(args.energy))))
0092
0093     # edep fraction by layers
0094     layers = np.asarray([
0095         [1, 5, 8,],
0096         [10, 15, 20],
0097     ])
0098     layers_bins = np.linspace(0, 0.5, 50)
0099
0100     fig, ax = plt.subplots(*layers.shape, figsize=(16, 9), dpi=160, sharex='col', sharey='all',
0101                            gridspec_kw=dict(hspace=0.05, wspace=0.05))
0102
0103     for ax, layer in zip(ax.flat, layers.flatten()):
0104         data = dfe[dfe['layer'] == layer]
0105         ax.hist(data['efrac'].values*100., weights=[1/float(len(data))]*len(data), bins=layers_bins,
0106                 ec='black', color=args.color)
0107         ax.tick_params(labelsize=24)
0108         ax.xaxis.set_minor_locator(MultipleLocator(1))
0109         ax.grid(linestyle=':', which='both')
0110         # ax.set_xlabel('Energy Deposit (MeV)', fontsize=26)
0111         # ax.set_ylabel('Normalized Counts', fontsize=26)
0112         mean, std = data.describe().loc[['mean', 'std'], 'edep'].values
0113         label = 'Layer {}'.format(layer)
0114     #            + '\n' + r'$\mu = {:.3f}$ MeV'.format(mean)
0115     #            + '\n' + r'$\sigma = {:.3f}$ MeV'.format(std)
0116         ax.text(*bpos, label, transform=ax.transAxes, fontsize=26, verticalalignment='top', bbox=bprops)
0117         ax.set_axisbelow(True)
0118         ax.set_yscale('log')
0119     fig.text(0.5, 0.02, 'Energy Deposit Percentage', fontsize=26, ha='center')
0120     fig.text(0.02, 0.5, 'Normalized Counts', fontsize=26, va='center', rotation=90)
0121     fig.savefig(os.path.join(args.outdir, 'efrac_layers_{}_{}.png'.format(args.type, int(args.energy))))
0122
0123
0124     if args.save:
0125         prof.loc[:, 'energy'] = args.energy*1000.
0126         prof.loc[:, 'type'] = args.type
0127         if os.path.exists(args.save):
0128             prev = pd.read_csv(args.save).set_index('layer', drop=True)
0129             prof = pd.concat([prof, prev])
0130         prof.to_csv(args.save)
0131