Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-11-26 02:34:40

0001 #!/usr/bin/env python3
0002 
0003 from builtins import range
0004 import ROOT
0005 from array import array
0006 
0007 from Validation.RecoTrack.plotting.ntuple import *
0008 import analysis
0009 
0010 from math import sqrt, copysign, sin, cos, pi
0011 
0012 class EventPlotter(object):
0013 
0014     def __init__(self):
0015         self.plots_2D = []
0016         self.plots_3D = []
0017         c = ROOT.TColor()
0018         self.colors_G = [c.GetColor(0,255,0), c.GetColor(0,185,0), c.GetColor(50,255,50), \
0019                        c.GetColor(0,100,0), c.GetColor(50,155,0), c.GetColor(0,70,155), \
0020                        c.GetColor(0,255,0), c.GetColor(0,255,0), c.GetColor(0,255,0)]
0021         self.colors_B = [c.GetColor(0,0,255), c.GetColor(0,0,155), c.GetColor(50,50,255), \
0022                        c.GetColor(0,0,80), c.GetColor(50,0,155), c.GetColor(0,70,155), \
0023                        c.GetColor(0,0,255), c.GetColor(0,0,255), c.GetColor(0,0,255)]
0024     
0025 
0026     def Reset(self):
0027         self.plots_2D = []
0028         self.plots_3D = []
0029  
0030     def PlotEvent3DHits(self, event, flag="PSG"):
0031         '''
0032         Plots the 3D hits of an event.
0033         flag is an string which defines which hit types to plot
0034         (p for pixel, s for strip, g for glued)
0035         '''        
0036         if('p' in flag or 'P' in flag):
0037             pixel_hits = event.pixelHits()
0038             pix_coords = []
0039             for hit in pixel_hits:
0040                 pix_coords.append(hit.x())
0041                 pix_coords.append(hit.y())
0042                 pix_coords.append(hit.z())
0043             pix_plot = ROOT.TPolyMarker3D(len(pix_coords)/3, array('f', pix_coords), 1)
0044             pix_plot.SetMarkerColor(4)
0045 
0046         if('s' in flag or 'S' in flag):
0047             strip_hits = event.stripHits()
0048             str_coords = []
0049             for hit in strip_hits:
0050                 str_coords.append(hit.x())
0051                 str_coords.append(hit.y())
0052                 str_coords.append(hit.z())
0053             str_plot = ROOT.TPolyMarker3D(len(str_coords)/3, array('f', str_coords), 1)
0054             str_plot.SetMarkerColor(2)
0055 
0056         if('g' in flag or 'G' in flag):
0057             glued_hits = event.gluedHits()
0058             glu_coords = []
0059             for hit in glued_hits:
0060                 glu_coords.append(hit.x())
0061                 glu_coords.append(hit.y())
0062                 glu_coords.append(hit.z())
0063             glu_plot = ROOT.TPolyMarker3D(len(glu_coords)/3, array('f', glu_coords), 1)
0064             glu_plot.SetMarkerColor(3)        
0065             
0066         if('p' in flag or 'P' in flag): self.plots_3D.append(pix_plot)
0067         if('s' in flag or 'S' in flag): self.plots_3D.append(str_plot)
0068         if('g' in flag or 'G' in flag): self.plots_3D.append(glu_plot)        
0069 
0070     def PlotXY(self, event, limits=[-1000,1000], flag="PSG"):
0071         '''
0072         Plots the hits of an event in an XY plane.
0073         flag is an string which defines which hit types to plot
0074         (p for pixel, s for strip, g for glued)
0075         '''
0076         #ntuple.tree().GetBranch("pix_x")
0077         if('p' in flag or 'P' in flag):
0078             pixel_hits = event.pixelHits()
0079             pix_x = []
0080             pix_y = []
0081             for hit in pixel_hits:
0082                 if(limits[0] < hit.z() < limits[1]):
0083                     pix_x.append(hit.x())
0084                     pix_y.append(hit.y())
0085             pix_plot = ROOT.TGraph(len(pix_x), array('f', pix_x), array('f', pix_y))
0086             pix_plot.SetMarkerColor(4)
0087 
0088         if('s' in flag or 'S' in flag):
0089             strip_hits = event.stripHits()
0090             str_x = []
0091             str_y = []
0092             for hit in strip_hits:
0093                 if(limits[0] < hit.z() < limits[1]):
0094                     str_x.append(hit.x())
0095                     str_y.append(hit.y())
0096             str_plot = ROOT.TGraph(len(str_x), array('f', str_x), array('f', str_y))
0097             str_plot.SetMarkerColor(2)
0098 
0099         if('g' in flag or 'G' in flag):
0100             glued_hits = event.gluedHits()
0101             glu_x = []
0102             glu_y = []
0103             for hit in glued_hits:
0104                 if(limits[0] < hit.z() < limits[1]):
0105                     glu_x.append(hit.x())
0106                     glu_y.append(hit.y())
0107             glu_plot = ROOT.TGraph(len(glu_x), array('f', glu_x), array('f', glu_y))
0108             glu_plot.SetMarkerColor(3)        
0109 
0110         plot = ROOT.TMultiGraph()
0111 
0112         if('p' in flag or 'P' in flag): plot.Add(pix_plot,"P")
0113         if('s' in flag or 'S' in flag): plot.Add(str_plot,"P")
0114         if('g' in flag or 'G' in flag): plot.Add(glu_plot,"P")
0115         self.plots_2D.append(plot)
0116 
0117     def PlotZY(self, event, limits=[-1000,1000], flag="PSG"):
0118         '''
0119         Plots the hits of an event in an ZR plane.
0120         flag is an string which defines which hit types to plot
0121         (p for pixel, s for strip, g for glued)
0122         '''
0123         #ntuple.tree().GetBranch("pix_x")
0124         if('p' in flag or 'P' in flag):
0125             pixel_hits = event.pixelHits()
0126             pix_z = []
0127             pix_y = []
0128             for hit in pixel_hits:
0129                 if(limits[0] < hit.z() < limits[1]):
0130                     pix_z.append(hit.z())
0131                     pix_y.append(hit.y())
0132             pix_plot = ROOT.TGraph(len(pix_z), array('f', pix_z), array('f', pix_y))
0133             pix_plot.SetMarkerColor(4)
0134 
0135         if('s' in flag or 'S' in flag):
0136             strip_hits = event.stripHits()
0137             str_z = []
0138             str_y = []
0139             for hit in strip_hits:
0140                 if(limits[0] < hit.z() < limits[1]):
0141                     str_z.append(hit.z())
0142                     str_y.append(hit.y())
0143             str_plot = ROOT.TGraph(len(str_z), array('f', str_z), array('f', str_y))
0144             str_plot.SetMarkerColor(2)
0145 
0146         if('g' in flag or 'G' in flag):
0147             glued_hits = event.gluedHits()
0148             glu_z = []
0149             glu_y = []
0150             for hit in glued_hits:
0151                 if(limits[0] < hit.z() < limits[1]):
0152                     glu_z.append(hit.z())
0153                     glu_y.append(hit.y())
0154             glu_plot = ROOT.TGraph(len(glu_z), array('f', glu_z), array('f', glu_y))
0155             glu_plot.SetMarkerColor(3)        
0156 
0157         plot = ROOT.TMultiGraph()
0158 
0159         if('p' in flag or 'P' in flag): plot.Add(pix_plot,"P")
0160         if('s' in flag or 'S' in flag): plot.Add(str_plot,"P")
0161         if('g' in flag or 'G' in flag): plot.Add(glu_plot,"P")
0162         self.plots_2D.append(plot)
0163 
0164     def PlotTracksXY(self, tracks):
0165         plot = ROOT.TMultiGraph()
0166 
0167         for track in tracks:
0168             X = []; Y = [];
0169             for hit in track.hits():
0170                 if(hit.isValid()):
0171                     X.append(hit.x())
0172                     Y.append(hit.y())
0173             plot.Add(ROOT.TGraph(len(X),array("f",X),array("f",Y)),"L")
0174 
0175         self.plots_2D.append(plot)
0176 
0177     def PlotTracksZY(self, tracks):
0178         plot = ROOT.TMultiGraph()
0179 
0180         for track in tracks:
0181             Y = []; Z = [];
0182             for hit in track.hits():
0183                 if(hit.isValid()):
0184                     Y.append(hit.y())
0185                     Z.append(hit.z())
0186             plot.Add(ROOT.TGraph(len(Z),array("f",Z),array("f",Y)),"L")
0187 
0188         self.plots_2D.append(plot)    
0189    
0190     def PlotTracks3D(self, tracks): 
0191         for track in tracks:
0192             X = []; Y = []; Z = [];
0193             for hit in track.hits():
0194                 if(hit.isValid()):
0195                     X.append(hit.x())
0196                     Y.append(hit.y())
0197                     Z.append(hit.z())        
0198             if(X):
0199                 self.plots_3D.append(ROOT.TPolyLine3D(len(X),array("f",X),array("f",Y),array("f",Z)))
0200         
0201     def PlotPixelTracks3D(self, tracks): 
0202         for track in tracks:
0203             X = []; Y = []; Z = [];
0204             for hit in track.pixelHits():
0205                 if(hit.isValid()):
0206                     X.append(hit.x())
0207                     Y.append(hit.y())
0208                     Z.append(hit.z())
0209             if(X):
0210                 self.plots_3D.append(ROOT.TPolyLine3D(len(X),array("f",X),array("f",Y),array("f",Z)))
0211 
0212     def PlotPixelGluedTracks3D(self, tracks):
0213         for track in tracks:
0214             X = []; Y = []; Z = [];
0215             for hit in track.hits():
0216                 if(hit.isValid() and hit.hitType != 1):
0217                     X.append(hit.x())
0218                     Y.append(hit.y())
0219                     Z.append(hit.z())
0220             if(X):
0221                 self.plots_3D.append(ROOT.TPolyLine3D(len(X),array("f",X),array("f",Y),array("f",Z)))
0222 
0223     def PlotTrack3D(self, track, color = 1): 
0224         '''Plots a single track and prints track info'''
0225         # Not so hasardous experimental edit:
0226         #hits = sorted([hit for hit in track.hits()], key = lambda hit: hit.index())
0227         #print [hit.index() for hit in hits]
0228         X = []; Y = []; Z = [];
0229         for hit in track.hits(): #hits: #track.hits():
0230             if(hit.isValid()):
0231                 X.append(hit.x())
0232                 Y.append(hit.y())
0233                 Z.append(hit.z())        
0234         if(not X):
0235             print("Track has no valid points")
0236             return
0237         plot = ROOT.TPolyLine3D(len(X),array("f",X),array("f",Y),array("f",Z))
0238         plot.SetLineColor(color)
0239         self.plots_3D.append(plot)
0240 
0241         '''
0242         print "Track parameters:"
0243         print "px  : " + str(track.px())
0244         print "py  : " + str(track.py())
0245         print "pz  : " + str(track.pz())
0246         print "pt  : " + str(track.pt())
0247         print "eta : " + str(track.eta())
0248         print "phi : " + str(track.phi())
0249         print "dxy : " + str(track.dxy())
0250         print "dz  : " + str(track.dz())
0251         print "q   : " + str(track.q())
0252         '''
0253 
0254     def TrackHelix(self, track, color = 1, style = 0):
0255         '''Creates a THelix object which can be plotted with Draw() method.'''
0256         if isinstance(track, TrackingParticle):
0257             phi = track.pca_phi()
0258             dxy = track.pca_dxy()
0259             dz = track.pca_dz()
0260         else:
0261             phi = track.phi()
0262             dxy = track.dxy()
0263             dz = track.dz()
0264         xyz = array("d", [-dxy*ROOT.TMath.Sin(phi), dxy*ROOT.TMath.Cos(phi), dz])
0265         v = array("d", [track.px(), track.py(), track.pz()])
0266         w = 0.3*3.8*track.q()*0.01 #Angular frequency = 0.3*B*q*hattuvakio, close enough
0267         z_last = dz
0268         for hit in track.hits(): 
0269             if(hit.isValidHit()): z_last = hit.z()
0270         
0271         helix = ROOT.THelix(xyz, v, w, array("d", [dz, z_last]))
0272         helix.SetLineColor(color)
0273         if style == 1: helix.SetLineStyle(9)
0274         if style == 2: helix.SetLineStyle(7)
0275         return helix
0276 
0277     def Plot3DHelixes(self, tracks, color = 1, style = 0):        
0278         for track in tracks:
0279             if(track.hits()):
0280                 self.plots_3D.append(self.TrackHelix(track, color, style))        
0281 
0282     def Plot3DHelix(self, track, color = 1, style = 0):        
0283         if(track.hits()):
0284             self.plots_3D.append(self.TrackHelix(track, color, style))
0285 
0286     def Plot3DHits(self, track, color = 1, style = 0):
0287         '''
0288         Plots the 3D hits from a track.
0289         '''        
0290         pix_coords = []
0291         for hit in track.pixelHits():    
0292             if hit.isValid():
0293                 pix_coords.append(hit.x())
0294                 pix_coords.append(hit.y())
0295                 pix_coords.append(hit.z())
0296             if pix_coords:
0297                 pix_plot = ROOT.TPolyMarker3D(len(pix_coords)/3, array('f', pix_coords), 2)
0298                 pix_plot.SetMarkerColor(color)
0299                 if style == 1: pix_plot.SetMarkerStyle(5)
0300                 if style == 2: pix_plot.SetMarkerStyle(4)
0301                 self.plots_3D.append(pix_plot)
0302         
0303         for hit in track.gluedHits():
0304             if hit.isValid():
0305                 x = hit.x(); y = hit.y(); z = hit.z()
0306                 if hit.isBarrel():
0307                     X = [x, x]
0308                     Y = [y, y]
0309                     Z = [z - sqrt(hit.zz()), z + sqrt(hit.zz())]
0310                 else:
0311                     X = [x - copysign(sqrt(hit.xx()),x), x + copysign(sqrt(hit.xx()),x)]
0312                     Y = [y - copysign(sqrt(hit.yy()),y), y + copysign(sqrt(hit.yy()),y)]
0313                     Z = [hit.z(), hit.z()]
0314                 glu_plot = ROOT.TPolyLine3D(len(X),array("f",X),array("f",Y),array("f",Z))
0315                 #glu_plot.SetLineStyle(2)    
0316                 if style == 1: glu_plot.SetLineStyle(2)
0317                 if style == 2: glu_plot.SetLineStyle(3)
0318                 glu_plot.SetLineColor(color) 
0319                 self.plots_3D.append(glu_plot)
0320 
0321         for hit in track.stripHits():
0322             if hit.isValid():
0323                 x = hit.x(); y = hit.y(); z = hit.z()
0324                 if hit.isBarrel():
0325                     X = [x, x]
0326                     Y = [y, y]
0327                     Z = [z - 1.5*sqrt(hit.zz()), z + 1.5*sqrt(hit.zz())]
0328                 else:
0329                     X = [x - 1.5*copysign(sqrt(hit.xx()),x), x + 1.5*copysign(sqrt(hit.xx()),x)]
0330                     Y = [y - 1.5*copysign(sqrt(hit.yy()),y), y + 1.5*copysign(sqrt(hit.yy()),y)]
0331                     Z = [hit.z(), hit.z()]
0332                 str_plot = ROOT.TPolyLine3D(len(X),array("f",X),array("f",Y),array("f",Z))
0333                 if style == 1: str_plot.SetLineStyle(2)
0334                 if style == 2: str_plot.SetLineStyle(3)
0335                 str_plot.SetLineColor(color) 
0336                 self.plots_3D.append(str_plot)
0337 
0338     def PlotVertex3D(self, vertex, color=1):
0339         plot = ROOT.TPolyMarker3D(1, array('f', [vertex.x(), vertex.y(), vertex.z()]),3)
0340         plot.SetMarkerColor(color)
0341         self.plots_3D.append(plot)
0342 
0343     def PlotPoint3D(self, point, color=1):
0344         plot = ROOT.TPolyMarker3D(1, array('f', [point[0], point[1], point[2]]),3)
0345         plot.SetMarkerColor(color)
0346         self.plots_3D.append(plot)
0347 
0348     def PlotTrackingFail(self, match):
0349         X = array('f', [match.last_loc[0], match.fail_loc[0]])
0350         Y = array('f', [match.last_loc[1], match.fail_loc[1]])
0351         Z = array('f', [match.last_loc[2], match.fail_loc[2]])
0352         plot = ROOT.TPolyLine3D(2, X, Y, Z)
0353         plot.SetLineWidth(3)
0354         plot.SetLineColor(2) 
0355         self.plots_3D.append(plot)
0356         self.PlotPoint3D(match.last_loc, 2)
0357 
0358     def PlotFakes_MatchedRecos(self, event, iterative = 1, reconstructed = 0):
0359         fakes = analysis.FindFakes(event)
0360         if iterative:
0361             # Plot fakes one by one
0362             for fake in fakes:
0363                 self.Reset()
0364                 self.Plot3DHelixes([fake],2)
0365                 self.Plot3DHits(fake, 2)
0366 
0367                 # Plot real particle tracks which include fake tracks hits
0368                 icol = 0
0369                 particle_inds = []
0370                 particles = []
0371                 reco_inds = []
0372                 recos = []
0373                 for hit in fake.hits():
0374                     if hit.isValid() and hit.nSimHits() >= 0:
0375                         for simHit in hit.simHits():
0376                             particle = simHit.trackingParticle()
0377                             if particle.index() not in particle_inds:
0378                                 particle_inds.append(particle.index())
0379                                 particles.append(particle)
0380 
0381                             '''
0382                             self.Plot3DHelix(particle, ROOT.TColor().GetColor(0,255,0), 1) # kAzure color, maybe ;)
0383                             self.Plot3DHits(particle, 3+icol, 1) 
0384                             icol += 1
0385 
0386                             print "Number of matched tracks to real particle: " + str(particle.nMatchedTracks())
0387                             '''
0388                             # Plot reconstructed tracks of these real particle tracks
0389                             if reconstructed and particle.nMatchedTracks() > 0:
0390                                 for info in particle.matchedTrackInfos():
0391                                     track = info.track()
0392                                     if track.index() not in reco_inds:
0393                                         reco_inds.append(track.index())
0394                                         recos.append(track)
0395 
0396                                     '''
0397                                     if particle.nMatchedTracks() == 1:
0398                                         self.Plot3DHelix(track,1,2)
0399                                         self.Plot3DHits(track,1,2)
0400                                     else:
0401                                         self.Plot3DHelix(track,5,2)
0402                                         self.Plot3DHits(track,5,2)
0403                                     '''
0404                 icol =  0
0405                 for particle in particles:
0406                     self.Plot3DHelix(particle, self.colors_G[icol], 1) # kAzure color, maybe ;)
0407                     self.Plot3DHits(particle, self.colors_G[icol], 1) 
0408                     icol += 1
0409                 for track in recos:
0410                     self.Plot3DHelix(track,1,2)
0411                     self.Plot3DHits(track,1,2)
0412                     '''
0413                     if track.trackingParticle().nMatchedTracks() == 1:
0414                         self.Plot3DHelix(track,1,2)
0415                         self.Plot3DHits(track,1,2)
0416                     else:
0417                         self.Plot3DHelix(track,5,2)
0418                         self.Plot3DHits(track,5,2)
0419                     '''
0420 
0421                 self.Draw()
0422             return
0423         # the following is useless by now
0424         # Plot all fakes at once (messy picture)
0425         '''
0426         self.Plot3DHelixes(fakes,2)
0427         for fake in fakes:
0428             self.Plot3DHits(fake, 2)
0429             # Plot real particle tracks which include fake tracks hits
0430             for hit in fake.hits():
0431                 if hit.isValid() and hit.nMatchedTrackingParticles() >= 0:
0432                     for info in hit.matchedTrackingParticleInfos():
0433                         particle = info.trackingParticle()
0434                         self.Plot3DHelix(particle,3,1)
0435                         self.Plot3DHits(particle,3,1)
0436  
0437                         print "Number of matched tracks to real particle: " + str(particle.nMatchedTracks())
0438                         # Plot reconstructed tracks of these real particle tracks
0439                         if reconstructed and particle.nMatchedTracks() > 0:
0440                             for info in particle.matchedTrackInfos():
0441                                 track = info.track()        
0442                                 if particle.nMatchedTracks() == 1:
0443                                     self.Plot3DHelix(track,1,2)
0444                                     self.Plot3DHits(track,1,2)
0445                                 else:
0446                                     self.Plot3DHelix(track,5,2)
0447                                     self.Plot3DHits(track,5,2)
0448         '''
0449 
0450     def PlotFakes(self, event, reconstructed = 1, fake_filter = None, end_filter = None, last_filter = None, particle_end_filter = None):
0451         iterative = 1
0452 
0453         fakes = analysis.FindFakes(event)
0454         if iterative:
0455             # Plot fakes one by one
0456             for fake in fakes:
0457                 # Check for filter
0458                 if fake_filter:
0459                     info = analysis.FakeInfo(fake)
0460                     if info.fake_class not in fake_filter:
0461                         continue
0462 
0463                 self.Reset()
0464                 self.Plot3DHelixes([fake],2)
0465                 self.Plot3DHits(fake, 2)
0466                 #self.PlotVertex3D(vertex,2)
0467                 fake_info = analysis.FakeInfo(fake)
0468                 analysis.PrintTrackInfo(fake, None, 0, fake_info)
0469                 if fake_info.matches:
0470                     for match in fake_info.matches:
0471                         continue #self.PlotTrackingFail(match)
0472                         #self.PlotPoint3D(fake_info.end_loc, 2)
0473 
0474                 # Find real particle tracks which include fake tracks hits
0475                 icol = 0
0476                 particle_inds = []
0477                 particles = []
0478                 reco_inds = []
0479                 recos = []
0480                 for hit in fake.hits():
0481                     if hit.isValid() and hit.nMatchedTrackingParticles() >= 0:        
0482                         for info in hit.matchedTrackingParticleInfos():
0483                             particle = info.trackingParticle()
0484                             if particle.index() not in particle_inds:
0485                                 particle_inds.append(particle.index())
0486                                 particles.append(particle)
0487                             if reconstructed and particle.nMatchedTracks() > 0:
0488                                 for info in particle.matchedTrackInfos():
0489                                     track = info.track()
0490                                     if track.index() not in reco_inds:
0491                                         reco_inds.append(track.index())
0492                                         recos.append(track)
0493 
0494 
0495                     # Find reconstructed tracks included in fakes hits 
0496                     if hit.isValid() and reconstructed and hit.ntracks() > 0:
0497                         for track in hit.tracks():
0498                             #track = info.track()
0499                             if (track.index() not in reco_inds) and track.index() != fake.index():
0500                                 reco_inds.append(track.index())
0501                                 recos.append(track)
0502  
0503                 # Plot the particles and reconstructed tracks
0504                 icol =  0
0505                 self.ParticleTest(particles, draw=False)
0506                 for particle in particles:
0507                     self.Plot3DHelix(particle, self.colors_G[icol], 1)
0508                     self.plots_3D[-1].SetLineStyle(5)
0509                     self.Plot3DHits(particle, self.colors_G[icol], 1) 
0510                     self.PlotVertex3D(particle.parentVertex(),self.colors_G[icol])
0511                     # EXPERIMENTAL LINE:
0512                     #self.PlotTrack3D(particle, self.colors_G[icol])
0513 
0514                     for decay in particle.decayVertices():
0515                         self.PlotVertex3D(decay, 5)
0516                     analysis.PrintTrackInfo(particle, fake)
0517                     icol += 1
0518                 icol = 0
0519                 for track in recos:
0520                     self.Plot3DHelix(track,self.colors_B[icol],2)
0521                     self.Plot3DHits(track,self.colors_B[icol],2)
0522                     #self.PlotVertex3D(vertex,self.colors_B[icol])
0523                     analysis.PrintTrackInfo(track, fake)
0524                     icol += 1
0525                     
0526                 if hit.isValid() and hit.z() >= 0: self.PlotDetectorRange("p",4)
0527                 elif hit.isValid() and hit.z() < 0: self.PlotDetectorRange("n",4)
0528                 else: self.PlotDetectorRange("b",4)
0529 
0530                 print("****************************\n")
0531                 self.Draw()
0532             return
0533 
0534     def DrawTrackTest_old(self, track):
0535         c3 = ROOT.TCanvas("3D Plot","3D Plot", 1200, 1200)
0536         self.PlotDetectorRange("b", 4)
0537         for i in range(8): self.plots_3D[i].Draw()
0538         axis = ROOT.TAxis3D()
0539         axis.Draw()
0540         for hit in track.hits():
0541             if hit.isValid(): 
0542                 point = ROOT.TPolyMarker3D(1, array('f', [hit.x(), hit.y(), hit.z()]), 2)
0543                 point.Draw()
0544                 input("continue")
0545         #axis.SetAxisRange(-278,278,"Z")
0546         
0547         axis.ToggleZoom()
0548 
0549     def DrawTrackTest(self, track): 
0550         self.PlotDetectorRange("b", 4)        
0551         self.PlotTrack3D(track)
0552         #self.Draw()
0553         '''
0554         for hit in track.hits():
0555             if hit.isValid(): 
0556                 
0557                 point = ROOT.TPolyMarker3D(1, array('f', [hit.x(), hit.y(), hit.z()]), 2)
0558                 self.plots_3D.append(point)        
0559                 self.Draw()
0560                 '''
0561 
0562     def ParticleTest(self, particles, draw=False):
0563         for particle in particles:
0564             print("\nPARTICLE " + str(particle.index()))
0565             for hit in particle.hits():
0566                 tof = -1
0567                 for info in hit.matchedTrackingParticleInfos():
0568                     if info.trackingParticle().index() == particle.index():
0569                         #if len(info.tof()): 
0570                         tof = info.tof()
0571                 print("Index: " + str(hit.index()) + ", Layer: " + str(hit.layerStr()) + ", TOF: " + str(tof) +\
0572                 "     XY distance: " + str(sqrt(hit.x()**2 + hit.y()**2)) + ", Z: " + str(hit.z()))
0573             self.DrawTrackTest(particle)
0574             if draw:
0575                 self.Draw()
0576                 self.Reset()
0577  
0578     def PlotDetectorRange(self, area="b", plates = 6):
0579         '''
0580         Plots the detector schematic layout.
0581         Area  means which part to plot (p = positive side, n = negative side, b = both)
0582         Plates means how many detector circle schemes to plot per side.
0583         '''
0584         r = [17, 17, 57, 57, 112, 112]
0585         z = [30, 70, 70, 120, 120, 280]
0586         for i in range(plates):
0587             X = []; Y = []; Z = []; ZN = []
0588             for t in range(0,101):
0589                 X.append(r[i]*cos(2*pi*t/100))
0590                 Y.append(r[i]*sin(2*pi*t/100))
0591                 Z.append(z[i]*1.0)
0592                 ZN.append(-1.0*z[i])
0593             plot = ROOT.TPolyLine3D(len(X),array("f",X),array("f",Y),array("f",Z),"S")
0594             nplot = ROOT.TPolyLine3D(len(X),array("f",X),array("f",Y),array("f",ZN),"S")
0595             plot.SetLineStyle(3)
0596             nplot.SetLineStyle(3)
0597             if area == "p": self.plots_3D.append(plot)
0598             if area == "n": self.plots_3D.append(nplot)
0599             if area == "b":
0600                 self.plots_3D.append(plot)
0601                 self.plots_3D.append(nplot)
0602 
0603     def ClassHistogram(self, data, name = "Histogram", labels = False, ready = False):
0604         '''
0605         Draws histogram with data. Data is in form of a dictionary,
0606         labels assigned to amount of hits. If labels is defined, it
0607         defines the labels assigned to data.
0608         '''
0609         if ready:
0610             if labels:
0611                 keys = [key for key in data]
0612                 hist_labels = [labels[key] for key in keys]
0613                 hist_data = [data[key] for key in keys]
0614             else:
0615                 hist_labels = [key for key in data]
0616                 hist_data = [data[key] for key in hist_labels]
0617         
0618         c1 = ROOT.TCanvas(name, name, 700, 500)
0619         c1.SetGrid()
0620         #c1.SetTopMargin(0.15)
0621         hist = ROOT.TH1F(name, name, 30,0,30)#, 3,0,3)
0622         hist.SetStats(0);
0623         hist.SetFillColor(38);
0624         hist.SetCanExtend(ROOT.TH1.kAllAxes);
0625         if ready: # cumulative results already calculated in data
0626             for i in range(len(hist_labels)):
0627                 #hist.Fill(hist_labels[i],hist_data[i])
0628                 [hist.Fill(hist_labels[i],1) for j in range(hist_data[i])]
0629         else: # data is in samples, not in cumulative results
0630             for entry in data: hist.Fill(entry,1)
0631         hist.LabelsDeflate()
0632         #hist.SetLabelSize(0.05)
0633         hist.Draw()
0634 
0635         input("Press any key to continue")
0636 
0637 
0638     def Draw(self):
0639         if self.plots_3D:
0640             c3 = ROOT.TCanvas("3D Plot","3D Plot", 1200, 1200)
0641             view = ROOT.TView3D()
0642             axis = ROOT.TAxis3D()
0643             axis.SetXTitle("x")
0644             axis.SetYTitle("y")
0645             axis.SetZTitle("z")
0646             #axis.GetZaxis().RotateTitle()
0647 
0648             for plot in self.plots_3D:
0649                 if plot: plot.Draw()
0650             
0651             #view.DefineViewDirection(array("d",[1,1,1]),array("d",[0,0,0]),1.0,0.0,0.0,1.0,0.0,1.0)#,array("d",[-1,-1,-1]),array("d",[1,1,1]))
0652             axis.Draw()
0653             #view.SetAxisNDC(array("d",[0,0,0]),array("d",[0,1,0]),array("d",[0,0,0]),array("d",[0,0,1]),array("d",[0,0,0]),array("d",[1,0,0]))
0654     
0655             #view.SetPsi(30)
0656             #view.SetView(0,0,pi/2,ROOT.Long())
0657             #view.DefineViewDirection(0,0,pi/2)
0658             #view.RotateView(pi/2,pi/2)
0659             axis.ToggleZoom()
0660 
0661         if self.plots_2D:
0662             c2 = ROOT.TCanvas("2D Plot","2D Plot", 1200, 1200)
0663 
0664             for plot in self.plots_2D: plot.Draw("A") 
0665 
0666         input("Press any key to continue")
0667 
0668