Belle II Software development
plots.py
1#!/usr/bin/python
2
3
10
11"""\
12Utilities for plotting
13"""
14from pxd.utils import getH1Sigma68WithError
15import basf2 as b2
16import ROOT
17import pandas as pd
18import numpy as np
19import matplotlib.pyplot as plt
20plt.style.use('belle2')
21
22
23def root_style():
24 """
25 ROOT style
26 """
27 # root style option
28
29 # canvas
30 ROOT.gStyle.SetCanvasBorderMode(0)
31 ROOT.gStyle.SetCanvasBorderSize(10)
32 ROOT.gStyle.SetCanvasColor(0)
33 ROOT.gStyle.SetCanvasDefH(450)
34 ROOT.gStyle.SetCanvasDefW(500)
35 ROOT.gStyle.SetCanvasDefX(10)
36 ROOT.gStyle.SetCanvasDefY(10)
37 # pad
38 ROOT.gStyle.SetPadBorderMode(0)
39 ROOT.gStyle.SetPadBorderSize(10)
40 ROOT.gStyle.SetPadColor(0)
41 ROOT.gStyle.SetPadBottomMargin(0.16)
42 ROOT.gStyle.SetPadTopMargin(0.10)
43 ROOT.gStyle.SetPadLeftMargin(0.17)
44 ROOT.gStyle.SetPadRightMargin(0.19)
45 ROOT.gStyle.SetPadGridX(1)
46 ROOT.gStyle.SetPadGridY(1)
47 ROOT.gStyle.SetGridStyle(2)
48 ROOT.gStyle.SetGridColor(ROOT.kGray + 1)
49 # frame
50 ROOT.gStyle.SetFrameFillStyle(0)
51 ROOT.gStyle.SetFrameFillColor(0)
52 ROOT.gStyle.SetFrameLineColor(1)
53 ROOT.gStyle.SetFrameLineStyle(0)
54 ROOT.gStyle.SetFrameLineWidth(2)
55 ROOT.gStyle.SetFrameBorderMode(0)
56 ROOT.gStyle.SetFrameBorderSize(10)
57 # axis
58 ROOT.gStyle.SetLabelFont(42, "xyz")
59 ROOT.gStyle.SetLabelOffset(0.015, "xyz")
60 ROOT.gStyle.SetLabelSize(0.048, "xyz")
61 ROOT.gStyle.SetNdivisions(505, "xyz")
62 ROOT.gStyle.SetTitleFont(42, "xyz")
63 ROOT.gStyle.SetTitleSize(0.050, "xyz")
64 ROOT.gStyle.SetTitleOffset(1.3, "x")
65 ROOT.gStyle.SetTitleOffset(1.2, "y")
66 ROOT.gStyle.SetTitleOffset(1.4, "z")
67 ROOT.gStyle.SetPadTickX(1)
68 ROOT.gStyle.SetPadTickY(1)
69 # title
70 ROOT.gStyle.SetTitleBorderSize(0)
71 ROOT.gStyle.SetTitleFillColor(10)
72 ROOT.gStyle.SetTitleAlign(12)
73 ROOT.gStyle.SetTitleFontSize(0.045)
74 ROOT.gStyle.SetTitleX(0.560)
75 ROOT.gStyle.SetTitleY(0.860)
76 ROOT.gStyle.SetTitleFont(42, "")
77 # stat
78 ROOT.gStyle.SetStatBorderSize(0)
79 ROOT.gStyle.SetStatColor(10)
80 ROOT.gStyle.SetStatFont(42)
81 ROOT.gStyle.SetStatX(0.94)
82 ROOT.gStyle.SetStatY(0.91)
83 # histo style
84 ROOT.gStyle.SetOptTitle(0)
85 ROOT.gStyle.SetOptStat(0)
86 ROOT.gStyle.SetHistLineWidth(3)
87 ROOT.TH2.SetDefaultSumw2()
88
89
90def df_plot_errorbar(df, x, y, yerr_low, yerr_up, ax=None, *args, **kwargs):
91 """
92 errorbar extension for pandas.DataFrame
93 Parameters:
94 df: pandas.DataFrame
95 x (str): column used as the x-axis
96 y (str): column used as the y-axis
97 yerr_low (str): column of the lower error of value y
98 yerr_up (str): column of the upper error of value y
99
100 ax: matplotlib.pyplot.axis
101 args: positional arguments
102 kwargs: named arguments
103 """
104 if ax is None:
105 ax = plt.gca()
106 array_x = df[x].to_numpy()
107 array_y = df[y].to_numpy()
108 yerr = [df[yerr_low].to_numpy(), df[yerr_up].to_numpy()]
109 default_kwargs = {"fmt": "o", "label": y, "alpha": 0.7}
110 kwargs = {**default_kwargs, **kwargs}
111
112 ax.errorbar(x=array_x, y=array_y, yerr=yerr, *args, **kwargs)
113 ax.legend()
114 ax.set_xlabel(x)
115 ax.set_ylabel(y)
116 return ax
117
118
119# Extend pandas.DataFrame
120if hasattr(pd.DataFrame, "errorbar"):
121 b2.B2WARNING("DataFrame.errorbar has already been defined. Replacing it!")
122pd.DataFrame.errorbar = df_plot_errorbar
123
124
125def plot_ip_resolutions(tree, cut="0<z0_err<1 & 0<d0_err<1", figsize=(16, 6), save_to=""):
126 """
127 Helper function to plot ip resolutions as a function of run number.
128 Parameters:
129 tree: a ROOT.TTree which contains run, hD0, hZ0 from PXD calibration validation.
130 figsize (tuple): figsize for matplotlib pyplot
131 save_to (str): Target file name to save a figure when it's not empty.
132 Return:
133 pandas.DataFrame with resolutions
134 """
135 plt.figure(figsize=figsize)
136 ip_res_dic = {"run": [], "d0": [], "d0_err": [], "z0": [], "z0_err": []}
137 for i in range(tree.GetEntries()):
138 tree.GetEntry(i)
139 # exp = getattr(tree, "exp")
140 run = getattr(tree, "run")
141 hD0 = getattr(tree, "hD0")
142 hZ0 = getattr(tree, "hZ0")
143 res_d0, res_d0_err = np.array(getH1Sigma68WithError(hD0, 1000)) / 2 * 1e4 # cm to um
144 res_z0, res_z0_err = np.array(getH1Sigma68WithError(hZ0, 1000)) / 2 * 1e4 # cm to um
145 # print(run,res_d0, res_z0)
146 ip_res_dic["run"].append(run)
147 ip_res_dic["d0"].append(res_d0)
148 ip_res_dic["d0_err"].append(res_d0_err)
149 ip_res_dic["z0"].append(res_z0)
150 ip_res_dic["z0_err"].append(res_z0_err)
151 df_res = pd.DataFrame(ip_res_dic)
152 df_res.query(cut).errorbar(y="z0", x="run", yerr_low="z0_err", yerr_up="z0_err")
153 df_res.query(cut).errorbar(y="d0", x="run", yerr_low="d0_err", yerr_up="d0_err")
154 plt.ylabel(r"$\sigma_{68}$ " + r"[$\mathrm{\mu m}$]")
155 plt.savefig(save_to)
156 return df_res
157
158
159def plot_efficiency_map(tree, exp_runs=[], num_name="hTrackClustersLayer1", den_name="hTrackPointsLayer1",
160 title="Efficiency of layer 1", canvas=None, save_to=""):
161 """
162 Helper function to plot efficiency map of modules for a given layer.
163 Color map can be set by ROOT.gStyle.SetPalette(ROOT.kDarkRainBow)
164 Parameters:
165 tree: ROOT.TTree
166 exp_runs (list): list of (experiment_number, run_number) or list of run_number
167 num_name (str): histogram name of the ROOT.TH2 used for numerator
168 den_name (str): histogram name of the ROOT.TH2 used for denominator
169 title (str): title of the plot
170 canvas (ROOT.TCanvas): TCanvas for plotting
171 save_to (str): File name to save the figure
172 Return:
173 the efficiency map (ROOT.TH2)
174 """
175 if canvas is None:
176 canvas = ROOT.TCanvas("c1", "c1", 900, 600)
177 use_exp_run_tuple = True
178 if not exp_runs:
179 return
180 elif not isinstance(exp_runs[0], tuple):
181 exp_runs = sorted(exp_runs)
182 use_exp_run_tuple = False
183 count = 0
184 h_eff, h_num, h_den = None, None, None
185 for i_evt in range(tree.GetEntries()):
186 tree.GetEntry(i_evt)
187 exp = getattr(tree, "exp")
188 run = getattr(tree, "run")
189 if not use_exp_run_tuple:
190 if run < exp_runs[0] or run > exp_runs[-1]:
191 continue
192 elif (exp, run) not in exp_runs:
193 continue
194 if count == 0:
195 h_num = getattr(tree, num_name).Clone()
196 h_den = getattr(tree, den_name).Clone()
197 else:
198 h_num.Add(getattr(tree, num_name))
199 h_den.Add(getattr(tree, den_name))
200 count += 1
201 if count:
202 h_temp = h_den.Clone()
203 h_temp.Divide(h_den) # h_temp bins filled to 1 if there is any counts in h_den
204 h_eff = h_num.Clone()
205 h_eff.Add(h_temp, 1e-9) # Added 1e-9 which shouldn't bias eff
206 h_eff.Divide(h_den)
207 h_eff.SetTitle(title)
208 h_eff.SetStats(True)
209 # default_opt_stat = ROOT.gStyle.GetOptStat()
210 ROOT.gStyle.SetOptStat(10)
211 h_eff.Draw("colz")
212 canvas.Draw()
213 s = h_eff.GetListOfFunctions().FindObject("stats")
214 # print(s.GetX1NDC(), s.GetX2NDC())
215 s.SetX1NDC(0.7)
216 s.SetX2NDC(0.9)
217 s.SetY1NDC(0.9)
218 s.SetY2NDC(0.95)
219 s.SetFillColorAlpha(0, 0.1)
220 s.SetLineColorAlpha(0, 0)
221 canvas.Update()
222 if save_to:
223 h_eff.GetZaxis().SetRangeUser(0.9, 1)
224 canvas.Update()
225 canvas.Print(save_to+".above90.png")
226 h_eff.GetZaxis().SetRangeUser(0, 1)
227 canvas.Update()
228 canvas.Print(save_to)
229
230 # ROOT.gStyle.SetOptStat(default_opt_stat)
231 return h_eff, h_num, h_den
232
233
234def plot_in_module_efficiency(df, pxdid=1052, figsize=(12, 16), alpha=0.7, save_to="",
235 y="eff_sel", x="run", cut="eff_sel>0&eff_sel_err_low<0.01",
236 yerr_low="eff_sel_err_low", yerr_up="eff_sel_err_up"):
237 """
238 Helper function to plot efficiencies of 4 x 6 ASIC regions in one module
239 Parameters:
240 df (pandas.DataFrame): pandas DataFrame for plotting
241 pxdid (int): pxdid, e.g., 1052
242 figsize (tuple): figsize for matplotlib pyplot
243 alpha (float): alpha color
244 save_to (str): File name to save the figure
245 y (str): column of the efficiency
246 x (str): column of the run
247 cut (str): additional cut
248 yerr_low (str): column of the lower error
249 yerr_up (str): column of the upper error
250 """
251
252 fig, axes = plt.subplots(4, 1, sharex=False, sharey=True, figsize=figsize)
253 for uBin in range(4):
254 ax = axes[uBin]
255 for vBin in range(6):
256 df.query(cut + "&" + f"pxdid=={pxdid}&uBin=={uBin}&vBin=={vBin}").errorbar(
257 y=y, x=x, yerr_low=yerr_low, yerr_up=yerr_up, ax=ax, label=f"{pxdid},u{uBin}v{vBin}", alpha=0.7)
258 ax.legend(bbox_to_anchor=(1, 1), loc="upper left")
259 ymin, ymax = ax.get_ylim()
260 ax.set_ylim(ymin, 1.)
261 if save_to:
262 fig.savefig(save_to, bbox_inches="tight")
263
264
265def plot_efficiency_vs_run(
266 df,
267 eff_var="eff",
268 eff_sel_var="eff_sel",
269 max_err=0.01,
270 figsize=(
271 12,
272 6),
273 save_to="pxd_efficiency_vs_run.png"):
274 """
275 Helper function to plot efficiency vs run
276 Parameters:
277 df (pandas.DataFrame): pandas DataFrame for plotting
278 eff_var (str): column of the efficiency
279 eff_sel_var (str): column of the efficiency of selected regions
280 max_err (float): the maximum error used to clean up points with large errors
281 figsize (tuple): figsize for matplotlib pyplot
282 save_to (str): file name to save the figure
283 """
284
285 plt.figure(figsize=figsize)
286 df.query(f"{eff_var}_err_low<{max_err}").errorbar(
287 y=eff_var,
288 x="run",
289 yerr_low=f"{eff_var}_err_low",
290 yerr_up=f"{eff_var}_err_up",
291 alpha=0.7,
292 label="All regions")
293 df.query(f"{eff_sel_var}_err_low<{max_err}").errorbar(
294 y=eff_sel_var,
295 x="run",
296 yerr_low=f"{eff_sel_var}_err_low",
297 yerr_up=f"{eff_sel_var}_err_up",
298 alpha=0.7,
299 label="Excluding hot/dead regions")
300 plt.ylabel("PXD efficiency")
301 ymin, ymax = plt.ylim()
302 plt.ylim(ymin, 1.0)
303 if save_to:
304 plt.savefig(save_to)
305
306
307def plot_module_efficiencies_in_DHHs(df, eff_var="eff_sel", phase="early_phase3", figsize=(12, 6), save_to="efficiency_vs_run.png"):
308 """
309 Helper function to plot efficiencies of modules in a DHH.
310 Parameters:
311 df (pandas.DataFrame): pandas DataFrame for plotting
312 eff_var (str): column of the efficiency
313 phase (str): belle2 phase, the default is earyly_phase3 with half PXD
314 figsize (tuple): figsize for matplotlib pyplot
315 save_to (str): file name to save the figure
316 """
317
318 dhh_modules_dic = {}
319 if phase == "early_phase3":
320 dhh_modules_dic = {
321 "H30": [1012, 1022, 1052, 1082, 2042],
322 "H40": [1042, 1062, 1072, 2052],
323 "H50": [1041, 1051, 1061, 1071, 2051],
324 "H60": [1011, 1021, 1031, 1081, 2041],
325 }
326 for dhh, pxdid_list in dhh_modules_dic.items():
327 plt.figure(figsize=figsize)
328 ymin = 1.0
329 for pxdid in pxdid_list:
330 df.query(f"{eff_var}>0&pxdid=={pxdid}&{eff_var}_err_low<0.01").errorbar(
331 y=eff_var, x="run", yerr_low=f"{eff_var}_err_low", yerr_up=f"{eff_var}_err_up", label=f"{pxdid}", alpha=0.7)
332 ymin = min(plt.ylim()[0], ymin)
333 plt.ylabel("Efficiency (selected)")
334 plt.title(dhh + " efficiency")
335 plt.ylim(ymin, 1.0)
336 if save_to:
337 plt.savefig(dhh + "_" + save_to)
338
339
340if __name__ == '__main__':
341 root_style()