Merge pull request #660 from efiring/colormap_alpha

efiring · efiring · commit 912ea58e0762 · 2012-01-08T11:13:06.000-08:00
Colormap alpha
diff --git a/examples/pylab_examples/custom_cmap.py b/examples/pylab_examples/custom_cmap.py
@@ -103,6 +103,16 @@
                    (1.0, 0.0, 0.0))
         }
 
+# Make a modified version of cdict3 with some transparency
+# in the middle of the range.
+cdict4 = cdict3.copy()
+cdict4['alpha'] = ((0.0, 1.0, 1.0),
+                #   (0.25,1.0, 1.0),
+                   (0.5, 0.3, 0.3),
+                #   (0.75,1.0, 1.0),
+                   (1.0, 1.0, 1.0))
+
+
 # Now we will use this example to illustrate 3 ways of
 # handling custom colormaps.
 # First, the most direct and explicit:
@@ -121,20 +131,27 @@
 # leave everything to register_cmap:
 
 plt.register_cmap(name='BlueRed3', data=cdict3) # optional lut kwarg
+plt.register_cmap(name='BlueRedAlpha', data=cdict4)
+
+# Make some illustrative fake data:
 
 x = np.arange(0, np.pi, 0.1)
 y = np.arange(0, 2*np.pi, 0.1)
 X, Y = np.meshgrid(x,y)
-Z = np.cos(X) * np.sin(Y)
+Z = np.cos(X) * np.sin(Y) * 10
+
+# Make the figure:
 
-plt.figure(figsize=(10,4))
-plt.subplots_adjust(wspace=0.3)
+plt.figure(figsize=(6,9))
+plt.subplots_adjust(left=0.02, bottom=0.06, right=0.95, top=0.94, wspace=0.05)
 
-plt.subplot(1,3,1)
+# Make 4 subplots:
+
+plt.subplot(2,2,1)
 plt.imshow(Z, interpolation='nearest', cmap=blue_red1)
 plt.colorbar()
 
-plt.subplot(1,3,2)
+plt.subplot(2,2,2)
 cmap = plt.get_cmap('BlueRed2')
 plt.imshow(Z, interpolation='nearest', cmap=cmap)
 plt.colorbar()
@@ -145,24 +162,32 @@
 
 plt.rcParams['image.cmap'] = 'BlueRed3'
 
-# Also see below for an alternative, particularly for
-# interactive use.
-
-plt.subplot(1,3,3)
+plt.subplot(2,2,3)
 plt.imshow(Z, interpolation='nearest')
 plt.colorbar()
+plt.title("Alpha = 1")
 
-# Or as yet another variation, we could replace the rcParams
+# Or as yet another variation, we can replace the rcParams
 # specification *before* the imshow with the following *after*
-# imshow:
-#
-# plt.set_cmap('BlueRed3')
-#
+# imshow.
 # This sets the new default *and* sets the colormap of the last
 # image-like item plotted via pyplot, if any.
+#
 
+plt.subplot(2,2,4)
+# Draw a line with low zorder so it will be behind the image.
+plt.plot([0, 10*np.pi], [0, 20*np.pi], color='c', lw=20, zorder=-1)
+
+plt.imshow(Z, interpolation='nearest')
+plt.colorbar()
+
+# Here it is: changing the colormap for the current image and its
+# colorbar after they have been plotted.
+plt.set_cmap('BlueRedAlpha')
+plt.title("Varying alpha")
+#
 
-plt.suptitle('Custom Blue-Red colormaps')
+plt.suptitle('Custom Blue-Red colormaps', fontsize=16)
 
 plt.show()
 
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -183,33 +183,59 @@ def set_colorbar(self, im, ax):
         self.colorbar = im, ax
 
     def to_rgba(self, x, alpha=None, bytes=False):
-        '''Return a normalized rgba array corresponding to *x*. If *x*
-        is already an rgb array, insert *alpha*; if it is already
-        rgba, return it unchanged. If *bytes* is True, return rgba as
-        4 uint8s instead of 4 floats.
-        '''
-        if alpha is None:
-            _alpha = 1.0
-        else:
-            _alpha = alpha
+        """
+        Return a normalized rgba array corresponding to *x*.
+
+        In the normal case, *x* is a 1-D or 2-D sequence of scalars, and
+        the corresponding ndarray of rgba values will be returned,
+        based on the norm and colormap set for this ScalarMappable.
+
+        There is one special case, for handling images that are already
+        rgb or rgba, such as might have been read from an image file.
+        If *x* is an ndarray with 3 dimensions,
+        and the last dimension is either 3 or 4, then it will be
+        treated as an rgb or rgba array, and no mapping will be done.
+        If the last dimension is 3, the *alpha* kwarg (defaulting to 1)
+        will be used to fill in the transparency.  If the last dimension
+        is 4, the *alpha* kwarg is ignored; it does not
+        replace the pre-existing alpha.  A ValueError will be raised
+        if the third dimension is other than 3 or 4.
+
+        In either case, if *bytes* is *False* (default), the rgba
+        array will be floats in the 0-1 range; if it is *True*,
+        the returned rgba array will be uint8 in the 0 to 255 range.
+
+        Note: this method assumes the input is well-behaved; it does
+        not check for anomalies such as *x* being a masked rgba
+        array, or being an integer type other than uint8, or being
+        a floating point rgba array with values outside the 0-1 range.
+        """
+        # First check for special case, image input:
         try:
             if x.ndim == 3:
                 if x.shape[2] == 3:
+                    if alpha is None:
+                        alpha = 1
                     if x.dtype == np.uint8:
-                        _alpha = np.array(_alpha*255, np.uint8)
+                        alpha = np.uint8(alpha * 255)
                     m, n = x.shape[:2]
                     xx = np.empty(shape=(m,n,4), dtype = x.dtype)
                     xx[:,:,:3] = x
-                    xx[:,:,3] = _alpha
+                    xx[:,:,3] = alpha
                 elif x.shape[2] == 4:
                     xx = x
                 else:
                     raise ValueError("third dimension must be 3 or 4")
                 if bytes and xx.dtype != np.uint8:
                     xx = (xx * 255).astype(np.uint8)
+                if not bytes and xx.dtype == np.uint8:
+                    xx = xx.astype(float) / 255
                 return xx
         except AttributeError:
+            # e.g., x is not an ndarray; so try mapping it
             pass
+
+        # This is the normal case, mapping a scalar array:
         x = ma.asarray(x)
         x = self.norm(x)
         x = self.cmap(x, alpha=alpha, bytes=bytes)
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -624,7 +624,7 @@ def __init__(self, name, segmentdata, N=256, gamma=1.0):
 
         segmentdata argument is a dictionary with a red, green and blue
         entries. Each entry should be a list of *x*, *y0*, *y1* tuples,
-        forming rows in a table.
+        forming rows in a table. Entries for alpha are optional.
 
         Example: suppose you want red to increase from 0 to 1 over
         the bottom half, green to do the same over the middle half,
@@ -680,6 +680,9 @@ def _init(self):
                 self._segmentdata['green'], self._gamma)
         self._lut[:-3, 2] = makeMappingArray(self.N,
                 self._segmentdata['blue'], self._gamma)
+        if 'alpha' in self._segmentdata:
+            self._lut[:-3, 3] = makeMappingArray(self.N,
+                    self._segmentdata['alpha'], 1)
         self._isinit = True
         self._set_extremes()
 
@@ -711,12 +714,13 @@ def from_list(name, colors, N=256, gamma=1.0):
         else:
             vals = np.linspace(0., 1., len(colors))
 
-        cdict = dict(red=[], green=[], blue=[])
+        cdict = dict(red=[], green=[], blue=[], alpha=[])
         for val, color in zip(vals, colors):
-            r,g,b = colorConverter.to_rgb(color)
+            r,g,b,a = colorConverter.to_rgba(color)
             cdict['red'].append((val, r, r))
             cdict['green'].append((val, g, g))
             cdict['blue'].append((val, b, b))
+            cdict['alpha'].append((val, a, a))
 
         return LinearSegmentedColormap(name, cdict, N, gamma)
 
@@ -733,7 +737,8 @@ def __init__(self, colors, name = 'from_list', N = None):
 
         *colors*
             a list of matplotlib color specifications,
-            or an equivalent Nx3 floating point array (*N* rgb values)
+            or an equivalent Nx3  or Nx4 floating point array
+            (*N* rgb or rgba values)
         *name*
             a string to identify the colormap
         *N*
@@ -774,11 +779,9 @@ def __init__(self, colors, name = 'from_list', N = None):
 
 
     def _init(self):
-        rgb = np.array([colorConverter.to_rgb(c)
-                    for c in self.colors], np.float)
+        rgba = colorConverter.to_rgba_array(self.colors)
         self._lut = np.zeros((self.N + 3, 4), np.float)
-        self._lut[:-3, :-1] = rgb
-        self._lut[:-3, -1] = 1
+        self._lut[:-3] = rgba
         self._isinit = True
         self._set_extremes()
 
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -199,7 +199,14 @@ def _get_unsampled_image(self, A, image_extents, viewlim):
                 im.is_grayscale = False
             else:
                 if self._rgbacache is None:
-                    x = self.to_rgba(self._A, bytes=True)
+                    x = self.to_rgba(self._A, bytes=False)
+                    # Avoid side effects: to_rgba can return its argument
+                    # unchanged.
+                    if np.may_share_memory(x, self._A):
+                        x = x.copy()
+                    # premultiply the colors
+                    x[...,0:3] *= x[...,3:4]
+                    x = (x * 255).astype(np.uint8)
                     self._rgbacache = x
                 else:
                     x = self._rgbacache
