PerformanceStudy/SinglePion/plot_single_resol.py

0001 import os
0002 from datetime import datetime
0003 from matplotlib.backends.backend_pdf import PdfPages
0004 import matplotlib.pyplot as plt
0005
0006 _PDF_BASENAME = 'plot_single_resol'
0007 _PDF_COVER_ADDED = False
0008
0009 def _cover_title_from_globals():
0010     for key in ('tag', 'fname'):
0011         if key in globals() and globals()[key] is not None:
0012             val = str(globals()[key])
0013             return os.path.basename(val)
0014     return _PDF_BASENAME
0015
0016 def _add_pdf_cover(fname, pdf_pages):
0017     fig = plt.figure(figsize=(8.5, 11))
0018     fig.text(0.5, 0.6, 'source file: '+fname+'.root', ha='center', va='center', fontsize=16)
0019     fig.text(0.5, 0.5, datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
0020              ha='center', va='center', fontsize=12)
0021     pdf_pages.savefig(fig)
0022     plt.close(fig)
0023
0024 _PDF_PATH = f'{_PDF_BASENAME}_figures.pdf'
0025 _pdf_pages = PdfPages(_PDF_PATH)
0026
0027 def _save_new_figs(figs_before):
0028     global _PDF_COVER_ADDED
0029     if not _PDF_COVER_ADDED:
0030         _add_pdf_cover(_cover_title_from_globals(), _pdf_pages)
0031         _PDF_COVER_ADDED = True
0032     for num in plt.get_fignums():
0033         if num not in figs_before:
0034             _pdf_pages.savefig(plt.figure(num))
0035
0036 # Minimal example for single particle analysis:
0037 # * pull distributions for theta, phi, and q/p (from qOverP). Definitions follow Tracking_Performances.C (cov[5], cov[9], cov[14]).
0038 # * resolutions (theta, phi, dp/p)
0039
0040
0041 import sys
0042 from pathlib import Path
0043 import numpy as np
0044 import pandas as pd
0045 import matplotlib.pyplot as plt
0046
0047 repo_root = Path('/global/homes/s/shujie/eic_dir/worksim/scripts/modules')
0048 if str(repo_root) not in sys.path:
0049     sys.path.insert(0, str(repo_root))
0050
0051 import epic_analysis_base as ana
0052 ana.configure_analysis_environment()
0053
0054 # ## load local file
0055 # fname = "/global/cfs/cdirs/m3763/shujie/worksim/matmap/rec__mom_0.5GeV_Central_2601_full_map.root"
0056 # ana.COL_TABLE = ana.get_col_table(fname)
0057 # tree = ana.read_ur(fname, 'events')
0058
0059 ## load simulation campaign file
0060 deg_list=['3to50deg']#, '45to135deg','130to177deg']
0061 mom_list=['500MeV', '1GeV', '2GeV', '5GeV', '10GeV', '20GeV'] ## omit 100 and 200 MeV
0062 deg_name = deg_list[0]
0063 mom_name = mom_list[0]
0064 s3_path = f"EPIC/RECO/25.12.0/epic_craterlake/SINGLE/pi+/{mom_name}/{deg_name}/"
0065 fname    = f"pi+_{mom_name}_{deg_name}.0001.eicrecon.edm4eic.root"
0066
0067 ana.COL_TABLE = ana.get_col_table(fname,s3_dir=s3_path)
0068 tree = ana.read_ur(fname, 'events',s3_dir=s3_path)
0069
0070 tag=fname.split('/')[-1]
0071 tag=tag.split('.root')[0]
0072
0073 track_params = 'CentralCKFTrackParameters'
0074 assoc_name = 'CentralCKFTrackAssociations'
0075 cov_branch = f'{track_params}/{track_params}.covariance.covariance[21]'
0076
0077 def _col_id(name):
0078     return next(k for k, v in ana.COL_TABLE.items() if v == name)
0079
0080 track_col_id = _col_id('CentralCKFTracks')
0081 mc_col_id = _col_id('MCParticles')
0082
0083 params = ana.get_branch_df(tree, track_params)
0084 cov_df = ana.get_branch_df(tree, cov_branch)
0085 cov_wide = cov_df.pivot_table(index=['entry', 'subentry'], columns='subsubentry', values='values')
0086 cov_wide.columns = [f'cov_{c}' for c in cov_wide.columns]
0087 cov_wide = cov_wide.reset_index()
0088
0089 params = params.merge(cov_wide, on=['entry', 'subentry'], how='left')
0090 params = params.rename(columns={'subentry': 'track_index', 'theta': 'theta_rec', 'phi': 'phi_rec', 'qOverP': 'qOverP_rec'})
0091
0092 rec_assoc = ana.get_branch_df(tree, f'_{assoc_name}_rec')
0093 sim_assoc = ana.get_branch_df(tree, f'_{assoc_name}_sim')
0094 assoc = rec_assoc[['entry', 'subentry', 'index', 'collectionID']].rename(
0095     columns={'subentry': 'assoc_index', 'index': 'track_index', 'collectionID': 'track_colID'}
0096 )
0097 assoc['mc_index'] = sim_assoc['index']
0098 assoc['mc_colID'] = sim_assoc['collectionID']
0099 assoc = assoc[(assoc['track_colID'] == track_col_id) & (assoc['mc_colID'] == mc_col_id)]
0100
0101 df = assoc.merge(params, on=['entry', 'track_index'], how='inner')
0102
0103 mc = ana.get_branch_df(tree, 'MCParticles')
0104 mc = mc.rename(columns={'subentry': 'mc_index'})
0105 mc['p_mc'] = np.sqrt(mc['momentum.x']**2 + mc['momentum.y']**2 + mc['momentum.z']**2)
0106 mc['theta_mc'] = np.arccos(mc['momentum.z'] / mc['p_mc'])
0107 mc['phi_mc'] = np.arctan2(mc['momentum.y'], mc['momentum.x'])
0108 mc['qOverP_true'] = mc['charge'] / mc['p_mc']
0109
0110 df = df.merge(mc[['entry', 'mc_index', 'theta_mc', 'phi_mc', 'qOverP_true', 'p_mc']], on=['entry', 'mc_index'], how='inner')
0111
0112 valid = (df['cov_5'] > 0) & (df['cov_9'] > 0) & (df['cov_14'] > 0)
0113 valid &= np.isfinite(df['theta_rec']) & np.isfinite(df['phi_rec']) & np.isfinite(df['qOverP_rec'])
0114 valid &= np.isfinite(df['theta_mc']) & np.isfinite(df['phi_mc']) & np.isfinite(df['qOverP_true'])
0115 df = df[valid]
0116
0117 df['pull_theta'] = (df['theta_rec'] - df['theta_mc']) / np.sqrt(df['cov_9'])
0118 df['pull_phi'] = (df['phi_rec'] - df['phi_mc']) / np.sqrt(df['cov_5'])
0119 df['pull_qoverp'] = (np.abs(df['qOverP_rec']) - np.abs(df['qOverP_true'])) / np.sqrt(df['cov_14'])
0120
0121
0122
0123 df['resol_theta'] = (df['theta_rec'] - df['theta_mc']) *1000
0124 df['resol_phi']   = (df['phi_rec'] - df['phi_mc']) *1000
0125 df['p_rec']     = (1./np.abs(df['qOverP_rec']))
0126 df['resol_dp']   = (df['p_rec'] - df['p_mc'])/df['p_mc']*100
0127
0128 _fig_nums_before = set(plt.get_fignums())
0129 bins = np.arange(-4, 4, 0.1)
0130 eta_mc_in_rec = ana.theta2eta(df['theta_mc'])
0131 eta_mc = ana.theta2eta(mc[mc.generatorStatus == 1]['theta_mc'])
0132
0133 ha_counts, bin_edges = np.histogram(eta_mc_in_rec, bins=bins)
0134 hb_counts, _ = np.histogram(eta_mc, bins=bins)
0135
0136 bin_centers = 0.5 * (bin_edges[:-1] + bin_edges[1:])
0137 eff = np.divide(ha_counts, hb_counts, out=np.zeros_like(ha_counts, dtype=float), where=hb_counts > 0)
0138 eff_err = np.zeros_like(eff)
0139 mask = hb_counts > 0
0140 eff_err[mask] = np.sqrt(eff[mask] * (1.0 - eff[mask]) / hb_counts[mask])
0141
0142 fig, (ax_top, ax_bot) = plt.subplots(2,1, figsize=(7, 6), sharex=True, gridspec_kw={'height_ratios': [3, 2]})
0143 ax_top.hist(eta_mc, bins=bins, histtype='step', color='black', label='Gen')
0144 ax_top.hist(eta_mc_in_rec, bins=bins, histtype='step', color='tab:blue', label='Reco matched')
0145 ax_top.set_ylabel('Entries')
0146 ax_top.legend(frameon=False)
0147 ax_top.grid(True, alpha=0.3)
0148
0149 ax_bot.errorbar(bin_centers, eff, yerr=eff_err, fmt='o', ms=3, lw=1, color='tab:blue', ecolor='gray', capsize=2)
0150 ax_bot.axhline(1.0, color='gray', lw=1, ls='--')
0151 ax_bot.set_xlabel(r'$\eta$')
0152 ax_bot.set_ylabel('Efficiency')
0153 ax_bot.set_xlim(bin_edges[0], bin_edges[-1])
0154 ax_bot.set_ylim(0, 1.05)
0155 ax_bot.grid(True, alpha=0.3)
0156 plt.tight_layout()
0157 ## save to pdf
0158 _save_new_figs(_fig_nums_before)
0159
0160 _fig_nums_before = set(plt.get_fignums())
0161 import math
0162 plots = [
0163     ('pull_theta' , 'Pull distribution(theta)'),
0164     ('pull_phi'   , 'Pull distribution(phi)'),
0165     ('pull_qoverp', 'Pull distribution(q/p)'),
0166     ('resol_theta', 'Resolution (theta[mrad])'),
0167     ('resol_phi'  , 'Resolution (phi[mrad])'),
0168     ('resol_dp'   , 'Resolution (dp/p[%])'),
0169     ('loc.a'   , 'Resolution (DCA$_r$[mm])'),
0170 ]
0171
0172 cols = 3
0173 rows = math.ceil(len(plots) / cols)
0174 fig, axes = plt.subplots(rows, cols, figsize=(4 * cols, 3 * rows))
0175
0176 def _label_from_title(title):
0177     ss=title.split("(")
0178     label1=ss[0]+":"
0179     label2 = ss[1].replace(')', '')
0180     return label1, label2
0181
0182 axes_flat = axes.flatten()
0183 for ax, (col, title) in zip(axes_flat, plots):
0184     bins = 101
0185     mean, sigma, sigma_err = ana.hist_gaus(df[col], ax, bins=bins)
0186     label1,label2 = _label_from_title(title)
0187     if np.isfinite(sigma) and np.isfinite(sigma_err):
0188         label2 = f'{label2} = {sigma:.3g} +/- {sigma_err:.3g}'
0189     ax.text(
0190         0.05,
0191         0.9,
0192         label1+"\n"+label2,
0193         transform=ax.transAxes,
0194         ha='left',
0195         va='center',
0196         fontsize=12,
0197         bbox=dict(boxstyle='round,pad=0.2', facecolor='white', alpha=0.7, edgecolor='none'),
0198     )
0199     ax.set_ylabel('Entries')
0200     ax.set_xlabel('')
0201     ax.grid(True, alpha=0.3)
0202
0203 for ax in axes_flat[len(plots):]:
0204     ax.axis('off')
0205
0206 plt.tight_layout()
0207
0208 ## save to pdf
0209 _save_new_figs(_fig_nums_before)
0210
0211
0212 _pdf_pages.close()