ATLAS Offline Software
fitman.py
Go to the documentation of this file.
1 #!/usr/bin/env python
2 
3 # Copyright (C) 2002-2017 CERN for the benefit of the ATLAS collaboration
4 """
5 Command line tool to extract desired vertices from a beam spot fitter
6 ntuple, store them in a standalone ntuple for use in RooFit, and run
7 different RooFit beam spot fits.
8 """
9 __author__ = 'Juerg Beringer'
10 __version__ = '$Id: fitman.py 721742 2016-02-03 23:34:44Z beringer $'
11 __usage__ = """%prog [options] [cmd ...]
12 
13 Commands are:
14 
15  extract Extract selected vertices into vertex ntuple
16  (if given, must be first command)
17  stdfit Standard beam spot fit (except for outlier removal)
18  stdfitvzz Dito, but includes z resolution
19  mypdf Similar to stdfit, but uses dedicated BeamSpotPdf class
20  (fast, but doesn't implement all integrals for plotting yet)
21  2gauss Double Gaussian fit - NOT YET WORKING
22  plot[xyz] Plot data and fit of x, y, and z, or single plot
23  data Histograms of x, y, z
24  cov Histograms of vertex error covariance
25  vtxerr Histograms of vertex errors in x, y, z
26  vtxerrxy Histograms of vertex errors in x, y
27 """
28 
29 import os, sys, re
30 from array import array
31 from math import sqrt, floor
32 import random
33 
34 
35 #
36 # Save properly quoted string of original command
37 #
38 qargv = [ ]
39 for s in sys.argv:
40  if re.search('\s|\*|\?',s): # any white space or special characters in word so we need quoting?
41  if "'" in s:
42  qargv.append('"%s"' % re.sub('"',"'",s))
43  else:
44  qargv.append("'%s'" % re.sub("'",'"',s))
45  else:
46  qargv.append(s)
47 qcmd = ' '.join(qargv)
48 
49 
50 #
51 # Argument parsing
52 #
53 from optparse import OptionParser
54 parser = OptionParser(usage=__usage__, version=__version__)
55 parser.add_option('-e', '--srcnt', dest='srcNtName', default=None, help='extract: file name of source beam spot fitter ntuple')
56 parser.add_option('', '--srctree', dest='srcTreeName', default='Beamspot/Vertices', help='extract: tree name of vertex ntuple (default: Beamspot/Vertices)')
57 parser.add_option('-l', '--lbl', dest='lbMin', type='int', default=0, help='extract: minimum LB number (inclusive)')
58 parser.add_option('-u', '--lbu', dest='lbMax', type='int', default=9999, help='extract: maximum LB number (inclusive)')
59 parser.add_option('', '--bcid', dest='bcid', type='int', default=None, help='extract: BCID')
60 parser.add_option('-t', '--vtype', dest='vType', type='int', default=None, help='extract: vertex type (1=PriVtx, 3=PileUp)')
61 parser.add_option('-n', '--ntracks', dest='nTracks', type='int', default=None, help='extract: min number of tracks per vertex')
62 parser.add_option('-p', '--passed', dest='passed', type='int', default=None, help='extract: passed flag (use only vertices passing InDetBeamSpotFinder selection)')
63 parser.add_option('', '--rndfrac', dest='rndFrac', type='float', default=None, help='extract: pick given fraction of random vertices after all other selections')
64 parser.add_option('', '--prescale', dest='prescale', type='int', default=None, help='extract: prescale selected vertices')
65 parser.add_option('', '--nblk', dest='nblk', type='int', default=None, help='extract: number of selected vertices per block (default: not used)')
66 parser.add_option('', '--iblk', dest='iblk', type='int', default=0, help='extract: block number to use (default: 0)')
67 parser.add_option('', '--chi2ndf', dest='chi2ndf', type='float', default=None, help='extract: max chi2/ndf')
68 parser.add_option('', '--chi2prob', dest='chi2prob', type='float', default=None, help='extract: min chi2 probability')
69 parser.add_option('', '--maxvtxerr', dest='maxvtxerr', type='float', default=None, help='extract: max allowed transverse vertex error')
70 parser.add_option('', '--minvtxerr', dest='minvtxerr', type='float', default=None, help='extract: min allowed transverse vertex error')
71 parser.add_option('', '--xmin', dest='xmin', type='float', default=None, help='extract: lower cut on x')
72 parser.add_option('', '--xmax', dest='xmax', type='float', default=None, help='extract: upper cut on x')
73 parser.add_option('', '--ymin', dest='ymin', type='float', default=None, help='extract: lower cut on y')
74 parser.add_option('', '--ymax', dest='ymax', type='float', default=None, help='extract: upper cut on y')
75 parser.add_option('', '--zmin', dest='zmin', type='float', default=None, help='extract: lower cut on z')
76 parser.add_option('', '--zmax', dest='zmax', type='float', default=None, help='extract: upper cut on z')
77 parser.add_option('-f', '--nt', dest='ntName', default='vtx.root', help='file name vertex ntuple (default: vtx.root)')
78 parser.add_option('-g', '--nttree', dest='ntTree', default="Vertices", help='name of TTree with vertices')
79 parser.add_option('', '--plot', dest='plot', default='ALL', help='what to plot (default: ALL)')
80 parser.add_option('', '--canvas', dest='canvas', default='default', help='canvas size: default, page, wide, extrawide or square')
81 parser.add_option('', '--lsize', dest='lsize', type='float', default=0.05, help='axis label size')
82 parser.add_option('', '--hnbins', dest='hnbins', type='int', default=None, help='number of histogram bins')
83 parser.add_option('', '--hmin', dest='hmin', type='float', default=None, help='histogram x axis minimum')
84 parser.add_option('', '--hmax', dest='hmax', type='float', default=None, help='histogram x axis maximum')
85 parser.add_option('', '--hrange', dest='hrange', type='float', default=None, help='histogram x axis range')
86 parser.add_option('', '--fixK', dest='fixK', type='float', default=None, help='stdfit: fix k factor in fit')
87 parser.add_option('', '--showbs', dest='showbs', action='store_true', default=False, help='show beam spot')
88 parser.add_option('', '--showratio', dest='showratio', action='store_true', default=False, help='show ratio data/fit')
89 parser.add_option('', '--minrange', dest='minrange', type='float', default=0.1, help='min range of x-axis on plots')
90 parser.add_option('', '--logy', dest='logy', action='store_true', default=False, help='log scale')
91 parser.add_option('-s', '--stats', dest='stats', action='store_true', default=False, help='show statistics box on plot')
92 parser.add_option('', '--optstat', dest='optstat', default='emruo', help='default OptStat value (Default: emruo)')
93 parser.add_option('', '--prelim', dest='prelim', action='store_true', default=False, help='add ATLAS Preliminary to figure')
94 parser.add_option('', '--approval', dest='approval', action='store_true', default=False, help='Label figure ATLAS for approval')
95 parser.add_option('', '--published', dest='published', action='store_true', default=False, help='add ATLAS to figure')
96 parser.add_option('', '--atlasx', dest='atlasx', type='float', default=0.2, help='x position for drawing ATLAS (Preliminary) label')
97 parser.add_option('', '--atlasy', dest='atlasy', type='float', default=0.85, help='y position for drawing ATLAS (Preliminary) label')
98 parser.add_option('', '--atlasdx', dest='atlasdx', type='float', default=0.115, help='x position offset for drawing Preliminary label')
99 parser.add_option('', '--comment', dest='comment', default='', help='additional text (use semicolon to indicate line breaks)')
100 parser.add_option('', '--commentx', dest='commentx', type='float', default=0.2, help='x position for additional text')
101 parser.add_option('', '--commenty', dest='commenty', type='float', default=0.69, help='y position for additional text')
102 parser.add_option('', '--legendx', dest='legendx', type='float', default=0.70, help='x position of legend')
103 parser.add_option('', '--legendy', dest='legendy', type='float', default=0.93, help='y position of legend')
104 parser.add_option('', '--name', dest='name', default='fitman', help='name for set of plots (default: fitman)')
105 parser.add_option('-o', '--output', dest='output', default='.gif', help='comma-separated list of output files or formats (default: .gif)')
106 parser.add_option('-c', '--cpu', dest='nCpu', type='int', default=2, help='number of CPUs to use for RooFit')
107 parser.add_option('-v', '--verbose', dest='verbose', action='store_true', default=False, help='verbose output')
108 parser.add_option('-i', '--interactive', dest='interactive', action='store_true', default=False, help='interactive')
109 parser.add_option('-b', '--batch', dest='batch', action='store_true', default=False, help='run in batch mode')
110 (options,args) = parser.parse_args()
111 if len(args)<1:
112  parser.error('wrong number of command line arguments')
113 if options.srcNtName and not 'extract' in args:
114  parser.error('must give extract command when using -e or --srcnt')
115 if options.interactive:
116  os.environ['PYTHONINSPECT'] = '1'
117 cmdList = args
118 
119 # Reset DISPLAY if in batch
120 if options.batch:
121  os.unsetenv('DISPLAY')
122 
123 # Place after option parsing to avoid conflict with ROOT option parsing
124 import ROOT
125 import ROOTUtils
126 ROOTUtils.setStyle()
127 ROOT.gStyle.SetOptStat(options.optstat)
128 
129 
130 #
131 # Utilities
132 #
133 def logresult(cmd='',fitresult=None):
134  log = open('%s.txt' % options.name,'w')
135  if cmd:
136  log.write(cmd)
137  log.write('\n')
138  if fitresult:
139  out = ROOT.std.stringstream()
140  fitresult.printStream(out,fitresult.defaultPrintContents(''),fitresult.defaultPrintStyle(''))
141  log.write(out.str())
142  del out
143  log.close()
144 
145 
146 def tmpHist(what,wmin=-1e10,wmax=+1e10):
147  vtxData.Draw(what,'%s > %s && %s < %s' % (what,wmin,what,wmax),'goff')
148  h = ROOT.gROOT.FindObject('htemp')
149  print 'Histogram for %s: mean = %7.4f rms = %1.4f' % (what,h.GetMean(),h.GetRMS())
150  return h.Clone('tmp-'+what)
151 
152 
153 #
154 # Definition of plots
155 #
156 class Plots(ROOTUtils.PlotLibrary):
157 
158  def __init__(self, name = options.name, whatList = []):
159  ROOTUtils.PlotLibrary.__init__(self,name)
160  self.whatList = whatList
161  self.allCanvasSize = 'landscape'
162  self.allCanvasDivs = (2,2)
163  self.singleCanvasSize = options.canvas
164 
165  def prepareMainPad(self):
166  mainPad = self.protect( ROOT.TPad("mainPad","mainPad",0.,0.3,1.,1.) )
167  mainPad.SetBottomMargin(0.03) # with space between the two pads
168  #mainPad.SetBottomMargin(0.00) # w/o space between the two pads
169  mainPad.Draw()
170  mainPad.cd()
171  return mainPad
172 
173  def drawRatioPad(self,frame,hmin,hmax):
174  ROOT.gROOT.ProcessLine(".L ratioPlot.C+")
175  axisTitle = frame.GetXaxis().GetTitle()
176  frame.GetXaxis().SetTitle('')
177  frame.GetXaxis().SetLabelSize(0.)
178  frame.GetYaxis().SetTitleOffset(1.0)
179  canvas = ROOT.gPad.GetCanvas()
180  canvas.cd()
181  ratioPad = self.protect( ROOT.TPad("ratioPad","ratioPad",0.,0.,1.,0.3) )
182  ratioPad.SetTopMargin(0.)
183  ratioPad.SetBottomMargin(0.3)
184  ratioPad.Draw()
185  ratioPad.cd()
186  dataHist = frame.getObject(0)
187  fitCurve = frame.getObject(1)
188  #global pull
189  #pull = dataHist.makeResidHist(fitCurve,True,True)
190  global ratio
191  ratio = ROOT.ratioPlot(dataHist,fitCurve)
192  ratio.Draw('AP')
193  # NOTE: SetRangeUser chooses bin edges, so ratio plot would not be perfectly aligned
194  ratio.GetXaxis().SetLimits(hmin,hmax)
195  scale = 2.5
196  size = scale*options.lsize
197  ratio.GetXaxis().SetLabelSize(size)
198  ratio.GetXaxis().SetTitle(axisTitle)
199  ratio.GetXaxis().SetTitleSize(size)
200  ratio.GetXaxis().SetTitleOffset(1.)
201  ratio.GetYaxis().SetRangeUser(0.3,1.7)
202  ratio.GetYaxis().SetLabelSize(size)
203  ratio.GetYaxis().SetTitle('Data / fit')
204  #ratio.GetYaxis().SetTitleSize(scale*ratio.GetYaxis().GetTitleSize())
205  ratio.GetYaxis().SetTitleSize(size)
206  ratio.GetYaxis().SetTitleOffset(1.0/scale)
207  ratio.GetYaxis().SetNdivisions(504)
208  ratio.SetMarkerSize(0.7)
209  line = self.protect( ROOT.TLine(hmin,1,hmax,1) )
210  line.Draw()
211 
212 
213  def labels(self):
214  # ATLAS labels and other text
215  if options.prelim:
216  ROOTUtils.atlasLabel(options.atlasx,options.atlasy,True,offset=options.atlasdx,energy=8,size=options.lsize)
217  if options.approval:
218  ROOTUtils.atlasLabel(options.atlasx,options.atlasy,False,offset=options.atlasdx,isForApproval=True,energy=8,size=options.lsize)
219  if options.published:
220  ROOTUtils.atlasLabel(options.atlasx,options.atlasy,False,offset=options.atlasdx,energy=8,size=options.lsize)
221  if options.comment:
222  self.protect( ROOTUtils.drawText(options.commentx,options.commenty,0.06,options.comment,font=42) )
223 
224  def legend(self,frame):
225  legendList = []
226  legendList.append([frame.getObject(0),'Data','P'])
227  legendList.append([frame.getObject(1),'Fit projection','L'])
228  if options.showbs:
229  legendList.append([frame.getObject(2),'Beam spot','L'])
230  legendMinY = max(options.legendy-options.lsize*1.4*len(legendList),0.2)
231  self.protect( ROOTUtils.drawLegend(options.legendx,legendMinY,0.92,options.legendy,legendList,textSize=options.lsize) )
232 
233  def hist(self,what):
234  if not what:
235  return
236  if options.hmin is not None or options.hmax is not None or options.hnbins is not None:
237  options.hmin = options.hmin if options.hmin is not None else 0.
238  options.hmax = options.hmax if options.hmax is not None else 1.
239  options.hnbins = options.hnbins if options.hnbins is not None else 40
240  h = ROOTUtils.protect( ROOT.TH1F('tmp','%s;%s' % (what,what),
241  options.hnbins,options.hmin,options.hmax))
242  vtxData.Draw('%s >> tmp' % what)
243  else:
244  vtxData.Draw(what)
245  if options.logy:
246  ROOT.gPad.SetLogy(options.logy)
247 
248  def x(self):
249  if options.showratio:
250  self.prepareMainPad()
251  dx = options.hrange/2 if options.hrange else max(options.minrange/2.,8.*sx.getVal())
252  hmin = options.hmin if options.hmin is not None else mx.getVal()-dx
253  hmax = options.hmax if options.hmax is not None else mx.getVal()+dx
254  global xframe
255  xframe = self.protect( x.frame(hmin,hmax) )
256  if options.hnbins is not None:
257  binning = ROOT.RooBinning(hmin, hmax)
258  binning.addUniform(options.hnbins, hmin, hmax)
259  data.plotOn(xframe,RooFit.Binning(binning))
260  xbinWidth = (hmax-hmin)/float(options.hnbins)
261  else:
262  data.plotOn(xframe)
263  xbinWidth = xframe.GetXaxis().GetBinWidth(1)
264  #xframe.GetYaxis().SetTitle("Number of primary vertices per %3.1f #mum" % round(1000.*xbinWidth,1))
265  xframe.GetYaxis().SetTitle("Events / %1.0f #mum" % round(1000.*xbinWidth,1))
266  xframe.GetYaxis().SetLabelSize(options.lsize)
267  xframe.GetYaxis().SetTitleSize(options.lsize)
268  if options.stats:
269  data.statOn(xframe)
270  cov = RooArgSet(vxx,vyy,vxy)
271  fitmodel.plotOn(xframe,RooFit.ProjWData(cov,data),
272  RooFit.NumCPU(options.nCpu))
273  if options.showbs:
274  k.setVal(0.0)
275  fitmodel.plotOn(xframe,RooFit.ProjWData(cov,data),
276  RooFit.NumCPU(options.nCpu),
277  RooFit.LineStyle(RooFit.kDashed))
278  xframe.Draw()
279  if options.logy:
280  ROOT.gPad.SetLogy(options.logy)
281  self.labels()
282  self.legend(xframe)
283  if options.showratio:
284  self.drawRatioPad(xframe,hmin,hmax)
285 
286  def y(self):
287  if options.showratio:
288  self.prepareMainPad()
289  dy = options.hrange/2 if options.hrange else max(options.minrange/2.,8.*sy.getVal())
290  hmin = options.hmin if options.hmin is not None else my.getVal()-dy
291  hmax = options.hmax if options.hmax is not None else my.getVal()+dy
292  global yframe
293  yframe = self.protect( y.frame(hmin,hmax) )
294  if options.hnbins is not None:
295  binning = ROOT.RooBinning(hmin, hmax)
296  binning.addUniform(options.hnbins, hmin, hmax)
297  data.plotOn(yframe,RooFit.Binning(binning))
298  ybinWidth = (hmax-hmin)/float(options.hnbins)
299  else:
300  data.plotOn(yframe)
301  ybinWidth = yframe.GetXaxis().GetBinWidth(1)
302  #yframe.GetYaxis().SetTitle("Number of primary vertices per %3.1f #mum" % round(1000.*ybinWidth,1))
303  yframe.GetYaxis().SetTitle("Events / %1.0f #mum" % round(1000.*ybinWidth,1))
304  yframe.GetYaxis().SetLabelSize(options.lsize)
305  yframe.GetYaxis().SetTitleSize(options.lsize)
306  if options.stats:
307  data.statOn(yframe)
308  cov = RooArgSet(vxx,vyy,vxy)
309  fitmodel.plotOn(yframe,RooFit.ProjWData(cov,data),
310  RooFit.NumCPU(options.nCpu))
311  if options.showbs:
312  k.setVal(0.0)
313  fitmodel.plotOn(yframe,RooFit.ProjWData(cov,data),
314  RooFit.NumCPU(options.nCpu),
315  RooFit.LineStyle(RooFit.kDashed))
316  yframe.Draw()
317  if options.logy:
318  ROOT.gPad.SetLogy(options.logy)
319  self.labels()
320  self.legend(yframe)
321  if options.showratio:
322  self.drawRatioPad(yframe,hmin,hmax)
323 
324  def z(self):
325  if options.showratio:
326  self.prepareMainPad()
327  dz = options.hrange/2 if options.hrange else max(options.minrange/2.,8.*sz.getVal())
328  hmin = options.hmin if options.hmin is not None else mz.getVal()-dz
329  hmax = options.hmax if options.hmax is not None else mz.getVal()+dz
330  global zframe
331  zframe = self.protect( z.frame(hmin,hmax) )
332  if options.hnbins is not None:
333  binning = ROOT.RooBinning(hmin, hmax)
334  binning.addUniform(options.hnbins, hmin, hmax)
335  data.plotOn(zframe,RooFit.Binning(binning))
336  zbinWidth = (hmax-hmin)/float(options.hnbins)
337  else:
338  data.plotOn(zframe)
339  zbinWidth = zframe.GetXaxis().GetBinWidth(1)
340  #zframe.GetYaxis().SetTitle("Number of primary vertices per %3.1f mm" % round(zbinWidth,1))
341  zframe.GetYaxis().SetLabelSize(options.lsize)
342  zframe.GetYaxis().SetTitleSize(options.lsize)
343  zframe.GetYaxis().SetTitle("Events / %1.0f mm" % round(zbinWidth,1))
344  if options.stats:
345  data.statOn(zframe)
346  cov = RooArgSet(vxx,vyy,vxy)
347  fitmodel.plotOn(zframe,RooFit.ProjWData(cov,data),
348  RooFit.NumCPU(options.nCpu)) # could also use e.g. RooFit.Range(-100.,100.)
349  if options.showbs:
350  k.setVal(0.0)
351  fitmodel.plotOn(zframe,RooFit.ProjWData(cov,data),
352  RooFit.NumCPU(options.nCpu),
353  RooFit.LineStyle(RooFit.kDashed))
354  zframe.Draw()
355  if options.logy:
356  ROOT.gPad.SetLogy(options.logy)
357  self.labels()
358  self.legend(zframe)
359  if options.showratio:
360  self.drawRatioPad(zframe,hmin,hmax)
361 
362 
363 #
364 # If extract command is given, must process it here, before setting up
365 # RooFit
366 #
367 if cmdList[0] == 'extract':
368  cmdList.pop(0)
369  if not options.srcNtName:
370  parser.error('must specify --srcnt for extract command')
371  if not os.path.exists(options.srcNtName):
372  sys.exit('ERROR: source ROOT file %s not found' % options.srcNtName)
373  srcFile = ROOT.TFile(options.srcNtName)
374  srcNt = srcFile.Get(options.srcTreeName)
375  nEntries = srcNt.GetEntries()
376  print '\nCuts for extracting vertices from: %s:' % options.srcNtName
377  print ' %i <= LB <= %i' % (options.lbMin, options.lbMax)
378  if options.bcid:
379  print ' BCID == %i' % options.bcid
380  if options.vType:
381  print ' vertex type == %i' % options.vType
382  if options.nTracks:
383  print ' nTracks >= %i' % options.nTracks
384  if options.passed:
385  print ' passed == %i' % options.passed
386  if options.chi2ndf:
387  print ' chi2/ndf <= %s' % options.chi2ndf
388  if options.chi2prob:
389  print ' chi2prob >= %s' % options.chi2prob
390  if options.minvtxerr:
391  print ' transverse vertex errors >= %s' % options.minvtxerr
392  if options.maxvtxerr:
393  print ' transverse vertex errors < %s' % options.maxvtxerr
394  if options.rndFrac is not None:
395  print ' using fraction %1.3f of selected vertices' % options.rndFrac
396  if options.nblk is not None:
397  print ' using block %i of blocks of %i selected vertices' % (options.iblk,options.nblk)
398  if options.prescale:
399  print ' prescaling selected vertices by factor %i' % options.prescale
400  if options.xmin is not None or options.xmax is not None:
401  print ' %s <= x <= %s' % (options.xmin,options.xmax)
402  if options.ymin is not None or options.ymax is not None:
403  print ' %s <= y <= %s' % (options.ymin,options.ymax)
404  if options.zmin is not None or options.zmax is not None:
405  print ' %s <= z <= %s' % (options.zmin,options.zmax)
406 
407  print '\nProcessing %i entries found in source ntuple ...' % nEntries
408  dstFile = ROOT.TFile(options.ntName,'recreate')
409  dstNt = ROOT.TTree(options.ntTree,'Vertices')
410  bx = array('d',[0]); dstNt.Branch('x',bx,'x/D')
411  by = array('d',[0]); dstNt.Branch('y',by,'y/D')
412  bz = array('d',[0]); dstNt.Branch('z',bz,'z/D')
413  bvxx = array('d',[0]); dstNt.Branch('vxx',bvxx,'vxx/D')
414  bvyy = array('d',[0]); dstNt.Branch('vyy',bvyy,'vyy/D')
415  bvxy = array('d',[0]); dstNt.Branch('vxy',bvxy,'vxy/D')
416  bvxz = array('d',[0]); dstNt.Branch('vxz',bvxz,'vxz/D')
417  bvyz = array('d',[0]); dstNt.Branch('vyz',bvyz,'vyz/D')
418  bvzz = array('d',[0]); dstNt.Branch('vzz',bvzz,'vzz/D')
419  nSelected = 0
420  nWritten = 0
421  if options.maxvtxerr:
422  maxVtxErr2 = options.maxvtxerr*options.maxvtxerr
423  if options.minvtxerr:
424  minVtxErr2 = options.minvtxerr*options.minvtxerr
425  for j in xrange(nEntries):
426  srcNt.GetEntry(j)
427  if srcNt.lb < options.lbMin:
428  continue
429  if srcNt.lb > options.lbMax:
430  continue
431  if options.bcid and srcNt.tck!=options.bcid:
432  continue
433  if options.vType and srcNt.vType!=options.vType:
434  continue
435  if options.nTracks and srcNt.nTracks<options.nTracks:
436  continue
437  if options.passed and srcNt.passed!=options.passed:
438  continue
439  if options.chi2ndf and srcNt.chi2/srcNt.ndf>options.chi2ndf:
440  continue
441  # NOTE: The strange Prob calculation mimicks what BeamSpotFinder does - see also
442  # FitQuality::numberDoF() and FitQuality::doubleNumberDoF() in
443  # Tracking/ TrkEvent/ TrkEventPrimitives/ TrkEventPrimitives/ FitQuality.h
444  if options.chi2prob and ROOT.TMath.Prob(srcNt.chi2,int(floor(srcNt.ndf+0.5)))<options.chi2prob:
445  continue
446  if options.maxvtxerr and (srcNt.vxx>=maxVtxErr2 or srcNt.vyy>=maxVtxErr2):
447  continue
448  if options.minvtxerr and (srcNt.vxx<minVtxErr2 or srcNt.vyy<minVtxErr2):
449  continue
450  if options.xmin is not None and srcNt.x<float(options.xmin):
451  continue
452  if options.xmax is not None and srcNt.x>float(options.xmax):
453  continue
454  if options.ymin is not None and srcNt.y<float(options.ymin):
455  continue
456  if options.ymax is not None and srcNt.y>float(options.ymax):
457  continue
458  if options.zmin is not None and srcNt.z<float(options.zmin):
459  continue
460  if options.zmax is not None and srcNt.z>float(options.zmax):
461  continue
462 
463  # No explicit Chi2/ndf cut - implicit in passed flag
464  nSelected += 1
465  if options.rndFrac is not None:
466  if random.random() > options.rndFrac:
467  continue
468  if options.prescale and (nSelected % options.prescale):
469  continue
470  if options.nblk is not None and not (options.iblk*options.nblk <= nSelected < (options.iblk+1)*options.nblk):
471  continue
472  nWritten += 1
473 
474  bx[0] = srcNt.x
475  by[0] = srcNt.y
476  bz[0] = srcNt.z
477  bvxx[0] = srcNt.vxx
478  bvyy[0] = srcNt.vyy
479  bvxy[0] = srcNt.vxy
480  bvxz[0] = srcNt.vxz
481  bvyz[0] = srcNt.vyz
482  bvzz[0] = srcNt.vzz
483  dstNt.Fill()
484  dstFile.Write()
485  dstFile.Close()
486  srcFile.Close()
487  print '%i vertices selected, %i vertices after prescales and block selection written to %s\n' % (nSelected,nWritten,options.ntName)
488 
489 
490 #
491 # RooFit setup common to all fits
492 #
493 from ROOT import kTRUE
494 from ROOT import RooFit, RooRealVar, RooDataSet, RooFormulaVar, RooArgList, RooArgSet
495 from ROOT import RooAddPdf
496 
497 # Observables and data
498 #x = RooRealVar("x","Primary vertex x [mm]",-3,3)
499 #y = RooRealVar("y","Primary vertex y [mm]",-3,3)
500 #z = RooRealVar("z","Primary vertex z [mm]",-300,300)
501 x = RooRealVar("x","x [mm]",-3,3)
502 y = RooRealVar("y","y [mm]",-3,3)
503 z = RooRealVar("z","z [mm]",-300,300)
504 vxx = RooRealVar("vxx","vxx",-3,3)
505 vyy = RooRealVar("vyy","vyy",-3,3)
506 vxy = RooRealVar("vxy","vxy",-3,3)
507 vxz = RooRealVar("vxz","vxz",-3,3)
508 vyz = RooRealVar("vyz","vyz",-3,3)
509 vzz = RooRealVar("vzz","vzz",-3,3)
510 
511 vtxDataFile = ROOT.TFile(options.ntName)
512 vtxData = vtxDataFile.Get(options.ntTree)
513 print '\nUsing %i input vertices from ntuple %s ...\n' % (vtxData.GetEntries(),options.ntName)
514 hx = tmpHist('x',-3.,3.)
515 hy = tmpHist('y',-3.,3.)
516 hz = tmpHist('z',-300.,300.)
517 data = RooDataSet("data","data", vtxData, RooArgSet(x,y,z,vxx,vyy,vxy,vzz))
518 
519 # Model (parameters and PDF)
520 mx = RooRealVar("mx","Mean x", hx.GetMean(), -2, 2)
521 sx = RooRealVar("sx","Sigma x", hx.GetRMS(), 0, 1)
522 ax = RooRealVar("ax","Tilt x", 0, -1e-3, 1e-3)
523 my = RooRealVar("my","Mean y", hy.GetMean(), -2, 2)
524 sy = RooRealVar("sy","Sigma y", hy.GetRMS(), 0, 1)
525 ay = RooRealVar("ay","Tilt y", 0, -1e-3, 1e-3)
526 mz = RooRealVar("mz","Mean z", hz.GetMean(), -300, 300)
527 sz = RooRealVar("sz","Sigma z", hz.GetRMS(), 0, 100)
528 k = RooRealVar("k","Error scale factor", 1, 0.0, 3)
529 if options.fixK:
530  k.setVal(options.fixK)
531  k.setConstant(kTRUE)
532 rho = RooRealVar("rho","Correlation coefficient", 0, -1, 1)
533 
534 # Double Gaussian parameters
535 f = RooRealVar("f", "Fraction of 1st Gauss", 0.9, 0., 1)
536 mx2 = RooRealVar("mx2","Mean x", -0.05, -2, 2)
537 sx2 = RooRealVar("sx2","Sigma x", 0.10, 0, 1)
538 ax2 = RooRealVar("ax2","Tilt x", 0, -1e-3, 1e-3)
539 my2 = RooRealVar("my2","Mean y", 1, -2, 2)
540 sy2 = RooRealVar("sy2","Sigma y", 0.10, 0, 1)
541 ay2 = RooRealVar("ay2","Tilt y", 0, -1e-3, 1e-3)
542 mz2 = RooRealVar("mz2","Mean z", -3, -300, 300)
543 sz2 = RooRealVar("sz2","Sigma z", 60, 0, 100)
544 rho2 = RooRealVar("rho2","Correlation coefficient", 0, -1, 1)
545 
546 # Intermediate parameters
547 mux = RooFormulaVar("mux","@0+@1*(@2-@3)",RooArgList(mx,ax,z,mz))
548 muy = RooFormulaVar("muy","@0+@1*(@2-@3)",RooArgList(my,ay,z,mz))
549 cxx = RooFormulaVar("cxx","@0*@0+@1*@1*@2",RooArgList(sx,k,vxx))
550 cxy = RooFormulaVar("cxy","@0*@1*@2+@3*@3*@4",RooArgList(rho,sx,sy,k,vxy))
551 cyy = RooFormulaVar("cyy","@0*@0+@1*@1*@2",RooArgList(sy,k,vyy))
552 czz = RooFormulaVar("czz","@0*@0",RooArgList(sz))
553 czzv = RooFormulaVar("czz","@0*@0+@1",RooArgList(sz,vzz))
554 
555 # Double Gaussian parameters
556 mux2 = RooFormulaVar("mux2","@0+@1*(@2-@3)",RooArgList(mx2,ax,z,mz))
557 muy2 = RooFormulaVar("muy2","@0+@1*(@2-@3)",RooArgList(my2,ay,z,mz))
558 cxx2 = RooFormulaVar("cxx2","@0*@0+@1*@1*@2",RooArgList(sx2,k,vxx))
559 cxy2 = RooFormulaVar("cxy2","@0*@1*@2+@3*@3*@4",RooArgList(rho,sx2,sy2,k,vxy))
560 cyy2 = RooFormulaVar("cyy2","@0*@0+@1*@1*@2",RooArgList(sy2,k,vyy))
561 czz2 = RooFormulaVar("czz","@0*@0",RooArgList(sz))
562 
563 
564 # Dummy matrix needed by some models
565 dummyMatrix = ROOT.TMatrixDSym(3)
566 dummyMatrix[0][0] = 1
567 dummyMatrix[1][1] = 1
568 dummyMatrix[2][2] = 1
569 
570 
571 #
572 # Command processing (except extract, see above)
573 #
574 for cmd in args:
575  cmdOk = False
576 
577  if cmd == 'none':
578  # Useful to end up in interactive mode
579  cmdOk = True
580 
581  if cmd == 'stdfit':
582  cmdOk = True
583  ROOT.gROOT.ProcessLine(".L GenGauss3D.cxx+")
584  fitmodel = ROOT.GenGauss3D("fitmodel","Full beam spot PDF",
585  RooArgList(x,y,z),
586  RooArgList(mux,muy,mz),
587  cxx,cxy,RooFit.RooConst(0),cyy,RooFit.RooConst(0),czz,
588  dummyMatrix)
589  # NOTE: need RooFit.Save() to get fit results
590  fitresult = fitmodel.fitTo(data,RooFit.ConditionalObservables(RooArgSet(vxx,vyy,vxy)),
591  RooFit.NumCPU(options.nCpu), RooFit.Timer(kTRUE), RooFit.Save())
592  logresult(qcmd,fitresult)
593 
594  if cmd == 'mypdf':
595  cmdOk = True
596  ROOT.gROOT.ProcessLine(".L BeamSpotPdf.cxx+")
597  fitmodel = ROOT.BeamSpotPdf("fitmodel","Full beam spot PDF",
598  x,y,z,vxx,vyy,vxy,mx,sx,ax,my,sy,ay,mz,sz,k,rho)
599  fitresult = fitmodel.fitTo(data,RooFit.ConditionalObservables(RooArgSet(vxx,vyy,vxy)),
600  RooFit.NumCPU(options.nCpu), RooFit.Timer(kTRUE), RooFit.Save())
601  logresult(qcmd,fitresult)
602 
603  if cmd == 'stdfitvzz':
604  cmdOk = True
605  ROOT.gROOT.ProcessLine(".L GenGauss3D.cxx+")
606  fitmodel = ROOT.GenGauss3D("fitmodel","Full beam spot PDF",
607  RooArgList(x,y,z),
608  RooArgList(mux,muy,mz),
609  cxx,cxy,RooFit.RooConst(0),cyy,RooFit.RooConst(0),czzv,
610  dummyMatrix)
611  fitresult = fitmodel.fitTo(data,RooFit.ConditionalObservables(RooArgSet(vxx,vyy,vxy,vzz)),
612  RooFit.NumCPU(options.nCpu), RooFit.Timer(kTRUE), RooFit.Save())
613  logresult(qcmd,fitresult)
614 
615  if cmd == '2gauss':
616  cmdOk = True
617  ROOT.gROOT.ProcessLine(".L GenGauss3D.cxx+")
618  ROOT.gROOT.ProcessLine(".L GenGauss3Dclone.cxx+")
619  g1 = ROOT.GenGauss3D("g1","Full beam spot PDF",
620  RooArgList(x,y,z),
621  RooArgList(mux,muy,mz),
622  cxx,cxy,RooFit.RooConst(0),cyy,RooFit.RooConst(0),czz,
623  dummyMatrix)
624  g2 = ROOT.GenGauss3Dclone("g2","Full beam spot PDF",
625  RooArgList(x,y,z),
626  RooArgList(mux,muy,mz),
627  cxx2,cxy2,RooFit.RooConst(0),cyy2,RooFit.RooConst(0),czz2,
628  dummyMatrix)
629  #fitmodel = RooAddPdf("fitmodel","sum",RooArgList(g1,g2),RooArgList(f))
630  fitmodel = RooAddPdf("fitmodel","sum",g1,g2,f)
631 
632  fitresult = fitmodel.fitTo(data,RooFit.ConditionalObservables(RooArgSet(vxx,vyy,vxy)),
633  RooFit.NumCPU(options.nCpu), RooFit.Timer(kTRUE), RooFit.Save())
634  logresult(qcmd,fitresult)
635 
636  if cmd == 'plot':
637  cmdOk = True
638  plots = Plots(whatList = ['x','y','z'])
639  plots.saveAsList = options.output.split(',')
640  plots.gPadSaveAsList = options.output.split(',')
641  plots.genPlot(options.plot)
642 
643  if cmd=='plotx' or cmd=='ploty' or cmd=='plotz':
644  cmdOk = True
645  plots = Plots()
646  plots.saveAsList = options.output.split(',')
647  plots.genPlot(cmd[-1])
648 
649  if cmd == 'data':
650  cmdOk = True
651  plots = Plots(whatList = ['x','y','z'])
652  plots.saveAsList = ['%s-data%s' % (options.name,ext) for ext in options.output.split(',')]
653  #plots.gPadSaveAsList = options.output.split(',')
654  plots.genPlot(options.plot,'hist')
655 
656  if cmd == 'cov':
657  cmdOk = True
658  plots = Plots(whatList = ['vxx','vxy','vxz',None,'vyy','vyz',None,None,'vzz'])
659  plots.allCanvasDivs = (3,3)
660  plots.saveAsList = ['%s-cov%s' % (options.name,ext) for ext in options.output.split(',')]
661  #plots.gPadSaveAsList = options.output.split(',')
662  plots.genPlot(options.plot,'hist')
663 
664  if cmd == 'vtxerr':
665  cmdOk = True
666  plots = Plots(whatList = ['sqrt(vxx)', 'sqrt(vyy)', 'sqrt(vzz)'])
667  plots.allCanvasSize = 'wide'
668  plots.allCanvasDivs = (3,1)
669  plots.saveAsList = ['%s-vtxerr%s' % (options.name,ext) for ext in options.output.split(',')]
670  #plots.gPadSaveAsList = options.output.split(',')
671  plots.genPlot(options.plot,'hist')
672 
673  if cmd == 'vtxerrxy':
674  cmdOk = True
675  plots = Plots(whatList = ['sqrt(vxx)', 'sqrt(vyy)'])
676  plots.allCanvasSize = 'wide'
677  plots.allCanvasDivs = (2,1)
678  plots.saveAsList = ['%s-vtxerr%s' % (options.name,ext) for ext in options.output.split(',')]
679  #plots.gPadSaveAsList = options.output.split(',')
680  plots.genPlot(options.plot,'hist')
681 
682  if cmd == 'debug':
683  cmdOk = True
684 
685  if not cmdOk:
686  print 'ERROR: unknown command: %s' % cmd
687  sys.exit(1)
fitman.Plots.z
def z(self)
Definition: fitman.py:324
xrange
void xrange(TH1 *h, bool symmetric)
Definition: computils.cxx:515
ROOTUtils.drawText
def drawText(x=0.74, y=0.87, dy=0.06, text='', font=62, color=1, align=11, linesep=';')
Definition: roofit/ROOTUtils.py:243
max
#define max(a, b)
Definition: cfImp.cxx:41
fitman.Plots.hist
def hist(self, what)
Definition: fitman.py:233
fitman.Plots.allCanvasDivs
allCanvasDivs
Definition: fitman.py:162
CaloCellPos2Ntuple.int
int
Definition: CaloCellPos2Ntuple.py:24
fitman.Plots.y
def y(self)
Definition: fitman.py:286
ROOTUtils.setStyle
def setStyle(style=None)
Definition: roofit/ROOTUtils.py:413
fitman.Plots
Definition: fitman.py:156
MuonGM::round
float round(const float toRound, const unsigned int decimals)
Definition: Mdt.cxx:27
ROOTUtils.protect
def protect(obj)
Definition: roofit/ROOTUtils.py:17
fitman.Plots.labels
def labels(self)
Definition: fitman.py:213
fitman.logresult
def logresult(cmd='', fitresult=None)
Definition: fitman.py:133
fitman.Plots.prepareMainPad
def prepareMainPad(self)
Definition: fitman.py:165
ROOTUtils.atlasLabel
def atlasLabel(x, y, isPreliminary=False, color=1, offset=0.115, isForApproval=False, energy=8, customstring="", size=0.05)
Definition: roofit/ROOTUtils.py:303
ROOTUtils.PlotLibrary.protect
def protect(self, obj)
Definition: roofit/ROOTUtils.py:101
fitman.tmpHist
def tmpHist(what, wmin=-1e10, wmax=+1e10)
Definition: fitman.py:146
fitman.Plots.allCanvasSize
allCanvasSize
Definition: fitman.py:161
fitman.Plots.drawRatioPad
def drawRatioPad(self, frame, hmin, hmax)
Definition: fitman.py:173
array
TCS::join
std::string join(const std::vector< std::string > &v, const char c=',')
Definition: Trigger/TrigT1/L1Topo/L1TopoCommon/Root/StringUtils.cxx:10
fitman.Plots.legend
def legend(self, frame)
Definition: fitman.py:224
fitman.Plots.singleCanvasSize
singleCanvasSize
Definition: fitman.py:163
ROOTUtils.PlotLibrary.__init__
def __init__(self, name='MyPlots', otherMethods=[])
Definition: roofit/ROOTUtils.py:85
ROOTUtils.PlotLibrary
Definition: roofit/ROOTUtils.py:77
Trk::open
@ open
Definition: BinningType.h:40
fitman.Plots.__init__
def __init__(self, name=options.name, whatList=[])
Definition: fitman.py:158
fitman.Plots.x
def x(self)
Definition: fitman.py:248
ROOTUtils.drawLegend
def drawLegend(x1, y1, x2, y2, legendList=[], fillColor=0, lineColor=0, textSize=None, protectLegend=True)
Definition: roofit/ROOTUtils.py:259
readCCLHist.float
float
Definition: readCCLHist.py:83
fitman.Plots.whatList
whatList
Definition: fitman.py:160
Generated on Thu Nov 7 2024 21:15:18 for ATLAS Offline Software by   doxygen 1.8.18