Add support for RGB and RGBA colors to pcolormesh

lib/cartopy/mpl/geoaxes.py

@@ -1795,7 +1795,7 @@ def _wrap_args(self, *args, **kwargs):
         kwargs['shading'] = 'flat'
         X = np.asanyarray(args[0])
         Y = np.asanyarray(args[1])
-        nrows, ncols = np.asanyarray(args[2]).shape
+        nrows, ncols, *_ = np.asanyarray(args[2]).shape
         Nx = X.shape[-1]
         Ny = Y.shape[0]
         if X.ndim != 2 or X.shape[0] == 1:
@@ -1839,15 +1839,6 @@ def _wrap_quadmesh(self, collection, **kwargs):
 
         # Get the quadmesh data coordinates
         coords = collection._coordinates
-        Ny, Nx, _ = coords.shape
-        if kwargs.get('shading') == 'gouraud':
-            # Gouraud shading has the same shape for coords and data
-            data_shape = Ny, Nx
-        else:
-            data_shape = Ny - 1, Nx - 1
-        # data array
-        C = collection.get_array().reshape(data_shape)
-
         transformed_pts = self.projection.transform_points(
             t, coords[..., 0], coords[..., 1])
 
@@ -1892,6 +1883,22 @@ def _wrap_quadmesh(self, collection, **kwargs):
             # or if there aren't any points to wrap
             return collection
 
+        Ny, Nx, _ = coords.shape
+        if kwargs.get('shading') == 'gouraud':
+            # Gouraud shading has the same shape for coords and data
+            data_shape = Ny, Nx
+        else:
+            data_shape = Ny - 1, Nx - 1
+        # Data array.  The -1 is so that we can give a helpful
+        # error message instead of a mysterious exception.
+        C = collection.get_array().reshape((*data_shape, -1))
+        if C.shape[-1] != 1:
+            # Later we call pcolor, which does not support RGB or RGBA data.
+            raise NotImplementedError("pcolormesh does not support wrapped"
+                                      " coordinates when RGB or RGBA data are"
+                                      " used.")
+        C = C.squeeze(axis=-1)
+
         # Wrapping with gouraud shading is error-prone. We will do our best,
         # but pcolor does not handle gouraud shading, so there needs to be
         # another way to handle the wrapped cells.

lib/cartopy/tests/mpl/test_mpl_integration.py

@@ -661,6 +661,73 @@ def test_pcolormesh_shading(shading, input_size, expected):
     assert coll._coordinates.shape == (expected, expected, 2)
 
 
+def make_pcolormesh_test_data(num_samples, x_bounds, y_bounds):
+    xbnds = np.linspace(*x_bounds, num_samples[0], endpoint=True)
+    ybnds = np.linspace(*y_bounds, num_samples[1], endpoint=True)
+
+    x, y = np.meshgrid(xbnds, ybnds)
+    data_r = 0.5 + 0.25 * (np.sin(np.deg2rad(x)) + np.cos(np.deg2rad(y)))
+    data_g = 0.5 + 0.25 * (np.sin(np.deg2rad(x + 120)) + np.cos(np.deg2rad(y)))
+    data_b = 0.5 + 0.25 * (np.sin(np.deg2rad(x + 240)) + np.cos(np.deg2rad(y)))
+    data = np.dstack((data_r, data_g, data_b))
+    return xbnds, ybnds, data
+
+
+@pytest.mark.mpl_image_compare(filename='pcolormesh_RGB_gouraud.png')
+def test_pcolormesh_RGB_gouraud():
+    num_samples = (36, 18)
+    x_bounds = (-180, 180)
+    y_bounds = (-90, 90)
+    xbnds, ybnds, data = make_pcolormesh_test_data(num_samples,
+                                                   x_bounds, y_bounds)
+
+    fig = plt.figure()
+    ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree())
+    ax.pcolormesh(xbnds, ybnds, data, transform=ccrs.PlateCarree(),
+                  snap=False, shading='gouraud')
+    ax.coastlines()
+    ax.set_global()  # make sure everything is visible
+
+    return fig
+
+
+@pytest.mark.mpl_image_compare(filename='pcolormesh_RGB_nearest.png')
+def test_pcolormesh_RGB_nearest():
+    num_samples = (36, 18)
+    x_bounds = (-175, 175)
+    y_bounds = (-90, 90)
+    xbnds, ybnds, data = make_pcolormesh_test_data(num_samples,
+                                                   x_bounds, y_bounds)
+
+    fig = plt.figure()
+    ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree())
+    ax.pcolormesh(xbnds, ybnds, data, transform=ccrs.PlateCarree(),
+                  snap=False, shading='nearest')
+    ax.coastlines()
+    ax.set_global()  # make sure everything is visible
+
+    return fig
+
+
+@pytest.mark.mpl_image_compare(filename='pcolormesh_RGB_flat.png')
+def test_pcolormesh_RGB_flat():
+    num_samples = (36, 18)
+    x_bounds = (-180, 180)
+    y_bounds = (-90, 90)
+    xbnds, ybnds, data = make_pcolormesh_test_data(num_samples,
+                                                   x_bounds, y_bounds)
+    data = data[:-1, :-1]
+
+    fig = plt.figure()
+    ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree())
+    ax.pcolormesh(xbnds, ybnds, data, transform=ccrs.PlateCarree(),
+                  snap=False, shading='flat')
+    ax.coastlines()
+    ax.set_global()  # make sure everything is visible
+
+    return fig
+
+
 @pytest.mark.natural_earth
 @pytest.mark.mpl_image_compare(filename='quiver_plate_carree.png')
 def test_quiver_plate_carree():