import sys sys.path.append('/home/mberger/fopiroot/fopiroot_dev/macro/tpc/FOPI/mberger') import ROOT,glob,math,argparse from functions import Divideh,set_palette,setRange,getColors from anaFile import anaFile class Cluster: def __init__(self): self.pos=ROOT.TVector3() self.posP=ROOT.TPolyMarker3D() self.posP_XY=ROOT.TPolyMarker() self.axes=[] self.axes_XY=[] self.len_XY=[] self.phi=[] self.cov_xy=ROOT.TMatrixD(2,2) self.cov_ell_xy=ROOT.TEllipse() self.axesProj=[] def setPos(self,pos,col=1): self.pos=pos self.posP.SetPoint(0,pos.X(),pos.Y(),pos.Z()) self.posP.SetMarkerStyle(20) self.posP.SetMarkerSize(0.6) self.posP.SetMarkerColor(col) def setPos_XY(self,pos,col): self.posP_XY.SetPoint(0,pos.X(),pos,Y()) self.posP_XY.SetMarkerStyle(20) self.posP_XY.SetMarkerSize(0.6) self.posP_XY.SetMarkerColor(col) def addAxis(self,vect,col=1): axis=ROOT.TPolyLine3D() axis.SetPoint(0,self.pos.X(),self.pos.Y(),self.pos.Z()) axis.SetPoint(1,self.pos.X()+vect.X(), self.pos.Y()+vect.Y(), self.pos.Z()+vect.Z()) axis.SetLineColor(col) self.axes.append(axis) self.addAxis_XY(vect,col) def setPlane(self,norm): ortho1=ROOT.TVector3(norm) ortho1.SetMag(1) ortho2=ortho1.Orthogonal() ortho2.SetMag(1) #print 'ortho1:' #ortho1.Print() #print 'ortho2:' #ortho2.Print() self.p1=ROOT.TVector3(ortho1) self.p2=ROOT.TVector3(ortho2) #print 'othogonal=',ortho1*ortho2 def setPlane2(self,p1,p2): self.p1=p1 self.p2=p2 def getProj(self,vect): proj1=self.p1*vect/self.p1.Mag() proj2=self.p2*vect/self.p1.Mag() return ROOT.TVector2(proj1,proj2) def addAxisProj(self,vect,col=1,): proj1=self.p1*vect/self.p1.Mag() proj2=self.p2*vect/self.p1.Mag() self.axesProj.append(ROOT.TLine(0,0,proj1,proj2)) self.axesProj[-1].SetLineColor(col) def setEllipseProj(self,cov): J=ROOT.TMatrixD(3,2) for i in range(3): J[i][0]=self.p1[i] J[i][1]=self.p2[i] cov_i=ROOT.TMatrixD(ROOT.TMatrix.kInverted,cov) Jt=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,J) if args.cutCov: tmp=ROOT.TMatrixD(cov_i,ROOT.TMatrixD.kMult,J) else: tmp=ROOT.TMatrixD(cov,ROOT.TMatrixD.kMult,J) cov_cut=ROOT.TMatrixD(Jt,ROOT.TMatrixD.kMult,tmp) #cov_cut.Print() if args.cutCov: cov_cut.Invert() #cov_cut.Print() eigenproblem=ROOT.TMatrixDEigen(cov_cut) eigenwert=eigenproblem.GetEigenValues() eigenvector=eigenproblem.GetEigenVectors() teigenvector=ROOT.TVector2(eigenvector[0][0],eigenvector[1][0]) teigenvector2=ROOT.TVector2(eigenvector[0][1],eigenvector[1][1]) #calc cov_cut scale proj=ROOT.TMatrixD(2,1) proj[0][0]=1 proj[1][0]=0 proj_t=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,proj) cov_cut_i=ROOT.TMatrixD(ROOT.TMatrixD.kInverted,cov_cut) if args.cutCov: temp2=ROOT.TMatrixD(cov_cut_i,ROOT.TMatrixD.kMult,proj) else: temp2=ROOT.TMatrixD(cov_cut,ROOT.TMatrixD.kMult,proj) sigma=ROOT.TMatrixD(proj_t,ROOT.TMatrixD.kMult,temp2) #print 'sigma=',sigma[0][0],1./math.sqrt(sigma[0][0]) proj2=ROOT.TMatrixD(2,1) proj2[0][0]=0 proj2[1][0]=1 proj2_t=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,proj2) temp3=ROOT.TMatrixD(cov_cut_i,ROOT.TMatrixD.kMult,proj2) sigma2=ROOT.TMatrixD(proj2_t,ROOT.TMatrixD.kMult,temp3) #print 'sigma2=',sigma2[0][0],1./math.sqrt(sigma2[0][0]) #self.cov_ell_proj=ROOT.TEllipse(0,0,math.sqrt(eigenwert[0][0]),math.sqrt(eigenwert[1][1]),0,360,ROOT.TMath.RadToDeg()*teigenvector.Phi()) self.cov_ell_proj=ROOT.TEllipse(0,0,math.sqrt(eigenwert[1][1]),math.sqrt(eigenwert[0][0]),0,360,ROOT.TMath.RadToDeg()*teigenvector2.Phi()) #self.cov_ell_proj=ROOT.TEllipse(0,0,(eigenwert[0][0]),(eigenwert[1][1]),0,360,ROOT.TMath.RadToDeg()*teigenvector.Phi()) self.cov_ell_proj.SetFillStyle(4000) def addAxis_XY(self,vect,col=1): axis=ROOT.TLine(self.pos.X(),self.pos.Y(), self.pos.X()+vect.X(),self.pos.Y()+vect.Y()) axis.SetLineColor(col) self.axes_XY.append(axis) self.len_XY.append(math.sqrt(vect.X()**2+vect.Y()**2)) self.phi.append(vect.Phi()*ROOT.TMath.RadToDeg()) def setEllipse_XY(self,cov): plane=ROOT.TMatrixD(3,3) for i in range(3): for j in range(3): if i==j and i!=2: plane[i][j]=1 else: plane[i][j]=0 #plane.Print() #self.cov_xy=ROOT.TMatrixD(plane,ROOT.TMatrixD.kMult,cov) #self.cov_xy.Print() for i in range(2): for j in range(2): self.cov_xy[i][j]=cov[i][j] eigenproblem=ROOT.TMatrixDEigen(self.cov_xy) eigenwert=eigenproblem.GetEigenValues() eigenvector=eigenproblem.GetEigenVectors() teigenvector=ROOT.TVector2(eigenvector[0][0],eigenvector[1][0]) teigenvector2=ROOT.TVector2(eigenvector[0][1],eigenvector[1][1]) #eigenwert.Print() self.cov_ell_xy=ROOT.TEllipse(self.pos.X(),self.pos.Y(),math.sqrt(eigenwert[0][0]),math.sqrt(eigenwert[1][1]),0,360,ROOT.TMath.RadToDeg()*teigenvector.Phi()) self.cov_ell_xy.SetFillStyle(4000) def Draw(self,canv): canv.cd() self.posP.Draw("same") for ax in self.axes: ax.Draw("same") def Draw_XY(self,canv): canv.cd() self.posP_XY.Draw("same") for ax in self.axes_XY: ax.Draw("same") #cov_ell=ROOT.TEllipse(self.pos.X(),self.pos.Y(),self.len_XY[0],self.len_XY[1],0,360,self.phi[0]) #cov_ell.Draw("same") self.cov_ell_xy.Draw("same") def Draw_Proj(self,canv): canv.cd() self.cov_ell_proj.Draw("same") for ax in self.axesProj: ax.Draw("same") def calcCovScale(cov,res): cov_i=ROOT.TMatrixD(ROOT.TMatrixD.kInverted,cov) # if cov[0][0]==0 or cov[1][1]==0 or cov[2][2]==0: # cov.Print() residual=ROOT.TMatrixD(3,1) for i in range(3): residual[i][0]=res[i] residual_t=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,residual) #temp2=ROOT.TMatrixD(cov_i,ROOT.TMatrixD.kMult,residual) temp2=ROOT.TMatrixD(cov,ROOT.TMatrixD.kMult,residual) lam=ROOT.TMatrixD(residual_t,ROOT.TMatrixD.kMult,temp2) lam=1./lam[0][0] if lam!=0: num=lam return num parser=argparse.ArgumentParser(description='plot a cluster and its axes') parser.add_argument('recofile',help='the reco file',type=str) parser.add_argument('digifile',help='the raw file with the digis',type=str) parser.add_argument('--padf',help='the padplane file to use',type=str,default='tpc/parfiles/padPlane_FOPI.dat') parser.add_argument('--ft0',help='ft for digi z calc',type=float,default=0) parser.add_argument('--ftbin',help='ftbin for digi z calc',type=float,default=64.3087) parser.add_argument('--vd',help='drift velocity for digi z calc',type=float,default=0.00237708) parser.add_argument('--fzGem',help='z gem for digi z calc',type=float,default=0) parser.add_argument('--sizecut',help='cut on size (inside window)',type=int,nargs=2,default=[0,10000]) parser.add_argument('--zcut',help='cut on zpos (inside window)',type=float,nargs=2,default=[-1000,1000]) parser.add_argument('--errVcut',help='cut on error V',type=float,nargs=2,default=[0,1000]) parser.add_argument('--hlp',help='display help',action='store_true') parser.add_argument('--cutCov',help='use cut cov instead of cov proj',action='store_true') args=parser.parse_args() if args.hlp: parser.print_help() exit(0) ROOT.gROOT.ProcessLine(".x rootlogon.C") ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")') ROOT.gROOT.LoadMacro("stlPYROOT.h+") set_palette() RFile=ROOT.TFile(args.recofile,'read') tree= RFile.Get("cbmsim") if tree==None: print 'no cbmsim tree found' if type(Rfile)==ROOT.TFile: Rfile.Close() exit() 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.AddFriend('cbmsim',args.digifile) tree.SetBranchStatus("*", 0) tree.SetBranchStatus("TpcCluster.*",1) tree.SetBranchStatus("TrackFitStat_0.*",1) tree.SetBranchStatus("TpcSPHit.*",1) tree.SetBranchStatus("TpcDigi.*",1) canvas=ROOT.TCanvas('c1','3D') canvas_XY=ROOT.TCanvas('c2','2D') ev=-1 for e in tree: ev+=1 print 'event:',ev for tr in e.TrackFitStat_0: print 'new track' posX=tr.GetHitPositionsX() posY=tr.GetHitPositionsY() posZ=tr.GetHitPositionsZ() MCposX=tr.GetMcRefPosX() MCposY=tr.GetMcRefPosY() MCposZ=tr.GetMcRefPosZ() axis1=tr.GetAxis1() axis2=tr.GetAxis2() axis3=tr.GetAxis3() sig=tr.GetSigResMC() errV=tr.GetErrV() planeU=tr.GetPlaneU() planeV=tr.GetPlaneV() errU=tr.GetErrU() errV=tr.GetErrV() size=tr.GetClSizes() PosMoms=tr.GetMcPosMom() clcount=-1 for clid in tr.GetHitIDs(): clcount+=1 SpHit=e.TpcSPHit.At(clid) cl=e.TpcCluster.At(SpHit.getClusterID()) # if cl.axis(2).Theta()len(size)-1 or clcount>len(posZ)-1: continue if size[clcount]args.sizecut[1]: continue if posZ[clcount]args.zcut[1]: continue if errV[clcount]args.errVcut[1]: continue #print 'size:',size[clcount] #print 'zpos:',posZ[clcount] #print 'amp:',cl.amp() clusP=Cluster() clusP.setPos(cl.pos()) eigenval=cl.sig(3) # clusP.addAxis(cl.axis(0),6) # clusP.addAxis(cl.axis(1),7) # clusP.addAxis(cl.axis(2),8) clusP.addAxis(axis1[clcount],6) clusP.addAxis(axis2[clcount],7) clusP.addAxis(axis3[clcount],8) #clusP.setEllipse_XY(cl.cov()) #print 'Theta: 0=',cl.axis(0).Theta()*ROOT.TMath.RadToDeg(),'1=',cl.axis(1).Theta()*ROOT.TMath.RadToDeg(),'2=',cl.axis(2).Theta()*ROOT.TMath.RadToDeg() pos=ROOT.TVector3(MCposX[clcount],MCposY[clcount],MCposZ[clcount]) mcP=Cluster() mcP.setPos(pos,2) PosMoms[clcount].SetMag(1) mcP.addAxis(PosMoms[clcount],2) mcRes=pos-clusP.pos #print 'scaling=',math.sqrt(calcCovScale(cl.cov(),mcRes)) #print 'sigresMC=',sig[clcount],mcRes.Mag()/sig[clcount] _3drange=1 apos=clusP.pos histo=ROOT.TH3D('htmp','Cluster', 100,apos.X()-_3drange,apos.X()+_3drange, 100,apos.Y()-_3drange,apos.Y()+_3drange, 100,apos.Z()-_3drange,apos.Z()+_3drange) histo.GetXaxis().SetTitle('X') histo.GetYaxis().SetTitle('Y') histo.GetZaxis().SetTitle('Z') histo_XY=ROOT.TH2D('htmp_xy','Cluster XY', 1000,-0.1,+0.1, 1000,-0.1,+0.1) histo_XY.GetXaxis().SetTitle('X') histo_XY.GetYaxis().SetTitle('Y') digis=ROOT.TPolyMarker3D() digimap=cl.getDigiMap() dcounter=-1 for dig in digimap: dcounter+=1 digi=e.TpcDigi.At(dig.first) dx=float(DigiMapper[digi.padId()][0]) dy=float(DigiMapper[digi.padId()][1]) dz=(digi.t()*args.ftbin+args.ft0)*args.vd+args.fzGem #print 'adding digi',dx,dy,dz digis.SetPoint(dcounter,dx,dy,dz) digis.SetMarkerStyle(24) digis.SetMarkerSize(0.6) #plane projected stuff #print 'mcres:' #mcRes.Print() #clusP.setPlane(mcRes) clusP.setPlane2(planeU[clcount],planeV[clcount]) clusP.addAxisProj(axis1[clcount],6) clusP.addAxisProj(axis2[clcount],7) clusP.addAxisProj(axis3[clcount],8) planeUscaled=ROOT.TVector3(planeU[clcount]) #planeUscaled.SetMag(planeUscaled.Mag()*errU[clcount]) planeUscaled.SetMag(errU[clcount]) planeVscaled=ROOT.TVector3(planeV[clcount]) #planeVscaled.SetMag(planeVscaled.Mag()*errV[clcount]) planeVscaled.SetMag(errV[clcount]) clusP.addAxisProj(planeUscaled,63) clusP.addAxisProj(planeVscaled,73) clusP.setEllipseProj(cl.cov()) mcResProj=clusP.getProj(mcRes) #print 'orig Mag:',mcRes.Mag(),'plane Mag:',math.sqrt(mcResProj.X()**2+mcResProj.Y()**2) mcResProjP=ROOT.TLine(0,0,mcResProj.X(),mcResProj.Y()) mcResProjP.SetLineColor(2) #sigVect=ROOT.TVector2(sig[clcount],0) sigVect3D=ROOT.TVector3(mcRes) sigVect3D.SetMag(sig[clcount]) sigVect=clusP.getProj(sigVect3D) #print 'sigvect:' #sigVect.Print() sigPoint=ROOT.TPolyMarker() sigPoint.SetPoint(0,sigVect.X(),sigVect.Y()) sigPoint.SetMarkerStyle(20) sigPoint.SetMarkerSize(.6) sigPoint.SetMarkerColor(30) #Uvec=getProj( canvas.cd() histo.Draw() mcP.Draw(canvas) clusP.Draw(canvas) digis.Draw("same") canvas.Update() canvas_XY.cd() histo_XY.Draw() #mcP.Draw_XY(canvas_XY) #clusP.Draw_XY(canvas_XY) clusP.Draw_Proj(canvas_XY) mcResProjP.Draw("same") sigPoint.Draw("same") canvas_XY.Update() u=raw_input("next?") if u=='q': exit() del histo del digis del histo_XY