Source Code for Module networkx.drawing.nx_vtk

  1  """ 
  2  Draw networks in 3d with vtk. 
  3   
  4  References: 
  5   - vtk:     http://www.vtk.org/ 
  6   
  7  """ 
  8  __author__ = """Aric Hagberg (hagberg@lanl.gov)""" 
  9  __date__ = "$Date: 2005-06-17 08:10:29 -0600 (Fri, 17 Jun 2005) $" 
 10  __credits__ = """""" 
 11  __revision__ = "$Revision: 1051 $" 
 12  #    Copyright (C) 2004,2005 by  
 13  #    Aric Hagberg <hagberg@lanl.gov> 
 14  #    Dan Schult <dschult@colgate.edu> 
 15  #    Pieter Swart <swart@lanl.gov> 
 16  #    Distributed under the terms of the GNU Lesser General Public License 
 17  #    http://www.gnu.org/copyleft/lesser.html 
 18  import vtk 
 19  from vtk.util.colors import banana, plum 
 20  import networkx 
 21   
 22  # this is a hack for drawing a network in 3d with vtk 
 23  # perhaps it will inspire someone to submit a polished version? 
 24  # much of the code borrowed from the vtk examples 
 25   
 26 -def draw_nxvtk(G, node_pos): 
 27      """ 
 28      Draw networkx graph in 3d with nodes at node_pos. 
 29   
 30      See layout.py for functions that compute node positions. 
 31   
 32      node_pos is a dictionary keyed by vertex with a three-tuple 
 33      of x-y positions as the value. 
 34   
 35      The node color is plum. 
 36      The edge color is banana. 
 37   
 38      All the nodes are the same size. 
 39   
 40      """ 
 41      # set node positions 
 42      np={} 
 43      for n in G.nodes(): 
 44          try: 
 45              np[n]=node_pos[n] 
 46          except KeyError: 
 47              raise networkx.NetworkXError, "node %s doesn't have position"%n 
 48   
 49      nodePoints = vtk.vtkPoints() 
 50   
 51      i=0 
 52      for (x,y,z) in np.values(): 
 53          nodePoints.InsertPoint(i, x, y, z) 
 54          i=i+1 
 55   
 56      # Create a polydata to be glyphed. 
 57      inputData = vtk.vtkPolyData() 
 58      inputData.SetPoints(nodePoints) 
 59   
 60      # Use sphere as glyph source. 
 61      balls = vtk.vtkSphereSource() 
 62      balls.SetRadius(.05) 
 63      balls.SetPhiResolution(20) 
 64      balls.SetThetaResolution(20) 
 65   
 66      glyphPoints = vtk.vtkGlyph3D() 
 67      glyphPoints.SetInput(inputData) 
 68      glyphPoints.SetSource(balls.GetOutput()) 
 69   
 70      glyphMapper = vtk.vtkPolyDataMapper() 
 71      glyphMapper.SetInput(glyphPoints.GetOutput()) 
 72   
 73      glyph = vtk.vtkActor() 
 74      glyph.SetMapper(glyphMapper) 
 75      glyph.GetProperty().SetDiffuseColor(plum) 
 76      glyph.GetProperty().SetSpecular(.3) 
 77      glyph.GetProperty().SetSpecularPower(30) 
 78   
 79      # Generate the polyline for the spline. 
 80      points = vtk.vtkPoints() 
 81      edgeData = vtk.vtkPolyData() 
 82   
 83      # Edges 
 84   
 85      lines = vtk.vtkCellArray() 
 86      i=0 
 87      for e in G.edges_iter(): 
 88          # The edge e can be a 2-tuple (Graph) or a 3-tuple (Xgraph) 
 89          u=e[0] 
 90          v=e[1] 
 91          if v in node_pos and u in node_pos: 
 92              lines.InsertNextCell(2) 
 93              for n in (u,v): 
 94                  (x,y,z)=node_pos[n] 
 95                  points.InsertPoint(i, x, y, z) 
 96                  lines.InsertCellPoint(i) 
 97                  i=i+1 
 98   
 99      edgeData.SetPoints(points) 
100      edgeData.SetLines(lines) 
101   
102      # Add thickness to the resulting line. 
103      Tubes = vtk.vtkTubeFilter() 
104      Tubes.SetNumberOfSides(16) 
105      Tubes.SetInput(edgeData) 
106      Tubes.SetRadius(.01) 
107      # 
108      profileMapper = vtk.vtkPolyDataMapper() 
109      profileMapper.SetInput(Tubes.GetOutput()) 
110   
111      #  
112      profile = vtk.vtkActor() 
113      profile.SetMapper(profileMapper) 
114      profile.GetProperty().SetDiffuseColor(banana) 
115      profile.GetProperty().SetSpecular(.3) 
116      profile.GetProperty().SetSpecularPower(30) 
117   
118      # Now create the RenderWindow, Renderer and Interactor 
119      ren = vtk.vtkRenderer() 
120      renWin = vtk.vtkRenderWindow() 
121      renWin.AddRenderer(ren) 
122   
123      iren = vtk.vtkRenderWindowInteractor() 
124      iren.SetRenderWindow(renWin) 
125   
126      # Add the actors 
127      ren.AddActor(glyph) 
128      ren.AddActor(profile) 
129   
130      renWin.SetSize(640, 640) 
131   
132      iren.Initialize() 
133      renWin.Render() 
134      iren.Start() 
135   
136 -def _test_suite(): 
137      import doctest 
138      suite = doctest.DocFileSuite('tests/drawing/nx_vtk.txt',package='networkx') 
139      return suite 
140   
141  if __name__ == "__main__": 
142      import os 
143      import sys 
144      import unittest 
145   
146      if sys.version_info[:2] < (2, 4): 
147          print "Python version 2.4 or later required for tests (%d.%d detected)." %  sys.version_info[:2] 
148          sys.exit(-1) 
149      # directory of networkx package (relative to this) 
150      nxbase=sys.path[0]+os.sep+os.pardir 
151      sys.path.insert(0,nxbase) # prepend to search path 
152      unittest.TextTestRunner().run(_test_suite()) 
153