Use Matplotlib colorbars in Mayavi

mayavi_mpl_colorbar

from mayavi import mlab
from pylab import plt
from numpy import array, linspace


def adjust_colorbar(pts, cb_title, vmin, vmax):
    pts.glyph.glyph_source.glyph_source.center = [0, 0, 0]
    pts.glyph.glyph.color_mode = 'color_by_scalar'
    pts.glyph.color_mode = 'color_by_scalar'
    pts.glyph.glyph.scale_mode = 'data_scaling_off'
    pts.glyph.glyph_source.glyph_source.center = [0, 0, 0]
    # use mpl colormap
    lut = pts.module_manager.scalar_lut_manager.lut.table.to_array()
    cmap = plt.cm.coolwarm  # @UndefinedVariable
    # cmap = plt.cm.PRGn  # @UndefinedVariable
    #     cmap = plt.cm.PuOr  # @UndefinedVariable
    #         cmap = plt.get_cmap('BrBG')
    cmaplist = array([cmap(i) for i in range(cmap.N)]) * 255
    lut = cmaplist
    pts.module_manager.scalar_lut_manager.lut.table = lut
    pts.module_manager.scalar_lut_manager.show_scalar_bar = True
    pts.module_manager.scalar_lut_manager.show_legend = True
    pts.module_manager.scalar_lut_manager.data_name = cb_title
    pts.module_manager.scalar_lut_manager.data_range = array([vmin, vmax])
    # make color bar smaller
    module_manager2 = pts.module_manager
    module_manager2.scalar_lut_manager.scalar_bar_representation.minimum_size = array([1, 1])
    module_manager2.scalar_lut_manager.scalar_bar_representation.position2 = array([0.12657727, 0.84742547])
    module_manager2.scalar_lut_manager.scalar_bar_representation.position = array([0.86853129, 0.02276423])
    module_manager2.scalar_lut_manager.scalar_bar_representation.maximum_size = array([100000, 100000])
    return pts


x = linspace(0, 10, 10)
y = linspace(0, 10, 10)
z = linspace(0, 10, 10)
s = x
pts = mlab.points3d(x, y, z, s)
adjust_colorbar(pts, 'test', 0, 10)
mlab.gcf().scene.z_plus_view()

mlab.show()