Repository: SciTools/cartopy

Issue:
Incompatible dimensions with pcolormesh when using 'nearest' shading
### Description
Pull request https://github.com/SciTools/cartopy/pull/1646 made large changes to how `pcolormesh` interpolates its arguments. The new code appears to recalculate the edges when `shading` is `nearest`, thus adding an element to both `x` and `y` dimensions. But in the process it does not update the `shading` argument accordingly, causing MPL to complain about incompatible dimensions.

#### Code to reproduce
A minimal example demonstrating the exception is
```python
import numpy as np
from matplotlib import pyplot as plt
import cartopy.crs as ccrs

ax = plt.axes(projection=ccrs.PlateCarree())

n=3
x = np.arange(n)+1
y = np.arange(n)+1
d = np.random.rand(n,n)

im = ax.pcolormesh(x, y, d, shading='nearest')
```

#### Traceback 
The traceback when running the code
```
Traceback (most recent call last):
  File "test.py", line 12, in <module>
    im = ax.pcolormesh(x, y, d, shading='nearest')
  File "/home/ec2-user/miniconda3/envs/arpav/lib/python3.8/site-packages/cartopy/mpl/geoaxes.py", line 318, in wrapper
    return func(self, *args, **kwargs)
  File "/home/ec2-user/miniconda3/envs/arpav/lib/python3.8/site-packages/cartopy/mpl/geoaxes.py", line 1797, in pcolormesh
    result = matplotlib.axes.Axes.pcolormesh(self, *args, **kwargs)
  File "/home/ec2-user/miniconda3/envs/arpav/lib/python3.8/site-packages/matplotlib/__init__.py", line 1361, in inner
    return func(ax, *map(sanitize_sequence, args), **kwargs)
  File "/home/ec2-user/miniconda3/envs/arpav/lib/python3.8/site-packages/matplotlib/axes/_axes.py", line 6183, in pcolormesh
    X, Y, C, shading = self._pcolorargs('pcolormesh', *args,
  File "/home/ec2-user/miniconda3/envs/arpav/lib/python3.8/site-packages/matplotlib/axes/_axes.py", line 5711, in _pcolorargs
    raise TypeError('Dimensions of C %s are incompatible with'
TypeError: Dimensions of C (3, 3) are incompatible with X (4) and/or Y (4); see help(pcolormesh)
```

### Environment
Conda environment on Amazon Linux with `cartopy` 0.20.0 from conda-forge and `matplotlib-base` 3.4.3.


Discussion:

Resolving patch:
diff --git a/lib/cartopy/mpl/geoaxes.py b/lib/cartopy/mpl/geoaxes.py
--- a/lib/cartopy/mpl/geoaxes.py
+++ b/lib/cartopy/mpl/geoaxes.py
@@ -1789,7 +1789,7 @@ def pcolormesh(self, *args, **kwargs):
         """
         # Add in an argument checker to handle Matplotlib's potential
         # interpolation when coordinate wraps are involved
-        args = self._wrap_args(*args, **kwargs)
+        args, kwargs = self._wrap_args(*args, **kwargs)
         result = super().pcolormesh(*args, **kwargs)
         # Wrap the quadrilaterals if necessary
         result = self._wrap_quadmesh(result, **kwargs)
@@ -1811,8 +1811,11 @@ def _wrap_args(self, *args, **kwargs):
         if not (kwargs.get('shading', default_shading) in
                 ('nearest', 'auto') and len(args) == 3 and
                 getattr(kwargs.get('transform'), '_wrappable', False)):
-            return args
+            return args, kwargs
 
+        # We have changed the shading from nearest/auto to flat
+        # due to the addition of an extra coordinate
+        kwargs['shading'] = 'flat'
         X = np.asanyarray(args[0])
         Y = np.asanyarray(args[1])
         nrows, ncols = np.asanyarray(args[2]).shape
@@ -1848,7 +1851,7 @@ def _interp_grid(X, wrap=0):
             X = _interp_grid(X.T, wrap=xwrap).T
             Y = _interp_grid(Y.T).T
 
-        return (X, Y, args[2])
+        return (X, Y, args[2]), kwargs
 
     def _wrap_quadmesh(self, collection, **kwargs):
         """
@@ -1864,8 +1867,13 @@ 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((Ny - 1, Nx - 1))
+        C = collection.get_array().reshape(data_shape)
 
         transformed_pts = self.projection.transform_points(
             t, coords[..., 0], coords[..., 1])
@@ -1894,6 +1902,23 @@ def _wrap_quadmesh(self, collection, **kwargs):
             # No wrapping needed
             return collection
 
+        # 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.
+        if kwargs.get('shading') == 'gouraud':
+            warnings.warn("Handling wrapped coordinates with gouraud "
+                          "shading is likely to introduce artifacts. "
+                          "It is recommended to remove the wrap manually "
+                          "before calling pcolormesh.")
+            # With gouraud shading, we actually want an (Ny, Nx) shaped mask
+            gmask = np.zeros(data_shape, dtype=bool)
+            # If any of the cells were wrapped, apply it to all 4 corners
+            gmask[:-1, :-1] |= mask
+            gmask[1:, :-1] |= mask
+            gmask[1:, 1:] |= mask
+            gmask[:-1, 1:] |= mask
+            mask = gmask
+
         # We have quadrilaterals that cross the wrap boundary
         # Now, we need to update the original collection with
         # a mask over those cells and use pcolor to draw those
@@ -1974,7 +1999,11 @@ def pcolor(self, *args, **kwargs):
         """
         # Add in an argument checker to handle Matplotlib's potential
         # interpolation when coordinate wraps are involved
-        args = self._wrap_args(*args, **kwargs)
+        args, kwargs = self._wrap_args(*args, **kwargs)
+        if matplotlib.__version__ < "3.3":
+            # MPL 3.3 introduced the shading option, and it isn't
+            # handled before that for pcolor calls.
+            kwargs.pop('shading', None)
         result = super().pcolor(*args, **kwargs)
 
         # Update the datalim for this pcolor.


Test patch:
diff --git a/lib/cartopy/tests/mpl/test_mpl_integration.py b/lib/cartopy/tests/mpl/test_mpl_integration.py
--- a/lib/cartopy/tests/mpl/test_mpl_integration.py
+++ b/lib/cartopy/tests/mpl/test_mpl_integration.py
@@ -612,7 +612,7 @@ def test_pcolormesh_diagonal_wrap():
     # and the bottom edge on the other gets wrapped properly
     xs = [[160, 170], [190, 200]]
     ys = [[-10, -10], [10, 10]]
-    zs = [[0, 1], [0, 1]]
+    zs = [[0]]
 
     ax = plt.axes(projection=ccrs.PlateCarree())
     mesh = ax.pcolormesh(xs, ys, zs)
@@ -684,6 +684,31 @@ def test_pcolormesh_wrap_set_array():
     return ax.figure
 
 
+@pytest.mark.parametrize('shading, input_size, expected', [
+    pytest.param('auto', 3, 4, id='auto same size'),
+    pytest.param('auto', 4, 4, id='auto input larger'),
+    pytest.param('nearest', 3, 4, id='nearest same size'),
+    pytest.param('nearest', 4, 4, id='nearest input larger'),
+    pytest.param('flat', 4, 4, id='flat input larger'),
+    pytest.param('gouraud', 3, 3, id='gouraud same size')
+])
+def test_pcolormesh_shading(shading, input_size, expected):
+    # Testing that the coordinates are all broadcast as expected with
+    # the various shading options
+    # The data shape is (3, 3) and we are changing the input shape
+    # based upon that
+    ax = plt.axes(projection=ccrs.PlateCarree())
+
+    x = np.arange(input_size)
+    y = np.arange(input_size)
+    d = np.zeros((3, 3))
+
+    coll = ax.pcolormesh(x, y, d, shading=shading)
+    # We can use coll.get_coordinates() once MPL >= 3.5 is required
+    # For now, we use the private variable for testing
+    assert coll._coordinates.shape == (expected, expected, 2)
+
+
 @pytest.mark.natural_earth
 @pytest.mark.mpl_image_compare(filename='quiver_plate_carree.png')
 def test_quiver_plate_carree():
diff --git a/lib/cartopy/tests/mpl/test_pseudo_color.py b/lib/cartopy/tests/mpl/test_pseudo_color.py
--- a/lib/cartopy/tests/mpl/test_pseudo_color.py
+++ b/lib/cartopy/tests/mpl/test_pseudo_color.py
@@ -14,7 +14,7 @@
 
 
 def test_pcolormesh_partially_masked():
-    data = np.ma.masked_all((40, 30))
+    data = np.ma.masked_all((39, 29))
     data[0:100] = 10
 
     # Check that a partially masked data array does trigger a pcolor call.
@@ -26,7 +26,7 @@ def test_pcolormesh_partially_masked():
 
 
 def test_pcolormesh_invisible():
-    data = np.zeros((3, 3))
+    data = np.zeros((2, 2))
 
     # Check that a fully invisible mesh doesn't fail.
     with mock.patch('cartopy.mpl.geoaxes.GeoAxes.pcolor') as pcolor: