Warning, /detector_benchmarks/benchmarks/backwards_ecal/backwards_ecal.org is written in an unsupported language. File is not indexed.
0001 #+PROPERTY: header-args:jupyter-python :session /jpy:localhost#8888:backwards_ecal :async yes :results drawer :exports both
0002
0003 #+TITLE: ePIC EEEMCal benchmark
0004 #+AUTHOR: Dmitry Kalinkin
0005 #+OPTIONS: d:t
0006
0007 #+LATEX_CLASS_OPTIONS: [9pt,letter]
0008 #+BIND: org-latex-image-default-width ""
0009 #+BIND: org-latex-image-default-option "scale=0.3"
0010 #+BIND: org-latex-images-centered nil
0011 #+BIND: org-latex-minted-options (("breaklines") ("bgcolor" "black!5") ("frame" "single"))
0012 #+LATEX_HEADER: \usepackage[margin=1in]{geometry}
0013 #+LATEX_HEADER: \setlength{\parindent}{0pt}
0014 #+LATEX: \sloppy
0015
0016 #+begin_src jupyter-python :results silent
0017 import os
0018 from pathlib import Path
0019
0020 import awkward as ak
0021 import boost_histogram as bh
0022 import dask_histogram as dh
0023 import numpy as np
0024 import uproot
0025 import vector
0026 #+end_src
0027
0028 #+begin_src jupyter-python :results silent
0029 from dask.distributed import Client
0030
0031 client = Client(os.environ.get("DASK_SCHEDULER", "localhost:8786"))
0032 #+end_src
0033
0034 #+begin_src jupyter-python
0035 vector.register_awkward()
0036
0037 from dask.distributed import WorkerPlugin
0038 class VectorImporter(WorkerPlugin):
0039 idempotent=True
0040
0041 def setup(self, worker):
0042 import vector
0043
0044 vector.register_awkward()
0045
0046 client.register_plugin(VectorImporter())
0047 #+end_src
0048
0049 * Plotting setup
0050
0051 #+begin_src jupyter-python :results silent
0052 import matplotlib as mpl
0053 import matplotlib.pyplot as plt
0054
0055 def setup_presentation_style():
0056 mpl.rcParams.update(mpl.rcParamsDefault)
0057 plt.style.use('ggplot')
0058 plt.rcParams.update({
0059 'axes.labelsize': 8,
0060 'axes.titlesize': 9,
0061 'figure.titlesize': 9,
0062 'figure.figsize': (4, 3),
0063 'legend.fontsize': 7,
0064 'xtick.labelsize': 8,
0065 'ytick.labelsize': 8,
0066 'xaxis.labellocation': 'right',
0067 'yaxis.labellocation': 'top',
0068 'pgf.rcfonts': False,
0069 })
0070
0071 setup_presentation_style()
0072 #+end_src
0073
0074 * Parameters
0075
0076 #+begin_src jupyter-python :results silent
0077 DETECTOR_CONFIG=os.environ.get("DETECTOR_CONFIG")
0078 PLOT_TITLE=os.environ.get("PLOT_TITLE")
0079 INPUT_PATH_FORMAT=os.environ.get("INPUT_PATH_FORMAT", "EPIC/RECO/24.04.0/epic_craterlake/SINGLE/{particle}/{energy}/130to177deg/{particle}_{energy}_130to177deg.{ix:04d}.eicrecon.edm4eic.root")
0080
0081 output_dir=Path(os.environ.get("OUTPUT_DIR", "./"))
0082 output_dir.mkdir(parents=True, exist_ok=True)
0083 #+end_src
0084
0085 * Analysis
0086
0087 First, we define a requirement on what phase we will consider for our
0088 analysis. The following function filters single-particle events that
0089 are thrown within $-3.5 < \eta < -2.0$:
0090
0091 #+begin_src jupyter-python
0092 def filter_pointing(events):
0093 part_momentum = ak.zip({
0094 "m": events["MCParticles.mass"],
0095 "px": events["MCParticles.momentum.x"],
0096 "py": events["MCParticles.momentum.y"],
0097 "pz": events["MCParticles.momentum.z"],
0098 }, with_name="Momentum4D")
0099 cond = (part_momentum.eta[:,0] > -3.5) & (part_momentum.eta[:,0] < -2.)
0100 return events[cond]
0101 #+end_src
0102
0103 #+begin_src jupyter-python
0104 energies = [
0105 "100MeV",
0106 "200MeV",
0107 "500MeV",
0108 "1GeV",
0109 "2GeV",
0110 "5GeV",
0111 "10GeV",
0112 "20GeV",
0113 ]
0114 filter_name = [
0115 "MCParticles.*",
0116 "EcalEndcapNClusters.energy",
0117 ]
0118
0119 pi_eval = {}
0120 e_eval = {}
0121
0122 def readlist(path):
0123 with open(path, "rt") as fp:
0124 paths = [line.rstrip() for line in fp.readlines()]
0125 return paths
0126
0127 for energy in energies:
0128 pi_eval[energy] = uproot.dask(
0129 {path: "events" for path in readlist(INPUT_PATH_FORMAT.format(particle="pi-", energy=energy))},
0130 filter_name=filter_name,
0131 steps_per_file=1,
0132 open_files=False,
0133 ).map_partitions(filter_pointing)
0134 e_eval[energy] = uproot.dask(
0135 {path: "events" for path in readlist(INPUT_PATH_FORMAT.format(particle="e-", energy=energy))},
0136 filter_name=filter_name,
0137 steps_per_file=1,
0138 open_files=False,
0139 ).map_partitions(filter_pointing)
0140 #+end_src
0141
0142 ** Energy resolution
0143
0144 The first step in the analysis is plotting distributions of cluster energies.
0145 Only electron distribution is needed, but we also do pi- for the pion rejection
0146 study.
0147
0148 #+begin_src jupyter-python
0149 e_over_p_hist = {}
0150 e_over_p_axis = bh.axis.Regular(110, 0., 1.10)
0151
0152 for ix, energy in enumerate(energies):
0153 for particle_name, dataset in [("e-", e_eval[energy]), ("pi-", pi_eval[energy])]:
0154 energy_value = float(energy.replace("GeV", "").replace("MeV", "e-3"))
0155 def clf(events):
0156 return ak.drop_none(ak.max(events["EcalEndcapNClusters.energy"], axis=-1)) / energy_value
0157 e_over_p_hist[(particle_name, energy)] = dh.factory(
0158 clf(dataset),
0159 axes=(e_over_p_axis,),
0160 )
0161
0162 e_over_p_hist = client.gather(client.compute(e_over_p_hist))
0163 #+end_src
0164
0165 The next step is to plot the histograms together with fits and produce a plot
0166 for resolution vs energy that summarizes them.
0167
0168 #+begin_src jupyter-python
0169 fig, axs = plt.subplots(2, 4, sharex=True, sharey=True, figsize=(15, 6))
0170
0171 axs = np.ravel(np.array(axs))
0172
0173 sigmas_rel_FWHM_cb = {}
0174 fractions_below = {}
0175
0176 def norm_by_sum(arr):
0177 return arr / np.sum(arr)
0178
0179 for ix, energy in enumerate(energies):
0180 plt.sca(axs[ix])
0181 hist = e_over_p_hist[("e-", energy)]
0182 patch = plt.stairs(norm_by_sum(hist.values()), hist.axes[0].edges, label=rf"$e^-$ {PLOT_TITLE}")
0183 plt.title(f"{energy}")
0184
0185 import scipy.stats
0186 def f(x, n, beta, m, loc, scale):
0187 return n * scipy.stats.crystalball.pdf(x, beta, m, loc, scale)
0188 n0 = np.sum(hist.values()[10:])
0189 loc0 = hist.axes[0].centers[5:][np.argmax(hist.values()[5:])]
0190 p0 = (n0, 2., 3., loc0, 0.05)
0191
0192 try:
0193 import scipy.optimize
0194 par, pcov = scipy.optimize.curve_fit(f, hist.axes[0].centers[5:], hist.values()[5:], p0=p0, maxfev=10000)
0195 except RuntimeError:
0196 print("Initial parameters:", p0)
0197 print(hist)
0198 raise
0199 plt.plot(hist.axes[0].centers, f(hist.axes[0].centers, *par), label=rf"Crystal Ball fit", color="tab:green", lw=0.8)
0200
0201 def summarize_fit(par):
0202 _, _, _, loc_cb, scale_cb = par
0203 # Calculate FWHM
0204 y_max = np.max(f(np.linspace(0., 1., 100), *par))
0205 f_prime = lambda x: f(x, *par) - y_max / 2
0206 x_plus, = scipy.optimize.root(f_prime, loc_cb + scale_cb).x
0207 x_minus, = scipy.optimize.root(f_prime, loc_cb - scale_cb).x
0208 plt.axvline(x_minus, ls="--", lw=0.75, color=patch.get_facecolor(), label=r"$\mu - $FWHM")
0209 plt.axvline(x_plus, ls=":", lw=0.75, color=patch.get_facecolor(), label=r"$\mu + $FWHM")
0210 fwhm = (x_plus - x_minus) / loc_cb
0211 sigma_rel_FWHM_cb = fwhm / 2 / np.sqrt(2 * np.log(2))
0212
0213 cutoff_x = loc_cb - fwhm
0214 fraction_below = np.sum(hist.values()[hist.axes[0].centers < cutoff_x]) / np.sum(hist.values())
0215
0216 return sigma_rel_FWHM_cb, fraction_below
0217
0218 sigma_rel_FWHM_cb, fraction_below = summarize_fit(par)
0219 sigmas_rel_FWHM_cb.setdefault(PLOT_TITLE, {})[energy] = sigma_rel_FWHM_cb
0220 fractions_below.setdefault(PLOT_TITLE, {})[energy] = fraction_below
0221
0222 plt.legend()
0223 plt.xlabel("$E/p$", loc="right")
0224 plt.ylabel("Event yield", loc="top")
0225
0226 fig.savefig(output_dir / f"resolution_plots.pdf", bbox_inches="tight")
0227 fig.savefig(output_dir / f"resolution_plots.png", bbox_inches="tight")
0228 plt.show()
0229 plt.close(fig)
0230
0231 plt.figure()
0232 energy_values = np.array([float(energy.replace("GeV", "").replace("MeV", "e-3")) for energy in energies])
0233
0234 for clf_label, sigma_rel_FWHM_cb in sigmas_rel_FWHM_cb.items():
0235 sigma_over_e = np.array([sigma_rel_FWHM_cb[energy] for energy in energies]) * 100 # convert to %
0236
0237 def f(energy, stochastic, constant):
0238 return np.sqrt((stochastic / np.sqrt(energy)) ** 2 + constant ** 2)
0239 cond = energy_values >= 0.5
0240 try:
0241 import scipy.optimize
0242 par, pcov = scipy.optimize.curve_fit(f, energy_values[cond], sigma_over_e[cond], maxfev=10000)
0243 except RuntimeError:
0244 print("energy_values", energy_values[cond])
0245 print("sigma_over_e", sigma_over_e[cond])
0246 raise
0247 stochastic, constant = par
0248
0249 plt.plot(
0250 energy_values,
0251 sigma_over_e,
0252 marker=".",
0253 ls="none",
0254 label=f"{clf_label}"
0255 )
0256 xmin = np.min(energy_values[cond])
0257 xmax = np.max(energy_values[cond])
0258 xs = np.arange(xmin, xmax, 0.1)
0259 plt.plot(
0260 xs,
0261 f(xs, *par),
0262 ls="--",
0263 lw=0.5,
0264 label=fr"Functional fit: ${np.ceil(stochastic * 10) / 10:.1f}\% / \sqrt{{E}} \oplus {np.ceil(constant * 10) / 10:.1f}\%$",
0265 )
0266 xmin = np.min(energy_values)
0267 xmax = np.max(energy_values) * 1.05
0268 xs = np.arange(xmin, xmax + 0.1, 0.1)
0269 plt.fill_between(
0270 xs,
0271 np.sqrt((2 / np.sqrt(xs)) ** 2 + 1 ** 2),
0272 np.sqrt((2 / np.sqrt(xs)) ** 2 + 3 ** 2),
0273 lw=0., alpha=0.2, color="black", label=r"YR requirement $2\% / \sqrt{E} \oplus (1-3)\%$",
0274 )
0275 plt.xlim(0., xmax)
0276 plt.ylim(top=6.)
0277 plt.legend()
0278 plt.xlabel("Energy, GeV", loc="right")
0279 plt.ylabel(r"$\sigma_{E} / E$ derived from FWHM, %", loc="top")
0280 plt.savefig(output_dir / f"resolution.pdf", bbox_inches="tight")
0281 plt.savefig(output_dir / f"resolution.png", bbox_inches="tight")
0282 plt.show()
0283 #+end_src
0284
0285 ** Pion rejection
0286
0287 #+begin_src jupyter-python
0288 fig, axs = plt.subplots(2, 4, sharex=True, sharey=True, figsize=(15, 6))
0289 fig_log, axs_log = plt.subplots(2, 4, sharex=True, sharey=True, figsize=(15, 6))
0290 fig_roc, axs_roc = plt.subplots(2, 4, sharex=True, sharey=True, figsize=(15, 6))
0291
0292 axs = np.ravel(np.array(axs))
0293 axs_log = np.ravel(np.array(axs_log))
0294 axs_roc = np.ravel(np.array(axs_roc))
0295
0296 rocs = {}
0297
0298 for ix, energy in enumerate(energies):
0299 energy_value = float(energy.replace("GeV", "").replace("MeV", "e-3"))
0300 def clf(events):
0301 return ak.drop_none(ak.max(events["EcalEndcapNClusters.energy"], axis=-1)) / energy_value
0302 e_pred = clf(e_eval[energy])
0303 pi_pred = clf(pi_eval[energy])
0304
0305 for do_log, ax in [(False, axs[ix]), (True, axs_log[ix])]:
0306 plt.sca(ax)
0307 patch = plt.stairs(norm_by_sum(e_over_p_hist[("e-", energy)].values()), e_over_p_hist[("e-", energy)].axes[0].edges, label=rf"$e^-$ {PLOT_TITLE}")
0308 patch = plt.stairs(norm_by_sum(e_over_p_hist[("pi-", energy)].values()), e_over_p_hist[("pi-", energy)].axes[0].edges, label=rf"$\pi^-$ {PLOT_TITLE}")
0309 plt.title(f"{energy}")
0310 plt.legend()
0311 plt.xlabel("Classifier output")
0312 plt.ylabel("Event yield")
0313 if do_log:
0314 plt.yscale("log")
0315
0316 plt.sca(axs_roc[ix])
0317 fpr = np.cumsum(e_over_p_hist[("pi-", energy)].values()[::-1])
0318 fpr /= fpr[-1]
0319 tpr = np.cumsum(e_over_p_hist[("e-", energy)].values()[::-1])
0320 tpr /= tpr[-1]
0321 cond = fpr != 0 # avoid infinite rejection (region of large uncertainty)
0322 cond &= tpr != 1 # avoid linear interpolation (region of large uncertainty)
0323 def mk_interp(tpr, fpr):
0324 def interp(eff):
0325 return np.interp(eff, tpr, fpr)
0326 return interp
0327 rocs.setdefault(PLOT_TITLE, {})[energy] = mk_interp(tpr, fpr)
0328 plt.plot(tpr[cond] * 100, 1 / fpr[cond], label=f"{PLOT_TITLE}")
0329 plt.yscale("log")
0330 plt.title(f"{energy}")
0331 plt.legend(loc="lower left")
0332 plt.xlabel("Electron efficiency, %")
0333 plt.ylabel("Pion rejection factor")
0334
0335 fig.savefig(output_dir / f"pred.pdf", bbox_inches="tight")
0336 fig.savefig(output_dir / f"pred.png", bbox_inches="tight")
0337 plt.close(fig)
0338 fig_log.savefig(output_dir / f"pred_log.pdf", bbox_inches="tight")
0339 fig_log.savefig(output_dir / f"pred_log.png", bbox_inches="tight")
0340 fig_log.show()
0341 fig_roc.savefig(output_dir / f"roc.pdf", bbox_inches="tight")
0342 fig_roc.savefig(output_dir / f"roc.png", bbox_inches="tight")
0343 fig_roc.show()
0344
0345 plt.figure()
0346 for clf_label, roc in rocs.items():
0347 plt.plot(
0348 [float(energy.replace("GeV", "").replace("MeV", "e-3")) for energy in energies],
0349 [1 / roc[energy](0.95) for energy in energies],
0350 marker=".",
0351 label=f"{clf_label}",
0352 )
0353 xmax = np.max(energy_values) * 1.05
0354 plt.xlim(0., xmax)
0355 plt.yscale("log")
0356 plt.legend()
0357 plt.xlabel("Energy, GeV")
0358 plt.ylabel("Pion rejection at 95%")
0359 plt.savefig(output_dir / f"pion_rej.pdf", bbox_inches="tight")
0360 plt.savefig(output_dir / f"pion_rej.png", bbox_inches="tight")
0361 plt.show()
0362 #+end_src