scripts/imaging_calorimeter/epi_separation.py

0001 '''
0002     A script to use the energy profile for e-pi separation
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 prepare_data(path, **kwargs):
0024     tmp = get_layers_data(path, **kwargs)
0025     tmp.loc[:, 'total_edep'] = tmp.groupby(['event', 'cluster'])['edep'].transform('sum')
0026     tmp.loc[:, 'efrac'] = tmp['edep']/tmp['total_edep']
0027     # tmp = tmp[tmp['cluster'] == 0]
0028     return tmp
0029
0030
0031 def calc_chi2(grp, pr, lmin=5, lmax=12):
0032     grp2 = grp[(grp['layer'] >= lmin) & (grp['layer'] <= lmax)]
0033     mean, std = pr.loc[grp2['layer'], ['mean', 'std']].values.T*pr['energy'].mean()
0034     edeps = grp2['edep'].values
0035     return np.sqrt(np.sum((edeps - mean)**2/std**2)/float(len(edeps)))
0036
0037
0038 # execute this script
0039 if __name__ == '__main__':
0040     parser = argparse.ArgumentParser(description='epi_separation')
0041     parser.add_argument('efile', type=str, help='path to root file (electrons)')
0042     parser.add_argument('pifile', type=str, help='path to root file (pions)')
0043     parser.add_argument('--prof', type=str, default='profile.csv', help='path to electron profile')
0044     parser.add_argument('--plot-dir', type=str, default='./plots', dest='outdir', help='output directory')
0045     parser.add_argument('-b', '--branch-name', type=str, default='EcalBarrelClustersLayers', dest='branch',
0046                         help='branch name in the root file (outputLayerCollection from ImagingClusterReco)')
0047     parser.add_argument('-m', '--macros', type=str, default='rootlogon.C', dest='macros',
0048                         help='root macros to load (accept multiple paths separated by \",\")')
0049     args = parser.parse_args()
0050
0051     os.makedirs(args.outdir, exist_ok=True)
0052
0053
0054     load_root_macros(args.macros)
0055     # prepare data
0056     dfe = prepare_data(args.efile, branch=args.branch)
0057     dfpi = prepare_data(args.pifile, branch=args.branch)
0058
0059     colors = ['royalblue', 'indianred', 'limegreen']
0060     prof = pd.concat([pd.read_csv(p.strip()) for p in args.prof.split(',')]).reset_index(drop=True)
0061
0062     # profile comparison
0063     fig, ax = plt.subplots(figsize=(16, 9), dpi=160)
0064     for title, color in zip(sorted(prof['type'].unique()), colors):
0065         pr = prof[prof['type'] == title]
0066         ax.plot(pr.index, pr['mean'], color=color, lw=2)
0067         ax.fill_between(pr.index, (pr['mean'] - pr['std']), (pr['mean'] + pr['std']),
0068                         color=color, alpha=0.3, label=title)
0069     ax.xaxis.set_major_locator(MaxNLocator(integer=True))
0070     ax.xaxis.set_minor_locator(MultipleLocator(1))
0071     ax.yaxis.set_minor_locator(MultipleLocator(1))
0072     ax.grid(linestyle=':', which='both')
0073     ax.tick_params(labelsize=24)
0074     ax.set_xlabel('Layer', fontsize=26)
0075     ax.set_ylabel('Energy Deposit Percentage', fontsize=26)
0076     ax.legend(fontsize=26, loc='upper left')
0077     fig.savefig(os.path.join(args.outdir, 'compare_prof.png'))
0078
0079
0080     # check profile
0081     layer_range = (4, 12)
0082     pre = prof[prof['type'].str.lower() == 'electron']
0083     fig, ax = plt.subplots(figsize=(16, 9), dpi=160)
0084     chi2 = dfe.groupby(['event', 'cluster']).apply(lambda x: calc_chi2(x, pre, *layer_range))
0085     ax.hist(chi2, bins=np.linspace(0, 5, 200), weights=[1/float(len(chi2))]*chi2,
0086             ec='royalblue', color='royalblue', alpha=0.5, label='electrons')
0087     chi2 = dfpi.groupby(['event', 'cluster']).apply(lambda x: calc_chi2(x, pre, *layer_range))
0088     # print(chi2[chi2 < 0.7])
0089     ax.hist(chi2, bins=np.linspace(0, 5, 200), weights=[1/float(len(chi2))]*chi2,
0090             ec='indianred', color='indianred', alpha=0.5, label='pions')
0091     ax.grid(linestyle=':', which='major')
0092     ax.set_axisbelow(True)
0093     ax.tick_params(labelsize=24)
0094     ax.set_xlabel(r'$\chi^2$', fontsize=26)
0095     ax.set_ylabel('Normalized Counts', fontsize=26)
0096     # ax.set_yscale('log')
0097     # ax.set_ylim(1e-3, 1.)
0098     ax.legend(title=r'$\chi^2$ for $E_{{dep}}$ in layer {:d} - {:d}'.format(*layer_range), title_fontsize=28, fontsize=26)
0099     fig.savefig(os.path.join(args.outdir, 'efrac_chi2.png'))
0100