scripts/imaging_calorimeter/draw_cluster.py

0001 '''
0002     A script to visualize the cluster
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 utils import *
0022 import sys
0023
0024
0025 # draw cluster in a 3d axis, expect a numpy array of (nhits, 4) shape with each row contains (x, y, z, E)
0026 # note z and x axes are switched
0027 def draw_hits3d(axis, data, cmap, units=('mm', 'mm', 'mm', 'MeV'), fontsize=24, **kwargs):
0028     # normalize energy to get colors
0029     x, y, z, edep = np.transpose(data)
0030     cvals = edep - min(edep) / (max(edep) - min(edep))
0031     cvals[np.isnan(cvals)] = 1.0
0032     colors = cmap(cvals)
0033
0034     # hits
0035     axis.scatter(z, y, x, c=colors, marker='o', **kwargs)
0036     axis.tick_params(labelsize=fontsize)
0037     axis.set_zlabel('x ({})'.format(units[2]), fontsize=fontsize + 2, labelpad=fontsize)
0038     axis.set_ylabel('y ({})'.format(units[1]), fontsize=fontsize + 2, labelpad=fontsize)
0039     axis.set_xlabel('z ({})'.format(units[0]), fontsize=fontsize + 2, labelpad=fontsize)
0040     cb = plt.colorbar(cm.ScalarMappable(norm=matplotlib.colors.Normalize(vmin=min(edep), vmax=max(edep)), cmap=cmap),
0041                       ax=axis, shrink=0.85)
0042     cb.ax.tick_params(labelsize=fontsize)
0043     cb.ax.get_yaxis().labelpad = fontsize
0044     cb.ax.set_ylabel('Energy Deposit ({})'.format(units[3]), rotation=90, fontsize=fontsize + 4)
0045     return axis
0046
0047
0048 # draw a cylinder in 3d axes
0049 # note z and x axes are switched
0050 def draw_cylinder3d(axis, r, z, order=['x', 'y', 'z'], rsteps=500, zsteps=500, **kwargs):
0051     x = np.linspace(-r, r, rsteps)
0052     z = np.linspace(-z, z, zsteps)
0053     Xc, Zc = np.meshgrid(x, z)
0054     Yc = np.sqrt(r**2 - Xc**2)
0055
0056     axis.plot_surface(Zc, Yc, Xc, alpha=0.1, **kwargs)
0057     axis.plot_surface(Zc, -Yc, Xc, alpha=0.1, **kwargs)
0058     return axis
0059
0060
0061 # fit the track of cluster and draw the fit
0062 def draw_track_fit(axis, dfh, length=200, stop_layer=8, scat_kw=dict(), line_kw=dict()):
0063     dfh = dfh[dfh['layer'] <= stop_layer]
0064     data = dfh.groupby('layer')[['z', 'y','x']].mean().values
0065     # data = dfh[['z', 'y', 'x']].values
0066     # ax.scatter(*data.T, **scat_kw)
0067     datamean = data.mean(axis=0)
0068     uu, dd, vv = np.linalg.svd(data - datamean)
0069     linepts = vv[0] * np.mgrid[-length:length:2j][:, np.newaxis]
0070     linepts += datamean
0071     axis.plot3D(*linepts.T, 'k:')
0072     return axis
0073
0074
0075 # color map
0076 def draw_heatmap(axis, x, y, weights, bins=1000, cmin=0., cmap=plt.get_cmap('rainbow'), pc_kw=dict()):
0077     w, xedg, yedg = np.histogram2d(x, y, weights=weights, bins=bins)
0078     xsz = np.mean(np.diff(xedg))
0079     ysz = np.mean(np.diff(yedg))
0080     wmin, wmax = w.min(), w.max()
0081     recs, clrs = [], []
0082     for i in np.arange(len(xedg) - 1):
0083         for j in np.arange(len(yedg) - 1):
0084             if w[i][j] > cmin:
0085                 recs.append(Rectangle((xedg[i], yedg[j]), xsz, ysz))
0086                 clrs.append(cmap((w[i][j] - wmin) / (wmax - wmin)))
0087     axis.add_collection(PatchCollection(recs, facecolor=clrs, **pc_kw))
0088     axis.set_xlim(xedg[0], xedg[-1])
0089     axis.set_ylim(yedg[0], yedg[-1])
0090     return axis, cm.ScalarMappable(norm=matplotlib.colors.Normalize(vmin=wmin, vmax=wmax), cmap=cmap)
0091
0092
0093 # execute this script
0094 if __name__ == '__main__':
0095     parser = argparse.ArgumentParser(description='Visualize the cluster from analysis')
0096     parser.add_argument('file', type=str, help='path to root file')
0097     parser.add_argument('-e', type=int, default=0, dest='iev', help='event number to plot')
0098     parser.add_argument('-c', type=int, default=0, dest='icl', help='cluster number to plot (0: all, -1: no cluster)')
0099     parser.add_argument('-s', type=int, default=8, dest='stop', help='stop layer for track fit')
0100     parser.add_argument('-o', type=str, default='./plots', dest='outdir', help='output directory')
0101     parser.add_argument('--compact', type=str, default='', dest='compact', help='compact file')
0102     parser.add_argument('-m', '--macros', type=str, default='rootlogon.C', dest='macros',
0103                         help='root macros to load (accept multiple paths separated by \",\")')
0104     parser.add_argument('-b', '--branch-name', type=str, default='RecoEcalBarrelHits', dest='branch',
0105                         help='branch name in the root file (outputLayerCollection from ImagingClusterReco)')
0106     parser.add_argument('--topo-size', type=float, default=2.0, dest='topo_size',
0107                         help='bin size for projection plot (mrad)')
0108     parser.add_argument('--topo-range', type=float, default=50.0, dest='topo_range',
0109                         help='half range for projection plot (mrad)')
0110     parser.add_argument('--zoom-factor', type=float, default=1.0, dest='zoom_factor',
0111                         help='factor for zoom-in')
0112     args = parser.parse_args()
0113
0114
0115     # we can read these values from xml file
0116     desc = compact_constants(args.compact, [
0117         'cb_ECal_RMin',
0118         'cb_ECal_ReadoutLayerThickness',
0119         'cb_ECal_ReadoutLayerNumber',
0120         'cb_ECal_Length'
0121     ])
0122     if not len(desc):
0123         # or define Ecal shapes
0124         rmin, thickness, length = 890, 20*(10. + 1.65), 860*2+500
0125     else:
0126         # convert cm to mm
0127         rmin = desc[0]*10.
0128         thickness = desc[1]*desc[2]*10.
0129         length = desc[3]*10.
0130
0131
0132     # read data
0133     load_root_macros(args.macros)
0134     df = get_hits_data(args.file, args.iev, branch=args.branch)
0135     if args.icl != 0:
0136         df = df[df['cluster'] == args.icl]
0137     if not len(df):
0138         print("Error: do not find any hits for cluster {:d} in event {:d}".format(args.icl, args.iev))
0139         exit(-1)
0140     # convert to polar coordinates (mrad), and stack all r values
0141     df['r'] = np.sqrt(df['x'].values**2 + df['y'].values**2 + df['z'].values**2)
0142     df['phi'] = np.arctan2(df['y'].values, df['x'].values)*1000.
0143     df['theta'] = np.arccos(df['z'].values/df['r'].values)*1000.
0144     df['eta'] = -np.log(np.tan(df['theta'].values/1000./2.))
0145
0146     # truth
0147     dfmcp = get_mcp_simple(args.file, args.iev, 'mcparticles2').iloc[0]
0148     pdgbase = ROOT.TDatabasePDG()
0149     inpart = pdgbase.GetParticle(int(dfmcp['pid']))
0150     print("Incoming particle = {}, pdgcode = {}, charge = {}, mass = {}"\
0151           .format(inpart.GetName(), inpart.PdgCode(), inpart.Charge(), inpart.Mass()))
0152     # neutral particle, no need to consider magnetic field
0153     if np.isclose(inpart.Charge(), 0., rtol=1e-5):
0154         vec = dfmcp[['px', 'py', 'pz']].values
0155     # charge particle, use the cluster center
0156     else:
0157         flayer = df[df['layer'] == df['layer'].min()]
0158         vec = flayer[['x', 'y', 'z']].mean().values
0159     vec = vec/np.linalg.norm(vec)
0160
0161     # particle line from (0, 0, 0) to the inner Ecal surface
0162     length = rmin/np.sqrt(vec[0]**2 + vec[1]**2)
0163     pline = np.transpose(vec*np.mgrid[0:length:2j][:, np.newaxis])
0164     cmap = truncate_colormap(plt.get_cmap('jet'), 0.1, 0.9)
0165
0166     os.makedirs(args.outdir, exist_ok=True)
0167     # cluster plot
0168     fig = plt.figure(figsize=(15, 12), dpi=160)
0169     ax = fig.add_subplot(111, projection='3d')
0170     # draw particle line
0171     ax.plot(*pline[[2, 1]], '--', zs=pline[0], color='green')
0172     # draw hits
0173     draw_hits3d(ax, df[['x', 'y', 'z', 'edep']].values, cmap, s=5.0)
0174     draw_cylinder3d(ax, rmin, length, rstride=10, cstride=10, color='royalblue')
0175     draw_cylinder3d(ax, rmin + thickness, length, rstride=10, cstride=10, color='forestgreen')
0176     ax.set_zlim(-(rmin + thickness), rmin + thickness)
0177     ax.set_ylim(-(rmin + thickness), rmin + thickness)
0178     ax.set_xlim(-length, length)
0179     fig.tight_layout()
0180     fig.savefig(os.path.join(args.outdir, 'e{}_cluster.png'.format(args.iev)))
0181
0182
0183     # zoomed-in cluster plot
0184     fig = plt.figure(figsize=(15, 12), dpi=160)
0185     ax = fig.add_subplot(111, projection='3d')
0186     # draw particle line
0187     ax.plot(*pline[[2, 1]], '--', zs=pline[0], color='green')
0188     # draw hits
0189     draw_hits3d(ax, df[['x', 'y', 'z', 'edep']].values, cmap, s=20.0)
0190     # draw_track_fit(ax, df, stop_layer=args.stop,
0191     #                scat_kw=dict(color='k', s=50.0), line_kw=dict(linestyle=':', color='k', lw=3))
0192     # view range
0193     center = (length + thickness/2.)*vec
0194     ranges = np.vstack([center - thickness/args.zoom_factor, center + thickness/args.zoom_factor]).T
0195     ax.set_zlim(*ranges[0])
0196     ax.set_ylim(*ranges[1])
0197     ax.set_xlim(*ranges[2])
0198
0199     fig.tight_layout()
0200     fig.savefig(os.path.join(args.outdir, 'e{}_cluster_zoom.png'.format(args.iev)))
0201
0202
0203     # projection plot
0204     # convert to mrad
0205     vecp = np.asarray([np.arccos(vec[2]), np.arctan2(vec[1], vec[0])])*1000.
0206     phi_rg = np.asarray([vecp[1] - args.topo_range, vecp[1] + args.topo_range])
0207     th_rg = np.asarray([vecp[0] - args.topo_range, vecp[0] + args.topo_range])
0208     eta_rg = np.resize(-np.log(np.tan(vecp[0]/1000./2.)), 2) + np.asarray([-args.topo_range, args.topo_range])/1000.
0209
0210     fig, axs = plt.subplots(1, 2, figsize=(13, 12), dpi=160, gridspec_kw={'wspace':0., 'width_ratios': [12, 1]})
0211     ax, sm = draw_heatmap(axs[0], df['eta'].values, df['phi'].values, weights=df['edep'].values,
0212                           bins=(np.arange(*eta_rg, step=args.topo_size/1000.), np.arange(*phi_rg, step=args.topo_size)),
0213                           cmap=cmap, cmin=0., pc_kw=dict(alpha=0.8, edgecolor='k'))
0214
0215     ax.set_ylabel(r'$\phi$ (mrad)', fontsize=32)
0216     ax.set_xlabel(r'$\eta$', fontsize=32)
0217     ax.tick_params(labelsize=28)
0218     ax.xaxis.set_minor_locator(MultipleLocator(5))
0219     ax.yaxis.set_minor_locator(MultipleLocator(5))
0220     ax.grid(linestyle=':', which='both')
0221     ax.set_axisbelow(True)
0222     cb = plt.colorbar(sm, cax=axs[1], shrink=0.85, aspect=1.2*20)
0223     cb.ax.tick_params(labelsize=28)
0224     cb.ax.get_yaxis().labelpad = 10
0225     cb.ax.set_ylabel('Energy Deposit (MeV)', rotation=90, fontsize=32)
0226     fig.savefig(os.path.join(args.outdir, 'e{}_topo.png'.format(args.iev)))
0227