import time
from os.path import join
import vtk
from vtk.util.misc import vtkGetDataRoot

r = vtk.vtkXMLStructuredGridReader()
r.SetFileName(join(vtkGetDataRoot(), 'Data/multicomb_1.vts'))
r.Update()
r.GetOutput().GetPointData().SetActiveScalars('StagnationEnergy')
m = vtk.vtkDataSetMapper()
m.SetInput(r.GetOutput())
a = vtk.vtkActor()
a.SetMapper(m)

ren = vtk.vtkRenderer()
ren.AddActor(a)
rw = vtk.vtkRenderWindow()
rw.AddRenderer(ren)

rw.Render()
time.sleep(2)
print "changing data attribute"
r.GetOutput().GetPointData().SetActiveScalars('Density')
# The following two don't help.
r.GetOutput().Modified()
m.Update()
print m.GetPolyDataMapper().GetInput().GetPointData().GetScalars().GetName()
# Uncomment the following to get the correct result.
#m = vtk.vtkDataSetMapper()
#m.SetInput(r.GetOutput())
#a.SetMapper(m)
rw.Render()
time.sleep(2)


