Improve colormap handling and default in add_data

Author: mwaskom

surfer/viz.py

@@ -15,7 +15,7 @@
 
 from . import utils, io
 from .config import config
-from .utils import Surface, verbose, _get_subjects_dir
+from .utils import Surface, verbose, create_color_lut, _get_subjects_dir
 
 
 import logging
@@ -757,7 +757,7 @@ def add_overlay(self, source, min=None, max=None, sign="abs", name=None,
         self._toggle_render(True, views)
 
     def add_data(self, array, min=None, max=None, thresh=None,
-                 colormap="blue-red", alpha=1,
+                 colormap="RdBu_r", alpha=1,
                  vertices=None, smoothing_steps=20, time=None,
                  time_label="time index=%d", colorbar=True,
                  hemi=None):
@@ -788,10 +788,10 @@ def add_data(self, array, min=None, max=None, thresh=None,
             max value in colormap (uses real max if None)
         thresh : None or float
             if not None, values below thresh will not be visible
-        colormap : str | array [256x4]
-            name of Mayavi colormap to use, or a custom look up table (a 256x4
-            array, with the columns representing RGBA (red, green, blue, alpha)
-            coded with integers going from 0 to 255).
+        colormap : string, list of colors, or array
+            name of matplotlib colormap to use, a list of matplotlib colors,
+            or a custom look up table (an n x 4 array coded with RBGA values
+            between 0 and 255).
         alpha : float in [0, 1]
             alpha level to control opacity
         vertices : numpy array
@@ -841,16 +841,9 @@ def add_data(self, array, min=None, max=None, thresh=None,
         # Copy and byteswap to deal with Mayavi bug
         mlab_plot = _prepare_data(array_plot)
 
-        # process colormap argument
-        if isinstance(colormap, basestring):
-            lut = None
-        else:
-            lut = np.asarray(colormap)
-            if lut.shape != (256, 4):
-                err = ("colormap argument must be mayavi colormap (string) or"
-                       " look up table (array of shape (256, 4))")
-                raise ValueError(err)
-            colormap = "blue-red"
+        # Process colormap argument into a lut
+        lut = create_color_lut(colormap)
+        colormap = "Greys"
 
         data = dict(array=array, smoothing_steps=smoothing_steps,
                     fmin=min, fmid=(min + max) / 2, fmax=max,