Minimal example of spheres and lines animated by changing colors based on scalar value

line_color_anim.py

# line_color_anim.py ---
# Author: Subhasis Ray
# Created: Thu Jul 19 16:58:58 2018 (-0400)
# Last-Updated: Fri Jul 20 11:18:41 2018 (-0400)
#           By: Subhasis Ray
# Version: $Id$

# Code:

"""Short example of animating line colors with color transfer function
and scalar value.

This shows use of color transfer function, points and lines in
vtkPolyData.

"""
from __future__ import print_function
import numpy as np
import vtk
from vtk.util import numpy_support as vtknp


def make_color_table():
    """Create a color transfer function for looking up colors
    corresponding to scalar value.

    """    
    colors = vtk.vtkNamedColors()
    ctf = vtk.vtkColorTransferFunction()
    ctf.SetColorSpaceToDiverging()
    ctf.AddRGBPoint(0.0, *list(colors.GetColor3d("DarkRed")))
    ctf.AddRGBPoint(1.0, *list(colors.GetColor3d("Red")))
    ctf.AddRGBPoint(2.0, *list(colors.GetColor3d("Orange")))
    ctf.AddRGBPoint(3.0, *list(colors.GetColor3d("Yellow")))
    return ctf


def make_lines(pre, post, ctf):
    """Create lines from positions in `pre` to positions in `post`.

    The colors are looked up from colortransfer function ctf.

    """
    nodes = vtk.vtkPoints()
    nodes.SetData(vtknp.numpy_to_vtk(
        np.concatenate((pre, post)),
        deep=True,
        array_type=vtk.VTK_FLOAT))
    lines = vtk.vtkCellArray()
    for ii in range(len(pre)):
        line = vtk.vtkLine()
        line.GetPointIds().SetId(0, ii)
        line.GetPointIds().SetId(1, len(pre)+ii)
        lines.InsertNextCell(line)
    colors = vtknp.numpy_to_vtk(np.ones(len(pre)),
                                deep=True,
                                array_type=vtk.VTK_FLOAT)
    polydata = vtk.vtkPolyData()
    polydata.SetPoints(nodes)
    polydata.SetLines(lines)
    polydata.GetCellData().SetScalars(colors)
    mapper = vtk.vtkPolyDataMapper()
    mapper.SetInputData(polydata)
    mapper.SetScalarRange(0, 2)
    mapper.SetLookupTable(ctf)
    actor = vtk.vtkActor()
    actor.SetMapper(mapper)
    return {'polydata': polydata,
            'mapper': mapper,
            'actor': actor}


def make_spheres(pos, ctf):
    """Create spheres at positions in `pos`. `ctf` is color transfer
    function for looking up colors corresponding to scalar value
    associated with the polydata"""
    sphere = vtk.vtkSphereSource()
    sphere.SetRadius(1.0)
    sphere.SetPhiResolution(10)
    sphere.SetThetaResolution(10)
    pts = vtk.vtkPoints()
    vtkpos = vtknp.numpy_to_vtk(pos, deep=True,
                                array_type=vtk.VTK_FLOAT)
    pts.SetData(vtkpos)
    polydata = vtk.vtkPolyData()
    polydata.SetPoints(pts)
    colors = vtknp.numpy_to_vtk(np.ones(len(pos)), deep=True,
                                array_type=vtk.VTK_FLOAT)
    polydata.GetPointData().SetScalars(colors)
    glyph = vtk.vtkGlyph3D()
    glyph.SetScaleModeToDataScalingOff()
    # comment out line above and uncomment line below to make the
    # sphere sizes change according to scalar value.
    
    # glyph.SetScaleFactor(1.0) 
    glyph.SetSourceConnection(sphere.GetOutputPort())
    glyph.SetInputData(polydata)
    mapper = vtk.vtkPolyDataMapper()
    mapper.SetInputConnection(glyph.GetOutputPort())
    mapper.SetScalarRange(0, 2)
    mapper.SetLookupTable(ctf)
    actor = vtk.vtkActor()
    actor.SetMapper(mapper)
    return {'source': sphere,
            'points': pts,
            'polydata': polydata,
            'glyph': glyph,
            'mapper': mapper,
            'actor': actor}


class UpdateCallback(object):
    def __init__(self, ids, polydata, line_polydata, size=3):
        """Callback object for updating data and rendering.
        
        ids: the indices of the points in polydata and lines in line_polydata

        polydata: vtkPolyData with point data for ids

        line_polydata: vtkPolyData with cell data for lines
        corresponding to ids.

        """
        self.ids = ids 
        self.polydata = polydata
        self.line_polydata = line_polydata
        self.size = size

    def execute(self, iren, event):
        """Change the color of a random set of ids"""
        indices = np.random.choice(self.ids, size=self.size, replace=False)
        colors = np.ones(len(self.ids))
        colors[indices] = 2.0
        colors = vtknp.numpy_to_vtk(colors, deep=True,
                                    array_type=vtk.VTK_FLOAT)
        self.polydata.GetPointData().SetScalars(colors)
        self.line_polydata.GetCellData().SetScalars(colors)
        iren.GetRenderWindow().Render()


def view3d(dt=100):
    """Put together a scene with spheres with attached lines that randomly
    change color every `dt` ms.

    """
    ctf = make_color_table()
    pre = np.random.sample(size=(10, 3)) * 10  # line start and spheres center position 
    post = pre.copy()       # line end position
    post[:, 0] += 100       # just make the lines 100 unit long, parallel to X axis
    spheres = make_spheres(pre, ctf)
    lines = make_lines(pre, post, ctf)
    renderer = vtk.vtkRenderer()
    renderer.AddActor(spheres['actor'])
    renderer.AddActor(lines['actor'])
    win = vtk.vtkRenderWindow()
    win.AddRenderer(renderer)
    win.Render()
    interactor = vtk.vtkRenderWindowInteractor()
    interactor.SetRenderWindow(win)
    interactor.Initialize()     # Must Initialize() before adding callback
    callback = UpdateCallback(np.arange(len(pre)), spheres['polydata'], lines['polydata'])
    interactor.AddObserver('TimerEvent', callback.execute)
    interactor.CreateRepeatingTimer(dt)     # Create a timer to generate TimerEvent every 100 ms
    interactor.Start()


if __name__ == '__main__':
    view3d()


#
# line_color_anim.py ends here