test/scripts/pv.py

0001 #! /usr/bin/env python
0002
0003 from __future__ import print_function
0004 import ROOT
0005 from builtins import range
0006 from ROOT import gROOT, gSystem, TFile, TH1I, TCanvas, TFitterMinuit
0007 #gROOT, TFile, TCanvas, TH1F, TH1I, TLegend, TH2F, gPad
0008
0009 #from ROOT import TCanvas, TH1F, TH1I
0010
0011 import sys,os, math
0012 sys.path.append( os.environ['HOME'])
0013
0014
0015 try:
0016     from RooAlias import *
0017     SetStyle()
0018 except:
0019     print("I cannot find RooAlias.py, ignore it and use plain style")
0020     gROOT.SetStyle('Plain')
0021
0022 gROOT.Reset()
0023
0024 gSystem.AddIncludePath(" -I$CMSSW_BASE/src/RecoVertex/BeamSpotProducer/interface");
0025 gSystem.AddIncludePath(" -I$CMSSW_BASE/src");
0026 #gSystem.AddLinkedLibs(" -L$CMSSW_BASE/lib/$SCRAM_ARCH -lRecoVertexBeamSpotProducer");
0027
0028 gROOT.SetMacroPath("$CMSSW_BASE/src/:.")
0029
0030 gROOT.ProcessLine(".L RecoVertex/BeamSpotProducer/interface/BeamSpotFitPVData.h+")
0031 gROOT.ProcessLine(".L RecoVertex/BeamSpotProducer/src/BeamSpotTreeData.cc+")
0032 gROOT.ProcessLine(".L BSVectorDict.h+")
0033 gROOT.ProcessLine(".L RecoVertex/BeamSpotProducer/src/FcnBeamSpotFitPV.cc+")
0034
0035 class BeamSpot:
0036     def __init__(self):
0037         self.Type = -1
0038         self.X = 0.
0039         self.Xerr = 0.
0040         self.Y = 0.
0041         self.Yerr = 0.
0042         self.Z = 0.
0043         self.Zerr = 0.
0044         self.sigmaZ = 0.
0045         self.sigmaZerr = 0.
0046         self.dxdz = 0.
0047         self.dxdzerr = 0.
0048         self.dydz = 0.
0049         self.dydzerr = 0.
0050         self.beamWidthX = 0.
0051         self.beamWidthXerr = 0.
0052         self.beamWidthY = 0.
0053         self.beamWidthYerr = 0.
0054         self.EmittanceX = 0.
0055         self.EmittanceY = 0.
0056         self.betastar = 0.
0057         self.IOVfirst = 0
0058         self.IOVlast = 0
0059         self.IOVBeginTime = 0
0060         self.IOVEndTime = 0
0061         self.Run = 0
0062     def Reset(self):
0063         self.Type = -1
0064         self.X = self.Y = self.Z = 0.
0065         self.Xerr = self.Yerr = self.Zerr = 0.
0066         self.sigmaZ = self.sigmaZerr = 0.
0067         self.dxdz = self.dydz = 0.
0068         self.dxdzerr = self.dydzerr = 0.
0069         self.beamWidthX = self.beamWidthY = 0.
0070         self.beamWidthXerr = self.beamWidthYerr = 0.
0071         self.EmittanceX = self.EmittanceY = self.betastar = 0.
0072         self.IOVfirst = self.IOVlast = 0
0073         self.Run = 0
0074
0075 # 3D fit
0076 def Fit3D( pvStore ):
0077
0078     errorScale_ = 0.9
0079     sigmaCut_ = 5.0
0080     fbeamspot = BeamSpot()
0081
0082     fcn = FcnBeamSpotFitPV(pvStore)
0083     minuitx = TFitterMinuit()
0084     minuitx.SetMinuitFCN(fcn)
0085
0086     # fit parameters: positions, widths, x-y correlations, tilts in xz and yz
0087     minuitx.SetParameter(0,"x",0.,0.02,-10.,10.)
0088     minuitx.SetParameter(1,"y",0.,0.02,-10.,10.)
0089     minuitx.SetParameter(2,"z",0.,0.20,-30.,30.)
0090     minuitx.SetParameter(3,"ex",0.015,0.01,0.,10.)
0091     minuitx.SetParameter(4,"corrxy",0.,0.02,-1.,1.)
0092     minuitx.SetParameter(5,"ey",0.015,0.01,0.,10.)
0093     minuitx.SetParameter(6,"dxdz",0.,0.0002,-0.1,0.1)
0094     minuitx.SetParameter(7,"dydz",0.,0.0002,-0.1,0.1)
0095     minuitx.SetParameter(8,"ez",1.,0.1,0.,30.)
0096     minuitx.SetParameter(9,"scale",errorScale_,errorScale_/10.,errorScale_/2.,errorScale_*2.)
0097
0098     # first iteration without correlations
0099     ierr = 0
0100     minuitx.FixParameter(4)
0101     minuitx.FixParameter(6)
0102     minuitx.FixParameter(7)
0103     minuitx.FixParameter(9)
0104     minuitx.SetMaxIterations(100)
0105     #       minuitx.SetPrintLevel(3)
0106     minuitx.SetPrintLevel(0)
0107     minuitx.CreateMinimizer()
0108     ierr = minuitx.Minimize()
0109     if ierr == 1:
0110     print("3D beam spot fit failed in 1st iteration")
0111     return (False, fbeamspot)
0112
0113     # refit with harder selection on vertices
0114
0115     fcn.setLimits(minuitx.GetParameter(0)-sigmaCut_*minuitx.GetParameter(3),
0116            minuitx.GetParameter(0)+sigmaCut_*minuitx.GetParameter(3),
0117            minuitx.GetParameter(1)-sigmaCut_*minuitx.GetParameter(5),
0118            minuitx.GetParameter(1)+sigmaCut_*minuitx.GetParameter(5),
0119            minuitx.GetParameter(2)-sigmaCut_*minuitx.GetParameter(8),
0120            minuitx.GetParameter(2)+sigmaCut_*minuitx.GetParameter(8));
0121     ierr = minuitx.Minimize();
0122     if ierr == 1:
0123     print("3D beam spot fit failed in 2nd iteration")
0124     return (False, fbeamspot)
0125
0126     # refit with correlations
0127
0128     minuitx.ReleaseParameter(4);
0129     minuitx.ReleaseParameter(6);
0130     minuitx.ReleaseParameter(7);
0131     ierr = minuitx.Minimize();
0132     if ierr == 1:
0133     print("3D beam spot fit failed in 3rd iteration")
0134     return (False, fbeamspot)
0135     # store results
0136
0137     fbeamspot.beamWidthX = minuitx.GetParameter(3);
0138     fbeamspot.beamWidthY = minuitx.GetParameter(5);
0139     fbeamspot.sigmaZ = minuitx.GetParameter(8);
0140     fbeamspot.beamWidthXerr = minuitx.GetParError(3);
0141     fbeamspot.beamWidthYerr = minuitx.GetParError(5);
0142     fbeamspot.sigmaZerr = minuitx.GetParError(8);
0143     fbeamspot.X = minuitx.GetParameter(0)
0144     fbeamspot.Y = minuitx.GetParameter(1)
0145     fbeamspot.Z = minuitx.GetParameter(2)
0146     fbeamspot.dxdz = minuitx.GetParameter(6)
0147     fbeamspot.dydz = minuitx.GetParameter(7)
0148     fbeamspot.Xerr = minuitx.GetParError(0)
0149     fbeamspot.Yerr = minuitx.GetParError(1)
0150     fbeamspot.Zerr = minuitx.GetParError(2)
0151     fbeamspot.dxdzerr = minuitx.GetParError(6)
0152     fbeamspot.dydzerr = minuitx.GetParError(7)
0153
0154     return (True, fbeamspot)
0155
0156
0157 def main():
0158
0159     # input files
0160     tfile = TFile("BeamFit_LumiBased_NewAlignWorkflow_1042_139407_1.root")
0161     tfile.cd()
0162
0163     fchain = ROOT.gDirectory.Get( 'PrimaryVertices' )
0164     entries = fchain.GetEntriesFast()
0165
0166     aData = BeamSpotTreeData()
0167
0168     aData.setBranchAddress(fchain);
0169
0170     histox = {}
0171     histoy = {}
0172     histoz = {}
0173     histobx = TH1I("bx","bx",5000,0,5000)
0174     histolumi = TH1I("lumi","lumi",3000,0,3000)
0175
0176 #pvStore = []
0177
0178     pvStore = ROOT.vector( BeamSpotFitPVData )(0)
0179
0180     for jentry in range( entries ):
0181     # get the next tree in the chain
0182     ientry = fchain.LoadTree(jentry)
0183     if ientry < 0:
0184         break
0185
0186     # verify file/tree/chain integrity
0187     nb = fchain.GetEntry( jentry )
0188     if nb <= 0 or not hasattr( fchain, 'lumi' ):
0189         continue
0190
0191
0192     #run = int( fchain.bunchCrossing )
0193     #lumi = int( fchain.lumi )
0194     #bx = int( fchain.bunchCrossing )
0195     run = aData.getRun()
0196     if ientry == 0 :
0197         therun = run
0198     if run != therun:
0199         print("FILES WITH DIFFERENT RUNS?? "+str(runt) + " and "+str(therun))
0200         break
0201
0202     pvdata = aData.getPvData()
0203
0204     bx = int( pvdata.bunchCrossing )
0205     histobx.Fill(bx)
0206     lumi = aData.getLumi()
0207     histolumi.Fill(lumi)
0208
0209     pvx = pvdata.position[0]
0210     pvy = pvdata.position[1]
0211     pvz = pvdata.position[2]
0212
0213     if (bx in histox) == False:
0214         print("bx: "+str(bx))
0215         histox[bx] = TH2F("x_"+str(bx),"x_"+str(bx),100,0,0.2,300,0,1500)#TH1F("x_"+str(bx),"x_"+str(bx),100,0,0.2)
0216         histoy[bx] = TH1F("y_"+str(bx),"y_"+str(bx),100,0,0.2)
0217         histoz[bx] = TH1F("z_"+str(bx),"z_"+str(bx),100,0,0.2)
0218
0219     histox[bx].Fill(pvx,lumi)
0220     histoy[bx].Fill(pvy)
0221     histoz[bx].Fill(pvz)
0222
0223         pvStore.push_back( pvdata )
0224         #if ientry > 10:
0225     #break
0226
0227     # fit
0228     results = Fit3D( pvStore )
0229
0230     if results[0]:
0231         print("Results:")
0232         print(" X = " +str(results[1].X))
0233         print("width X = " +str(results[1].beamWidthX))
0234
0235     # plots
0236     cvbx = TCanvas("bx","bx",700,700)
0237
0238     histobx.Draw()
0239     histobx.SetXTitle("bunch ID")
0240
0241     cvlumi = TCanvas("lumi","lumi",700,700)
0242
0243     histolumi.Draw()
0244     histolumi.SetXTitle("lumi section")
0245
0246     cvx = TCanvas("cvx","cvx",700,700)
0247
0248     i = 0
0249     for ibx in histox.keys():
0250     if i==0:
0251         t1 = histox[ibx].ProjectionX("xall_"+str(ibx))
0252         t1.Draw()
0253     else:
0254         t1 = histox[ibx].ProjectionX("xall_"+str(ibx))
0255         t1.Draw("same")
0256
0257     i =+ 1
0258
0259     raw_input ("Enter to quit:")
0260
0261
0262 if __name__ == "__main__":
0263     main()
0264
0265
0266
0267