User-generated colormaps are handled more easily.

efiring · efiring · commit 75ede62e8b85 · 2009-07-30T19:32:15.000Z
svn path=/trunk/matplotlib/; revision=7309
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,3 +1,7 @@
+2009-07-30 Add set_cmap and register_cmap, and improve get_cmap,
+           to provide convenient handling of user-generated
+           colormaps. - EF
+
 2009-07-28 Quiver speed improved, thanks to tip by Ray Speth. -EF
 
 2009-07-27 Simplify argument handling code for plot method. -EF
diff --git a/doc/api/api_changes.rst b/doc/api/api_changes.rst
@@ -21,6 +21,11 @@ list may help describe what changes may be necessary in your code.
 Changes beyond 0.98.x
 =====================
 
+* User-generated colormaps can now be added to the set recognized
+  by :func:`matplotlib.cm.get_cmap`.  Colormaps can be made the
+  default and applied to the current image using
+  :func:`matplotlib.pyplot.set_cmap`.
+
 * changed use_mrecords default to False in mlab.csv2rec since this is
   partially broken
 
diff --git a/examples/pylab_examples/custom_cmap.py b/examples/pylab_examples/custom_cmap.py
@@ -103,10 +103,24 @@
                    (1.0, 0.0, 0.0))
         }
 
+# Now we will use this example to illustrate 3 ways of
+# handling custom colormaps.
+# First, the most direct and explicit:
 
 blue_red1 = LinearSegmentedColormap('BlueRed1', cdict1)
+
+# Second, create the map explicitly and register it.
+# Like the first method, this method works with any kind
+# of Colormap, not just
+# a LinearSegmentedColormap:
+
 blue_red2 = LinearSegmentedColormap('BlueRed2', cdict2)
-blue_red3 = LinearSegmentedColormap('BlueRed3', cdict3)
+plt.register_cmap(cmap=blue_red2)
+
+# Third, for LinearSegmentedColormap only,
+# leave everything to register_cmap:
+
+plt.register_cmap(name='BlueRed3', data=cdict3) # optional lut kwarg
 
 x = np.arange(0, np.pi, 0.1)
 y = np.arange(0, 2*np.pi, 0.1)
@@ -121,13 +135,33 @@
 plt.colorbar()
 
 plt.subplot(1,3,2)
-plt.imshow(Z, interpolation='nearest', cmap=blue_red2)
+cmap = plt.get_cmap('BlueRed2')
+plt.imshow(Z, interpolation='nearest', cmap=cmap)
 plt.colorbar()
 
+# Now we will set the third cmap as the default.  One would
+# not normally do this in the middle of a script like this;
+# it is done here just to illustrate the method.
+
+plt.rcParams['image.cmap'] = 'BlueRed3'
+
+# Also see below for an alternative, particularly for
+# interactive use.
+
 plt.subplot(1,3,3)
-plt.imshow(Z, interpolation='nearest', cmap=blue_red3)
+plt.imshow(Z, interpolation='nearest')
 plt.colorbar()
 
+# Or as yet another variation, we could replace the rcParams
+# specification *before* the imshow with the following *after*
+# imshow:
+#
+# plt.set_cmap('BlueRed3')
+#
+# This sets the new default *and* sets the colormap of the last
+# image-like item plotted via pyplot, if any.
+
+
 plt.suptitle('Custom Blue-Red colormaps')
 
 plt.show()
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -9,16 +9,79 @@
 import matplotlib.cbook as cbook
 from matplotlib._cm import *
 
+# Dictionary for user-registered colormaps:
+cmap_d = dict()
+
+# Using this second dictionary allows us to handle any
+# Colormap instance; the built-in datad is only for
+# LinearSegmentedColormaps.  The advantage of keeping
+# datad is that it delays the generation of the Colormap
+# instance until it is actually needed. Generating the
+# instance is fast enough, and typically done few enough
+# times, that there is no need to cache the result.
+
+def register_cmap(name=None, cmap=None, data=None, lut=None):
+    """
+    Add a colormap to the set recognized by :func:`get_cmap`.
+
+    It can be used in two ways::
+
+        register_cmap(name='swirly', cmap=swirly_cmap)
+
+        register_cmap(name='choppy', data=choppydata, lut=128)
+
+    In the first case, *cmap* must be a :class:`colors.Colormap`
+    instance.  The *name* is optional; if absent, the name will
+    be the :attr:`name` attribute of the *cmap*.
+
+    In the second case, the three arguments are passed to
+    the :class:`colors.LinearSegmentedColormap` initializer,
+    and the resulting colormap is registered.
+
+    """
+    if name is None:
+        try:
+            name = cmap.name
+        except AttributeError:
+            raise ValueError("Arguments must include a name or a Colormap")
+
+    if not cbook.is_string_like(name):
+        raise ValueError("Colormap name must be a string")
+
+    if isinstance(cmap, colors.Colormap):
+        cmap_d[name] = cmap
+        return
 
+    # For the remainder, let exceptions propagate.
+    if lut is None:
+        lut = mpl.rcParams['image.lut']
+    cmap = colors.LinearSegmentedColormap(name, data, lut)
+    cmap_d[name] = cmap
 
 def get_cmap(name=None, lut=None):
     """
-    Get a colormap instance, defaulting to rc values if *name* is None
+    Get a colormap instance, defaulting to rc values if *name* is None.
+
+    Colormaps added with :func:`register_cmap` take precedence over
+    builtin colormaps.
+
+    If *name* is a :class:`colors.Colormap` instance, it will be
+    returned.
     """
-    if name is None: name = mpl.rcParams['image.cmap']
-    if lut is None: lut = mpl.rcParams['image.lut']
+    if name is None:
+        name = mpl.rcParams['image.cmap']
 
-    assert(name in datad.keys())
+    if isinstance(name, colors.Colormap):
+        return name
+
+    if name in cmap_d:
+        return cmap_d[name]
+
+    if name not in datad:
+        raise ValueError("%s is not a known colormap name" % name)
+
+    if lut is None:
+        lut = mpl.rcParams['image.lut']
     return colors.LinearSegmentedColormap(name,  datad[name], lut)
 
 class ScalarMappable:
@@ -116,9 +179,9 @@ def set_cmap(self, cmap):
         """
         set the colormap for luminance data
 
-        ACCEPTS: a colormap
+        ACCEPTS: a colormap or registered colormap name
         """
-        if cmap is None: cmap = get_cmap()
+        cmap = get_cmap(cmap)
         self.cmap = cmap
         self.changed()
 
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -512,7 +512,7 @@ def set_norm(self, norm):
     def set_cmap(self, cmap):
         if self._A is not None:
             raise RuntimeError('Cannot change colors after loading data')
-        cm.ScalarMappable.set_cmap(self, norm)
+        cm.ScalarMappable.set_cmap(self, cmap)
 
 class PcolorImage(martist.Artist, cm.ScalarMappable):
     '''
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -17,7 +17,7 @@
 from matplotlib.scale import get_scale_docs, get_scale_names
 
 from matplotlib import cm
-from matplotlib.cm import get_cmap
+from matplotlib.cm import get_cmap, register_cmap
 
 import numpy as np
 
@@ -1396,6 +1396,24 @@ def clim(vmin=None, vmax=None):
     im.set_clim(vmin, vmax)
     draw_if_interactive()
 
+def set_cmap(cmap):
+    '''
+    set the default colormap to *cmap* and apply to current image if any.
+    See help(colormaps) for more information.
+
+    *cmap* must be a :class:`colors.Colormap` instance, or
+    the name of a registered colormap.
+
+    See :func:`register_cmap` and :func:`get_cmap`.
+    '''
+    cmap = cm.get_cmap(cmap)
+
+    rc('image', cmap=cmap.name)
+    im = gci()
+
+    if im is not None:
+        im.set_cmap(cmap)
+    draw_if_interactive()
 
 
 def imread(*args, **kwargs):
@@ -6327,12 +6345,12 @@ def legend(*args, **kwargs):
 *bbox_to_anchor* keyword argument. bbox_to_anchor can be an instance
 of BboxBase(or its derivatives) or a tuple of 2 or 4 floats.
 For example, ::
-        
+
  loc = 'upper right', bbox_to_anchor = (0.5, 0.5)
 
 will place the legend so that the upper right corner of the legend at
 the center of the axes.
-        
+
 The legend location can be specified in other coordinate, by using the
 *bbox_transform* keyword.
 
@@ -6365,7 +6383,7 @@ def legend(*args, **kwargs):
 
   *fancybox*: [ None | False | True ]
     if True, draw a frame with a round fancybox.  If None, use rc
-    
+
   *shadow*: [ None | False | True ]
     If *True*, draw a shadow behind legend. If *None*, use rc settings.
 
