Belle II Software  release-05-01-25
plotPXDPositionEstimator.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 """
5 Created on Wed Jan 31 20:54:00 2018
6 Plot PXD position estimator payload
7 @author: benni
8 """
9 
10 import numpy as np
11 import matplotlib.pyplot as plt
12 import pandas as pd
13 import seaborn as sns
14 
15 from basf2 import *
16 from ROOT import Belle2
17 import ROOT
18 
19 
20 class PlotClusterPositionEstimatorPayload(Module):
21  """Plot the PXDClusterPositionEstimator playload """
22 
23  def __init__(self, resultdir):
24  """Initialize"""
25  super().__init__() # don't forget to call parent constructor
26 
27  self.position_estimator = Belle2.PyDBObj('PXDClusterPositionEstimatorPar')
28 
29  self.resultdir = resultdir
30 
31  self.counter = 0
32 
33  def event(self):
34  """Plot position payload in case it has changed"""
35 
36  # We plot the payload whenever it changes
37  if self.position_estimator.hasChanged():
38  B2INFO("PXDClusterPositionEstimator payload has changed. Plot the new payload.")
39  self.plot()
40  self.counter += 1
41 
42  def average_covariance(self, shape_classifier):
43  """Compute the average covariance for a shape classifier"""
44  flat_covs = []
45  weights = []
46 
47  offsetMap = shape_classifier.getOffsetMap()
48  likelyhoodMap = shape_classifier.getLikelyhoodMap()
49 
50  for item in offsetMap:
51  shape_index = item.first
52  offsets = item.second
53 
54  for eta_index, offset in enumerate(offsets):
55  likelyhood = likelyhoodMap[shape_index][eta_index]
56  flat_covs.append([offset.getUSigma2(), offset.getVSigma2(), offset.getUVCovariance()])
57  weights.append(likelyhood)
58 
59  weights = np.array(weights)
60  flat_covs = np.array(flat_covs)
61  flat_average = np.average(flat_covs, axis=0, weights=weights)
62  return np.array([[flat_average[0], flat_average[2]], [flat_average[2], flat_average[1]]])
63 
64  def plot(self):
65  """Plot position estimator payload"""
66  for pair in self.position_estimator.getGridMap():
67  pixelkind = pair.first
68  grid = pair.second
69 
70  # Keep some numbers to gain overview
71  summary_dict = {'shapes': [],
72  'corrections': [],
73  'coverage': [],
74  'sigma_u': [],
75  'sigma_v': [],
76  'corr_uv': [],
77  'thetaU': [],
78  'thetaV': [],
79  'pixelkind': [],
80  }
81 
82  # Loop over angular grid and plot shape classifiers
83  for uBin in range(1, grid.GetXaxis().GetNbins() + 1):
84  for vBin in range(1, grid.GetYaxis().GetNbins() + 1):
85 
86  # Shape classifier for angle bin
87  shape_classifier = self.position_estimator.getShapeClassifier(uBin, vBin, pixelkind)
88 
89  # Bin is centered around angles
90  thetaU = grid.GetXaxis().GetBinCenter(uBin)
91  thetaV = grid.GetYaxis().GetBinCenter(vBin)
92 
93  # Create plots for classifier
94  stats = self.get_classifier_stats(shape_classifier, pixelkind)
95 
96  # Fill summary dict
97  summary_dict['thetaU'].append(thetaU)
98  summary_dict['thetaV'].append(thetaV)
99  summary_dict['pixelkind'].append(pixelkind)
100  summary_dict['shapes'].append(stats['shapes'])
101  summary_dict['corrections'].append(stats['corrections'])
102  summary_dict['coverage'].append(stats['coverage'])
103  summary_dict['sigma_u'].append(stats['sigma_u'] * 10000) # cm -> um
104  summary_dict['sigma_v'].append(stats['sigma_v'] * 10000) # cm -> um
105  summary_dict['corr_uv'].append(stats['corr_uv'])
106 
107  # Create heatmaps to gain overview
108  df = pd.DataFrame(summary_dict)
109 
110  pivot_table = df.pivot(index='thetaU', columns='thetaV', values='corrections')
111  fig = plt.figure(figsize=(12, 12))
112  ax = fig.add_subplot(111)
113  ax.set_xlabel('thetaU / degree', size=20)
114  ax.set_ylabel('thetaV / degree', size=20)
115  ax.set_title('Number of corrections kind={:d}'.format(pixelkind), size=20)
116  ax = sns.heatmap(
117  pivot_table,
118  mask=pivot_table.isnull(),
119  annot=True,
120  fmt="d",
121  linewidths=.5,
122  square=True,
123  cmap='Blues_r',
124  cbar_kws={
125  'label': '#corrections'})
126  ax.invert_yaxis()
127  fig.savefig(self.resultdir + '/Corrections_Heatmap_kind_{:d}.png'.format(pixelkind), dpi=100)
128  fig.clf()
129  plt.close(fig)
130 
131  pivot_table = df.pivot(index='thetaU', columns='thetaV', values='shapes')
132  fig = plt.figure(figsize=(12, 12))
133  ax = fig.add_subplot(111)
134  ax.set_xlabel('thetaU / degree', size=20)
135  ax.set_ylabel('thetaV / degree', size=20)
136  ax.set_title('Number of shapes kind={:d}'.format(pixelkind), size=20)
137  ax = sns.heatmap(
138  pivot_table,
139  mask=pivot_table.isnull(),
140  annot=True,
141  fmt="d",
142  linewidths=.5,
143  square=True,
144  cmap='Blues_r',
145  cbar_kws={
146  'label': '#shapes'})
147  ax.invert_yaxis()
148  fig.savefig(self.resultdir + '/Shapes_Heatmap_kind_{:d}.png'.format(pixelkind), dpi=100)
149  fig.clf()
150  plt.close(fig)
151 
152  pivot_table = df.pivot(index='thetaU', columns='thetaV', values='coverage')
153  fig = plt.figure(figsize=(12, 12))
154  ax = fig.add_subplot(111)
155  ax.set_xlabel('thetaU / degree', size=20)
156  ax.set_ylabel('thetaV / degree', size=20)
157  ax.set_title('Coverage kind={:d}'.format(pixelkind), size=20)
158  ax = sns.heatmap(
159  pivot_table,
160  mask=pivot_table.isnull(),
161  annot=True,
162  fmt=".1f",
163  linewidths=.5,
164  square=True,
165  cmap='Blues_r',
166  cbar_kws={
167  'label': 'coverage / %'})
168  ax.invert_yaxis()
169  fig.savefig(self.resultdir + '/Coverage_Heatmap_kind_{:d}.png'.format(pixelkind), dpi=100)
170  fig.clf()
171  plt.close(fig)
172 
173  pivot_table = df.pivot(index='thetaU', columns='thetaV', values='sigma_u')
174  fig = plt.figure(figsize=(12, 12))
175  ax = fig.add_subplot(111)
176  ax.set_xlabel('thetaU / degree', size=20)
177  ax.set_ylabel('thetaV / degree', size=20)
178  ax.set_title('Average cluster sigma u kind={:d}'.format(pixelkind), size=20)
179  ax = sns.heatmap(
180  pivot_table,
181  mask=pivot_table.isnull(),
182  annot=True,
183  fmt=".1f",
184  linewidths=.5,
185  square=True,
186  cmap='Blues_r',
187  cbar_kws={
188  'label': 'sigma u / um'})
189  ax.invert_yaxis()
190  fig.savefig(self.resultdir + '/SigmaU_Heatmap_kind_{:d}.png'.format(pixelkind), dpi=100)
191  fig.clf()
192  plt.close(fig)
193 
194  pivot_table = df.pivot(index='thetaU', columns='thetaV', values='sigma_v')
195  fig = plt.figure(figsize=(12, 12))
196  ax = fig.add_subplot(111)
197  ax.set_xlabel('thetaU / degree', size=20)
198  ax.set_ylabel('thetaV / degree', size=20)
199  ax.set_title('Average cluster sigma v kind={:d}'.format(pixelkind), size=20)
200  ax = sns.heatmap(
201  pivot_table,
202  mask=pivot_table.isnull(),
203  annot=True,
204  fmt=".1f",
205  linewidths=.5,
206  square=True,
207  cmap='Blues_r',
208  cbar_kws={
209  'label': 'sigma v / um'})
210  ax.invert_yaxis()
211  fig.savefig(self.resultdir + '/SigmaV_Heatmap_kind_{:d}.png'.format(pixelkind), dpi=100)
212  fig.clf()
213  plt.close(fig)
214 
215  pivot_table = df.pivot(index='thetaU', columns='thetaV', values='corr_uv')
216  fig = plt.figure(figsize=(12, 12))
217  ax = fig.add_subplot(111)
218  ax.set_xlabel('thetaU / degree', size=20)
219  ax.set_ylabel('thetaV / degree', size=20)
220  ax.set_title('Average uv correlation kind={:d}'.format(pixelkind), size=20)
221  ax = sns.heatmap(
222  pivot_table,
223  mask=pivot_table.isnull(),
224  annot=True,
225  fmt=".1f",
226  linewidths=.5,
227  square=True,
228  cmap='Blues_r',
229  cbar_kws={
230  'label': 'correlation'})
231  ax.invert_yaxis()
232  fig.savefig(self.resultdir + '/CorrelationUV_Heatmap_kind_{:d}.png'.format(pixelkind), dpi=100)
233  fig.clf()
234  plt.close(fig)
235 
236  def get_classifier_stats(self, shape_classifier, pixelkind):
237  """Compute some statistics for a shape classifier"""
238  # Read corrections data
239  offsetMap = shape_classifier.getOffsetMap()
240  likelyhoodMap = shape_classifier.getLikelyhoodMap()
241  shapeLikelyhoodMap = shape_classifier.getShapeLikelyhoodMap()
242 
243  # Some counters
244  nCorrections = 0
245  nShapes = 0
246  coverage = 0.0
247 
248  for item in offsetMap:
249  shape_index = item.first
250  offsets_array = item.second
251  shape_likelyhood = shapeLikelyhoodMap[shape_index]
252 
253  nShapes += 1
254 
255  for eta_index, offset in enumerate(offsets_array):
256  coverage += likelyhoodMap[shape_index][eta_index]
257  nCorrections += 1
258 
259  if nShapes > 0:
260  # Average covariance of all hits in classifier
261  cov = self.average_covariance(shape_classifier)
262  sigma_u = np.sqrt(cov[0, 0])
263  sigma_v = np.sqrt(cov[1, 1])
264  corr_uv = cov[0, 1] / (np.sqrt(cov[0, 0]) * np.sqrt(cov[1, 1]))
265  else:
266  sigma_u = 0.0
267  sigma_v = 0.0
268  corr_uv = 0.0
269 
270  # Create some summary statistics
271  stats = {}
272  stats['shapes'] = nShapes
273  stats['corrections'] = nCorrections
274  stats['coverage'] = 100 * coverage
275  stats['sigma_u'] = sigma_u
276  stats['sigma_v'] = sigma_v
277  stats['corr_uv'] = corr_uv
278 
279  # return some statistics
280  return stats
281 
282 
283 if __name__ == "__main__":
284 
285  import argparse
286  import os
287  import shutil
288  parser = argparse.ArgumentParser(description="Plot summary of hit estimator")
289  parser.add_argument('--resultdir', default="results", type=str, help='Put all plots in this directory')
290  args = parser.parse_args()
291 
292  # Remove old stuff
293  if os.path.isdir(os.getcwd() + '/' + args.resultdir):
294  shutil.rmtree(os.getcwd() + '/' + args.resultdir)
295 
296  # Create nice clean folder for results
297  os.mkdir(os.getcwd() + '/' + args.resultdir)
298 
299  # Now let's create a path to simulate our events.
300  main = create_path()
301  main.add_module("EventInfoSetter", evtNumList=[1])
302  main.add_module(PlotClusterPositionEstimatorPayload(args.resultdir))
303  main.add_module("Progress")
304 
305  process(main)
306  print(statistics)
plotPXDPositionEstimator.PlotClusterPositionEstimatorPayload.event
def event(self)
Definition: plotPXDPositionEstimator.py:33
plotPXDPositionEstimator.PlotClusterPositionEstimatorPayload.average_covariance
def average_covariance(self, shape_classifier)
Definition: plotPXDPositionEstimator.py:42
plotPXDPositionEstimator.PlotClusterPositionEstimatorPayload.__init__
def __init__(self, resultdir)
Definition: plotPXDPositionEstimator.py:23
plotPXDPositionEstimator.PlotClusterPositionEstimatorPayload.counter
counter
Counter for number of different payloads.
Definition: plotPXDPositionEstimator.py:31
plotPXDPositionEstimator.PlotClusterPositionEstimatorPayload.position_estimator
position_estimator
Position estimator payload.
Definition: plotPXDPositionEstimator.py:27
Belle2::PyDBObj
Class to access a DBObjPtr from Python.
Definition: PyDBObj.h:50
plotPXDPositionEstimator.PlotClusterPositionEstimatorPayload.get_classifier_stats
def get_classifier_stats(self, shape_classifier, pixelkind)
Definition: plotPXDPositionEstimator.py:236
plotPXDPositionEstimator.PlotClusterPositionEstimatorPayload.plot
def plot(self)
Definition: plotPXDPositionEstimator.py:64
plotPXDPositionEstimator.PlotClusterPositionEstimatorPayload.resultdir
resultdir
Directory to put all plots.
Definition: plotPXDPositionEstimator.py:29
plotPXDPositionEstimator.PlotClusterPositionEstimatorPayload
Definition: plotPXDPositionEstimator.py:20