#
# volume.py
# (c) Neil Gershenfeld  2/10/11
# volume render example
#

import vtk
from numpy import *

N = 100
i = ones((N,N,N))
i = cumsum(i,axis=0)
i = i - N/2 - .5
i = i/15.
x = transpose(i,(0,1,2))
y = transpose(i,(2,0,1))
z = transpose(i,(1,2,0))
r = (x**2 + y**2 + z**2)
a = sin(r)/r
a = 255*(1+a)/2.
a = a.astype(uint8)
a[N/2:N,N/2:N,N/2:N] = 0 

dataImporter = vtk.vtkImageImport()
string = a.tostring()
dataImporter.CopyImportVoidPointer(string,len(string))
dataImporter.SetDataScalarTypeToUnsignedChar()
dataImporter.SetNumberOfScalarComponents(1)
dataImporter.SetDataExtent(0,N-1,0,N-1,0,N-1)
dataImporter.SetWholeExtent(0,N-1,0,N-1,0,N-1)
alphaChannelFunc = vtk.vtkPiecewiseFunction()
alphaChannelFunc.AddPoint(0, 0.0)
alphaChannelFunc.AddPoint(100, 0.1)
alphaChannelFunc.AddPoint(255, 0.2)
colorFunc = vtk.vtkColorTransferFunction()
colorFunc.AddRGBPoint(0, 1.0, 1.0, 1.0)
colorFunc.AddRGBPoint(100, 1.0, 1.0, 1.0)
colorFunc.AddRGBPoint(255, 0.1, 0.1, 0.1)
volumeProperty = vtk.vtkVolumeProperty()
volumeProperty.SetColor(colorFunc)
volumeProperty.SetScalarOpacity(alphaChannelFunc)
compositeFunction = vtk.vtkVolumeRayCastCompositeFunction()
volumeMapper = vtk.vtkVolumeRayCastMapper()
volumeMapper.SetVolumeRayCastFunction(compositeFunction)
volumeMapper.SetInputConnection(dataImporter.GetOutputPort())
volume = vtk.vtkVolume()
volume.SetMapper(volumeMapper)
volume.SetProperty(volumeProperty)
renderer = vtk.vtkRenderer()
renderWin = vtk.vtkRenderWindow()
renderWin.AddRenderer(renderer)
renderInteractor = vtk.vtkRenderWindowInteractor()
renderInteractor.SetRenderWindow(renderWin)
renderer.AddVolume(volume)
renderer.SetBackground(1,1,1)
renderWin.SetSize(500, 500)
renderWin.SetWindowName("sin(r)/r (q to quit)")
renderInteractor.Initialize()
renderWin.Render()
renderInteractor.Start()