Source Code for Module pyllar.Plot3

  1  """ 3d plot class. 
  2  Use Plot3 to draw points and lines. 
  3  """ 
  4   
  5   
  6  from Plottable import Plottable 
  7  import vtk 
  8  from numpy import * 
  9  import re 
 10   
 11 -class Plot3(Plottable): 
 12       
 13 -    def __init__(self, data, style='b-'): 
 14          """ Create a new Plot3 object. 
 15           
 16          Plot3(data, style) 
 17          data is a Nx3 NumPy array that defines the 
 18          N points of the plot.  
 19          style is a string and can be any combination of color 
 20          symbols (r, g, b, y, m, c, k) and line styles (-, --, -). 
 21          """ 
 22          # create plottable defaults 
 23          self.pos = array([0,0,0]) 
 24          self.direction = array([1.0,0.0,0.0]) 
 25   
 26          # move data into a vtk.vtkDataArray 
 27          numRows = data.shape[0] 
 28          PlotData = vtk.vtkFloatArray() 
 29          PlotData.SetNumberOfComponents(3) 
 30          PlotData.SetNumberOfTuples(numRows) 
 31   
 32          # create a cell array to define the points 
 33          PointArray = vtk.vtkCellArray() 
 34          PointArray.InsertNextCell(numRows) 
 35       
 36          # create a cell array to define the line 
 37          LineArray = vtk.vtkCellArray() 
 38          LineArray.InsertNextCell(numRows) 
 39           
 40          # loop through rows and fill everything 
 41          for i in arange(numRows): 
 42              PlotData.SetTuple3(i, data[i,0], data[i,1], data[i,2]) 
 43              PointArray.InsertCellPoint(i) 
 44              LineArray.InsertCellPoint(i) 
 45       
 46          # set the data as points 
 47          Points = vtk.vtkPoints() 
 48          Points.SetNumberOfPoints(numRows) 
 49          Points.SetData(PlotData) 
 50       
 51          # create the polydata 
 52          vertData = vtk.vtkPolyData() 
 53          vertData.SetPoints(Points) 
 54          vertData.SetVerts(PointArray) 
 55          lineData = vtk.vtkPolyData() 
 56          lineData.SetPoints(Points) 
 57          lineData.SetLines(LineArray) 
 58           
 59          # create mapper 
 60          self.vertMapper = vtk.vtkPolyDataMapper() 
 61          self.vertMapper.SetInput(vertData) 
 62          self.lineMapper = vtk.vtkPolyDataMapper() 
 63          self.lineMapper.SetInput(lineData) 
 64           
 65          # create actor 
 66          self.vertActor = vtk.vtkActor() 
 67          self.vertActor.SetMapper(self.vertMapper)      
 68          self.lineActor = vtk.vtkActor() 
 69          self.lineActor.SetMapper(self.lineMapper)         
 70           
 71          # set line style 
 72          self.SetStyle(style)         
 73           
 74          # add to the axes 
 75          self.actor = vtk.vtkAssembly() 
 76          self.actor.AddPart(self.vertActor) 
 77          self.actor.AddPart(self.lineActor) 
 78          self.AddToAxes() 
 79           
 80           
 81      #def AddToAxes(self): 
 82      #    """ Add this plottable to the current axes. 
 83      #    """ 
 84      #    self.axes = gca() 
 85      #    self.axes.ren.AddActor(self.actor) 
 86      #    self.axes.AddChild(self) 
 87   
 88           
 89           
 90 -    def SetStyle(self, style): 
 91          """ Set the plot style (color and marker) from a style string. 
 92          """ 
 93                   
 94          # find colors 
 95          if re.search("b", style): 
 96              color = [0, 0, 1] 
 97          elif re.search("g", style): 
 98              color = [0, 1, 0] 
 99          elif re.search("r", style): 
100              color = [1, 0, 0] 
101          else: # default color 
102              color = [0, 0, 1] 
103               
104          # find line style 
105          if re.search("--", style): 
106              stipple = 0xFF00 
107          elif re.search("-", style): 
108              stipple = 0xFFFF  
109          else: # default line style 
110              stipple = 0xFFFF  
111               
112          # set line style 
113          prop = self.lineActor.GetProperty() # get the current properties 
114       
115          # set the color of the line 
116          prop.SetColor(color) 
117       
118               
119          # find point style       
120          if re.search('\.', style): 
121              pointSize = 2 
122          else: # default point size   
123              pointSize = 0       
124               
125          # find line size 
126          if re.search('\-', style): 
127              lineSize = 2 
128              prop.SetRepresentationToWireframe() 
129          else: # default point size   
130              lineSize = 0       
131              prop.SetRepresentationToPoints() 
132               
133          # set the thickness of the points and lines 
134          prop.SetLineWidth(lineSize) 
135       
136          # set the culling so that the points can be seen from all angles 
137          prop.FrontfaceCullingOff() 
138          prop.BackfaceCullingOff() 
139   
140          # set the line pattern 
141          prop.SetLineStipplePattern(stipple)  
142       
143          self.lineActor.SetProperty(prop)       
144               
145          # set point style 
146          prop = self.vertActor.GetProperty() # get the current properties 
147       
148          # set the color of the line 
149          prop.SetColor(color) 
150       
151          # set the thickness of the points and lines 
152          prop.SetPointSize(pointSize) 
153          if pointSize == 0: 
154              prop.SetOpacity(0.0) 
155       
156          # set the culling so that the points can be seen from all angles 
157          prop.SetRepresentationToPoints() 
158          prop.FrontfaceCullingOff() 
159          prop.BackfaceCullingOff() 
160       
161          self.vertActor.SetProperty(prop) 
162           
163       
164      # set the line width 
165 -    def SetLineWidth(self, width): 
166          prop = self.lineActor.GetProperty() # get the current properties         
167          prop.SetLineWidth(width) # set the thickness of the line 
168          self.lineActor.SetProperty(prop) 
169           
170      # set the marker size 
171 -    def SetMarkerSize(self, size): 
172          prop = self.vertActor.GetProperty() # get the current properties         
173          prop.SetPointSize(size) # set the thickness of the line 
174          self.vertActor.SetProperty(prop) 
175           
176           
177  ## End Plot3 class definition ## 
178   
179   
180  ## Convenience Matlab-style functions ## 
181 -def plot3(x,y,z,style='b-'): 
182      """ Create a line plot. 
183       
184      plot3(x, y, z, style) 
185      x, y, and z are Nx1 NumPy arrays that define the 
186      N points of the plot. 
187      style is a string and can be any combination of color 
188      symbols (r, g, b, y, m, c, k) and line styles (-, --, -). 
189      """ 
190           
191      if x.shape[0]==1: 
192          x = x.T 
193           
194      if y.shape[0]==1: 
195          y = y.T 
196           
197      if z.shape[0]==1: 
198          y = z.T 
199           
200      data = c_[x, y, z] 
201           
202      return Plot3(data, style) 
203