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 extention 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. Return:
132 pandas.DataFrame with resolutions
133 """
134 plt.figure(figsize=figsize)
135 ip_res_dic = {"run": [], "d0": [], "d0_err": [], "z0": [], "z0_err": []}
136 for i in range(tree.GetEntries()):
137 tree.GetEntry(i)
138 # exp = getattr(tree, "exp")
139 run = getattr(tree, "run")
140 hD0 = getattr(tree, "hD0")
141 hZ0 = getattr(tree, "hZ0")
142 res_d0, res_d0_err = np.array(getH1Sigma68WithError(hD0, 1000)) / 2 * 1e4 # cm to um
143 res_z0, res_z0_err = np.array(getH1Sigma68WithError(hZ0, 1000)) / 2 * 1e4 # cm to um
144 # print(run,res_d0, res_z0)
145 ip_res_dic["run"].append(run)
146 ip_res_dic["d0"].append(res_d0)
147 ip_res_dic["d0_err"].append(res_d0_err)
148 ip_res_dic["z0"].append(res_z0)
149 ip_res_dic["z0_err"].append(res_z0_err)
150 df_res = pd.DataFrame(ip_res_dic)
151 df_res.query(cut).errorbar(y="z0", x="run", yerr_low="z0_err", yerr_up="z0_err")
152 df_res.query(cut).errorbar(y="d0", x="run", yerr_low="d0_err", yerr_up="d0_err")
153 plt.ylabel(r"$\sigma_{68}$ " + r"[$\mathrm{\mu m}$]")
154 plt.savefig(save_to)
155 return df_res
156
157
158def plot_efficiency_map(tree, exp_runs=[], num_name="hTrackClustersLayer1", den_name="hTrackPointsLayer1",
159 title="Efficiency of layer 1", canvas=None, save_to=""):
160 """
161 Helper function to plot efficiency map of modules for a given layer.
162 Color map can be set by ROOT.gStyle.SetPalette(ROOT.kDarkRainBow)
163 Parameters:
164 tree: ROOT.TTree
165 exp_runs (list): list of (experiment_number, run_number) or list of run_number
166 num_name (str): histogram name of the ROOT.TH2 used for numerator
167 den_name (str): histogram name of the ROOT.TH2 used for denominator
168 title (str): title of the plot
169 canvas (ROOT.TCanvas): TCanvas for plotting
170 save_to (str): File name to save the figure
171 Return:
172 the efficiency map (ROOT.TH2)
173 """
174 if canvas is None:
175 canvas = ROOT.TCanvas("c1", "c1", 900, 600)
176 use_exp_run_tuple = True
177 if not exp_runs:
178 return
179 elif not isinstance(exp_runs[0], tuple):
180 exp_runs = sorted(exp_runs)
181 use_exp_run_tuple = False
182 count = 0
183 h_eff, h_num, h_den = None, None, None
184 for i_evt in range(tree.GetEntries()):
185 tree.GetEntry(i_evt)
186 exp = getattr(tree, "exp")
187 run = getattr(tree, "run")
188 if not use_exp_run_tuple:
189 if run < exp_runs[0] or run > exp_runs[-1]:
190 continue
191 elif (exp, run) not in exp_runs:
192 continue
193 if count == 0:
194 h_num = getattr(tree, num_name).Clone()
195 h_den = getattr(tree, den_name).Clone()
196 else:
197 h_num.Add(getattr(tree, num_name))
198 h_den.Add(getattr(tree, den_name))
199 count += 1
200 if count:
201 h_temp = h_den.Clone()
202 h_temp.Divide(h_den) # h_temp bins filled to 1 if there is any counts in h_den
203 h_eff = h_num.Clone()
204 h_eff.Add(h_temp, 1e-9) # Added 1e-9 which shouldn't bias eff
205 h_eff.Divide(h_den)
206 h_eff.SetTitle(title)
207 h_eff.SetStats(True)
208 # default_opt_stat = ROOT.gStyle.GetOptStat()
209 ROOT.gStyle.SetOptStat(10)
210 h_eff.Draw("colz")
211 canvas.Draw()
212 s = h_eff.GetListOfFunctions().FindObject("stats")
213 # print(s.GetX1NDC(), s.GetX2NDC())
214 s.SetX1NDC(0.7)
215 s.SetX2NDC(0.9)
216 s.SetY1NDC(0.9)
217 s.SetY2NDC(0.95)
218 s.SetFillColorAlpha(0, 0.1)
219 s.SetLineColorAlpha(0, 0)
220 canvas.Update()
221 if save_to:
222 h_eff.GetZaxis().SetRangeUser(0.9, 1)
223 canvas.Update()
224 canvas.Print(save_to+".above90.png")
225 h_eff.GetZaxis().SetRangeUser(0, 1)
226 canvas.Update()
227 canvas.Print(save_to)
228
229 # ROOT.gStyle.SetOptStat(default_opt_stat)
230 return h_eff, h_num, h_den
231
232
233def plot_in_module_efficiency(df, pxdid=1052, figsize=(12, 16), alpha=0.7, save_to="",
234 y="eff_sel", x="run", cut="eff_sel>0&eff_sel_err_low<0.01",
235 yerr_low="eff_sel_err_low", yerr_up="eff_sel_err_up"):
236 """
237 Helper function to plot efficiencies of 4 x 6 ASIC regions in one module
238 Parameters:
239 df (pandas.DataFrame): pandas DataFrame for plotting
240 pxdid (int): pxdid, e.g., 1052
241 figsize (tuple): figsize for matplotlib pyplot
242 alpha (float): alpha color
243 save_to (str): File name to save the figure
244 y (str): column of the efficiency
245 x (str): column of the run
246 cut (str): additional cut
247 yerr_low (str): column of the lower error
248 yerr_up (str): column of the upper error
249 """
250
251 fig, axes = plt.subplots(4, 1, sharex=False, sharey=True, figsize=figsize)
252 for uBin in range(4):
253 ax = axes[uBin]
254 for vBin in range(6):
255 df.query(cut + "&" + f"pxdid=={pxdid}&uBin=={uBin}&vBin=={vBin}").errorbar(
256 y=y, x=x, yerr_low=yerr_low, yerr_up=yerr_up, ax=ax, label=f"{pxdid},u{uBin}v{vBin}", alpha=0.7)
257 ax.legend(bbox_to_anchor=(1, 1), loc="upper left")
258 ymin, ymax = ax.get_ylim()
259 ax.set_ylim(ymin, 1.)
260 if save_to:
261 fig.savefig(save_to, bbox_inches="tight")
262
263
264def plot_efficiency_vs_run(
265 df,
266 eff_var="eff",
267 eff_sel_var="eff_sel",
268 max_err=0.01,
269 figsize=(
270 12,
271 6),
272 save_to="pxd_efficiency_vs_run.png"):
273 """
274 Helper function to plot efficiency vs run
275 Parameters:
276 df (pandas.DataFrame): pandas DataFrame for plotting
277 eff_var (str): column of the efficiency
278 eff_sel_var (str): column of the efficiency of selected regions
279 max_err (float): the maximum error used to clean up points with large errors
280 figsize (tuple): figsize for matplotlib pyplot
281 save_to (str): file name to save the figure
282 """
283
284 plt.figure(figsize=figsize)
285 df.query(f"{eff_var}_err_low<{max_err}").errorbar(
286 y=eff_var,
287 x="run",
288 yerr_low=f"{eff_var}_err_low",
289 yerr_up=f"{eff_var}_err_up",
290 alpha=0.7,
291 label="All regions")
292 df.query(f"{eff_sel_var}_err_low<{max_err}").errorbar(
293 y=eff_sel_var,
294 x="run",
295 yerr_low=f"{eff_sel_var}_err_low",
296 yerr_up=f"{eff_sel_var}_err_up",
297 alpha=0.7,
298 label="Excluding hot/dead regions")
299 plt.ylabel("PXD efficiency")
300 ymin, ymax = plt.ylim()
301 plt.ylim(ymin, 1.0)
302 if save_to:
303 plt.savefig(save_to)
304
305
306def plot_module_efficiencies_in_DHHs(df, eff_var="eff_sel", phase="early_phase3", figsize=(12, 6), save_to="efficiency_vs_run.png"):
307 """
308 Helper function to plot efficiencies of modules in a DHH.
309 Parameters:
310 df (pandas.DataFrame): pandas DataFrame for plotting
311 eff_var (str): column of the efficiency
312 phase (str): belle2 phase, the default is earyly_phase3 with half PXD
313 figsize (tuple): figsize for matplotlib pyplot
314 save_to (str): file name to save the figure
315 """
316
317 dhh_modules_dic = {}
318 if phase == "early_phase3":
319 dhh_modules_dic = {
320 "H30": [1012, 1022, 1052, 1082, 2042],
321 "H40": [1042, 1062, 1072, 2052],
322 "H50": [1041, 1051, 1061, 1071, 2051],
323 "H60": [1011, 1021, 1031, 1081, 2041],
324 }
325 for dhh, pxdid_list in dhh_modules_dic.items():
326 plt.figure(figsize=figsize)
327 ymin = 1.0
328 for pxdid in pxdid_list:
329 df.query(f"{eff_var}>0&pxdid=={pxdid}&{eff_var}_err_low<0.01").errorbar(
330 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)
331 ymin = min(plt.ylim()[0], ymin)
332 plt.ylabel("Efficiency (selected)")
333 plt.title(dhh + " efficiency")
334 plt.ylim(ymin, 1.0)
335 if save_to:
336 plt.savefig(dhh + "_" + save_to)
337
338
339if __name__ == '__main__':
340 root_style()
341