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) files = glob.glob(dir + "/TpcFopiRecoPersist_run"+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 "+path+"/cand_workaround.C+") 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 #initialize TClonesArrays used tcaCombTracks=ROOT.TClonesArray("GFTrack") tcaCdcTracks=ROOT.TClonesArray("CdcTrack") tcaTpcTracks=ROOT.TClonesArray("GFTrack") tcaTpcClusters=ROOT.TClonesArray("TpcCluster") tcaTpcPoints=ROOT.TClonesArray("TpcSPHit") tcaCdcPoints=ROOT.TClonesArray("CdcHit") tcaCdcPointsGF=ROOT.TClonesArray("PseudoSpacePoint") tree.SetBranchAddress("CdcTrack",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("CdcPseudoSPHit",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=0 #canvas for 2D event display c=ROOT.TCanvas("All tracks 2D","All tracks 2D") c2=ROOT.TCanvas("Matched tracks 2D","Matched tracks 2D") 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 nComb>nCdc: 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,-60,60,1000,-60,60) #plotting the TPC-tracks graphs = {} #dictionary for graphs with the points of the tpc tracks trackNb=0 hist.Draw() 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() graphs[trackNb].SetPoint(hitNb,coord[0][0],coord[1][0]) #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 = {} for track in tcaCdcTracks: graphsCdc[trackNbCdc]=ROOT.TGraph() graphsCdc[trackNbCdc].SetMarkerColor(trackNbCdc+1) hitNb=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()) print "Drawing cdc graph nb:"+str(trackNbCdc)+" .GetN() "+ str(graphsCdc[trackNbCdc].GetN()) graphsCdc[trackNbCdc].Draw("*") graphsCdc[trackNbCdc].SetMarkerStyle(24) 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,-60,60,1000,-60,60) hist2.Draw() 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][0],coord[1][0]) x.append(coord[0][0]) y.append(coord[1][0]) z.append(coord[2][0]) else: hit=tcaCdcPointsGF.At(hitId) coord=hit.getRawHitCoord() graphsTot[trackNbTot].SetPoint(hitNb,coord[0][0],coord[1][0]) x.append(coord[0][0]) y.append(coord[1][0]) z.append(coord[2][0]) # 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()