Merge pull request #1892 from greglucas/pcolor-shading-nearest

FIX: pcolor shading with nearest

lib/cartopy/mpl/geoaxes.py

@@ -1793,7 +1793,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 = matplotlib.axes.Axes.pcolormesh(self, *args, **kwargs)
         # Wrap the quadrilaterals if necessary
         result = self._wrap_quadmesh(result, **kwargs)
@@ -1815,8 +1815,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
@@ -1852,7 +1855,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):
         """
@@ -1868,8 +1871,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])
@@ -1898,6 +1906,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
@@ -1978,7 +2003,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 = matplotlib.axes.Axes.pcolor(self, *args, **kwargs)
 
         # Update the datalim for this pcolor.

lib/cartopy/tests/mpl/test_mpl_integration.py

@@ -583,7 +583,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)
@@ -652,6 +652,31 @@ def test_pcolormesh_wrap_set_array():
     coll.set_array(Z.ravel())
 
 
+@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
 @ImageTesting(['quiver_plate_carree'])
 def test_quiver_plate_carree():

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.
@@ -27,7 +27,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: