Merge pull request #15090 from kmader/patch-1

tacaswell · web-flow · commit 02c8199a76f6 · 2019-10-31T20:51:51.000-04:00
ENH: Coerce MxNx1 images into MxN images for imshow
diff --git a/doc/users/next_whats_new/imshow_m_n_1.rst b/doc/users/next_whats_new/imshow_m_n_1.rst
@@ -0,0 +1,5 @@
+Imshow now coerces 3D arrays with depth 1 to 2D
+------------------------------------------------
+Starting from this version arrays of size MxNx1 will be coerced into MxN 
+for displaying. This means commands like ``plt.imshow(np.random.rand(3, 3, 1))`` 
+will no longer return an error message that the image shape is invalid.
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -684,6 +684,10 @@ def set_data(self, A):
             raise TypeError("Image data of dtype {} cannot be converted to "
                             "float".format(self._A.dtype))
 
+        if (self._A.ndim == 3 and self._A.shape[-1] == 1):
+            # If just one dimension assume scalar and apply colormap
+            self._A = self._A[:, :, 0]
+
         if not (self._A.ndim == 2
                 or self._A.ndim == 3 and self._A.shape[-1] in [3, 4]):
             raise TypeError("Invalid shape {} for image data"
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
@@ -467,6 +467,35 @@ def test_imshow():
     ax.set_ylim(0, 3)
 
 
+@check_figures_equal(extensions=['png'])
+def test_imshow_10_10_1(fig_test, fig_ref):
+    # 10x10x1 should be the same as 10x10
+    arr = np.arange(100).reshape((10, 10, 1))
+    ax = fig_ref.subplots()
+    ax.imshow(arr[:, :, 0], interpolation="bilinear", extent=(1, 2, 1, 2))
+    ax.set_xlim(0, 3)
+    ax.set_ylim(0, 3)
+
+    ax = fig_test.subplots()
+    ax.imshow(arr, interpolation="bilinear", extent=(1, 2, 1, 2))
+    ax.set_xlim(0, 3)
+    ax.set_ylim(0, 3)
+
+
+def test_imshow_10_10_2():
+    fig, ax = plt.subplots()
+    arr = np.arange(200).reshape((10, 10, 2))
+    with pytest.raises(TypeError):
+        ax.imshow(arr)
+
+
+def test_imshow_10_10_5():
+    fig, ax = plt.subplots()
+    arr = np.arange(500).reshape((10, 10, 5))
+    with pytest.raises(TypeError):
+        ax.imshow(arr)
+
+
 @image_comparison(['no_interpolation_origin'], remove_text=True)
 def test_no_interpolation_origin():
     fig, axs = plt.subplots(2)
