FIX: fix twoslope norm for colorbars

jklymak · jklymak · commit 9ef1caa0d5ad · 2021-07-07T06:55:46.000-07:00
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -1300,9 +1300,13 @@ def __call__(self, value, clip=None):
 
         if not self.vmin <= self.vcenter <= self.vmax:
             raise ValueError("vmin, vcenter, vmax must increase monotonically")
+        # must linearly extrapolate for ticks on colorbars that go
+        # beyond vmin and vmax:
+        min = self.vmin - (self.vcenter - self.vmin)
+        max = self.vmax + (self.vmax - self.vcenter)
         result = np.ma.masked_array(
-            np.interp(result, [self.vmin, self.vcenter, self.vmax],
-                      [0, 0.5, 1.]), mask=np.ma.getmask(result))
+            np.interp(result, [min, self.vcenter, max],
+                      [-0.5, 0.5, 1.5]), mask=np.ma.getmask(result))
         if is_scalar:
             result = np.atleast_1d(result)[0]
         return result
@@ -1313,8 +1317,9 @@ def inverse(self, value):
         (vmin,), _ = self.process_value(self.vmin)
         (vmax,), _ = self.process_value(self.vmax)
         (vcenter,), _ = self.process_value(self.vcenter)
-
-        result = np.interp(value, [0, 0.5, 1.], [vmin, vcenter, vmax])
+        min = vmin - (vcenter - vmin)
+        max = vmax + (vmax - vcenter)
+        result = np.interp(value, [-0.5, 0.5, 1.5], [min, vcenter, max])
         return result
 
 
diff --git a/lib/matplotlib/tests/baseline_images/test_colorbar/colorbar_twoslope.png b/lib/matplotlib/tests/baseline_images/test_colorbar/colorbar_twoslope.png
diff --git a/lib/matplotlib/tests/test_colorbar.py b/lib/matplotlib/tests/test_colorbar.py
@@ -776,11 +776,12 @@ def test_inset_colorbar_layout():
 @image_comparison(['colorbar_twoslope.png'], remove_text=True,
                   style='mpl20')
 def test_twoslope_colorbar():
-    # Note that the first tick = 20, and should be in the middle
+    # Note that the second tick = 20, and should be in the middle
     # of the colorbar (white)
+    # There should be no tick right at the bottom, nor at the top.
     fig, ax = plt.subplots()
 
-    norm = mcolors.TwoSlopeNorm(20, 0, 100)
+    norm = mcolors.TwoSlopeNorm(20, 5, 95)
     pc = ax.pcolormesh(np.arange(1, 11), np.arange(1, 11),
                        np.arange(100).reshape(10, 10),
                        norm=norm, cmap='RdBu_r')
diff --git a/tutorials/colors/colormapnorms.py b/tutorials/colors/colormapnorms.py
@@ -277,7 +277,8 @@
 # longitude depends on latitude.
 ax.set_aspect(1 / np.cos(np.deg2rad(49)))
 ax.set_title('TwoSlopeNorm(x)')
-fig.colorbar(pcm, shrink=0.6)
+cb = fig.colorbar(pcm, shrink=0.6)
+cb.set_ticks([-500, 0, 1000, 2000, 3000, 4000])
 plt.show()
 
 
@@ -312,7 +313,8 @@ def _inverse(x):
 # ----------------------------------------------------------
 #
 # The `.TwoSlopeNorm` described above makes a useful example for
-# defining your own norm.
+# defining your own norm.  Note for the colorbar to work, you must
+# define an inverse for your norm:
 
 
 class MidpointNormalize(colors.Normalize):
@@ -323,9 +325,18 @@ def __init__(self, vmin=None, vmax=None, vcenter=None, clip=False):
     def __call__(self, value, clip=None):
         # I'm ignoring masked values and all kinds of edge cases to make a
         # simple example...
-        x, y = [self.vmin, self.vcenter, self.vmax], [0, 0.5, 1]
+        # Note also that we must extrapolate linearly beyond vmin/vmax
+        min = self.vmin - (self.vcenter - self.vmin)
+        max = self.vmax + (self.vmax - self.vcenter)
+        x, y = [min, self.vcenter, max], [-0.5, 0.5, 1.5]
         return np.ma.masked_array(np.interp(value, x, y))
 
+    def inverse(self, value):
+        min = self.vmin - (self.vcenter - self.vmin)
+        max = self.vmax + (self.vmax - self.vcenter)
+        y, x = [min, self.vcenter, max], [-0.5, 0.5, 1.5]
+        return np.interp(value, x, y)
+
 
 fig, ax = plt.subplots()
 midnorm = MidpointNormalize(vmin=-500., vcenter=0, vmax=4000)
@@ -334,5 +345,7 @@ def __call__(self, value, clip=None):
                     cmap=terrain_map, shading='auto')
 ax.set_aspect(1 / np.cos(np.deg2rad(49)))
 ax.set_title('Custom norm')
-fig.colorbar(pcm, shrink=0.6, extend='both')
+cb = fig.colorbar(pcm, shrink=0.6, extend='both')
+cb.set_ticks([-500, 0, 1000, 2000, 3000, 4000])
+
 plt.show()
