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