Merge pull request #2050 from pelson/colormap_api

mdboom · mdboom · commit ebb7a9ddcbe4 · 2013-05-28T12:03:48.000-04:00
Added the from_levels_and_colors function.
diff --git a/doc/api/api_changes.rst b/doc/api/api_changes.rst
@@ -104,6 +104,10 @@ Changes in 1.3.x
   Deep copying a `Path` always creates an editable (i.e. non-readonly)
   `Path`.
 
+* matplotlib.colors.normalize and matplotlib.colors.no_norm have been
+  deprecated in favour of matplotlib.colors.Normalize and
+  matplotlib.colors.NoNorm respectively.
+
 * The `font.*` rcParams now affect only text objects created after the
   rcParam has been set, and will not retroactively affect already
   existing text objects.  This brings their behavior in line with most
diff --git a/doc/users/whats_new.rst b/doc/users/whats_new.rst
@@ -88,6 +88,14 @@ rcParam has been set, and will not retroactively affect already
 existing text objects.  This brings their behavior in line with most
 other rcParams.
 
+Easier creation of colormap and normalizer for levels with colors
+-----------------------------------------------------------------
+Phil Elson added the :func:`matplotlib.colors.from_levels_and_colors`
+function to easily create a colormap and normalizer for representation
+of discrete colors for plot types such as
+:func:`matplotlib.pyplot.pcolormesh`, with a similar interface to that of
+contourf.
+
 Catch opening too many figures using pyplot
 -------------------------------------------
 Figures created through `pyplot.figure` are retained until they are
diff --git a/lib/matplotlib/colorbar.py b/lib/matplotlib/colorbar.py
@@ -852,8 +852,8 @@ def __init__(self, ax, mappable, **kw):
         mappable.autoscale_None()
 
         self.mappable = mappable
-        kw['cmap'] = mappable.cmap
-        kw['norm'] = mappable.norm
+        kw['cmap'] = cmap = mappable.cmap
+        kw['norm'] = norm = mappable.norm
 
         if isinstance(mappable, contour.ContourSet):
             CS = mappable
@@ -868,6 +868,9 @@ def __init__(self, ax, mappable, **kw):
             if not CS.filled:
                 self.add_lines(CS)
         else:
+            if getattr(cmap, 'colorbar_extend', False) is not False:
+                kw.setdefault('extend', cmap.colorbar_extend)
+            
             if isinstance(mappable, martist.Artist):
                 kw['alpha'] = mappable.get_alpha()
 
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -508,6 +508,12 @@ def __init__(self, name, N=256):
         self._i_bad = N + 2
         self._isinit = False
 
+        #: When this colormap exists on a scalar mappable and colorbar_extend
+        #: is not False, colorbar creation will pick up ``colorbar_extend`` as
+        #: the default value for the ``extend`` keyword in the
+        #: :class:`matplotlib.colorbar.Colorbar` constructor.
+        self.colorbar_extend = False
+
     def __call__(self, X, alpha=None, bytes=False):
         """
         Parameters
@@ -832,7 +838,7 @@ def _init(self):
 class Normalize(object):
     """
     A class which, when called, can normalize data into
-    the ``[0, 1]`` interval.
+    the ``[0.0, 1.0]`` interval.
 
     """
     def __init__(self, vmin=None, vmax=None, clip=False):
@@ -1212,8 +1218,12 @@ def inverse(self, value):
         return value
 
 # compatibility with earlier class names that violated convention:
-normalize = Normalize
-no_norm = NoNorm
+normalize = cbook.deprecated('1.3', alternative='Normalize',
+                             name='normalize',
+                             obj_type='class alias')(Normalize)
+no_norm = cbook.deprecated('1.3', alternative='NoNorm',
+                           name='no_norm',
+                           obj_type='class alias')(NoNorm)
 
 
 def rgb_to_hsv(arr):
@@ -1405,3 +1415,71 @@ def shade_rgb(self, rgb, elevation, fraction=1.):
         hsv[:, :, 1:] = np.where(hsv[:, :, 1:] > 1., 1, hsv[:, :, 1:])
         # convert modified hsv back to rgb.
         return hsv_to_rgb(hsv)
+
+
+def from_levels_and_colors(levels, colors, extend='neither'):
+    """
+    A helper routine to generate a cmap and a norm instance which
+    behave similar to contourf's levels and colors arguments.
+
+    Parameters
+    ----------
+    levels : sequence of numbers
+        The quantization levels used to construct the :class:`BoundaryNorm`.
+        Values ``v`` are quantizized to level ``i`` if
+        ``lev[i] <= v < lev[i+1]``.
+    colors : sequence of colors
+        The fill color to use for each level. If `extend` is "neither" there
+        must be ``n_level - 1`` colors. For an `extend` of "min" or "max" add
+        one extra color, and for an `extend` of "both" add two colors.
+    extend : {'neither', 'min', 'max', 'both'}, optional
+        The behaviour when a value falls out of range of the given levels.
+        See :func:`~matplotlib.pyplot.contourf` for details.
+
+    Returns
+    -------
+    (cmap, norm) : tuple containing a :class:`Colormap` and a \
+                   :class:`Normalize` instance
+    """
+    colors_i0 = 0
+    colors_i1 = None
+
+    if extend == 'both':
+        colors_i0 = 1
+        colors_i1 = -1
+        extra_colors = 2
+    elif extend == 'min':
+        colors_i0 = 1
+        extra_colors = 1
+    elif extend == 'max':
+        colors_i1 = -1
+        extra_colors = 1
+    elif extend == 'neither':
+        extra_colors = 0
+    else:
+        raise ValueError('Unexpected value for extend: {0!r}'.format(extend))
+
+    n_data_colors = len(levels) - 1
+    n_expected_colors = n_data_colors + extra_colors
+    if len(colors) != n_expected_colors:
+        raise ValueError('With extend == {0!r} and n_levels == {1!r} expected'
+                         ' n_colors == {2!r}. Got {3!r}.'
+                         ''.format(extend, len(levels), n_expected_colors,
+                                   len(colors)))
+
+    cmap = ListedColormap(colors[colors_i0:colors_i1], N=n_data_colors)
+
+    if extend in ['min', 'both']:
+        cmap.set_under(colors[0])
+    else:
+        cmap.set_under('none')
+
+    if extend in ['max', 'both']:
+        cmap.set_over(colors[-1])
+    else:
+        cmap.set_over('none')
+
+    cmap.colorbar_extend = extend
+
+    norm = BoundaryNorm(levels, ncolors=n_data_colors)
+    return cmap, norm
diff --git a/lib/matplotlib/tests/baseline_images/test_colors/levels_and_colors.png b/lib/matplotlib/tests/baseline_images/test_colors/levels_and_colors.png
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -1,12 +1,14 @@
-"""
-Tests for the colors module.
-"""
-
 from __future__ import print_function
+from nose.tools import assert_raises
 import numpy as np
 from numpy.testing.utils import assert_array_equal, assert_array_almost_equal
+
+
 import matplotlib.colors as mcolors
 import matplotlib.cm as cm
+import matplotlib.pyplot as plt
+from matplotlib.testing.decorators import image_comparison
+
 
 def test_colormap_endian():
     """
@@ -23,6 +25,7 @@ def test_colormap_endian():
         #print(anative.dtype.isnative, aforeign.dtype.isnative)
         assert_array_equal(cmap(anative), cmap(aforeign))
 
+
 def test_BoundaryNorm():
     """
     Github issue #1258: interpolation was failing with numpy
@@ -36,7 +39,8 @@ def test_BoundaryNorm():
     ncolors = len(boundaries)
     bn = mcolors.BoundaryNorm(boundaries, ncolors)
     assert_array_equal(bn(vals), expected)
-    
+
+
 def test_LogNorm():
     """
     LogNorm igornoed clip, now it has the same
@@ -46,6 +50,7 @@ def test_LogNorm():
     ln = mcolors.LogNorm(clip=True, vmax=5)
     assert_array_equal(ln([1, 6]), [0, 1.0])
 
+
 def test_Normalize():
     norm = mcolors.Normalize()
     vals = np.arange(-10, 10, 1, dtype=np.float)
@@ -74,6 +79,7 @@ def _inverse_tester(norm_instance, vals):
     """
     assert_array_almost_equal(norm_instance.inverse(norm_instance(vals)), vals)
 
+
 def _scalar_tester(norm_instance, vals):
     """
     Checks if scalars and arrays are handled the same way.
@@ -82,10 +88,88 @@ def _scalar_tester(norm_instance, vals):
     scalar_result = [norm_instance(float(v)) for v in vals]
     assert_array_almost_equal(scalar_result, norm_instance(vals))
 
+
 def _mask_tester(norm_instance, vals):
     """
     Checks mask handling
     """
     masked_array = np.ma.array(vals)
     masked_array[0] = np.ma.masked
     assert_array_equal(masked_array.mask, norm_instance(masked_array).mask)
+
+
+@image_comparison(baseline_images=['levels_and_colors'],
+                  extensions=['png'])
+def test_cmap_and_norm_from_levels_and_colors():
+    data = np.linspace(-2, 4, 49).reshape(7, 7)
+    levels = [-1, 2, 2.5, 3]
+    colors = ['red', 'green', 'blue', 'yellow', 'black']
+    extend = 'both'
+    cmap, norm = mcolors.from_levels_and_colors(levels, colors, extend=extend)
+
+    ax = plt.axes()
+    m = plt.pcolormesh(data, cmap=cmap, norm=norm)
+    plt.colorbar(m)
+
+    # Hide the axes labels (but not the colorbar ones, as they are useful)
+    for lab in ax.get_xticklabels() + ax.get_yticklabels():
+        lab.set_visible(False)
+
+
+def test_cmap_and_norm_from_levels_and_colors2():
+    levels = [-1, 2, 2.5, 3]
+    colors = ['red', (0, 1, 0), 'blue', (0.5, 0.5, 0.5), (0.0, 0.0, 0.0, 1.0)]
+    clr = mcolors.colorConverter.to_rgba_array(colors)
+    bad = (0.1, 0.1, 0.1, 0.1)
+    no_color = (0.0, 0.0, 0.0, 0.0)
+
+    # Define the test values which are of interest.
+    # Note: levels are lev[i] <= v < lev[i+1]
+    tests = [('both', None, {-2: clr[0],
+                             -1: clr[1],
+                             2: clr[2],
+                             2.25: clr[2],
+                             3: clr[4],
+                             3.5: clr[4],
+                             np.ma.array(1, mask=True): bad}),
+
+             ('min', -1, {-2: clr[0],
+                          -1: clr[1],
+                          2: clr[2],
+                          2.25: clr[2],
+                          3: no_color,
+                          3.5: no_color,
+                          np.ma.array(1, mask=True): bad}),
+
+             ('max', -1, {-2: no_color,
+                          -1: clr[0],
+                          2: clr[1],
+                          2.25: clr[1],
+                          3: clr[3],
+                          3.5: clr[3],
+                          np.ma.array(1, mask=True): bad}),
+
+             ('neither', -2, {-2: no_color,
+                              -1: clr[0],
+                              2: clr[1],
+                              2.25: clr[1],
+                              3: no_color,
+                              3.5: no_color,
+                              np.ma.array(1, mask=True): bad}),
+             ]
+
+    for extend, i1, cases in tests:
+        cmap, norm = mcolors.from_levels_and_colors(levels, colors[0:i1],
+                                                    extend=extend)
+        cmap.set_bad(bad)
+        for d_val, expected_color in sorted(cases.items()):
+            assert_array_equal(expected_color, cmap(norm([d_val]))[0],
+                               'Wih extend={0!r} and data '
+                               'value={1!r}'.format(extend, d_val))
+
+    assert_raises(ValueError, mcolors.from_levels_and_colors, levels, colors)
+
+
+if __name__ == '__main__':
+    import nose
+    nose.runmodule(argv=['-s', '--with-doctest'], exit=False)
