Merge pull request #16937 from greglucas/colormap_reverse_extend

anntzer · web-flow · commit c29282737bf0 · 2020-03-28T01:11:07.000+01:00
Making sure to keep over/under/bad in cmap resample/reverse.
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -790,7 +790,12 @@ def from_list(name, colors, N=256, gamma=1.0):
 
     def _resample(self, lutsize):
         """Return a new color map with *lutsize* entries."""
-        return LinearSegmentedColormap(self.name, self._segmentdata, lutsize)
+        new_cmap = LinearSegmentedColormap(self.name, self._segmentdata,
+                                           lutsize)
+        new_cmap._rgba_over = self._rgba_over
+        new_cmap._rgba_under = self._rgba_under
+        new_cmap._rgba_bad = self._rgba_bad
+        return new_cmap
 
     # Helper ensuring picklability of the reversed cmap.
     @staticmethod
@@ -821,7 +826,12 @@ def reversed(self, name=None):
                         [(1.0 - x, y1, y0) for x, y0, y1 in reversed(data)])
                   for key, data in self._segmentdata.items()}
 
-        return LinearSegmentedColormap(name, data_r, self.N, self._gamma)
+        new_cmap = LinearSegmentedColormap(name, data_r, self.N, self._gamma)
+        # Reverse the over/under values too
+        new_cmap._rgba_over = self._rgba_under
+        new_cmap._rgba_under = self._rgba_over
+        new_cmap._rgba_bad = self._rgba_bad
+        return new_cmap
 
 
 class ListedColormap(Colormap):
@@ -885,7 +895,12 @@ def _init(self):
     def _resample(self, lutsize):
         """Return a new color map with *lutsize* entries."""
         colors = self(np.linspace(0, 1, lutsize))
-        return ListedColormap(colors, name=self.name)
+        new_cmap = ListedColormap(colors, name=self.name)
+        # Keep the over/under values too
+        new_cmap._rgba_over = self._rgba_over
+        new_cmap._rgba_under = self._rgba_under
+        new_cmap._rgba_bad = self._rgba_bad
+        return new_cmap
 
     def reversed(self, name=None):
         """
@@ -906,7 +921,12 @@ def reversed(self, name=None):
             name = self.name + "_r"
 
         colors_r = list(reversed(self.colors))
-        return ListedColormap(colors_r, name=name, N=self.N)
+        new_cmap = ListedColormap(colors_r, name=name, N=self.N)
+        # Reverse the over/under values too
+        new_cmap._rgba_over = self._rgba_under
+        new_cmap._rgba_under = self._rgba_over
+        new_cmap._rgba_bad = self._rgba_bad
+        return new_cmap
 
 
 class Normalize:
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -39,13 +39,25 @@ def test_resample():
     colorlist[:, 3] = 0.7
     lsc = mcolors.LinearSegmentedColormap.from_list('lsc', colorlist)
     lc = mcolors.ListedColormap(colorlist)
+    # Set some bad values for testing too
+    for cmap in [lsc, lc]:
+        cmap.set_under('r')
+        cmap.set_over('g')
+        cmap.set_bad('b')
     lsc3 = lsc._resample(3)
     lc3 = lc._resample(3)
     expected = np.array([[0.0, 0.2, 1.0, 0.7],
                          [0.5, 0.2, 0.5, 0.7],
                          [1.0, 0.2, 0.0, 0.7]], float)
     assert_array_almost_equal(lsc3([0, 0.5, 1]), expected)
     assert_array_almost_equal(lc3([0, 0.5, 1]), expected)
+    # Test over/under was copied properly
+    assert_array_almost_equal(lsc(np.inf), lsc3(np.inf))
+    assert_array_almost_equal(lsc(-np.inf), lsc3(-np.inf))
+    assert_array_almost_equal(lsc(np.nan), lsc3(np.nan))
+    assert_array_almost_equal(lc(np.inf), lc3(np.inf))
+    assert_array_almost_equal(lc(-np.inf), lc3(-np.inf))
+    assert_array_almost_equal(lc(np.nan), lc3(np.nan))
 
 
 def test_register_cmap():
@@ -818,6 +830,10 @@ def test_colormap_reversing(name):
         cmap._init()
         cmap_r._init()
     assert_array_almost_equal(cmap._lut[:-3], cmap_r._lut[-4::-1])
+    # Test the bad, over, under values too
+    assert_array_almost_equal(cmap(-np.inf), cmap_r(np.inf))
+    assert_array_almost_equal(cmap(np.inf), cmap_r(-np.inf))
+    assert_array_almost_equal(cmap(np.nan), cmap_r(np.nan))
 
 
 def test_cn():
