File indexing completed on 2026-09-01 09:34:06
0001 """Reproduce the single uncalibrated CALOROC energy-response plot."""
0002
0003 import argparse
0004
0005 import matplotlib.pyplot as plt
0006 import numpy as np
0007 from scipy.optimize import curve_fit
0008
0009 from script.common import (
0010 LOAD_CHUNK_SIZE,
0011 SliceResult,
0012 applyPlotStyle,
0013 boundsCrystalBallPositive,
0014 buildSliceHistogramFromChunks,
0015 crystal_ball,
0016 curveFitReducer,
0017 loadData,
0018 maskPositiveAboveThreshold,
0019 nERecoBins,
0020 nETruthBins,
0021 p0CrystalBall,
0022 plotThetaDistribution,
0023 reduceSlices,
0024 )
0025
0026
0027 applyPlotStyle()
0028
0029
0030 def linear(x, a, b):
0031 return a * x + b
0032
0033
0034 def main(chunk_iter, output, title, e_truth_min, e_truth_max, e_reco_min,
0035 e_reco_max, txt=None, ytitle=''):
0036 """Accumulate response histograms, fit each truth-energy slice, and plot."""
0037 e_truth_bins = np.linspace(e_truth_min, e_truth_max, nETruthBins)
0038 e_reco_bins = np.linspace(e_reco_min, e_reco_max, nERecoBins)
0039 theta_bins = np.linspace(0, 180, 180)
0040
0041 hist = np.zeros((len(e_truth_bins) - 1, len(e_reco_bins) - 1))
0042 theta_counts = np.zeros(len(theta_bins) - 1)
0043 for chunk in chunk_iter:
0044 h, _, _ = np.histogram2d(
0045 chunk['ETruth'], chunk['EReco'], bins=(e_truth_bins, e_reco_bins)
0046 )
0047 hist += h
0048 tc, _ = np.histogram(chunk['theta'], bins=theta_bins)
0049 theta_counts += tc
0050
0051 err = hist / np.sqrt(hist)
0052 xc = 0.5 * (e_truth_bins[1:] + e_truth_bins[:-1])
0053 yc = 0.5 * (e_reco_bins[1:] + e_reco_bins[:-1])
0054 mean, _, std, _ = reduceSlices(
0055 xc, yc, hist, err,
0056 reducer=curveFitReducer(
0057 crystal_ball, p0CrystalBall, boundsCrystalBallPositive
0058 ),
0059 sliceMask=maskPositiveAboveThreshold(0.05, 0.4),
0060 )
0061
0062 x, y = np.meshgrid(e_truth_bins, e_reco_bins)
0063 fig, (ax, ax2) = plt.subplots(ncols=2, figsize=(14, 7))
0064 pc = ax.pcolormesh(x, y, hist.T, cmap='magma_r')
0065 ax.errorbar(xc, mean, yerr=np.abs(std), marker='o')
0066 try:
0067 finite = np.isfinite(mean) & np.isfinite(std)
0068 popt, _ = curve_fit(
0069 linear, xc[finite], mean[finite], sigma=np.abs(std[finite]),
0070 absolute_sigma=True, p0=[1, 0]
0071 )
0072 fine_bins = np.linspace(e_truth_bins[0], e_truth_bins[-1], 100)
0073 ax.plot(
0074 fine_bins, linear(fine_bins, *popt),
0075 label='y = %.3e x + %.3e' % (popt[0], popt[1])
0076 )
0077 except Exception:
0078 pass
0079 ax.legend(frameon=False)
0080
0081 if txt is not None:
0082 np.savetxt(txt, np.column_stack((xc, mean, np.abs(std))), fmt='%.6f')
0083 ax.set_ylim(e_reco_min, e_reco_max)
0084 try:
0085 fig.colorbar(pc, ax=ax)
0086 except Exception:
0087 pass
0088 ax.set_xlabel(r'$E_{truth}$ (GeV)')
0089 ax.set_ylabel(ytitle)
0090 plotThetaDistribution(ax2, theta_counts, theta_bins)
0091 fig.suptitle(title)
0092 plt.tight_layout(rect=[0, 0, 1, 0.95])
0093 plt.savefig(output)
0094
0095
0096 if __name__ == '__main__':
0097 parser = argparse.ArgumentParser()
0098 parser.add_argument('output')
0099 parser.add_argument('data')
0100 parser.add_argument('pdg', type=int)
0101 parser.add_argument('title')
0102 parser.add_argument('--eTruthMin', type=float, required=True)
0103 parser.add_argument('--eTruthMax', type=float, required=True)
0104 parser.add_argument('--eRecoMin', type=float, required=True)
0105 parser.add_argument('--eRecoMax', type=float, required=True)
0106 parser.add_argument('--ytitle')
0107 parser.add_argument('--txt', default=None)
0108 args = parser.parse_args()
0109 columns = [
0110 'mcId', 'PDG', 'theta', 'EReco', 'ETruth', 'PhiReco', 'EtaReco',
0111 'PhiTruth', 'EtaTruth', 'status', 'nevent'
0112 ]
0113 main(
0114 loadData(args.data, args.pdg, columns), args.output, args.title,
0115 args.eTruthMin, args.eTruthMax, args.eRecoMin, args.eRecoMax,
0116 args.txt, args.ytitle
0117 )
