Merge pull request #2002 from pelson/contour_levels_with_colors

mdboom · mdboom · commit 355b0567f839 · 2013-05-21T07:14:30.000-07:00
Added support for providing 1 or 2 extra colours to the contour routines to easily specify the under and over colors.
diff --git a/doc/api/api_changes.rst b/doc/api/api_changes.rst
@@ -122,16 +122,25 @@ Changes in 1.3.x
   value of this kwarg is ``head_width = 20 * width``.
 
 * Removed call of :meth:`~matplotlib.axes.Axes.grid` in
-  :meth:`~matplotlib.pyplot.plotfile`. To draw the axes grid, set to *True*
-  matplotlib.rcParams['axes.grid'] or ``axes.grid`` in ``.matplotlibrc`` or
-  explicitly call :meth:`~matplotlib.axes.Axes.grid`
+  :meth:`~matplotlib.pyplot.plotfile`. To draw the axes grid, set the
+  ``axes.grid`` rcParam to *True*, or explicitly call
+  :meth:`~matplotlib.axes.Axes.grid`.
+
+* It is now posible to provide ``number of levels + 1`` colors in the case of
+  `extend='both'` for contourf (or just ``number of levels`` colors for an
+  extend value ``min`` or ``max``) such that the resulting colormap's
+  ``set_under`` and ``set_over`` are defined appropriately. Any other number
+  of colors will continue to behave as before (if more colors are provided
+  than levels, the colors will be unused). A similar change has been applied
+  to contour, where ``extend='both'`` would expect ``number of levels + 2``
+  colors.
 
 * A new keyword *extendrect* in :meth:`~matplotlib.pyplot.colorbar` and
   :class:`~matplotlib.colorbar.ColorbarBase` allows one to control the shape
   of colorbar extensions.
 
 * The `~matplotlib.mpl` module is now deprecated. Those who relied on this
-  module should transition to simply using `import matplotlib as mpl`.
+  module should transition to simply using ``import matplotlib as mpl``.
 
 * The extension of :class:`~matplotlib.widgets.MultiCursor` to both vertical
   (default) and/or horizontal cursor implied that ``self.line`` is replaced
@@ -142,8 +151,8 @@ Changes in 1.3.x
   raises :class:`NotImplementedError` instead of :class:`OSError` if the
   :command:`ps` command cannot be run.
 
-* The :func:`~matplotlib.cbook.check_output` function has been moved to
-  `~matplotlib.compat.subprocess`.
+* The :func:`matplotlib.cbook.check_output` function has been moved to
+  :func:`matplotlib.compat.subprocess`.
 
 * :class:`~matplotlib.patches.Patch` now fully supports using RGBA values for
   its ``facecolor`` and ``edgecolor`` attributes, which enables faces and
diff --git a/lib/matplotlib/contour.py b/lib/matplotlib/contour.py
@@ -864,7 +864,32 @@ def __init__(self, ax, *args, **kwargs):
             ncolors = len(self.levels)
             if self.filled:
                 ncolors -= 1
-            cmap = colors.ListedColormap(self.colors, N=ncolors)
+            i0 = 0
+
+            # Handle the case where colors are given for the extended
+            # parts of the contour.
+            extend_min = self.extend in ['min', 'both']
+            extend_max = self.extend in ['max', 'both']
+            use_set_under_over = False
+            # if we are extending the lower end, and we've been given enough
+            # colors then skip the first color in the resulting cmap. For the
+            # extend_max case we don't need to worry about passing more colors
+            # than ncolors as ListedColormap will clip.
+            total_levels = ncolors + int(extend_min) + int(extend_max)
+            if (len(self.colors) == total_levels and
+                    any([extend_min, extend_max])):
+                use_set_under_over = True
+                if extend_min:
+                    i0 = 1
+
+            cmap = colors.ListedColormap(self.colors[i0:None], N=ncolors)
+
+            if use_set_under_over:
+                if extend_min:
+                    cmap.set_under(self.colors[0])
+                if extend_max:
+                    cmap.set_over(self.colors[-1])
+
         if self.filled:
             self.collections = cbook.silent_list('mcoll.PathCollection')
         else:
@@ -1172,16 +1197,16 @@ def _process_levels(self):
         # (Colorbar needs this even for line contours.)
         self._levels = list(self.levels)
 
-        if not self.filled:
-            self.layers = self.levels
-            return
-
         if self.extend in ('both', 'min'):
             self._levels.insert(0, min(self.levels[0], self.zmin) - 1)
         if self.extend in ('both', 'max'):
             self._levels.append(max(self.levels[-1], self.zmax) + 1)
         self._levels = np.asarray(self._levels)
 
+        if not self.filled:
+            self.layers = self.levels
+            return
+
         # layer values are mid-way between levels
         self.layers = 0.5 * (self._levels[:-1] + self._levels[1:])
         # ...except that extended layers must be outside the
@@ -1526,9 +1551,7 @@ def _check_xyz(self, args, kwargs):
             if y.shape != z.shape:
                 raise TypeError("Shape of y does not match that of z: found "
                                 "{0} instead of {1}.".format(y.shape, z.shape))
-
         else:
-
             raise TypeError("Inputs x and y must be 1D or 2D.")
 
         return x, y, z
diff --git a/lib/matplotlib/tests/baseline_images/test_contour/contour_manual_colors_and_levels.png b/lib/matplotlib/tests/baseline_images/test_contour/contour_manual_colors_and_levels.png
diff --git a/lib/matplotlib/tests/test_contour.py b/lib/matplotlib/tests/test_contour.py
@@ -149,3 +149,34 @@ def test_contour_manual_labels():
     cs = plt.contour(x,y,z)
     pts = np.array([(1.5, 3.0), (1.5, 4.4), (1.5, 6.0)])
     plt.clabel(cs, manual=pts)
+
+
+@image_comparison(baseline_images=['contour_manual_colors_and_levels'],
+                  extensions=['png'], remove_text=True)
+def test_given_colors_levels_and_extends():
+    _, axes = plt.subplots(2, 4)
+
+    data = np.arange(12).reshape(3, 4)
+    
+    colors = ['red', 'yellow', 'pink', 'blue', 'black']
+    levels = [2, 4, 8, 10]
+    
+    for i, ax in enumerate(axes.flatten()):
+        plt.sca(ax)
+        
+        filled = i % 2 == 0.
+        extend = ['neither', 'min', 'max', 'both'][i // 2]
+        
+        if filled:
+            last_color = -1 if extend in ['min', 'max'] else None
+            plt.contourf(data, colors=colors[:last_color], levels=levels, extend=extend)
+        else:
+            last_level = -1 if extend == 'both' else None
+            plt.contour(data, colors=colors, levels=levels[:last_level], extend=extend)
+    
+        plt.colorbar()
+
+
+if __name__ == '__main__':
+    import nose
+    nose.runmodule(argv=['-s', '--with-doctest'], exit=False)
