vtkparticle.py

#     
# Data Type:
#   x  y  z  vx  vy  vz  particle_id
# 
# References:
# http://www.vtk.org/Wiki/VTK/Examples/Python/GeometricObjects/Display/Point
# http://stackoverflow.com/questions/7591204/how-to-display-point-cloud-in-vtk-in-different-colors
# http://stackoverflow.com/questions/519633/lazy-method-for-reading-big-file-in-python

import vtk

def readline_in_chunks(file_object, chunk_size=4096):
  buf = ''
  while True:
    data = file_object.read(chunk_size)
    if not data:
      break
    buf = buf + data
    while True:
      pos = buf.find('\n')
      if pos < 0:
        break
      yield buf[:pos]
      buf = buf[pos+1:]

def main():
  f = open('data/particle_sp3_rnk3_step4000', 'r')

  points = vtk.vtkPoints()
  vertices = vtk.vtkCellArray()
  depth = vtk.vtkDoubleArray()
  depth.SetName('DepthArray')

  e_min = e_max = 0
  for line in readline_in_chunks(f):
    pos = [ float(v) for v in line.strip().split()[:6] ]
    e = pos[3]**2 + pos[4]**2 + pos[5]**2
    if e > e_max:
      e_max = e
    elif e < e_min:
      e_min = e
    pid = points.InsertNextPoint(pos[:3])
    depth.InsertNextValue(e)
    vertices.InsertNextCell(1)
    vertices.InsertCellPoint(pid)

  point = vtk.vtkPolyData()
  point.GetPointData().SetScalars(depth)
  point.GetPointData().SetActiveScalars('DepthArray')
  point.SetPoints(points)
  point.SetVerts(vertices)

  mapper = vtk.vtkPolyDataMapper()
  mapper.SetInput(point)
  mapper.SetColorModeToDefault()
  mapper.SetScalarRange(e_min, e_max)
  mapper.SetScalarVisibility(1)

  actor = vtk.vtkActor()
  actor.SetMapper(mapper)
  actor.GetProperty().SetPointSize(1)

  renderer = vtk.vtkRenderer()
  renderer.AddActor(actor)
  renderer.SetBackground(.2, .3, .4)
  renderer.ResetCamera()

  renderWindow = vtk.vtkRenderWindow()
  renderWindow.AddRenderer(renderer)
  renderWindowInteractor = vtk.vtkRenderWindowInteractor()
  renderWindowInteractor.SetRenderWindow(renderWindow)
  renderWindow.Render()
  renderWindowInteractor.Start()

if __name__ == '__main__':
  main()