Using perceptually uniform colormaps in Mayavi volumetric visualization

Question

AFAIK Mayavi does not come with any perceptually uniform colormaps. I tried naively to just pass it one of Matplotlib's colormaps but it failed:

from mayavi import mlab
import multiprocessing
import matplotlib.pyplot as plt

plasma = plt.get_cmap('plasma')

...
mlab.pipeline.volume(..., colormap=plasma)

TraitError: Cannot set the undefined 'colormap' attribute of a 'VolumeFactory' object.

---

Edit: I found a guide to convert Matplotlib colormaps to Mayavi colormaps. However, it unfortunately doesn't work since I am trying to use a volume using a perceptually uniform colormap.

from matplotlib.cm import get_cmap
import numpy as np
from mayavi import mlab

values = np.linspace(0., 1., 256)
lut_dict = {}
lut_dict['plasma'] = get_cmap('plasma')(values.copy())

x, y, z = np.ogrid[-10:10:20j, -10:10:20j, -10:10:20j]
s = np.sin(x*y*z)/(x*y*z)

mlab.pipeline.volume(mlab.pipeline.scalar_field(s), vmin=0, vmax=0.8, colormap=lut_dict['plasma'])  # still getting the same error
mlab.axes()
mlab.show()

...

Answer 1

Instead of setting it as the colormap argument, if you set it as the ColorTransferFunction of the volume, it works as expected.

import numpy as np
from mayavi import mlab
from tvtk.util import ctf
from matplotlib.pyplot import cm

values = np.linspace(0., 1., 256)
x, y, z = np.ogrid[-10:10:20j, -10:10:20j, -10:10:20j]
s = np.sin(x*y*z)/(x*y*z)

volume = mlab.pipeline.volume(mlab.pipeline.scalar_field(s), vmin=0, vmax=0.8)
# save the existing colormap
c = ctf.save_ctfs(volume._volume_property)
# change it with the colors of the new colormap
# in this case 'plasma'
c['rgb']=cm.get_cmap('plasma')(values.copy())
# load the color transfer function to the volume
ctf.load_ctfs(c, volume._volume_property)
# signal for update
volume.update_ctf = True

mlab.show()

Answer 2

While the previous answer by like444 helped me partially with a similar problem, it leads to incorrect translation between colormaps. This is because the format in which matplotlib and tvtk store color information is slightly different: Matplotlib uses RGBA, while ColorTransferFunction uses VRGB, where V is the value in the shown data that this part of the colormap is assigned to. So by doing a 1-to-1 copy, green becomes red, blue becomes green and alpha becomes blue. The following code snippet fixes that:

def cmap_to_ctf(cmap_name):
    values = list(np.linspace(0, 1, 256))
    cmap = cm.get_cmap(cmap_name)(values)
    transfer_function = ctf.ColorTransferFunction()
    for i, v in enumerate(values):
        transfer_function.add_rgb_point(v, cmap[i, 0], cmap[i, 1], cmap[i, 2])
    return transfer_function