Analysis Software
Documentation for sPHENIX simulation software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
RecoPV_optimization.py
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file RecoPV_optimization.py
1 #! /usr/bin/env python
2 from optparse import OptionParser
3 import sys
4 import os
5 import datetime
6 from array import *
7 import ROOT
8 import numpy as np
9 import pandas as pd
10 import math
11 import glob
12 import ctypes
13 from plotUtil import Draw_2Dhist
14 from sigmaEff import minimum_size_range
15 
16 ROOT.gROOT.LoadMacro('./sPHENIXStyle/sPhenixStyle.C')
17 ROOT.gROOT.ProcessLine('SetsPhenixStyle()')
18 ROOT.gROOT.SetBatch(True)
19 
20 TickSize = 0.03
21 AxisTitleSize = 0.05
22 AxisLabelSize = 0.04
23 LeftMargin = 0.15
24 RightMargin = 0.08
25 TopMargin = 0.08
26 BottomMargin = 0.13
27 
28 global nbinsdphi, nbinsdca
29 global dphiinterval, dcainterval
30 
31 def Draw_1Dhist_wTF1(hist, histdata, norm1, logy, ymaxscale, XaxisName, Ytitle_unit, outname):
32  xmin, xmax = minimum_size_range(histdata, 68.5)
33  effsigma = xmax - xmin
34 
35  hist.Sumw2()
36  binwidth = hist.GetXaxis().GetBinWidth(1)
37  histmax = hist.GetMaximum()
38  c = ROOT.TCanvas('c', 'c', 800, 700)
39  if norm1:
40  hist.Scale(1. / hist.Integral(-1, -1))
41  if logy:
42  c.SetLogy()
43  c.cd()
44  ROOT.gPad.SetRightMargin(RightMargin)
45  ROOT.gPad.SetTopMargin(TopMargin)
46  ROOT.gPad.SetLeftMargin(LeftMargin)
47  ROOT.gPad.SetBottomMargin(BottomMargin)
48  if norm1:
49  if Ytitle_unit == '':
50  hist.GetYaxis().SetTitle(
51  'Normalized entries / ({:g})'.format(binwidth))
52  else:
53  hist.GetYaxis().SetTitle(
54  'Normalized entries / ({:g} {unit})'.format(binwidth, unit=Ytitle_unit))
55  else:
56  if Ytitle_unit == '':
57  hist.GetYaxis().SetTitle('Entries / ({:g})'.format(binwidth))
58  else:
59  hist.GetYaxis().SetTitle(
60  'Entries / ({:g} {unit})'.format(binwidth, unit=Ytitle_unit))
61 
62  # hist.GetXaxis().SetRangeUser(hist.GetBinLowEdge(1)-binwidth, hist.GetBinLowEdge(hist.GetNbinsX())+2*binwidth)
63  if logy:
64  hist.GetYaxis().SetRangeUser(hist.GetMinimum(0)*0.5, histmax * ymaxscale)
65  else:
66  hist.GetYaxis().SetRangeUser(0., histmax * ymaxscale)
67  hist.GetXaxis().SetTitle(XaxisName)
68  hist.GetXaxis().SetTickSize(TickSize)
69  hist.GetXaxis().SetTitleSize(AxisTitleSize)
70  hist.GetXaxis().SetLabelSize(AxisLabelSize)
71  hist.GetYaxis().SetTickSize(TickSize)
72  hist.GetYaxis().SetTitleSize(AxisTitleSize)
73  hist.GetYaxis().SetLabelSize(AxisLabelSize)
74  hist.GetXaxis().SetTitleOffset(1.1)
75  hist.GetYaxis().SetTitleOffset(1.35)
76  hist.SetLineColor(1)
77  hist.SetLineWidth(2)
78  hist.Draw('hist')
79  # Gaussian fit
80  # maxbincenter = hist.GetBinCenter(hist.GetMaximumBin())
81  # f1 = ROOT.TF1('f1', 'gaus', maxbincenter - hist.GetStdDev(), maxbincenter + hist.GetStdDev())
82  # f1.SetParameter(1,hist.GetMean())
83  # f1.SetParLimits(1,-1,1)
84  # f1.SetParameter(2,hist.GetStdDev())
85  # f1.SetParLimits(2,0,5)
86  # f1.SetLineColorAlpha(ROOT.TColor.GetColor('#F54748'), 0.9)
87  # hist.Fit(f1,'NQ')
88  # f1.Draw('same')
89  # Draw lines
90  l1 = ROOT.TLine(xmin, 0, xmin, histmax)
91  l1.SetLineColor(ROOT.kRed)
92  l1.SetLineStyle(ROOT.kDashed)
93  l1.SetLineWidth(2)
94  l1.Draw()
95  l2 = ROOT.TLine(xmax, 0, xmax, histmax)
96  l2.SetLineColor(ROOT.kRed)
97  l2.SetLineStyle(ROOT.kDashed)
98  l2.SetLineWidth(2)
99  l2.Draw()
100  leg = ROOT.TLegend((1-RightMargin)-0.45, (1-TopMargin)-0.1,
101  (1-RightMargin)-0.1, (1-TopMargin)-0.03)
102  leg.SetTextSize(0.045)
103  leg.SetFillStyle(0)
104  leg.AddEntry('', '#it{#bf{sPHENIX}} Simulation', '')
105  # leg.AddEntry('','Au+Au #sqrt{s_{NN}}=200 GeV','')
106  leg.Draw()
107  leg2 = ROOT.TLegend(0.45, 0.67, 0.88, 0.8)
108  leg2.SetTextSize(0.033)
109  leg2.SetFillStyle(0)
110  # leg2.AddEntry(f1,'Gaussian fit','l')
111  # leg2.AddEntry('', '#mu={0:.2e}#pm{1:.2e}'.format(f1.GetParameter(1), f1.GetParError(1)), '')
112  # leg2.AddEntry('', '#sigma={0:.2e}#pm{1:.2e}'.format(f1.GetParameter(2), f1.GetParError(2)), '')
113  leg2.AddEntry('', '#LT\Deltaz#GT={0:.2e}#pm{1:.2e} {2}'.format(hist.GetMean(), hist.GetMeanError(), Ytitle_unit), '')
114  # leg2.AddEntry('', '#sigma={0:.2e}#pm{1:.2e}'.format(hist.GetStdDev(), hist.GetStdDevError()), '')
115  leg2.AddEntry('', '#sigma_{{eff}} (68.5%) ={0:.4g} {1}'.format(effsigma, Ytitle_unit), '')
116  leg2.Draw()
117  c.RedrawAxis()
118  c.Draw()
119  c.SaveAs(outname+'.pdf')
120  c.SaveAs(outname+'.png')
121  if(c):
122  c.Close()
123  ROOT.gSystem.ProcessEvents()
124  del c
125  c = 0
126  # return f1.GetParameter(1), f1.GetParError(1), f1.GetParameter(2), f1.GetParError(2)
127  # return hist.GetMean(), hist.GetMeanError(), hist.GetStdDev(), hist.GetStdDevError()
128  return hist.GetMean(), hist.GetMeanError(), effsigma, 0.0
129 
130 
131 def Draw2Dhisttable(hist, XaxisName, YaxisName, ZaxisName, DrawOpt, outname):
132  c = ROOT.TCanvas('c', 'c', 4000, 3700)
133  c.cd()
134  ROOT.gPad.SetRightMargin(0.185)
135  ROOT.gPad.SetTopMargin(TopMargin)
136  ROOT.gPad.SetLeftMargin(LeftMargin)
137  ROOT.gPad.SetBottomMargin(BottomMargin)
138  ROOT.gStyle.SetPaintTextFormat('1.5f')
139  hist.SetMarkerSize(0.4)
140  hist.GetXaxis().SetTitle(XaxisName)
141  hist.GetYaxis().SetTitle(YaxisName)
142  hist.GetZaxis().SetTitle(ZaxisName)
143  hist.GetXaxis().SetTickSize(TickSize)
144  hist.GetYaxis().SetTickSize(TickSize)
145  hist.GetXaxis().SetTitleSize(AxisTitleSize)
146  hist.GetYaxis().SetTitleSize(AxisTitleSize)
147  hist.GetZaxis().SetTitleSize(AxisTitleSize)
148  hist.GetXaxis().SetLabelSize(AxisLabelSize)
149  hist.GetYaxis().SetLabelSize(AxisLabelSize)
150  hist.GetXaxis().SetTitleOffset(1.1)
151  hist.GetYaxis().SetTitleOffset(1.3)
152  hist.GetZaxis().SetTitleOffset(1.3)
153  hist.GetZaxis().SetLabelSize(AxisLabelSize)
154  hist.GetZaxis().SetRangeUser(hist.GetMinimum(), hist.GetMaximum())
155  hist.LabelsOption("v")
156  hist.SetContour(10000)
157  hist.Draw(DrawOpt)
158  # bx,by,bz = (ctypes.c_int(),ctypes.c_int(),ctypes.c_int()) # Ref: https://github.com/svenkreiss/decouple/blob/master/Decouple/plot_utils.py#L44-L52
159  # hist.GetBinXYZ(hist.GetMinimumBin(),bx,by,bz)
160  # binxCenter = hist.GetXaxis().GetBinCenter(bx.value)
161  # binyCenter = hist.GetYaxis().GetBinCenter(by.value)
162  # el = ROOT.TEllipse(binxCenter,binyCenter,.1,.1)
163  # el.SetFillColor(2)
164  # el.Draw()
165  c.RedrawAxis()
166  c.Draw()
167  c.SaveAs(outname+'.pdf')
168  c.SaveAs(outname+'.png')
169  if(c):
170  c.Close()
171  ROOT.gSystem.ProcessEvents()
172  del c
173  c = 0
174 
175 def doFitSaveFitresult(reshist, resdata, fitresname, histwTF1name):
176  print ('In doFitSaveFitresult()')
177  out = ROOT.TFile(fitresname, 'recreate')
178  tree = ROOT.TTree('tree', 'tree')
179  arrmean = array('f', [0])
180  arrmeanerr = array('f', [0])
181  arrsigma = array('f', [0])
182  arrsigmaerr = array('f', [0])
183  tree.Branch('mean', arrmean, 'mean/F')
184  tree.Branch('mean_err', arrmeanerr, 'mean_err/F')
185  tree.Branch('sigma', arrsigma, 'sigma/F')
186  tree.Branch('sigma_err', arrsigmaerr, 'sigma_err/F')
187  arrmean[0], arrmeanerr[0], arrsigma[0], arrsigmaerr[0] = Draw_1Dhist_wTF1(reshist, resdata, False, True, 50, '#Deltaz(PV_{Truth}, PV_{Reco}) [cm]', 'cm', histwTF1name)
188  tree.Fill()
189  tree.Write('', ROOT.TObject.kOverwrite)
190  out.Close()
191  del reshist, out, tree, arrmean, arrmeanerr, arrsigma, arrsigmaerr
192 
193 def main(dphitxt, dcatxt, infiledir, outfiledir):
194  print ('In main()')
195 
196  hM_DiffVtxZ = ROOT.TH1F('hM_DiffVtxZ', 'hM_DiffVtxZ_altrange', 200, -5, 5)
197  hM_DiffVtxZ_NClusLayer1 = ROOT.TH2F('hM_DiffVtxZ_NClusLayer1', 'hM_DiffVtxZ_NClusLayer1', 200, -5, 5, 200, 0, 4000)
198 
199  df = ROOT.RDataFrame('minitree', '{}/dphi{}deg_dca{}cm/INTTVtxZ.root'.format(infiledir, dphitxt, dcatxt))
200  data = df.AsNumpy(columns=['NTruthVtx','PV_z','TruthPV_z','NClusLayer1'])
201 
202  for i, ntruthvtx in enumerate(data['NTruthVtx']):
203  if ntruthvtx == 1:
204  hM_DiffVtxZ.Fill((data['PV_z'][i] - data['TruthPV_z'][i]))
205  hM_DiffVtxZ_NClusLayer1.Fill((data['PV_z'][i] - data['TruthPV_z'][i]), data['NClusLayer1'][i])
206 
207  diffarray = data['PV_z'] - data['TruthPV_z']
208  doFitSaveFitresult(hM_DiffVtxZ, diffarray, outfiledir+'fitresult_dPhi{}deg_dca{}cm.root'.format(dphitxt, dcatxt), outfiledir+'DiffVtxZ_wTF1')
209  Draw_2Dhist(hM_DiffVtxZ_NClusLayer1, False, False, False, 0.13, '#Deltaz(PV_{Truth}, PV_{Reco}) [cm]', 'Number of clusters (inner)', 'colz', outfiledir+'DiffVtxZ_NClusLayer1_dphi{}deg_dca{}cm'.format(dphitxt, dcatxt))
210 
211  del df, data, hM_DiffVtxZ, hM_DiffVtxZ_NClusLayer1
212 
213 def get2DTable(fitresdir):
214  print('In get2DTable')
215  hM_mean_dPhidcaCut = ROOT.TH2F('hM_mean_dPhidcaCut', 'hM_mean_dPhidcaCut', nbinsdphi, 0, nbinsdphi, nbinsdca, 0, nbinsdca)
216  hM_sigma_dPhidcaCut = ROOT.TH2F('hM_sigma_dPhidcaCut', 'hM_sigma_dPhidcaCut', nbinsdphi, 0, nbinsdphi, nbinsdca, 0, nbinsdca)
217 
218  l_sigma = []
219  l_mean = []
220  l_dPhi = []
221  l_dca = []
222  for i in range(0, nbinsdphi):
223  for j in range(0, nbinsdca):
224  dphicut_deg = dphiinterval * (i+1)
225  dcacut = dcainterval * (j+1)
226  dphitext = '{:.3f}'.format(dphicut_deg).replace('.', 'p')
227  dcadtext = '{:.3f}'.format(dcacut).replace('.', 'p')
228  try:
229  fitrs = ROOT.RDataFrame('tree', '{}/dPhi{}deg_dca{}cm/fitresult_dPhi{}deg_dca{}cm.root'.format(fitresdir, dphitext, dcadtext, dphitext, dcadtext))
230  res = fitrs.AsNumpy(columns=['mean','sigma'])
231  l_dPhi.append(dphicut_deg)
232  l_dca.append(dcacut)
233  l_mean.append(res['mean'][0])
234  l_sigma.append(res['sigma'][0])
235  hM_mean_dPhidcaCut.SetBinContent(i + 1, j + 1, res['mean'][0])
236  hM_sigma_dPhidcaCut.SetBinContent(i + 1, j + 1, res['sigma'][0])
237  hM_mean_dPhidcaCut.GetXaxis().SetBinLabel(i+1, '{:.2f}'.format((i+1)*dphiinterval))
238  hM_mean_dPhidcaCut.GetYaxis().SetBinLabel(j+1, '{:.3f}'.format((j+1)*dcainterval))
239  hM_sigma_dPhidcaCut.GetXaxis().SetBinLabel(i+1, '{:.2f}'.format((i+1)*dphiinterval))
240  hM_sigma_dPhidcaCut.GetYaxis().SetBinLabel(j+1, '{:.3f}'.format((j+1)*dcainterval))
241  del fitrs, res
242 
243  except Exception as e:
244  print(e)
245 
246  df_opt = pd.DataFrame(list(zip(l_dPhi, l_dca, l_mean, l_sigma)), columns =['dPhiBin', 'dZBin', 'mean', 'sigma'])
247  df_opt = df_opt.sort_values(by='sigma', ascending=True, ignore_index=True)
248 
249  del l_dPhi, l_dca, l_mean, l_sigma
250  return hM_mean_dPhidcaCut, hM_sigma_dPhidcaCut, df_opt
251 
252 if __name__ == '__main__':
253  parser = OptionParser(usage="usage: %prog ver [options -h]")
254  parser.add_option("-i", "--infiledir", dest="infiledir", default='/sphenix/user/hjheng/TrackletAna/minitree/INTT/VtxEvtMap_ana382_zvtx-20cm_dummyAlignParams', help="Input file directory")
255  parser.add_option("-o", "--outdirprefix", dest="outdirprefix", default='ana382_zvtx-20cm_dummyAlignParams', help="Output file directory")
256  parser.add_option("-f", "--dofit", action="store_true", dest="dofit", default=False, help="Do fit")
257  parser.add_option("-p", "--nbinsdphi", dest="nbinsdphi", default=20, help="Number of bins in dPhi")
258  parser.add_option("-d", "--nbinsdca", dest="nbinsdca", default=20, help="Number of bins in dca")
259 
260  (opt, args) = parser.parse_args()
261  print('opt: {}'.format(opt))
262 
263  infiledir = opt.infiledir
264  outdirprefix = opt.outdirprefix
265  dofit = opt.dofit
266  nbinsdphi = int(opt.nbinsdphi)
267  nbinsdca = int(opt.nbinsdca)
268 
269  plotpath = './RecoPV_optimization/'
270  os.makedirs('{}/{}'.format(plotpath,outdirprefix), exist_ok=True)
271 
272  dphiinterval = 0.1
273  dcainterval = 0.025
274 
275  for i in range(0, nbinsdphi):
276  for j in range(0, nbinsdca):
277  dphicut_deg = dphiinterval * (i+1)
278  dphicut_rad = dphicut_deg * (np.pi / 180.)
279  dcacut = dcainterval * (j+1)
280  dphitext = '{:.3f}'.format(dphicut_deg).replace('.', 'p')
281  dcadtext = '{:.3f}'.format(dcacut).replace('.', 'p')
282 
283  outpath = './RecoPV_optimization/{}/dPhi{}deg_dca{}cm/'.format(outdirprefix,dphitext,dcadtext)
284  os.makedirs(outpath, exist_ok=True)
285  try:
286  if dofit == True:
287  main(dphitext,dcadtext,infiledir,outpath)
288  except Exception as e:
289  print(e)
290 
291  hM_mean_dPhidcaCut, hM_sigma_dPhidcaCut, df_opt = get2DTable('./RecoPV_optimization/{}'.format(outdirprefix))
292 
293  ROOT.gStyle.SetPalette(ROOT.kThermometer)
294  ROOT.TGaxis.SetMaxDigits(3)
295  # set the number of divisions to show the bin labels
296  # hM_mean_dPhidcaCut.GetYaxis().SetMaxDigits(4)
297  # hM_sigma_dPhidcaCut.GetYaxis().SetMaxDigits(4)
298  # hM_mean_dPhidcaCut.GetYaxis().SetDecimals()
299  # hM_sigma_dPhidcaCut.GetYaxis().SetDecimals()
300 
301  Draw2Dhisttable(hM_mean_dPhidcaCut, '#Delta#phi [degree]', 'DCA [cm]', '#LT\Deltaz#GT [cm]', 'colztext45', '{}/{}/Optimization_GaussianMean_Table'.format(plotpath,outdirprefix))
302  Draw2Dhisttable(hM_sigma_dPhidcaCut, '#Delta#phi [degree]', 'DCA [cm]', '#sigma_{eff} (68.5%) [cm]', 'colztext45', '{}/{}/Optimization_GaussianSigma_Table'.format(plotpath,outdirprefix))
303 
304  for i in range(5):
305  sigmainfo = df_opt.iloc[i]
306  print(sigmainfo)
307 
308