File indexing completed on 2025-01-18 09:12:12
0001
0002 from pathlib import Path
0003 from typing import List
0004 import re
0005
0006 import uproot
0007 import typer
0008 import numpy
0009 import matplotlib.pyplot
0010
0011
0012 def main(files: List[Path], output: str, title: str = ""):
0013 mus = []
0014
0015 for file in files:
0016 m = re.match(".*mu(\d+).*", file.name)
0017 mu = int(m.group(1))
0018 mus.append(mu)
0019
0020 mus = list(sorted(mus))
0021
0022 fig, (ax, tax) = matplotlib.pyplot.subplots(
0023 2, 1, gridspec_kw={"hspace": 0, "height_ratios": [2, 1]}
0024 )
0025
0026 for fitter, color in zip(("Iterative", "AMVF"), ["tab:blue", "tab:orange"]):
0027 mean = numpy.array([])
0028 stddev = numpy.array([])
0029 ntrue_mean = numpy.array([])
0030 ntrue_stddev = numpy.array([])
0031 time = numpy.array([])
0032
0033 for mu in mus:
0034 for file in files:
0035 if f"mu{mu}" in file.name and fitter in file.name:
0036 break
0037
0038 time_file = file.parent / f"{file.stem}_time.txt"
0039 if time_file.exists():
0040 time = numpy.append(time, float(time_file.read_text()))
0041 else:
0042 time.append(float("nan"))
0043
0044 rf = uproot.open(f"{file}:vertexing")
0045
0046 nreco = rf["nRecoVtx"].array(library="np")
0047 ntrue = rf["nTrueVtx"].array(library="np")
0048
0049 nreco_mean = nreco.mean()
0050
0051 nreco_std = nreco.std()
0052
0053 mean = numpy.append(mean, nreco_mean)
0054 stddev = numpy.append(stddev, nreco_std)
0055
0056 ntrue_mean = numpy.append(ntrue_mean, ntrue.mean())
0057 ntrue_stddev = numpy.append(ntrue_stddev, ntrue.std())
0058
0059 ax.fill_between(mus, mean - stddev, mean + stddev, fc=color, alpha=0.2)
0060 ax.plot(mus, mean, label=fitter, c=color)
0061 tax.plot(mus, time, label=f"{fitter} time", c=color)
0062
0063 ax.plot(mus, mus, label="truth", c="gray", ls="--")
0064
0065 tax.set_xlabel("$\mu$")
0066 ax.set_ylabel("Number of vertices")
0067 tax.set_ylabel("wall time [s]")
0068
0069 fig.align_labels()
0070
0071 ax.legend()
0072 tax.legend()
0073 ax.set_title(title)
0074
0075 fig.tight_layout()
0076 fig.savefig(output)
0077
0078
0079 if __name__ == "__main__":
0080 typer.run(main)
