scripts/imaging_calorimeter/draw_cluster_layers.py

0001 '''
0002     A script to visualize the cluster layer-by-layer
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 argparse
0014 import matplotlib
0015 from matplotlib import cm
0016 from matplotlib import pyplot as plt
0017 from matplotlib.ticker import MultipleLocator
0018 from matplotlib.collections import PatchCollection
0019 from matplotlib.patches import Rectangle
0020 from mpl_toolkits.axes_grid1 import make_axes_locatable
0021 from matplotlib.backends.backend_pdf import PdfPages
0022 from utils import *
0023 import sys
0024 import imageio
0025
0026
0027 # color map
0028 def draw_heatmap(axis, x, y, weights, bins=1000, vmin=None, vmax=None, cmap=plt.get_cmap('rainbow'), pc_kw=dict()):
0029     w, xedg, yedg = np.histogram2d(x, y, weights=weights, bins=bins)
0030     xsz = np.mean(np.diff(xedg))
0031     ysz = np.mean(np.diff(yedg))
0032     if vmin == None:
0033         vmin = w.min()
0034     if vmax == None:
0035         vmax = w.max()
0036     recs, clrs = [], []
0037     for i in np.arange(len(xedg) - 1):
0038         for j in np.arange(len(yedg) - 1):
0039             if w[i][j] > vmin:
0040                 recs.append(Rectangle((xedg[i], yedg[j]), xsz, ysz))
0041                 clrs.append(cmap((w[i][j] - vmin) / (vmax - vmin)))
0042     axis.add_collection(PatchCollection(recs, facecolor=clrs, **pc_kw))
0043     axis.set_xlim(xedg[0], xedg[-1])
0044     axis.set_ylim(yedg[0], yedg[-1])
0045     return axis, cm.ScalarMappable(norm=matplotlib.colors.Normalize(vmin=vmin, vmax=vmax), cmap=cmap)
0046
0047
0048 # execute this script
0049 if __name__ == '__main__':
0050     parser = argparse.ArgumentParser(description='Visualize the cluster (layer-wise) from analysis')
0051     parser.add_argument('file', type=str, help='path to root file')
0052     parser.add_argument('-e', type=int, default=0, dest='iev', help='event number to plot')
0053     parser.add_argument('-c', type=int, default=0, dest='icl', help='cluster number to plot (0: all, -1: no cluster)')
0054     parser.add_argument('-s', type=int, default=8, dest='stop', help='stop layer for track fit')
0055     parser.add_argument('-o', type=str, default='./plots', dest='outdir', help='output directory')
0056     parser.add_argument('-d', type=float, default=1.0, dest='dura', help='duration of gif')
0057     parser.add_argument('--compact', type=str, default='', dest='compact', help='compact file')
0058     parser.add_argument('-m', '--macros', type=str, default='rootlogon.C', dest='macros',
0059                         help='root macros to load (accept multiple paths separated by \",\")')
0060     parser.add_argument('-b', '--branch-name', type=str, default='RecoEcalBarrelHits', dest='branch',
0061                         help='branch name in the root file (outputLayerCollection from ImagingClusterReco)')
0062     parser.add_argument('--topo-size', type=float, default=2.0, dest='topo_size',
0063                         help='bin size for projection plot (mrad)')
0064     parser.add_argument('--topo-range', type=float, default=50.0, dest='topo_range',
0065                         help='half range for projection plot (mrad)')
0066     args = parser.parse_args()
0067
0068
0069     # we can read these values from xml file
0070     desc = compact_constants(args.compact, [
0071         'cb_ECal_RMin',
0072         'cb_ECal_ReadoutLayerThickness',
0073         'cb_ECal_ReadoutLayerNumber',
0074         'cb_ECal_Length'
0075     ])
0076     if not len(desc):
0077         # or define Ecal shapes
0078         rmin, thickness, length = 890, 20*(10. + 1.65), 860*2+500
0079     else:
0080         # convert cm to mm
0081         rmin = desc[0]*10.
0082         thickness = desc[1]*desc[2]*10.
0083         length = desc[3]*10.
0084
0085
0086     # read data
0087     load_root_macros(args.macros)
0088     df = get_hits_data(args.file, args.iev, args.branch)
0089     if args.icl != 0:
0090         df = df[df['cluster'] == args.icl]
0091     if not len(df):
0092         print("Error: do not find any hits for cluster {:d} in event {:d}".format(args.icl, args.iev))
0093         exit(-1)
0094     # convert to polar coordinates (mrad), and stack all r values
0095     df['r'] = np.sqrt(df['x'].values**2 + df['y'].values**2 + df['z'].values**2)
0096     df['phi'] = np.arctan2(df['y'].values, df['x'].values)*1000.
0097     df['theta'] = np.arccos(df['z'].values/df['r'].values)*1000.
0098     df['eta'] = -np.log(np.tan(df['theta'].values/1000./2.))
0099
0100     # truth
0101     dfmcp = get_mcp_simple(args.file, args.iev, 'mcparticles2').iloc[0]
0102     pdgbase = ROOT.TDatabasePDG()
0103     inpart = pdgbase.GetParticle(int(dfmcp['pid']))
0104     print("Incoming particle = {}, pdgcode = {}, charge = {}, mass = {}"\
0105           .format(inpart.GetName(), inpart.PdgCode(), inpart.Charge(), inpart.Mass()))
0106     # neutral particle, no need to consider magnetic field
0107     if np.isclose(inpart.Charge(), 0., rtol=1e-5):
0108         vec = dfmcp[['px', 'py', 'pz']].values
0109     # charge particle, use the cluster center
0110     else:
0111         flayer = df[df['layer'] == df['layer'].min()]
0112         vec = flayer[['x', 'y', 'z']].mean().values
0113     vec = vec/np.linalg.norm(vec)
0114
0115     # particle line from (0, 0, 0) to the inner Ecal surface
0116     length = rmin/np.sqrt(vec[0]**2 + vec[1]**2)
0117     pline = np.transpose(vec*np.mgrid[0:length:2j][:, np.newaxis])
0118     cmap = truncate_colormap(plt.get_cmap('jet'), 0.1, 0.9)
0119
0120     # convert truth to mrad
0121     vecp = np.asarray([np.arccos(vec[2]), np.arctan2(vec[1], vec[0])])*1000.
0122     phi_rg = np.asarray([vecp[1] - args.topo_range, vecp[1] + args.topo_range])
0123     th_rg = np.asarray([vecp[0] - args.topo_range, vecp[0] + args.topo_range])
0124     eta_rg = np.resize(-np.log(np.tan(vecp[0]/1000./2.)), 2) + np.asarray([-args.topo_range, args.topo_range])/1000.
0125
0126     os.makedirs(args.outdir, exist_ok=True)
0127     # cluster plot by layers (rebinned)
0128     pdf = PdfPages(os.path.join(args.outdir, 'e{}_c{}_layers.pdf'.format(args.iev, args.icl)))
0129     bpos, bprops = (0.5, 0.95), dict(boxstyle='round', facecolor='wheat', alpha=0.2)
0130     frames = []
0131     for i in np.arange(1, df['layer'].max() + 1, dtype=int):
0132         data = df[df['layer'] == i]
0133         fig, axs = plt.subplots(1, 2, figsize=(17, 16), dpi=160, gridspec_kw={'wspace':0., 'width_ratios': [16, 1]})
0134         ax, sm = draw_heatmap(axs[0], data['eta'].values, data['phi'].values, weights=data['edep'].values,
0135                               bins=(np.arange(*eta_rg, step=args.topo_size/1000.), np.arange(*phi_rg, step=args.topo_size)),
0136                               cmap=cmap, vmin=0.0, vmax=1.0, pc_kw=dict(alpha=0.8, edgecolor='k'))
0137         ax.set_ylabel(r'$\phi$ (mrad)', fontsize=28)
0138         ax.set_xlabel(r'$\eta$', fontsize=28)
0139         ax.tick_params(labelsize=24)
0140         ax.xaxis.set_minor_locator(MultipleLocator(5))
0141         ax.yaxis.set_minor_locator(MultipleLocator(5))
0142         ax.grid(linestyle=':', which='both')
0143         ax.set_axisbelow(True)
0144         ax.text(*bpos, 'Layer {:d}'.format(i), transform=ax.transAxes,
0145                 fontsize=26, va='top', ha='center', bbox=bprops)
0146         cb = plt.colorbar(sm, cax=axs[1], shrink=0.85)
0147         cb.ax.tick_params(labelsize=24)
0148         cb.ax.get_yaxis().labelpad = 24
0149         cb.ax.set_ylabel('Energy Deposit (MeV)', rotation=90, fontsize=28)
0150         pdf.savefig(fig)
0151         # gif frames
0152         fig.savefig('ltmp.png')
0153         plt.close(fig)
0154         frames.append(imageio.imread('ltmp.png'))
0155     pdf.close()
0156     os.remove('ltmp.png')
0157
0158     # build GIF
0159     imageio.mimsave(os.path.join(args.outdir, 'e{:d}_c{:d}_layers.gif'.format(args.iev, args.icl)),
0160                     frames, 'GIF', duration=args.dura)
0161
0162
0163     # cluster plot by layers (scatters)
0164     pdf = PdfPages(os.path.join(args.outdir, 'e{}_c{}_layers2.pdf'.format(args.iev, args.icl)))
0165     bpos, bprops = (0.5, 0.95), dict(boxstyle='round', facecolor='wheat', alpha=0.2)
0166     frames = []
0167     for i in np.arange(1, df['layer'].max() + 1, dtype=int):
0168         data = df[df['layer'] == i]
0169         fig, axs = plt.subplots(1, 2, figsize=(17, 16), dpi=160, gridspec_kw={'wspace':0., 'width_ratios': [16, 1]})
0170         ax = axs[0]
0171         colors = cmap(data['edep']/1.0)
0172         ax.scatter(data['eta'].values, data['phi'].values, c=colors, marker='s', s=15.0)
0173         ax.set_xlim(*eta_rg)
0174         ax.set_ylim(*phi_rg)
0175         ax.set_ylabel(r'$\phi$ (mrad)', fontsize=28)
0176         ax.set_xlabel(r'$\eta$', fontsize=28)
0177         ax.tick_params(labelsize=24)
0178         ax.xaxis.set_minor_locator(MultipleLocator(5))
0179         ax.yaxis.set_minor_locator(MultipleLocator(5))
0180         ax.grid(linestyle=':', which='both')
0181         ax.set_axisbelow(True)
0182         ax.text(*bpos, 'Layer {:d}'.format(i), transform=ax.transAxes,
0183                 fontsize=26, va='top', ha='center', bbox=bprops)
0184         cb = plt.colorbar(sm, cax=axs[1], shrink=0.85)
0185         cb.ax.tick_params(labelsize=24)
0186         cb.ax.get_yaxis().labelpad = 24
0187         cb.ax.set_ylabel('Energy Deposit (MeV)', rotation=90, fontsize=28)
0188         pdf.savefig(fig)
0189         # gif frames
0190         fig.savefig('ltmp.png')
0191         plt.close(fig)
0192         frames.append(imageio.imread('ltmp.png'))
0193     pdf.close()
0194     os.remove('ltmp.png')
0195
0196     # build GIF
0197     imageio.mimsave(os.path.join(args.outdir, 'e{:d}_c{:d}_layers2.gif'.format(args.iev, args.icl)),
0198                     frames, 'GIF', duration=args.dura)
0199