Merge pull request #13986 from anntzer/pcolorfast-quadmesh-rgba

efiring · web-flow · commit c05537d310bf · 2019-04-23T07:51:35.000-10:00
Support RGBA for quadmesh mode of pcolorfast.
diff --git a/doc/api/api_changes/2019-02-11-PGE.rst b/doc/api/api_changes/2019-02-11-PGE.rst
@@ -1,6 +1,4 @@
 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
+pcolorfast now accepts 3D images (RGB or RGBA) arrays.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -6191,12 +6191,18 @@ def pcolorfast(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
         Parameters
         ----------
         C : array-like(M, N)
-            A 2D array or masked array. The values will be color-mapped.
-            This argument can only be passed positionally.
+            The image data. Supported array shapes are:
 
-            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.
+            - (M, N): an image with scalar data. The data is visualized
+              using a colormap.
+            - (M, N, 3): an image with RGB values (0-1 float or 0-255 int).
+            - (M, N, 4): an image with RGBA values (0-1 float or 0-255 int),
+              i.e. including transparency.
+
+            The first two dimensions (M, N) define the rows and columns of
+            the image.
+
+            This parameter can only be passed positionally.
 
         X, Y : tuple or array-like, default: ``(0, N)``, ``(0, M)``
             *X* and *Y* are used to specify the coordinates of the
@@ -6221,9 +6227,9 @@ def pcolorfast(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
             - Use 2D arrays *X*, *Y* if you need an *arbitrary quadrilateral
               grid* (i.e. if the quadrilaterals are not rectangular).
 
-              In this case *X* and *Y* are 2D arrays with shape (M, N),
+              In this case *X* and *Y* are 2D arrays with shape (M + 1, N + 1),
               specifying the x and y coordinates of the corners of the colored
-              quadrilaterals. See `~.Axes.pcolormesh` for details.
+              quadrilaterals.
 
               This is the most general, but the slowest to render.  It may
               produce faster and more compact output using ps, pdf, and
@@ -6292,10 +6298,6 @@ 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")
@@ -6305,9 +6307,15 @@ def pcolorfast(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
         if style == "quadmesh":
             # data point in each cell is value at lower left corner
             coords = np.stack([x, y], axis=-1)
+            if np.ndim(C) == 2:
+                qm_kwargs = {"array": np.ma.ravel(C)}
+            elif np.ndim(C) == 3:
+                qm_kwargs = {"color": np.ma.reshape(C, (-1, C.shape[-1]))}
+            else:
+                raise ValueError("C must be 2D or 3D")
             collection = mcoll.QuadMesh(
-                nc, nr, coords,
-                array=np.ma.ravel(C), alpha=alpha, cmap=cmap, norm=norm,
+                nc, nr, coords, **qm_kwargs,
+                alpha=alpha, cmap=cmap, norm=norm,
                 antialiased=False, edgecolors="none")
             self.add_collection(collection, autolim=False)
             xl, xr, yb, yt = x.min(), x.max(), y.min(), y.max()
@@ -6331,7 +6339,7 @@ def pcolorfast(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
 
         if vmin is not None or vmax is not None:
             ret.set_clim(vmin, vmax)
-        else:
+        elif np.ndim(C) == 2:  # C.ndim == 3 is RGB(A) so doesn't need scaling.
             ret.autoscale_None()
 
         ret.sticky_edges.x[:] = [xl, xr]
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -5151,27 +5151,14 @@ def test_no_None():
          mpl.collections.QuadMesh),
     ]
 )
-def test_pcolorfast_colormapped(xy, cls):
+@pytest.mark.parametrize(
+    "data", [np.arange(12).reshape((3, 4)), np.random.rand(3, 4, 3)]
+)
+def test_pcolorfast(xy, data, cls):
     fig, ax = plt.subplots()
-    data = np.arange(12).reshape((3, 4))
     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)
 
