test/scripts/PlotLumiScan.py

0001 #!/usr/bin/env python
0002 #____________________________________________________________
0003 #
0004 #  PlotLumiScan
0005 #
0006 # A very simple way to make lumiscan plots from beamfit.txt files
0007 # Needed files: A txt file specifying lumi section ranges eg. RunLumiScan.txt
0008 #   All the beam fit txt files in <data_dir> created after running AnalyzeLumiScan.py
0009 #
0010 # Geng-yuan Jeng
0011 # Geng-yuan.Jeng@cern.ch
0012 #
0013 # Fermilab, 2009
0014 #
0015 #____________________________________________________________
0016
0017 from __future__ import print_function
0018 from builtins import range
0019 import sys,os
0020 import string
0021 from array import array
0022 from math import sqrt
0023 from ROOT import gROOT,gPad,TH1F,TLegend,TFile,TStyle,TAxis
0024
0025 def get_list_files(directory,pattern="txt"):
0026     dir = []
0027     dir = os.listdir(directory)
0028     dir.sort(cmp)
0029     lfiles = []
0030
0031     for f in dir:
0032         if f.find(pattern) != -1:
0033         #print f
0034             lfiles.append(f)
0035     return lfiles;
0036
0037
0038 def plot_trending(type,label,x,xe):
0039     h = TH1F(type+'_lumi',type+'_lumi',len(x),0.5,0.5+len(x))
0040     h.SetStats(0)
0041     h.GetXaxis().SetTitle("Lumisection")
0042     h.GetYaxis().SetTitle(type[:len(type)-1]+"_{0} (cm)")
0043     h.GetYaxis().SetLabelSize(0.03)
0044     h.SetTitleOffset(1.1)
0045     h.SetOption("e1")
0046
0047     for i in range(len(x)):
0048         h.SetBinContent(i+1, x[i])
0049         h.SetBinError(i+1, xe[i])
0050         h.GetXaxis().SetBinLabel(i+1,label[i])
0051     return h
0052
0053 def main():
0054
0055     if len(sys.argv) < 4:
0056         print("\n [Usage] python PlotLumiScan.py <LumiScanLists.txt> <data dir> <verbose:True/False>")
0057         sys.exit()
0058
0059     lumilistfile = sys.argv[1]
0060     runinfofile = open(lumilistfile,"r")
0061     runinfolist = runinfofile.readlines()
0062     runsinfo = {}
0063
0064     for line in runinfolist:
0065         npos=0
0066         for i in line.split():
0067             npos+=1
0068             if npos == 1:
0069                 run="Run"+str(i)+"/"
0070             else:
0071                 runsinfo.setdefault(run,[]).append(int(i))
0072 ##    print runsinfo
0073
0074     infiledir = sys.argv[2]
0075     if infiledir.endswith("/") != 1:
0076         infiledir+="/"
0077     verbose = sys.argv[3]
0078     files = get_list_files(infiledir,"txt")
0079     nfiles = len(files)-1
0080     labels = []
0081     x0=[]
0082     y0=[]
0083     z0=[]
0084     sigZ=[]
0085     x0Err=[]
0086     y0Err=[]
0087     z0Err=[]
0088     sigZErr=[]
0089
0090     ## Read files and put values into data containers
0091     ## Labels:
0092
0093     lumilist = runsinfo.get(infiledir)
0094
0095     for j in range((len(lumilist)+1)/2):
0096         labelName=str(lumilist[j*2])+"-"+str(lumilist[j*2+1])
0097         labels.append(labelName)
0098 ##    print labels
0099
0100     for f in files:
0101         readfile = open(infiledir+f)
0102         for line in readfile:
0103             if line.find("X") != -1 and not "BeamWidth" in line and not "Emittance" in line:
0104                 count=0
0105                 for val in line.split():
0106                     count+=1
0107                     if count > 1:
0108                         x0.append(float(val))
0109             if line.find("Cov(0,j)") != -1:
0110                 count=0
0111                 for val in line.split():
0112                     count+=1
0113                     if count == 2:
0114                         valErr=sqrt(float(val))
0115                         x0Err.append(valErr)
0116             if line.find("Y") != -1 and not "BeamWidth" in line and not "Emittance" in line:
0117                 count=0
0118                 for val in line.split():
0119                     count+=1
0120                     if count > 1:
0121                         y0.append(float(val))
0122             if line.find("Cov(1,j)") != -1:
0123                 count=0
0124                 for val in line.split():
0125                     count+=1
0126                     if count == 3:
0127                         valErr=sqrt(float(val))
0128                         y0Err.append(valErr)
0129             if line.find("Z") != -1 and not "sigma" in line:
0130                 count=0
0131                 for val in line.split():
0132                     count+=1
0133                     if count > 1:
0134                         z0.append(float(val))
0135             if line.find("Cov(2,j)") != -1:
0136                 count=0
0137                 for val in line.split():
0138                     count+=1
0139                     if count == 4:
0140                         valErr=sqrt(float(val))
0141                         z0Err.append(valErr)
0142             if line.find("sigmaZ") != -1:
0143                 count=0
0144                 for val in line.split():
0145                     count+=1
0146                     if count > 1:
0147                         sigZ.append(float(val))
0148             if line.find("Cov(3,j)") != -1:
0149                 count=0
0150                 for val in line.split():
0151                     count+=1
0152                     if count == 5:
0153                         valErr=sqrt(float(val))
0154                         sigZErr.append(valErr)
0155
0156     if verbose == "True":
0157         for i in range(len(x0)):
0158             print("     x0 = "+str(x0[i])+" +/- %1.8f (stats) [cm]" % (x0Err[i]))
0159             print("     y0 = "+str(y0[i])+" +/- %1.8f (stats) [cm]" % (y0Err[i]))
0160             print("     z0 = "+str(z0[i])+" +/- %1.6f (stats) [cm]" % (z0Err[i]))
0161             print("sigmaZ0 = "+str(sigZ[i])+" +/- %1.6f (stats) [cm]" % (sigZErr[i]))
0162
0163     ## Make plots and save to root file
0164     rootFile = TFile("Summary.root","RECREATE");
0165     gROOT.SetStyle("Plain")
0166
0167     hx0_lumi=plot_trending("x0",labels,x0,x0Err)
0168     hx0_lumi.SetTitle("x coordinate of beam spot vs. lumi")
0169
0170     hy0_lumi=plot_trending("y0",labels,y0,y0Err)
0171     hy0_lumi.SetTitle("y coordinate of beam spot vs. lumi")
0172
0173     hz0_lumi=plot_trending("z0",labels,z0,z0Err)
0174     hz0_lumi.SetTitle("z coordinate of beam spot vs. lumi")
0175
0176     hsigZ_lumi=plot_trending("sigmaZ0",labels,sigZ,sigZErr)
0177     hsigZ_lumi.SetTitle("sigma z_{0} of beam spot vs. lumi")
0178
0179     rootFile.Write();
0180     rootFile.Close();
0181
0182
0183 #_________________________________
0184 if __name__ =='__main__':
0185     sys.exit(main())
0186
0187