Scripts/MaterialMapping/material_comparison.py

0001 import matplotlib.pyplot as plt
0002 import numpy as np
0003 import uproot
0004 import math
0005 from collections import namedtuple
0006
0007 FileRecord = namedtuple("FileRecord", ["name", "tree", "label", "color", "marker_size"])
0008 PlotRecord = namedtuple(
0009     "PlotRecord", ["x_axis", "y_axis", "x_range", "x_bins", "x_label", "saveAs"]
0010 )
0011
0012 # The file records
0013 fileRecords = [
0014     FileRecord("geant4_material_tracks.root", "material-tracks", "Geant4", "blue", 3),
0015     FileRecord("acts_material_tracks.root", "material-tracks", "Acts", "orange", 4),
0016 ]
0017
0018
0019 # The plot records
0020 plotRecords = [
0021     PlotRecord("v_eta", "t_X0", (-4.0, 4.0), 80, "η", "tX0_vs_eta.svg"),
0022     PlotRecord("v_phi", "t_X0", (-math.pi, math.pi), 72, "φ", "tX0_vs_phi.svg"),
0023 ]
0024
0025 # Different plot records
0026 for pr in plotRecords:
0027
0028     fig, axs = plt.subplots(2, 1, sharex=True)
0029     fig.subplots_adjust(hspace=0.05)
0030
0031     # Prepare limit & ratios
0032     y_lim = 0
0033     y_ratio_values = []
0034     y_ratio_errors = [0.0 for i in range(pr.x_bins)]
0035
0036     # Loop over the file records
0037     for ifr, fr in enumerate(fileRecords):
0038
0039         # Load the three
0040         tree = uproot.open(fr.name + ":" + fr.tree)
0041
0042         x_arr = tree[pr.x_axis].array(library="np")
0043         y_arr = tree[pr.y_axis].array(library="np")
0044         y_max = y_arr.max()
0045         y_lim = y_max if y_max > y_lim else y_lim
0046
0047         # Generate the central bin values
0048         x_step = (pr.x_range[1] - pr.x_range[0]) / pr.x_bins
0049         x_vals = [pr.x_range[0] + (ix + 0.5) * x_step for ix in range(pr.x_bins)]
0050
0051         # Prepare the min /max
0052         y_min_vals = [1000.0] * pr.x_bins
0053         y_max_vals = [0.0] * pr.x_bins
0054         y_vals_sorted = [np.array([])] * pr.x_bins
0055
0056         for iv in range(len(x_arr)):
0057             x_b = int((x_arr[iv] - pr.x_range[0]) / x_step)
0058             y_v = y_arr[iv]
0059             # Take min / max
0060             y_min_vals[x_b] = y_v if y_v < y_min_vals[x_b] else y_min_vals[x_b]
0061             y_max_vals[x_b] = y_v if y_v > y_max_vals[x_b] else y_max_vals[x_b]
0062             # Regulate the x value
0063             y_vals_sorted[x_b] = np.append(y_vals_sorted[x_b], y_v)
0064
0065         axs[0].fill_between(
0066             x=x_vals,
0067             y1=y_min_vals,
0068             y2=y_max_vals,
0069             alpha=0.1,
0070             label=fr.label + " spread",
0071             color=fr.color,
0072         )
0073         axs[0].grid(axis="x")
0074         y_vals_mean = [y_bin_vals.mean() for y_bin_vals in y_vals_sorted]
0075
0076         y_ratio_values += [y_vals_mean]
0077         y_vals_mse = [
0078             y_bin_vals.std() ** 2 / len(y_bin_vals) for y_bin_vals in y_vals_sorted
0079         ]
0080
0081         axs[0].errorbar(
0082             x=x_vals,
0083             y=y_vals_mean,
0084             yerr=y_vals_mse,
0085             markersize=fr.marker_size,
0086             marker="o",
0087             mfc=fr.color if ifr == 0 else "none",
0088             linestyle="none",
0089             label=fr.label + " mean",
0090             color=fr.color,
0091         )
0092
0093         if ifr > 0:
0094             y_ratios = [
0095                 y_ratio_values[ifr][ib] / y_ratio_values[0][ib]
0096                 for ib in range(pr.x_bins)
0097             ]
0098             axs[1].errorbar(
0099                 x=x_vals,
0100                 y=y_ratios,
0101                 yerr=y_ratio_errors,
0102                 markersize=fr.marker_size,
0103                 marker="o",
0104                 mfc="none",
0105                 linestyle="none",
0106                 color=fr.color,
0107                 label=fr.label,
0108             )
0109             axs[1].set_ylabel("Ratio to " + fileRecords[0].label)
0110
0111     # Some final cosmetics
0112     axs[0].set_ylim(0.0, y_lim)
0113     axs[0].grid(axis="x", linestyle="dotted")
0114
0115     axs[1].set_ylim(0.9, 1.1)
0116     axs[1].grid(axis="x", linestyle="dotted")
0117     axs[1].axhline(y=1.0, color="black", linestyle="-")
0118
0119     axs[0].legend(loc="upper center")
0120     axs[1].legend(loc="upper center")
0121
0122     # Set the range of x-axis
0123     plt.xlabel(pr.x_label)
0124     fig.savefig(pr.saveAs)
0125
0126
0127 plt.show()