test/scripts/pv.py

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