import ROOT, glob, math, sys, os from ROOT import std from mpl_toolkits.mplot3d import Axes3D import matplotlib.pyplot as plt dummy=False for iarg in range(len(sys.argv)) : arg = sys.argv[iarg] if arg == "-path": dir = sys.argv[iarg+1]; if arg == "-run": run = sys.argv[iarg+1]; if arg =="-dummy": dummy=True #files = glob.glob(dir + "/*"+str(run)+"*.smoothed.reco.root") #getting list of files runStr=str(run) if runStr=="0": runStr="" files = glob.glob(dir + "/*"+runStr+"*smoothed.reco.root") files.sort() ROOT.gROOT.ProcessLine(".x $VMCWORKDIR/gconfig/rootlogon.C") #Getting absolute path of script path=os.path.dirname(sys.argv[0]) path=os.path.abspath(path) #load workaround to return by reference used in Genfit ROOT.gROOT.ProcessLine(".L $SCRIPTPATH/cand_workaround.C+") def drawCdcSectorBorders(): ll=[] lc=[] for iB in range(0,16): iX=20*math.cos(iB*22.5/180*math.pi) iY=20*math.sin(iB*22.5/180*math.pi) oX=80*math.cos(8./180*math.pi+iB*22.5/180*math.pi) oY=80*math.sin(8./180*math.pi+iB*22.5/180*math.pi) iXc=20*math.cos(iB*22.5/180*math.pi+11.25/180*math.pi) iYc=20*math.sin(iB*22.5/180*math.pi+11.25/180*math.pi) oXc=80*math.cos(8./180*math.pi+iB*22.5/180*math.pi+11.25/180*math.pi) oYc=80*math.sin(8./180*math.pi+iB*22.5/180*math.pi+11.25/180*math.pi) ### print iX,iY,oX,oY ll.append(ROOT.TLine(iX,iY,oX,oY)) ll[iB].SetLineColor(ROOT.kRed) ll[iB].Draw() lc.append(ROOT.TLine(iXc,iYc,oXc,oYc)) lc[iB].SetLineColor(ROOT.kBlack) lc[iB].Draw() return ll,lc for file in files: print file if dummy==True: continue gfile=ROOT.TFile(file) if gfile.IsZombie() : print "zooooombie" continue tree=gfile.Get("cbmsim") if tree == None: print "no tree" continue #dummy reco: fRun=ROOT.FairRunAna() fRun.SetInputFile(file) fMagField = ROOT.FOPIField(0.616); fMagField.SetTargetOffset(-61); fRun.SetField(fMagField) ROOT.GFFieldManager.getInstance().init(ROOT.PndFieldAdaptor(fRun.GetField())); ROOT.GFMaterialEffects.getInstance().init(ROOT.GFTGeoMaterialInterface()); taroff=-61 targetpos=taroff+41 geoFile="/tmp/FopiGeom_s339"+str(targetpos)+".root" fRun.SetGeomFile(geoFile) rtdb = fRun.GetRuntimeDb(); parInput1=ROOT.FairParRootFileIo() parInput1.open(file.replace("reco","param")) rtdb.setFirstInput(parInput1); fRun.SetOutputFile("dummy.root"); fRun.Init() #initialize TClonesArrays used tcaCombTracks=ROOT.TClonesArray("GFTrack") tcaCdcTracks=ROOT.TClonesArray("GFTrack") tcaTpcTracks=ROOT.TClonesArray("GFTrack") # tcaTpcClusters=ROOT.TClonesArray("TpcCluster") tcaTpcPoints=ROOT.TClonesArray("TpcSPHit") # tcaCdcPoints=ROOT.TClonesArray("CdcHit") tcaCdcPointsGF=ROOT.TClonesArray("PseudoSpacePoint") tree.SetBranchAddress("CdcTrackPostFit",ROOT.AddressOf(tcaCdcTracks)) tree.CdcTrack=tcaCdcTracks tree.SetBranchAddress("TpcTrackPostFit",ROOT.AddressOf(tcaTpcTracks)) tree.TpcTrackPostFit=tcaTpcTracks tree.SetBranchAddress("FopiTrackPostFit",ROOT.AddressOf(tcaCombTracks)) tree.FopiTrackPostFit=tcaCombTracks # tree.SetBranchAddress("TpcCluster",ROOT.AddressOf(tcaTpcClusters)) # tree.TpcCluster=tcaTpcClusters tree.SetBranchAddress("TpcSPHit",ROOT.AddressOf(tcaTpcPoints)) tree.TpcSPHit=tcaTpcPoints # tree.SetBranchAddress("CdcHit",ROOT.AddressOf(tcaCdcPoints)) # tree.CdcHit=tcaCdcPoints tree.SetBranchAddress("CdcSpacepoint",ROOT.AddressOf(tcaCdcPointsGF)) tree.CdcPseudoSPHit=tcaCdcPointsGF #simple colorMap for matplotlib colors=[] colors.append('r') colors.append('b') colors.append('g') colors.append('c') colors.append('m') colors.append('y') colors.append('k') colors.append('w') colors.append('#FF6600')#orange totentries=tree.GetEntries() print totentries i=1 #canvas for 2D event display c=ROOT.TCanvas("All tracks 2D","All tracks 2D",1921,0,700,600) c2=ROOT.TCanvas("Matched tracks 2D","Matched tracks 2D",2700,000,700,600) while i < totentries: tree.GetEntry(i) # if i%500==0: print i i=i+1 if tcaCombTracks == None or tcaCdcTracks == None or tcaTpcTracks == None : continue nTpc=tcaTpcTracks.GetEntriesFast() nCdc=tcaCdcTracks.GetEntriesFast() nComb=tcaCombTracks.GetEntriesFast() #condition can be removed if needed if True: print "Event: "+str(i)+"\t nComb: "+str(nComb) +"\t nCdc: "+str(nCdc)+ "\t nTpc: " + str(nTpc) #Plot all TPC and CDC tracks, matched and unmatched in 2D c.cd() hist=ROOT.TH2D("All_tracks_2D","All tracks ;x (cm);y (cm)",1000,-90,90,1000,-90,90) #plotting the TPC-tracks graphs = {} #dictionary for graphs with the points of the tpc tracks trackNb=0 hist.Draw() a,b= drawCdcSectorBorders() # hist.GetXaxis().SetRangeUser(-25,25) # hist.GetYaxis().SetRangeUser(-25,25) tpcDetId=0 for track in tcaTpcTracks: cand=track.getCand() graphs[trackNb]=ROOT.TGraph() graphs[trackNb].SetMarkerColor(trackNb+1) for hitNb in range (cand.getNHits()): hitId=ROOT.workAroundGetHitId(cand,hitNb) tpcDetId=ROOT.workAroundGetDetId(cand,hitNb) hit=tcaTpcPoints.At(hitId) coord=hit.getRawHitCoord() # print coord[0],coord[1] graphs[trackNb].SetPoint(hitNb,coord[0],coord[1]) #print str(coord[0][0])+", "+str(coord[1][0])+", "+str(coord[2][0]) print "Drawing tpc graph nb:"+str(trackNb)+" .GetN() "+ str(graphs[trackNb].GetN()) graphs[trackNb].Draw("*") trackNb+=1 #plotting the CDC-tracks trackNbCdc=0 graphsCdc = {} graphsCdcTracks={} for track in tcaCdcTracks: graphsCdc[trackNbCdc]=ROOT.TGraph() graphsCdc[trackNbCdc].SetMarkerColor(trackNbCdc+1) graphsCdcTracks[trackNbCdc]=ROOT.TGraph() graphsCdcTracks[trackNbCdc].SetMarkerColor(trackNbCdc+1) graphsCdcTracks[trackNbCdc].SetLineColor(trackNbCdc+1) hitNb=0 cand=track.getCand() for hitNb in range (cand.getNHits()): hitId=ROOT.workAroundGetHitId(cand,hitNb) cdcDetId=ROOT.workAroundGetDetId(cand,hitNb) hit=tcaCdcPointsGF.At(hitId) coord=hit.getRawHitCoord() graphsCdc[trackNbCdc].SetPoint(hitNb,coord[0],coord[1]) #print str(coord[0][0])+", "+str(coord[1][0])+", "+str(coord[2][0]) # for hitId in track.GetCdcIdHits(): # hit=tcaCdcPoints.At(hitId) # coord=hit.GetHitPos() # graphsCdc[trackNbCdc].SetPoint(hitNb,coord.X(),coord.Y()) # hitNb+=1 # print str(coord.X())+", "+str(coord.Y())+", "+str(coord.Z())+" bla "+str( graphsCdc[trackNbCdc].GetN()) rep=track.getCardinalRep() dummy=ROOT.TVector3() if rep.getMom(rep.getReferencePlane()).Perp()>0.05 and rep.getStatusFlag()==0 and rep.getPos(rep.getReferencePlane()).Perp()<40: # rep.Print() rep.getPos(rep.getReferencePlane()).Print() try: for iStep in range(300): pos=rep.stepalongChange(-0.1) # pos.Print() graphsCdcTracks[trackNbCdc].SetPoint(iStep,pos.X(),pos.Y()) except Exception: print "failed" else: print "not extrapolating track: statusFlag:",rep.getStatusFlag(),", pt: ", rep.getMom(rep.getReferencePlane()).Perp(), ", trackposR: ", rep.getPos(rep.getReferencePlane()).Perp() print "Drawing cdc graph nb:"+str(trackNbCdc)+" .GetN() "+ str(graphsCdc[trackNbCdc].GetN()) graphsCdc[trackNbCdc].Draw("*") graphsCdc[trackNbCdc].SetMarkerStyle(24) graphsCdcTracks[trackNbCdc].Draw("C") graphsCdc[trackNbCdc].SetMarkerStyle(32) trackNbCdc+=1 c.Draw() c.Update() #Plot all combined tracks, in 2D and 3D c2.cd() hist2=ROOT.TH2D("comb_tracks_2D","combined tracks 2d ;x (cm);y (cm)",1000,-90,90,1000,-90,90) hist2.Draw() hist2.GetXaxis().SetRangeUser(-25,25) hist2.GetYaxis().SetRangeUser(-25,25) graphsTot = {} graphsTot3d = {} trackNbTot=0 #initializing 3D-matplotlib-figure fig = plt.figure() ax = Axes3D(fig) for track in tcaCombTracks: #arrays of points used by matplotlib x=[] y=[] z=[] cand=track.getCand() graphsTot[trackNbTot]=ROOT.TGraph() graphsTot[trackNbTot].SetMarkerColor(trackNbTot+1) for hitNb in range (cand.getNHits()): detId=ROOT.workAroundGetDetId(cand,hitNb) hitId=ROOT.workAroundGetHitId(cand,hitNb) if detId==tpcDetId: hit=tcaTpcPoints.At(hitId) coord=hit.getRawHitCoord() graphsTot[trackNbTot].SetPoint(hitNb,coord[0],coord[1]) x.append(coord[0]) y.append(coord[1]) z.append(coord[2]) else: hit=tcaCdcPointsGF.At(hitId) coord=hit.getRawHitCoord() graphsTot[trackNbTot].SetPoint(hitNb,coord[0],coord[1]) x.append(coord[0]) y.append(coord[1]) z.append(coord[2]) # print "detId "+str(detId) #only plot tracks for wich one has color !!! if trackNbTot < len(colors): ax.scatter(x,y,z,c=colors[trackNbTot]) graphsTot[trackNbTot].Draw("*") trackNbTot+=1 ax.set_xlabel('X (cm)') ax.set_ylabel('Y (cm)') ax.set_zlabel('Z (cm)') ax.set_xlim3d([-60,60]) ax.set_ylim3d([-60,60]) c2.Draw() c2.Update() plt.show() #show 3d plot # bah = raw_input("enter to continue") input()