import ROOT, argparse, math, sys, os, copy pandapath=os.environ.get('PANDAPATH') sys.path.append(pandapath+'/macro/tpc/FOPI/mberger') from functions import Divideh parser=argparse.ArgumentParser(description='plot correlation of 2pad cluster') parser.add_argument('recofile',help='recofile',type=str) parser.add_argument('--events',help='how many events to process',type=int,default=-1) parser.add_argument('--padf',help='the padplane file to use',type=str,default='tpc/parfiles/padPlane_FOPI.dat') parser.add_argument('--rfile',help='the rootfile to save the shit into',type=str,default='None') parser.add_argument('--pfile',help='the pdf file to save the shit into',type=str,default='None') parser.add_argument('--read',help='Read the histograms from file',action='store_true') parser.add_argument('--data',help='read real data, not sim',action='store_true') parser.add_argument('--moreData',help='additional data files',type=str,default=['None'],nargs='+') args=parser.parse_args() ROOT.gROOT.ProcessLine(".x rootlogon.C") ROOT.gROOT.ProcessLine("gStyle->SetPalette(1)") ROOT.gROOT.LoadMacro("stlPYROOT.h+") DigiMapper = [] file = open(args.padf, "r") lines=file.readlines() file.close() lines.reverse() for line in lines: koord = line.split(" ") DigiMapper.append([koord[3], koord[4]]) tree=ROOT.TChain("cbmsim","cbmsim") tree.Add(args.recofile) if args.moreData[0]!='None': # tree.Add(args.recofile) for addfile in args.moreData: tree.Add(addfile) #else: # Infile=ROOT.TFile(args.recofile,"read") # tree=Infile.Get("cbmsim") if not args.data: tree.AddFriend("cbmsim",args.digifile) tree.SetBranchStatus("CutCosmics.*",1) tree.SetBranchStatus("TpcSPHit.*",1) tree.SetBranchStatus("TpcCluster.*",1) if not args.data: tree.SetBranchStatus("TpcDigi.*",1) ev=-1 c1=ROOT.TCanvas("c1","c1",1000,1400) c1.Divide(1,2) hresVsAmpfrac=ROOT.TH2D("hresVsAmpfrac","MC Residual vs Amp Fraction",1200,-0.1,1.1,1000,-0.3,0.3) haxphi=ROOT.TH1D("haxphi","axis phi size2",720,-360,360) c2=ROOT.TCanvas("c2","amp fracs",1000,1000) c2.Divide(3,3) c3=ROOT.TCanvas("c3","amp fracs profiles",1000,1000) c3.Divide(3,3) c4=ROOT.TCanvas("c4","amp frac res norm dir",1000,1000) c4.Divide(3,3) resVsAmpfracHists=[] resAmpfrac=[] Ampfrac=[] for i in range(9): resVsAmpfracHists.append(ROOT.TH2D("hresVsAmpfrac"+str(i), "MC Residual vs Amp Fraction (Size "+str(i+2)+")", 1200,-0.1,1.1,1000,-0.3,0.3)) resAmpfrac.append(ROOT.TH2D("hresAmpfrac"+str(i),"Ampfrac dir res (Size "+str(i+2)+")", 22,-1.1,1.1,22,-1.1,1.1)) Ampfrac.append(ROOT.TH2D("hAmpfrac"+str(i),"Ampfrac dir (Size "+str(i+2)+")", 22,-1.1,1.1,22,-1.1,1.1)) tree.SetBranchStatus("*",0) tree.SetBranchStatus("CutCosmics.*",1) if not args.read: for e in tree: ev+=1 if ev%100==0: print "Event:",ev if args.events!=-1 and ev>=args.events: break for tr in e.CutCosmics: if args.data and tr.GetLength()<3: continue size2D=tr.Get2DClSizes() size3D=tr.GetClSizes() if args.data: mcres=tr.GetRes() else: mcres=tr.GetResMC() amps=tr.GetAmps() xpos=tr.GetHitPositionsX() ids=tr.GetHitIDs() for icl in range(len(mcres)): sphit=e.TpcSPHit.At(ids[icl]) cluster=e.TpcCluster.At(sphit.getClusterID()) digimap=cluster.getDigiMap() eta=-1 leftamp=-1 rightamp=-1 leftpos=[80,80] rightpos=[-80,-80] bigpos=[0,0] bigamp=0 smallpos=[0,0] smallamp=9999999999999999999 digicol={} for dig in digimap: digi=e.TpcDigi.At(dig.first) padId=str(digi.padId()) dx=float(DigiMapper[int(padId)][0]) dy=float(DigiMapper[int(padId)][1]) amp=digi.amp() if digicol.get(padId,None)!=None: digicol[padId]['amp']+=amp else: digicol[padId]={} digicol[padId]['amp']=amp digicol[padId]['dx']=dx digicol[padId]['dy']=dy for padId in digicol: dx=digicol[padId]['dx'] dy=digicol[padId]['dy'] amp=digicol[padId]['amp'] if dxrightpos[0]: rightpos=[dx,dy] rightamp=amp if ampbigamp: bigamp=amp bigpos=[dx,dy] if (size2D[icl]==2 and size3D[icl]==2): eta=leftamp/amps[icl] axis=ROOT.TVector3(rightpos[0]-leftpos[0],rightpos[1]-leftpos[1],0) #print leftamp,rightamp,smallamp,bigamp if axis.X()<0: axis*=-1 if axis.Mag()!=0: axis.SetMag(1) hresVsAmpfrac.Fill(eta,mcres[icl]*axis) if size2D[icl]<11 and size2D[icl]>1: eta=smallamp/(smallamp+bigamp) axis=ROOT.TVector3(smallpos[0]-bigpos[0],smallpos[1]-bigpos[1],0) if size2D[icl]==2: haxphi.Fill(axis.Phi()*ROOT.TMath.RadToDeg()) if axis.Mag()!=0: axis.SetMag(1) resVsAmpfracHists[size2D[icl]-2].Fill(eta,mcres[icl]*axis) resVsAmpfracHists[size2D[icl]-2].Fill(1-eta,-1*(mcres[icl]*axis)) weight=mcres[icl]*axis resAmpfrac[size2D[icl]-2].Fill(axis.X()*(1-eta),axis.Y()*(1-eta),weight) Ampfrac[size2D[icl]-2].Fill(axis.X()*(1-eta),axis.Y()*(1-eta)) if mcres[icl]*axis < -0.05 and eta<0.1 and (size2D[icl]==2 or size2D[icl]==3): print 'event',ev,'cluster:',sphit.getClusterID(),icl else: infile=ROOT.TFile(args.rfile,'read') hresVsAmpfrac=copy.deepcopy(infile.Get("hresVsAmpfrac")) for i in range(9): resVsAmpfracHists[i]=copy.deepcopy(infile.Get("hresVsAmpfrac"+str(i))) infile.Close() c1.cd(1) hresVsAmpfrac.Draw("colz") c1.cd(2) haxphi.Draw() profiles=[] fits=[] fitfunc=ROOT.TF1("func","pol4",-0.1,0.6) resAmpfracNorm=[] for i in range(9): c2.cd(i+1) resVsAmpfracHists[i].Draw("colz") profiles.append(resVsAmpfracHists[i].ProfileX()) c3.cd(i+1) profiles[-1].Draw() profiles[-1].Fit(fitfunc) c4.cd(i+1) print "drawing",i resAmpfracNorm.append(Divideh(resAmpfrac[i],Ampfrac[i],"hresAmpfracNorm"+str(i),"res amp frac norm (Size "+str(i+2)+")")) # for idi in range(resAmpfracNorm[-1].GetBinsX()): # for ibi in range(resAmpfracNorm[-1].GetBinsY()): # if resAmpfracNorm[-1].GetBinsContent(idi,ibi)==0: # resAmpfracNorm[-1].SetBinsContent(idi,ibi-) resAmpfracNorm[-1].Draw("colz") #resAmpfrac[i].Draw("colz") print "updating" c1.Update() print "updating 2" c2.Update() print "updating 3" c3.Update() print "updating 4" c4.Update() if args.pfile!='None': c1.SaveAs(args.pfile+'(') c2.SaveAs(args.pfile) c3.SaveAs(args.pfile) c4.SaveAs(args.pfile+')') if args.rfile!='None': outfile=ROOT.TFile(args.rfile,'recreate') c1.Write() c2.Write() c3.Write() c4.Write() hresVsAmpfrac.Write() for i in range(9): resVsAmpfracHists[i].Write() profiles[i].Write() resAmpfrac[i].Write() Ampfrac[i].Write() u='n' while u!='q': u=raw_input('quit?')