FIX: Update pcolor(mesh) to create grid edges

MPL 3.3 added a new interpolation routine to automatically create
edges for the grids when the input X/Y are the same shape as the
color data. This causes issues with map data due to the possiblity
of jumps in the coordinate arrays. This adds a modulo wrap to
the input arguments to do the interpolation in Cartopy before
sending the arguments on to Matplotlib.

Author: greglucas

lib/cartopy/mpl/geoaxes.py

@@ -70,6 +70,16 @@
 CARTOPY_USER_BACKGROUNDS environment variable.
 """
 
+# A list of projections that can have wrapped coordinates.
+_WRAP_PROJECTIONS = (ccrs._RectangularProjection,
+                     ccrs._WarpedRectangularProjection,
+                     ccrs.InterruptedGoodeHomolosine,
+                     ccrs.Mercator,
+                     ccrs.LambertAzimuthalEqualArea,
+                     ccrs.AzimuthalEquidistant,
+                     ccrs.TransverseMercator,
+                     ccrs.Stereographic)
+
 
 # XXX call this InterCRSTransform
 class InterProjectionTransform(mtransforms.Transform):
@@ -1649,6 +1659,9 @@ def pcolormesh(self, *args, **kwargs):
             A :class:`~cartopy.crs.Projection`.
 
         """
+        # Add in an argument checker to handle Matplotlib's potential
+        # interpolation when coordinate wraps are involved
+        args = self._wrap_args(*args, **kwargs)
         result = matplotlib.axes.Axes.pcolormesh(self, *args, **kwargs)
         # Wrap the quadrilaterals if necessary
         result = self._wrap_quadmesh(result, **kwargs)
@@ -1659,22 +1672,62 @@ def pcolormesh(self, *args, **kwargs):
         self.autoscale_view()
         return result
 
+    def _wrap_args(self, *args, **kwargs):
+        """
+        Handle the interpolation when a wrap could be involved with
+        the data coordinates before passing on to Matplotlib.
+        """
+        if (kwargs.get('shading', 'auto') in ('nearest', 'auto') and
+                len(args) == 3 and
+                isinstance(kwargs.get('transform'), _WRAP_PROJECTIONS)):
+            X = np.asanyarray(args[0])
+            Y = np.asanyarray(args[1])
+            nrows, ncols = np.asanyarray(args[2]).shape
+            Nx = X.shape[-1]
+            Ny = Y.shape[0]
+            if X.ndim != 2 or X.shape[0] == 1:
+                x = X.reshape(1, Nx)
+                X = x.repeat(Ny, axis=0)
+            if Y.ndim != 2 or Y.shape[1] == 1:
+                y = Y.reshape(Ny, 1)
+                Y = y.repeat(Nx, axis=1)
+
+            def _interp_grid(X, wrap=0):
+                # helper for below
+                if np.shape(X)[1] > 1:
+                    dX = np.diff(X, axis=1)
+                    # account for the wrap
+                    if wrap:
+                        dX = (dX + wrap/2) % wrap - wrap/2
+                    dX = dX/2
+                    X = np.hstack((X[:, [0]] - dX[:, [0]],
+                                   X[:, :-1] + dX,
+                                   X[:, [-1]] + dX[:, [-1]]))
+                else:
+                    # This is just degenerate, but we can't reliably guess
+                    # a dX if there is just one value.
+                    X = np.hstack((X, X))
+                return X
+            t = kwargs.get('transform')
+            xwrap = abs(t.x_limits[1] - t.x_limits[0])
+            if ncols == Nx:
+                X = _interp_grid(X, wrap=xwrap)
+                Y = _interp_grid(Y)
+            if nrows == Ny:
+                X = _interp_grid(X.T, wrap=xwrap).T
+                Y = _interp_grid(Y.T).T
+
+            args = (X, Y, args[2])
+        return args
+
     def _wrap_quadmesh(self, collection, **kwargs):
         """
         Handles the Quadmesh collection when any of the quadrilaterals
         cross the boundary of the projection.
         """
         t = kwargs.get('transform', None)
-        wrap_proj_types = (ccrs._RectangularProjection,
-                           ccrs._WarpedRectangularProjection,
-                           ccrs.InterruptedGoodeHomolosine,
-                           ccrs.Mercator,
-                           ccrs.LambertAzimuthalEqualArea,
-                           ccrs.AzimuthalEquidistant,
-                           ccrs.TransverseMercator,
-                           ccrs.Stereographic)
-        if not (isinstance(t, wrap_proj_types) and
-                isinstance(self.projection, wrap_proj_types)):
+        if not (isinstance(t, _WRAP_PROJECTIONS) and
+                isinstance(self.projection, _WRAP_PROJECTIONS)):
             # Nothing to do
             return collection
 
@@ -1785,6 +1838,9 @@ def pcolor(self, *args, **kwargs):
             A :class:`~cartopy.crs.Projection`.
 
         """
+        # Add in an argument checker to handle Matplotlib's potential
+        # interpolation when coordinate wraps are involved
+        args = self._wrap_args(*args, **kwargs)
         result = matplotlib.axes.Axes.pcolor(self, *args, **kwargs)
 
         # Update the datalim for this pcolor.

lib/cartopy/tests/mpl/test_pseudo_color.py

@@ -54,3 +54,34 @@ def test_savefig_tight():
     buf = io.BytesIO()
     plt.savefig(buf, format='png', bbox_inches='tight')
     plt.close()
+
+
+def test_pcolormesh_arg_interpolation():
+    # Test that making the edges from center point inputs
+    # properly accounts for the wrapped coordinates appearing
+    # in the middle of an array
+    x = [359, 1, 3]
+    y = [-10, 10]
+
+    xs, ys = np.meshgrid(x, y)
+    # Z with the same shape as X/Y to force the interpolation
+    z = np.zeros(xs.shape)
+
+    ax = plt.subplot(211, projection=ccrs.PlateCarree())
+    coll = ax.pcolormesh(xs, ys, z, shading='auto',
+                         transform=ccrs.PlateCarree())
+
+    # Compare the output coordinates of the generated mesh
+    expected = np.array([[[358, -20],
+                          [360, -20],
+                          [2, -20],
+                          [4, -20]],
+                         [[358, 0],
+                          [360, 0],
+                          [2, 0],
+                          [4, 0]],
+                         [[358, 20],
+                          [360, 20],
+                          [2, 20],
+                          [4, 20]]])
+    np.testing.assert_array_almost_equal(expected, coll._coordinates)