import ROOT, glob, math, sys, os from ROOT import std from mpl_toolkits.mplot3d import Axes3D import matplotlib.pyplot as plt import matplotlib import random from matplotlib.patches import FancyArrowPatch from mpl_toolkits.mplot3d import proj3d #small class to draw an arrow, used to draw the beam axis class Arrow3D(FancyArrowPatch): def __init__(self, xs, ys, zs, *args, **kwargs): FancyArrowPatch.__init__(self, (0,0), (0,0), *args, **kwargs) self._verts3d = xs, ys, zs def draw(self, renderer): xs3d, ys3d, zs3d = self._verts3d xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M) self.set_positions((xs[0],ys[0]),(xs[1],ys[1])) FancyArrowPatch.draw(self, renderer) class EventDisplay: def __init__(self,fileName, mcFileName): ROOT.gROOT.ProcessLine(".x $VMCWORKDIR/gconfig/rootlogon.C") 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+") self.rootFile=ROOT.TFile(fileName) if self.rootFile.IsZombie() : print "zooooombie" raise TypeError self.tree=self.rootFile.Get("cbmsim") if self.tree == None: print "no tree" raise TypeError self.tcaTpcTracks=ROOT.TClonesArray("GFTrack") self.tcaTpcClusters=ROOT.TClonesArray("TpcCluster") self.tcaTpcPoints=ROOT.TClonesArray("TpcSPHit") self.tree.SetBranchAddress("TpcSPHit",ROOT.AddressOf(self.tcaTpcPoints)) self.tree.TpcSPHit=self.tcaTpcPoints self.tree.SetBranchAddress("TpcTrackPostFit",ROOT.AddressOf(self.tcaTpcTracks)) self.tree.TpcTrackPostFit=self.tcaTpcTracks self.tree.SetBranchAddress("TpcCluster",ROOT.AddressOf(self.tcaTpcClusters)) self.tree.TpcCluster=self.tcaTpcClusters self.totentries=self.tree.GetEntries() self.colors=[] self.colors.append('r') self.colors.append('b') self.colors.append('g') self.colors.append('c') self.colors.append('m') self.colors.append('y') self.colors.append('k') self.colors.append('w') self.colors.append('#FF6600')#orange self.colors.append('#000000')#black self.c=ROOT.TCanvas("Hits from found tracks","Hits from found tracks",1800,400,700,700) self.hist=ROOT.TH2D("hits from tracks","hits from tracks ;x (cm);y (cm)",1000,-60,60,1000,-60,60) self.hist.Draw() self.c2=ROOT.TCanvas("All clusters (mcid)","All clusters (mcid)",1100,400,700,700) self.hist2=ROOT.TH2D("hits from mctracks","hits from mctracks ;x (cm);y (cm)",1000,-60,60,1000,-60,60) self.hist2.Draw() #event number and the gui which should show the eventDisplay def gotoEvent(self, number,gui) : print "goint to event "+ str(number) self.tcaTpcClusters.Delete() self.tcaTpcTracks.Delete() self.tcaTpcPoints.Delete() self.tree.GetEntry(number) print "nClusters: "+str(self.tcaTpcClusters.GetEntriesFast()) +" nTracks: "+str(self.tcaTpcTracks.GetEntriesFast()) self.c.cd() self.graphsTracks = {} #dictionary for graphs with the points of the tpc tracks trackNb=0 tpcDetId=0 gui.axClusterMcId.clear() gui.axClustersFromTrack.clear() gui.axClusterMcId.mouse_init() gui.axClustersFromTrack.mouse_init() # ax = Axes3D(self.fig) # gui.canvas2.draw() #show 3d plot # gui.axClustersFromTrack = self.fig.add_subplot(1, 2, 2, projection='3d') for track in self.tcaTpcTracks: #arrays used by matplotlib x=[] y=[] z=[] cand=track.getCand() self.graphsTracks[trackNb]=ROOT.TGraph() self.graphsTracks[trackNb].SetMarkerColor(trackNb+1) for hitNb in range (cand.getNHits()): hitId=ROOT.workAroundGetHitId(cand,hitNb) tpcDetId=ROOT.workAroundGetDetId(cand,hitNb) hit=self.tcaTpcPoints.At(hitId) coord=hit.getRawHitCoord() self.graphsTracks[trackNb].SetPoint(hitNb,coord[0][0],coord[1][0]) x.append(coord[0][0]) y.append(coord[1][0]) z.append(coord[2][0]) self.graphsTracks[trackNb].Draw("*") if trackNb < len(self.colors): gui.axClustersFromTrack.scatter(x,y,z,c=self.colors[trackNb]) else: gui.axClustersFromTrack.scatter(x,y,z,c="#000000") trackNb+=1 self.c.Update() print trackNb self.c2.cd() self.graphsTracks2 = {} #dictionary for graphs with the points of the tpc tracks test = {} test2 = {} for cluster in self.tcaTpcClusters: mcId = cluster.mcId().DominantID().mctrackID() if mcId not in self.graphsTracks2: test[mcId]=(len(test)) self.graphsTracks2[mcId]=ROOT.TGraph() self.graphsTracks2[mcId].SetMarkerColor(len(test)) test2[mcId]=[] pos=cluster.pos() self.graphsTracks2[mcId].SetPoint(self.graphsTracks2[mcId].GetN(),pos.X(),pos.Y()) test2[mcId].append(pos) for graph in self.graphsTracks2: self.graphsTracks2[graph].Draw("*") x=[] y=[] z=[] for pos in test2[graph]: x.append(pos.X()) y.append(pos.Y()) z.append(pos.Z()) if test[graph] < len(self.colors): gui.axClusterMcId.scatter(x,y,z,c=self.colors[test[graph]]) else: gui.axClusterMcId.scatter(x,y,z,c="#000000") self.c2.Update() gui.axClustersFromTrack.set_ylabel('Y (cm)') gui.axClustersFromTrack.set_zlabel('Z (cm)') gui.axClustersFromTrack.set_xlim3d([-60,60]) gui.axClustersFromTrack.set_ylim3d([-60,60]) gui.axClusterMcId.set_ylabel('Y (cm)') gui.axClusterMcId.set_zlabel('Z (cm)') gui.axClusterMcId.set_xlim3d([-60,60]) gui.axClusterMcId.set_ylim3d([-60,60]) a = Arrow3D([0,0],[0,0],[0,30], mutation_scale=20, lw=1, arrowstyle="-|>", color="r", linestyle="dashed") b = Arrow3D([0,0],[0,0],[0,-40], mutation_scale=20, lw=1, arrowstyle="-|>", color="b", linestyle="dashed") gui.axClusterMcId.add_artist(a) gui.axClusterMcId.add_artist(b) gui.canvasClusterMcId.draw() gui.canvasClustersFromTrack.draw()