Merge pull request #1958 from greglucas/pcolormesh-bbox

FIX: Calculate bbox for pcolormesh

lib/cartopy/mpl/geoaxes.py

@@ -1794,10 +1794,6 @@ def _wrap_quadmesh(self, collection, **kwargs):
         cross the boundary of the projection.
         """
         t = kwargs.get('transform', None)
-        if not (getattr(t, '_wrappable', False) and
-                getattr(self.projection, '_wrappable', False)):
-            # Nothing to do
-            return collection
 
         # Get the quadmesh data coordinates
         coords = collection._coordinates
@@ -1833,8 +1829,25 @@ def _wrap_quadmesh(self, collection, **kwargs):
                     np.isnan(diagonal1_lengths) |
                     (diagonal1_lengths > size_limit))
 
-        if not np.any(mask):
-            # No wrapping needed
+        # Update the data limits based on the corners of the mesh
+        # in transformed coordinates, ignoring nan values
+        with warnings.catch_warnings():
+            warnings.filterwarnings('ignore', 'All-NaN slice encountered')
+            # If we have all nans, that is OK and will be handled by the
+            # Bbox calculations later, so suppress that warning from the user
+            corners = ((np.nanmin(xs), np.nanmin(ys)),
+                       (np.nanmax(xs), np.nanmax(ys)))
+        collection._corners = mtransforms.Bbox(corners)
+        self.update_datalim(collection._corners)
+
+        # We need to keep the transform/projection check after
+        # update_datalim to make sure we are getting the proper
+        # datalims on the returned collection
+        if (not (getattr(t, '_wrappable', False) and
+                 getattr(self.projection, '_wrappable', False)) or
+                not np.any(mask)):
+            # If both projections are unwrappable
+            # or if there aren't any points to wrap
             return collection
 
         # Wrapping with gouraud shading is error-prone. We will do our best,

lib/cartopy/mpl/geocollection.py

@@ -52,3 +52,8 @@ def set_clim(self, vmin=None, vmax=None):
 
         # Update color limits for the rest of the cells.
         super().set_clim(vmin, vmax)
+
+    def get_datalim(self, transData):
+        # Return the corners that were calculated in
+        # the pcolormesh routine.
+        return self._corners

lib/cartopy/tests/mpl/test_pseudo_color.py

@@ -81,3 +81,46 @@ def test_pcolormesh_arg_interpolation():
                           [2, 20],
                           [4, 20]]])
     np.testing.assert_array_almost_equal(expected, coll._coordinates)
+
+
+def test_pcolormesh_datalim():
+    # Test that wrapping the coordinates still produces proper data limits
+    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(2, 1, 1, projection=ccrs.PlateCarree())
+    coll = ax.pcolormesh(xs, ys, z, shading='auto',
+                         transform=ccrs.PlateCarree())
+
+    coll_bbox = coll.get_datalim(ax.transData)
+    np.testing.assert_array_equal(coll_bbox, [[-2, -20], [4, 20]])
+
+    # Non-wrapped coordinates
+    x = [-80, 0, 80]
+    y = [-10, 10]
+
+    xs, ys = np.meshgrid(x, y)
+    ax = plt.subplot(2, 1, 1, projection=ccrs.PlateCarree())
+    coll = ax.pcolormesh(xs, ys, z, shading='auto',
+                         transform=ccrs.PlateCarree())
+
+    coll_bbox = coll.get_datalim(ax.transData)
+    np.testing.assert_array_equal(coll_bbox, [[-120, -20], [120, 20]])
+
+    # A projection that doesn't support wrapping
+    x = [-10, 0, 10]
+    y = [-10, 10]
+
+    xs, ys = np.meshgrid(x, y)
+    ax = plt.subplot(2, 1, 1, projection=ccrs.Orthographic())
+    coll = ax.pcolormesh(xs, ys, z, shading='auto',
+                         transform=ccrs.PlateCarree())
+
+    coll_bbox = coll.get_datalim(ax.transData)
+    expected = [[-1650783.327873, -2181451.330891],
+                [1650783.327873, 2181451.330891]]
+    np.testing.assert_array_almost_equal(coll_bbox, expected)