How to directly set RGB/RGBA colors in mayavi

Question

As far as I am aware, there is no documentation for doing this, but I have found a way to do it with only a minimum amount of hacking around. Here is a minimal example, which might require a little tinkering for different kinds of sources:

from tvtk.api import tvtk; from mayavi import mlab; import numpy as np

x,y,z=np.random.random((3,nr_points)) #some data
colors=np.random.randint(256,size=(100,3)) #some RGB or RGBA colors

pts=mlab.points3d(x,y,z)
sc=tvtk.UnsignedCharArray()
sc.from_array(colors)

pts.mlab_source.dataset.point_data.scalars=sc
pts.mlab_source.dataset.modified()

It also looks like sometimes you have to ensure that the mapper points to the right thing. This is not necessary for the above example, but it may be for other sources

pts.actor.mapper.input=pts.mlab_source.dataset

At some point the mayavi API should be fixed better so that there is an API exposed to just do this for all the pipeline functions, but that turns out to be a rather complicated and sweeping set of changes which I don't currently have time to finish.

Edit: User eqzx posted an answer to another question (Specify absolute colour for 3D points in MayaVi) which may be simpler, especially for certain source types that are hard to get to work with tvtk.UnsignedCharArray.

His idea is to create a LUT spanning the entire range of 256x256x256 RGB values. Note that this LUT therefore has 16,777,216 entries. Which, if you wanted to use it in many vtk objects, may waste quite a lot of memory if you are not careful.

#create direct grid as 256**3 x 4 array 
def create_8bit_rgb_lut():
    xl = numpy.mgrid[0:256, 0:256, 0:256]
    lut = numpy.vstack((xl[0].reshape(1, 256**3),
                        xl[1].reshape(1, 256**3),
                        xl[2].reshape(1, 256**3),
                        255 * numpy.ones((1, 256**3)))).T
    return lut.astype('int32')

# indexing function to above grid
def rgb_2_scalar_idx(r, g, b):
    return 256**2 *r + 256 * g + b

#N x 3 colors
colors = numpy.array([_.color for _ in points])

#N scalars
scalars = numpy.zeros((colors.shape[0],))

for (kp_idx, kp_c) in enumerate(colors):
    scalars[kp_idx] = rgb_2_scalar_idx(kp_c[0], kp_c[1], kp_c[2])

rgb_lut = create_8bit_rgb_lut()

points_mlab = mayavi.mlab.points3d(x, y, z
                                   keypoint_scalars,
                                   mode = 'point')

#magic to modify lookup table 
points_mlab.module_manager.scalar_lut_manager.lut._vtk_obj.SetTableRange(0, rgb_lut.shape[0])
points_mlab.module_manager.scalar_lut_manager.lut.number_of_colors = rgb_lut.shape[0]
points_mlab.module_manager.scalar_lut_manager.lut.table = rgb_lut