Add a helper to copy a colormap and set its extreme colors.

anntzer · anntzer · commit 625bdc74e60b · 2019-06-28T18:48:11.000+02:00
See changelog.  See also the number of explicit copies in the examples,
which suggest that the old setter-based API is a bit of a footgun (as
forgetting to copy seems easy).
diff --git a/doc/users/next_whats_new/2019-06-28-AL.rst b/doc/users/next_whats_new/2019-06-28-AL.rst
@@ -0,0 +1,13 @@
+`Colormap.set_extremes` and `Colormap.with_extremes`
+````````````````````````````````````````````````````
+
+Because the `Colormap.set_bad`, `Colormap.set_under` and `Colormap.set_over`
+methods modify the colormap in place, the user must be careful to first make a
+copy of the colormap if setting the extreme colors e.g. for a builtin colormap.
+
+The new ``Colormap.with_extremes(bad=..., under=..., over=...)`` can be used to
+first copy the colormap and set the extreme colors on that copy.
+
+The new `Colormap.set_extremes` method is provided for API symmetry with
+`Colormap.with_extremes`, but note that it suffers from the same issue as the
+earlier individual setters.
diff --git a/examples/images_contours_and_fields/image_masked.py b/examples/images_contours_and_fields/image_masked.py
@@ -8,7 +8,6 @@
 The second subplot illustrates the use of BoundaryNorm to
 get a filled contour effect.
 """
-from copy import copy
 
 import numpy as np
 import matplotlib.pyplot as plt
@@ -25,11 +24,7 @@
 Z = (Z1 - Z2) * 2
 
 # Set up a colormap:
-# use copy so that we do not mutate the global colormap instance
-palette = copy(plt.cm.gray)
-palette.set_over('r', 1.0)
-palette.set_under('g', 1.0)
-palette.set_bad('b', 1.0)
+palette = plt.cm.gray.with_extremes(over='r', under='g', bad='b')
 # Alternatively, we could use
 # palette.set_bad(alpha = 0.0)
 # to make the bad region transparent.  This is the default.
diff --git a/examples/images_contours_and_fields/quadmesh_demo.py b/examples/images_contours_and_fields/quadmesh_demo.py
@@ -9,8 +9,6 @@
 This demo illustrates a bug in quadmesh with masked data.
 """
 
-import copy
-
 from matplotlib import cm, pyplot as plt
 import numpy as np
 
@@ -30,10 +28,8 @@
 axs[0].pcolormesh(Qx, Qz, Z, shading='gouraud')
 axs[0].set_title('Without masked values')
 
-# You can control the color of the masked region. We copy the default colormap
-# before modifying it.
-cmap = copy.copy(cm.get_cmap(plt.rcParams['image.cmap']))
-cmap.set_bad('y', 1.0)
+# You can control the color of the masked region.
+cmap = cm.get_cmap(plt.rcParams['image.cmap']).with_extremes(bad='y')
 axs[1].pcolormesh(Qx, Qz, Zm, shading='gouraud', cmap=cmap)
 axs[1].set_title('With masked values')
 
diff --git a/examples/specialty_plots/leftventricle_bulleye.py b/examples/specialty_plots/leftventricle_bulleye.py
@@ -176,9 +176,8 @@ def bullseye_plot(ax, data, seg_bold=None, cmap=None, norm=None):
 # The second example illustrates the use of a ListedColormap, a
 # BoundaryNorm, and extended ends to show the "over" and "under"
 # value colors.
-cmap3 = mpl.colors.ListedColormap(['r', 'g', 'b', 'c'])
-cmap3.set_over('0.35')
-cmap3.set_under('0.75')
+cmap3 = (mpl.colors.ListedColormap(['r', 'g', 'b', 'c'])
+         .with_extremes(over='0.35', under='0.75'))
 
 # If a ListedColormap is used, the length of the bounds array must be
 # one greater than the length of the color list.  The bounds must be
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -60,6 +60,7 @@
 """
 
 from collections.abc import Sized
+import copy
 import itertools
 import re
 
@@ -606,8 +607,7 @@ def __copy__(self):
         return cmapobject
 
     def set_bad(self, color='k', alpha=None):
-        """Set color to be used for masked values.
-        """
+        """Set color to be used for masked values."""
         self._rgba_bad = to_rgba(color, alpha)
         if self._isinit:
             self._set_extremes()
@@ -628,6 +628,28 @@ def set_over(self, color='k', alpha=None):
         if self._isinit:
             self._set_extremes()
 
+    def set_extremes(self, *, bad=None, under=None, over=None):
+        """
+        Set the colors for masked (*bad*) values and, when ``norm.clip =
+        False``, low (*under*) and high (*over*) out-of-range values.
+        """
+        if bad is not None:
+            self.set_bad(bad)
+        if under is not None:
+            self.set_under(under)
+        if over is not None:
+            self.set_over(over)
+
+    def with_extremes(self, *, bad=None, under=None, over=None):
+        """
+        Return a copy of the colormap, for which the colors for masked (*bad*)
+        values and, when ``norm.clip = False``, low (*under*) and high (*over*)
+        out-of-range values, have been set accordingly.
+        """
+        new_cm = copy.copy(self)
+        new_cm.set_extremes(bad=bad, under=under, over=over)
+        return new_cm
+
     def _set_extremes(self):
         if self._rgba_under:
             self._lut[self._i_under] = self._rgba_under
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
@@ -1,5 +1,4 @@
 from contextlib import ExitStack
-from copy import copy
 import io
 import os
 import sys
@@ -769,10 +768,7 @@ def test_mask_image_over_under():
           (2 * np.pi * 0.5 * 1.5))
     Z = 10*(Z2 - Z1)  # difference of Gaussians
 
-    palette = copy(plt.cm.gray)
-    palette.set_over('r', 1.0)
-    palette.set_under('g', 1.0)
-    palette.set_bad('b', 1.0)
+    palette = plt.cm.gray.with_extremes(over='r', under='g', bad='b')
     Zm = ma.masked_where(Z > 1.2, Z)
     fig, (ax1, ax2) = plt.subplots(1, 2)
     im = ax1.imshow(Zm, interpolation='bilinear',
@@ -830,10 +826,7 @@ def test_imshow_endianess():
                   remove_text=True, style='mpl20')
 def test_imshow_masked_interpolation():
 
-    cm = copy(plt.get_cmap('viridis'))
-    cm.set_over('r')
-    cm.set_under('b')
-    cm.set_bad('k')
+    cm = plt.get_cmap('viridis').with_extremes(over='r', under='b', bad='k')
 
     N = 20
     n = colors.Normalize(vmin=0, vmax=N*N-1)
diff --git a/tutorials/colors/colorbar_only.py b/tutorials/colors/colorbar_only.py
@@ -64,9 +64,8 @@
 fig, ax = plt.subplots(figsize=(6, 1))
 fig.subplots_adjust(bottom=0.5)
 
-cmap = mpl.colors.ListedColormap(['red', 'green', 'blue', 'cyan'])
-cmap.set_over('0.25')
-cmap.set_under('0.75')
+cmap = (mpl.colors.ListedColormap(['red', 'green', 'blue', 'cyan'])
+        .with_extremes(over='0.25', under='0.75'))
 
 bounds = [1, 2, 4, 7, 8]
 norm = mpl.colors.BoundaryNorm(bounds, cmap.N)
@@ -91,10 +90,8 @@
 fig, ax = plt.subplots(figsize=(6, 1))
 fig.subplots_adjust(bottom=0.5)
 
-cmap = mpl.colors.ListedColormap(['royalblue', 'cyan',
-                                  'yellow', 'orange'])
-cmap.set_over('red')
-cmap.set_under('blue')
+cmap = (mpl.colors.ListedColormap(['royalblue', 'cyan', 'yellow', 'orange'])
+        .with_extremes(over='red', under='blue'))
 
 bounds = [-1.0, -0.5, 0.0, 0.5, 1.0]
 norm = mpl.colors.BoundaryNorm(bounds, cmap.N)
