CMS 3D CMS Logo

/afs/cern.ch/work/a/aaltunda/public/www/CMSSW_6_2_5/src/RecoVertex/BeamSpotProducer/scripts/ntuplemaker.py

Go to the documentation of this file.
00001 #!/usr/bin/env python
00002 #____________________________________________________________
00003 #
00004 #
00005 # A very simple way to make plots with ROOT via an XML file
00006 #
00007 # Francisco Yumiceva
00008 # yumiceva@fnal.gov
00009 #
00010 # Fermilab, 2010
00011 #
00012 #____________________________________________________________
00013 
00014 """
00015    ntuplemaker
00016 
00017    A very simple script to plot the beam spot data stored in condDB
00018 
00019    usage: %prog -t <tag name>
00020    -a, --auth    = AUTH: DB authorization path. online(/nfshome0/popcondev/conddb).
00021    -b, --batch : Run ROOT in batch mode.
00022    -c, --create  = CREATE: name for beam spot data file.
00023    -d, --data    = DATA: input beam spot data file.
00024    -D, --destDB  = DESTDB: destination DB string. online(oracle://cms_orcon_prod/CMS_COND_31X_BEAMSPOT).
00025    -i, --initial = INITIAL: First IOV. Options: run number, or run:lumi, eg. \"133200:21\"
00026    -f, --final   = FINAL: Last IOV. Options: run number, or run:lumi
00027    -o, --output  = OUTPUT: filename of ROOT file with plots.
00028    -x, --xcrossing = XCROSSING : Bunch crossing number.
00029    
00030    Francisco Yumiceva (yumiceva@fnal.gov)
00031    Fermilab 2010
00032    
00033 """
00034 
00035 
00036 import os, string, re, sys, math
00037 import commands, time
00038 from BeamSpotObj import BeamSpot
00039 from IOVObj import IOV
00040 from CommonMethods import *
00041 
00042 try:
00043     import ROOT
00044 except:
00045     print "\nCannot load PYROOT, make sure you have setup ROOT in the path"
00046     print "and pyroot library is also defined in the variable PYTHONPATH, try:\n"
00047     if (os.getenv("PYTHONPATH")):
00048         print " setenv PYTHONPATH ${PYTHONPATH}:$ROOTSYS/lib\n"
00049     else:
00050         print " setenv PYTHONPATH $ROOTSYS/lib\n"
00051         sys.exit()
00052 
00053 from ROOT import *
00054 from array import array
00055 
00056 def getFill( json, run ):
00057 
00058     thefill = 0
00059     run = int(run)
00060     keys = json.keys()
00061     
00062     for i in keys:
00063 
00064         run0 = int(json[i][0])
00065         run1 = int(json[i][1])
00066         if run>= run0 and run<=run1:
00067             thefill = i
00068     
00069     return int(thefill)
00070     
00071 if __name__ == '__main__':
00072 
00073 
00074 
00075     # fill and runs
00076     FillList = {}
00077     runsfile = open("FillandRuns.txt")
00078     for line in runsfile:
00079         if line.find('fill:') != -1:
00080             aline = line.split()
00081             afill = aline[1]
00082             run0 = aline[3]
00083             run1 = aline[5]
00084             FillList[int(afill)] = [int(run0),int(run1)]
00085 
00086     #print FillList
00087 
00088     # create ntuple
00089     gROOT.ProcessLine(
00090         "struct spot {\
00091         Float_t   position[3];\
00092         Float_t   posError[3];\
00093         Float_t   width[3];\
00094         Float_t   widthError[3];\
00095         Float_t   slope[2];\
00096         Float_t   slopeError[2];\
00097         Float_t   time[2];\
00098         Int_t     run;\
00099         Int_t     lumi[2];\
00100         Int_t     fill;\
00101         };" );
00102     
00103     bntuple = spot()
00104     fntuple = TFile( 'bntuple.root', 'RECREATE' )
00105     tbylumi = TTree( 'bylumi', 'beam spot data lumi by lumi' )
00106     tbylumi.Branch('fill', AddressOf( bntuple, 'fill'), 'fill/I' )
00107     tbylumi.Branch('run', AddressOf( bntuple, 'run'), 'run/I' )
00108     tbylumi.Branch('lumi', AddressOf( bntuple, 'lumi'), 'lumi[2]/I' )
00109     tbylumi.Branch('position', AddressOf( bntuple, 'position'),'position[3]/F')
00110     tbylumi.Branch('posErr', AddressOf( bntuple, 'posError'),'posError[3]/F')
00111     tbylumi.Branch('width', AddressOf( bntuple, 'width'),'width[3]/F')
00112     tbylumi.Branch('widthErr', AddressOf( bntuple, 'widthError'),'widthError[3]/F')
00113     tbylumi.Branch('slope', AddressOf( bntuple, 'slope'),'slope[2]/F')
00114     tbylumi.Branch('slopeErr', AddressOf( bntuple, 'slopeError'),'slopeError[2]/F')
00115     tbylumi.Branch('time', AddressOf( bntuple, 'time'),'time[2]/F')
00116             
00117     tbyIOV = TTree( 'byIOV', 'beam spot data by IOV' )
00118     tbyIOV.Branch('fill', AddressOf( bntuple, 'fill'), 'fill/I' )
00119     tbyIOV.Branch('run', AddressOf( bntuple, 'run'), 'run/I' )
00120     tbyIOV.Branch('lumi', AddressOf( bntuple, 'lumi'), 'lumi[2]/I' )
00121     tbyIOV.Branch('position', AddressOf( bntuple, 'position'),'position[3]/F')
00122     tbyIOV.Branch('posErr', AddressOf( bntuple, 'posError'),'posError[3]/F')
00123     tbyIOV.Branch('width', AddressOf( bntuple, 'width'),'width[3]/F')
00124     tbyIOV.Branch('widthErr', AddressOf( bntuple, 'widthError'),'widthError[3]/F')
00125     tbyIOV.Branch('slope', AddressOf( bntuple, 'slope'),'slope[2]/F')
00126     tbyIOV.Branch('slopeErr', AddressOf( bntuple, 'slopeError'),'slopeError[2]/F')
00127     tbyIOV.Branch('time', AddressOf( bntuple, 'time'),'time[2]/F')
00128     
00129     tbyrun = TTree( 'byrun', 'beam spot data by run' )
00130     tbyrun.Branch('fill', AddressOf( bntuple, 'fill'), 'fill/I' )
00131     tbyrun.Branch('run', AddressOf( bntuple, 'run'), 'run/I' )
00132     tbyrun.Branch('lumi', AddressOf( bntuple, 'lumi'), 'lumi[2]/I' )
00133     tbyrun.Branch('position', AddressOf( bntuple, 'position'),'position[3]/F')
00134     tbyrun.Branch('posErr', AddressOf( bntuple, 'posError'),'posError[3]/F')
00135     tbyrun.Branch('width', AddressOf( bntuple, 'width'),'width[3]/F')
00136     tbyrun.Branch('widthErr', AddressOf( bntuple, 'widthError'),'widthError[3]/F')
00137     tbyrun.Branch('slope', AddressOf( bntuple, 'slope'),'slope[2]/F')
00138     tbyrun.Branch('slopeErr', AddressOf( bntuple, 'slopeError'),'slopeError[2]/F')
00139     tbyrun.Branch('time', AddressOf( bntuple, 'time'),'time[2]/F')
00140                     
00141 
00142     # COMMAND LINE OPTIONS
00143     #################################
00144     option,args = parse(__doc__)
00145     if not args and not option: exit()
00146     
00147     if not option.data: 
00148         print " need to provide beam spot data file"
00149         exit()
00150     
00151     if option.batch:
00152         ROOT.gROOT.SetBatch()
00153         
00154     datafilename = "tmp_beamspot.dat"
00155     if option.create:
00156         datafilename = option.create
00157     
00158     getDBdata = True
00159     if option.data:
00160         getDBdata = False
00161    
00162     IOVbase = 'lumibase'
00163     firstRun = "0:0"
00164     lastRun = "4999999999:4999999999"
00165     
00166     if option.initial:
00167         firstRun = option.initial
00168     if option.final:
00169         lastRun = option.final
00170         
00171     # GET IOVs
00172     ################################
00173 
00174     if getDBdata:
00175 
00176         print " read DB to get list of IOVs for the given tag"
00177         acommand = 'cmscond_list_iov -c frontier://PromptProd/CMS_COND_31X_BEAMSPOT -P /afs/cern.ch/cms/DB/conddb -t '+ tag
00178         tmpstatus = commands.getstatusoutput( acommand )
00179         tmplistiov = tmpstatus[1].split('\n')
00180         #print tmplistiov
00181 
00182         iovlist = []
00183         passline = False
00184         iline = jline = 0
00185         totlines = len(tmplistiov)
00186         for line in tmplistiov:
00187         
00188             if line.find('since') != -1:
00189                 passline = True
00190                 jline = iline
00191             if passline and iline > jline and iline < totlines-1:
00192                 linedata = line.split()
00193                 #print linedata
00194                 aIOV = IOV()
00195                 aIOV.since = int(linedata[0])
00196                 aIOV.till = int(linedata[1])
00197                 iovlist.append( aIOV )
00198             iline += 1
00199     
00200         print " total number of IOVs = " + str(len(iovlist))
00201 
00202 
00203         #  GET DATA
00204         ################################
00205 
00206         otherArgs = ''
00207         if option.destDB:
00208             otherArgs = " -d " + option.destDB
00209             if option.auth:
00210                 otherArgs = otherArgs + " -a "+ option.auth
00211         
00212         print " get beam spot data from DB for IOVs. This can take a few minutes ..."
00213 
00214         tmpfile = open(datafilename,'w')
00215 
00216         for iIOV in iovlist:
00217             passiov = False
00218             tmprunfirst = firstRun
00219             tmprunlast = lastRun
00220             tmplumifirst = 1
00221             tmplumilast = 9999999
00222             if IOVbase=="lumibase":
00223                 #tmprunfirst = int(firstRun.split(":")[0])
00224                 #tmprunlast  = int(lastRun.split(":")[0])
00225                 #tmplumifirst = int(firstRun.split(":")[1])
00226                 #tmplumilast  = int(lastRun.split(":")[1])
00227                 tmprunfirst = pack( int(firstRun.split(":")[0]) , int(firstRun.split(":")[1]) )
00228                 tmprunlast  = pack( int(lastRun.split(":")[0]) , int(lasstRun.split(":")[1]) )
00229             #print "since = " + str(iIOV.since) + " till = "+ str(iIOV.till)
00230             if iIOV.since >= int(tmprunfirst) and int(tmprunlast) < 0 and iIOV.since <= int(tmprunfirst):
00231                 print " IOV: " + str(iIOV.since)
00232                 passiov = True
00233             if iIOV.since >= int(tmprunfirst) and int(tmprunlast) > 0 and iIOV.till <= int(tmprunlast):
00234                 print " IOV: " + str(iIOV.since) + " to " + str(iIOV.till)
00235                 passiov = True
00236             if iIOV.since >= int(tmprunlast) and iIOV.till >= 4294967295:
00237                 print " IOV: " + str(iIOV.since) + " to " + str(iIOV.till)
00238                 passiov = True                
00239             if passiov:
00240                 acommand = 'getBeamSpotDB.py -t '+ tag + " -r " + str(iIOV.since) +otherArgs
00241                 if IOVbase=="lumibase":
00242                     tmprun = unpack(iIOV.since)[0]
00243                     tmplumi = unpack(iIOV.since)[1]
00244                     acommand = 'getBeamSpotDB.py -t '+ tag + " -r " + str(tmprun) +" -l "+tmplumi +otherArgs
00245                 status = commands.getstatusoutput( acommand )
00246                 tmpfile.write(status[1])
00247     
00248         print " beam spot data collected and stored in file " + datafilename
00249     
00250         tmpfile.close()
00251 
00252     
00253     # PROCESS DATA
00254     ###################################
00255 
00256     # check if input data exists if given
00257     if option.data:
00258         if os.path.isdir(option.data):
00259             tmp = commands.getstatusoutput("ls "+option.data)
00260             files = tmp[1].split()
00261             datafilename = "combined_all.txt"
00262             output = open(datafilename,"w")
00263             
00264             for f in files:
00265                 if os.path.isdir(option.data+"/"+f) is False:
00266                     input = open(option.data +"/"+f)
00267                     output.writelines(input.readlines())
00268             output.close()
00269             print " data files have been collected in "+datafilename
00270             
00271         elif os.path.exists(option.data):
00272             datafilename = option.data
00273         else:
00274             print " input beam spot data DOES NOT exist, file " + option.data
00275             exit()
00276 
00277     listbeam = []
00278 
00279     if option.xcrossing:
00280         listmap = readBeamSpotFile(datafilename,listbeam,IOVbase,firstRun,lastRun)
00281         # bx
00282         print "List of bunch crossings in the file:"
00283         print listmap.keys()
00284         listbeam = listmap[option.Xrossing]
00285     else:
00286         readBeamSpotFile(datafilename,listbeam,IOVbase,firstRun,lastRun)
00287     
00288     sortAndCleanBeamList(listbeam,IOVbase)
00289             
00290     
00291     ###################################    
00292     
00293     for ii in range(0,len(listbeam)):
00294             
00295         ibeam = listbeam[ii]
00296         
00297         bntuple.position = array('f', [float(ibeam.X), float(ibeam.Y), float(ibeam.Z)])
00298         bntuple.posError = array('f', [float(ibeam.Xerr),float(ibeam.Yerr),float(ibeam.Zerr)])
00299         bntuple.width = array('f', [float(ibeam.beamWidthX), float(ibeam.beamWidthY), float(ibeam.sigmaZ)])
00300         bntuple.widthError = array('f',[float(ibeam.beamWidthXerr),float(ibeam.beamWidthYerr),float(ibeam.sigmaZerr)])
00301         bntuple.run = int(ibeam.Run)
00302         bntuple.fill = int( getFill( FillList, int(ibeam.Run) ) )
00303         bntuple.lumi = array('i', [int(ibeam.IOVfirst),int(ibeam.IOVlast)])
00304         line = ibeam.IOVBeginTime
00305         begintime = time.mktime( time.strptime(line.split()[0] +  " " + line.split()[1] + " " + line.split()[2],"%Y.%m.%d %H:%M:%S %Z") )
00306         line = ibeam.IOVEndTime
00307         endtime = time.mktime( time.strptime(line.split()[0] +  " " + line.split()[1] + " " + line.split()[2],"%Y.%m.%d %H:%M:%S %Z") )
00308         bntuple.time = array('f', [begintime, endtime])
00309         tbylumi.Fill()
00310 
00311 
00312     iovlist = listbeam
00313     iovlist = createWeightedPayloads("tmp.txt",iovlist,False)
00314 
00315     for ii in range(0,len(iovlist)):
00316 
00317         ibeam = iovlist[ii]
00318         
00319         bntuple.position = array('f', [float(ibeam.X), float(ibeam.Y), float(ibeam.Z)])
00320         bntuple.posError = array('f', [float(ibeam.Xerr),float(ibeam.Yerr),float(ibeam.Zerr)])
00321         bntuple.width = array('f', [float(ibeam.beamWidthX), float(ibeam.beamWidthY), float(ibeam.sigmaZ)])
00322         bntuple.widthError = array('f',[float(ibeam.beamWidthXerr),float(ibeam.beamWidthYerr),float(ibeam.sigmaZerr)])
00323         bntuple.run = int(ibeam.Run)
00324         bntuple.fill = int( getFill( FillList, int(ibeam.Run) ) )
00325         bntuple.lumi = array('i', [int(ibeam.IOVfirst),int(ibeam.IOVlast)])
00326         line = ibeam.IOVBeginTime
00327         begintime = time.mktime( time.strptime(line.split()[0] +  " " + line.split()[1] + " " + line.split()[2],"%Y.%m.%d %H:%M:%S %Z") )
00328         line = ibeam.IOVEndTime
00329         endtime = time.mktime( time.strptime(line.split()[0] +  " " + line.split()[1] + " " + line.split()[2],"%Y.%m.%d %H:%M:%S %Z") )
00330         bntuple.time = array('f', [begintime, endtime])
00331                                         
00332         tbyIOV.Fill()
00333                                                                                 
00334     weightedlist = listbeam
00335     weightedlist = createWeightedPayloads("tmp.txt",weightedlist,True)
00336     
00337     for ii in range(0,len(weightedlist)):
00338 
00339         ibeam = weightedlist[ii]
00340 
00341         bntuple.position = array('f', [float(ibeam.X), float(ibeam.Y), float(ibeam.Z)])
00342         bntuple.posError = array('f', [float(ibeam.Xerr),float(ibeam.Yerr),float(ibeam.Zerr)])
00343         bntuple.width = array('f', [float(ibeam.beamWidthX), float(ibeam.beamWidthY), float(ibeam.sigmaZ)])
00344         bntuple.widthError = array('f',[float(ibeam.beamWidthXerr),float(ibeam.beamWidthYerr),float(ibeam.sigmaZerr)])
00345         bntuple.run = int(ibeam.Run)
00346         bntuple.fill = int( getFill( FillList, int(ibeam.Run) ) )
00347         bntuple.lumi = array('i', [int(ibeam.IOVfirst),int(ibeam.IOVlast)])
00348         line = ibeam.IOVBeginTime
00349         begintime = time.mktime( time.strptime(line.split()[0] +  " " + line.split()[1] + " " + line.split()[2],"%Y.%m.%d %H:%M:%S %Z") )
00350         line = ibeam.IOVEndTime
00351         endtime = time.mktime( time.strptime(line.split()[0] +  " " + line.split()[1] + " " + line.split()[2],"%Y.%m.%d %H:%M:%S %Z") )
00352         bntuple.time = array('f', [begintime, endtime])
00353                                         
00354         tbyrun.Fill()
00355 
00356 
00357     os.system('rm tmp.txt')
00358     fntuple.cd()
00359     tbylumi.Write()
00360     tbyIOV.Write()
00361     tbyrun.Write()
00362     fntuple.Close()
00363             
00364     # CLEAN temporal files
00365     ###################################
00366     #os.system('rm tmp_beamspotdata.log')