Merge pull request #8690 from paalge/pcolorfast

efiring · web-flow · commit 94b0bca18446 · 2019-03-18T09:39:42.000-10:00
Adds support for rgba and rgb images to pcolorfast
diff --git a/doc/api/api_changes/2019-02-11-PGE.rst b/doc/api/api_changes/2019-02-11-PGE.rst
@@ -0,0 +1,6 @@
+Added support for RGB(A) images in pcolorfast
+`````````````````````````````````````````````
+
+pcolorfast now accepts 3D images (RGB or RGBA) arrays if the X and Y
+specifications allow image or pcolorimage rendering; they remain unsupported by
+the more general quadmesh rendering
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -6172,6 +6172,10 @@ def pcolorfast(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
             A 2D array or masked array. The values will be color-mapped.
             This argument can only be passed positionally.
 
+            C can in some cases be 3D with the last dimension as rgb(a).
+            This is available when C qualifies for image or pcolorimage type,
+            will throw a TypeError if C is 3D and quadmesh.
+
         X, Y : tuple or array-like, default: ``(0, N)``, ``(0, M)``
             *X* and *Y* are used to specify the coordinates of the
             quadrilaterals. There are different ways to do this:
@@ -6245,7 +6249,7 @@ def pcolorfast(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
                 "'norm' must be an instance of 'mcolors.Normalize'")
 
         C = args[-1]
-        nr, nc = np.shape(C)
+        nr, nc = np.shape(C)[:2]
         if len(args) == 1:
             style = "image"
             x = [0, nc]
@@ -6266,6 +6270,10 @@ def pcolorfast(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
                     else:
                         style = "pcolorimage"
             elif x.ndim == 2 and y.ndim == 2:
+                if C.ndim > 2:
+                    raise ValueError(
+                        'pcolorfast needs to use quadmesh, '
+                        'which is not supported when x and y are 2D and C 3D')
                 style = "quadmesh"
             else:
                 raise TypeError("arguments do not match valid signatures")
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -5190,6 +5190,21 @@ def test_pcolorfast_colormapped(xy, cls):
     assert type(ax.pcolorfast(*xy, data)) == cls
 
 
+def test_pcolor_fast_RGB():
+
+    fig, ax = plt.subplots(1, 1)
+
+    np.random.seed(19680801)
+    C = np.random.rand(10, 10, 3)  # RGB image [0,1]
+    x = np.arange(11, dtype=np.float)
+    y = np.arange(11, dtype=np.float)
+
+    xv, yv = np.meshgrid(x, y)
+
+    with pytest.raises(ValueError):
+        ax.pcolorfast(xv, yv, C)
+
+
 def test_shared_scale():
     fig, axs = plt.subplots(2, 2, sharex=True, sharey=True)
 
