Add function to create lut from many sources

Author: mwaskom

surfer/tests/test_utils.py

@@ -1,6 +1,7 @@
 from os.path import join as pjoin
 
 import numpy as np
+import nose.tools as nt
 from numpy.testing import assert_array_almost_equal, assert_array_equal
 
 from surfer import utils
@@ -64,3 +65,29 @@ def test_huge_cross():
     z = np.cross(x, y)
     zz = utils._fast_cross_3d(x, y)
     assert_array_equal(z, zz)
+
+
+def test_create_color_lut():
+    """Test various ways of making a colormap."""
+    # Test valid lut
+    cmap_in = (np.random.rand(256, 4) * 255).astype(np.int)
+    cmap_out = utils.create_color_lut(cmap_in)
+    assert_array_equal(cmap_in, cmap_out)
+
+    # Test mostly valid lut
+    cmap_in = cmap_in[:, :3]
+    cmap_out = utils.create_color_lut(cmap_in)
+    assert_array_equal(cmap_in, cmap_out[:, :3])
+    assert_array_equal(cmap_out[:, 3], np.ones(256, np.int) * 255)
+
+    # Test matplotlib lut
+    cmap_out = utils.create_color_lut("BuGn_r")
+    nt.assert_equal(cmap_out.shape, (256, 4))
+
+    # Test list of colors lut
+    cmap_out = utils.create_color_lut(["purple", "pink", "white"])
+    nt.assert_equal(cmap_out.shape, (256, 4))
+
+    # Test that we can ask for a specific number of colors
+    cmap_out = utils.create_color_lut("Reds", 12)
+    nt.assert_equal(cmap_out.shape, (12, 4))

surfer/utils.py

@@ -10,6 +10,8 @@
 import nibabel as nib
 from scipy import sparse
 from scipy.spatial.distance import cdist
+import matplotlib as mpl
+from matplotlib import cm
 
 from .config import config
 
@@ -413,6 +415,60 @@ def mesh_edges(faces):
     return edges
 
 
+def create_color_lut(cmap, n_colors=256):
+    """Return a colormap suitable for setting as a Mayavi LUT.
+
+    Parameters
+    ----------
+    cmap : string, list of colors, n x 3 or n x 4 array
+        Input colormap definition. This can be the name of a matplotlib
+        colormap, a list of valid matplotlib colors, or a suitable
+        mayavi LUT (possibly missing the alpha channel).
+    n_colors : int, optional
+        Number of colors in the resulting LUT. This is ignored if cmap
+        is a 2d array.
+    Returns
+    -------
+    lut : n_colors x 4 integer array
+        Color LUT suitable for passing to mayavi
+    """
+    if isinstance(cmap, np.ndarray):
+        if np.ndim(cmap) == 2:
+            if cmap.shape[1] == 4:
+                # This looks likes a LUT that's ready to go
+                lut = cmap.astype(np.int)
+            elif cmap.shape[1] == 3:
+                # This looks like a LUT, but it's missing the alpha channel
+                alpha = np.ones(len(cmap), np.int) * 255
+                lut = np.c_[cmap, alpha]
+
+            return lut
+
+    # Otherwise, we're going to try and use matplotlib to create it
+
+    if cmap in dir(cm):
+        # This is probably a matplotlib colormap, so build from that
+        # The matplotlib colormaps are a superset of the mayavi colormaps
+        # except for one or two cases (i.e. blue-red, which is a crappy
+        # rainbow colormap and shouldn't be used for anything, although in
+        # its defense it's better than "Jet")
+        cmap = getattr(cm, cmap)
+
+    elif np.iterable(cmap):
+        # This looks like a list of colors? Let's try that.
+        colors = list(map(mpl.colors.colorConverter.to_rgb, cmap))
+        cmap = mpl.colors.LinearSegmentedColormap.from_list("_", colors)
+
+    else:
+        # If we get here, it's a bad input
+        raise ValueError("Input %s was not valid for making a lut" % cmap)
+
+    # Convert from a matplotlib colormap to a lut array
+    lut = (cmap(np.linspace(0, 1, n_colors)) * 255).astype(np.int)
+
+    return lut
+
+
 @verbose
 def smoothing_matrix(vertices, adj_mat, smoothing_steps=20, verbose=None):
     """Create a smoothing matrix which can be used to interpolate data defined