Calorimetery/ClusterVisualizer/draw_clusters.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='RecoEcalBarrelImagingHits', 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 #        'EcalBarrel_rmin',
0072 #        'EcalBarrel_ReadoutLayerThickness',
0073 #        'EcalBarrel_ReadoutLayerNumber',
0074 #        'EcalBarrel_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, 'MCParticles').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         dfc = df[(df['eta'] <= df['eta'].quantile(0.95)) & (df['eta'] >= df['eta'].quantile(0.05)) &
0112                  (df['phi'] <= df['phi'].quantile(0.95)) & (df['phi'] >= df['phi'].quantile(0.05))]
0113         vec = np.average(dfc[['x', 'y', 'z']].values, axis=0, weights= dfc['energy'].values if np.sum(dfc['energy'].values) > 0 else None)
0114     vec = vec/np.linalg.norm(vec)
0115
0116     # particle line from (0, 0, 0) to the inner Ecal surface
0117     length = rmin/np.sqrt(vec[0]**2 + vec[1]**2)
0118     pline = np.transpose(vec*np.mgrid[0:length:2j][:, np.newaxis])
0119     cmap = truncate_colormap(plt.get_cmap('jet'), 0.1, 0.9)
0120
0121     # convert truth to mrad
0122     vecp = np.asarray([np.arccos(vec[2]), np.arctan2(vec[1], vec[0])])*1000.
0123     phi_rg = np.asarray([vecp[1] - args.topo_range, vecp[1] + args.topo_range])
0124     th_rg = np.asarray([vecp[0] - args.topo_range, vecp[0] + args.topo_range])
0125     eta_rg = np.resize(-np.log(np.tan(vecp[0]/1000./2.)), 2) + np.asarray([-args.topo_range, args.topo_range])/1000.
0126
0127     os.makedirs(args.outdir, exist_ok=True)
0128
0129     # cluster plot by layers (rebinned grids)
0130     pdf = PdfPages(os.path.join(args.outdir, 'e{}_c{}_layers.pdf'.format(args.iev, args.icl)))
0131     bpos, bprops = (0.5, 0.95), dict(boxstyle='round', facecolor='wheat', alpha=0.2)
0132     frames = []
0133     for i in np.arange(1, df['layer'].max() + 1, dtype=int):
0134         data = df[df['layer'] == i]
0135         fig, axs = plt.subplots(1, 2, figsize=(17, 16), dpi=160, gridspec_kw={'wspace':0., 'width_ratios': [16, 1]})
0136         ax, sm = draw_heatmap(axs[0], data['eta'].values, data['phi'].values, weights=data['energy'].values,
0137                               bins=(np.arange(*eta_rg, step=args.topo_size/1000.), np.arange(*phi_rg, step=args.topo_size)),
0138                               cmap=cmap, vmin=0.0, vmax=1.0, pc_kw=dict(alpha=0.8, edgecolor='k'))
0139         ax.set_ylabel(r'$\phi$ (mrad)', fontsize=28)
0140         ax.set_xlabel(r'$\eta$', fontsize=28)
0141         ax.tick_params(labelsize=24)
0142         ax.xaxis.set_minor_locator(MultipleLocator(5))
0143         ax.yaxis.set_minor_locator(MultipleLocator(5))
0144         ax.grid(linestyle=':', which='both')
0145         ax.set_axisbelow(True)
0146         ax.text(*bpos, 'Layer {:d}'.format(i), transform=ax.transAxes,
0147                 fontsize=26, va='top', ha='center', bbox=bprops)
0148         cb = plt.colorbar(sm, cax=axs[1], shrink=0.85)
0149         cb.ax.tick_params(labelsize=24)
0150         cb.ax.get_yaxis().labelpad = 24
0151         cb.ax.set_ylabel('Energy Deposit (MeV)', rotation=90, fontsize=28)
0152         pdf.savefig(fig)
0153         # gif frames
0154         fig.savefig('ltmp.png')
0155         plt.close(fig)
0156         frames.append(imageio.v2.imread('ltmp.png'))
0157     pdf.close()
0158     os.remove('ltmp.png')
0159
0160     # build GIF
0161     imageio.mimsave(os.path.join(args.outdir, 'e{:d}_c{:d}_layers.gif'.format(args.iev, args.icl)),
0162                     frames, 'GIF', duration=args.dura)
0163
0164
0165     # cluster plot by layers (scatters)
0166     pdf = PdfPages(os.path.join(args.outdir, 'e{}_c{}_layers2.pdf'.format(args.iev, args.icl)))
0167     bpos, bprops = (0.5, 0.95), dict(boxstyle='round', facecolor='wheat', alpha=0.2)
0168     frames = []
0169     for i in np.arange(1, df['layer'].max() + 1, dtype=int):
0170         data = df[df['layer'] == i]
0171         fig, axs = plt.subplots(1, 2, figsize=(17, 16), dpi=160, gridspec_kw={'wspace':0., 'width_ratios': [16, 1]})
0172         ax = axs[0]
0173         colors = cmap(data['energy']/1.0)
0174         ax.scatter(data['eta'].values, data['phi'].values, c=colors, marker='s', s=15.0)
0175         ax.set_xlim(*eta_rg)
0176         ax.set_ylim(*phi_rg)
0177         ax.set_ylabel(r'$\phi$ (mrad)', fontsize=28)
0178         ax.set_xlabel(r'$\eta$', fontsize=28)
0179         ax.tick_params(labelsize=24)
0180         ax.xaxis.set_minor_locator(MultipleLocator(5))
0181         ax.yaxis.set_minor_locator(MultipleLocator(5))
0182         ax.grid(linestyle=':', which='both')
0183         ax.set_axisbelow(True)
0184         ax.text(*bpos, 'Layer {:d}'.format(i), transform=ax.transAxes,
0185                 fontsize=26, va='top', ha='center', bbox=bprops)
0186         cb = plt.colorbar(sm, cax=axs[1], shrink=0.85)
0187         cb.ax.tick_params(labelsize=24)
0188         cb.ax.get_yaxis().labelpad = 24
0189         cb.ax.set_ylabel('Energy Deposit (MeV)', rotation=90, fontsize=28)
0190         pdf.savefig(fig)
0191         # gif frames
0192         fig.savefig('ltmp.png')
0193         plt.close(fig)
0194         frames.append(imageio.v2.imread('ltmp.png'))
0195     pdf.close()
0196     os.remove('ltmp.png')
0197
0198     # build GIF
0199     imageio.mimsave(os.path.join(args.outdir, 'e{:d}_c{:d}_layers2.gif'.format(args.iev, args.icl)),
0200                     frames, 'GIF', duration=args.dura)
0201