Search code examples
pythonmultidimensional-arrayplot3dmayavi

Plotting 3D points with different colors in Mayavi (Python)


Is there any way to give mayavi a list of tuples, or maybe some numpy array of number_of_points x 3 size, such that I can specify different colour for each point?

So, I have the following data:

x of size Nx1 (contains x coordinates of N points)

y of size Nx1 (contains y coordinates of N points)

z of size Nx1 (contains z coordinates of N points)

R of size Nx1 (contains the values for the R channel of N points)

G of size Nx1 (contains the values for the G channel of N points)

B of size Nx1 (contains the values for the B channel of N points)

I want somehow to give this RGB data to mayavi so it will use the actual colour of the point, so I would like something like this:

from mayavi import mlab
plt = mlab.points3d(x, y, z, color = (R, G, B))

This works if N = 1, or with other words, only if I gave Mayavi a single point, otherwise it doesn't. Thus, I can iterate it, but it is very slow and hard on memory for some reason.

I've tried many things, but I can't seem to find a single approach (apart from doing it in a loop) that does what I need. Any ideas on how to do it?


Solution

  • One way is to put your RGB arrays into a lookup table that you then tell your points3d object to use. For example:

    import numpy as np
    import mayavi.mlab as mlab
    
    # Fake data from:
    # http://docs.enthought.com/mayavi/mayavi/auto/mlab_helper_functions.html#points3d
    t = np.linspace(0, 2 * np.pi, 20)
    
    x = np.sin(2 * t)
    y = np.cos(t)
    z = np.cos(2 * t)
    
    # Create a [0..len(t)) index that we'll pass as 's'
    s = np.arange(len(t))
    
    # Create and populate lookup table (the integer index in s corresponding
    #   to the point will be used as the row in the lookup table
    lut = np.zeros((len(s), 4))
    
    # A simple lookup table that transitions from red (at index 0) to
    #   blue (at index len(data)-1)
    for row in s:
        f = (row/len(s))
        lut[row,:] = [255*(1-f),0,255*f,255]
    
    # Plot the points, update its lookup table
    p3d = mlab.points3d(x, y, z, s, scale_mode='none')
    p3d.module_manager.scalar_lut_manager.lut.number_of_colors = len(s)
    p3d.module_manager.scalar_lut_manager.lut.table = lut
    
    mlab.draw()
    mlab.show()
    

    Produces

    enter image description here

    Reference: