ENH: Adding a contour fast-path to speed up contouring

lib/cartopy/mpl/geoaxes.py

@@ -321,6 +321,50 @@ def wrapper(self, *args, **kwargs):
     return wrapper
 
 
+def _add_transform_first(func):
+    """
+    A decorator that adds and validates the transform_first keyword argument.
+
+    This handles a fast-path optimization that projects the points before
+    creating any patches or lines. This means that the lines/patches will be
+    calculated in projected-space, not data-space. It requires the first
+    three arguments to be x, y, and z and all must be two-dimensional to use
+    the fast-path option.
+
+    This should be added after the _add_transform wrapper so that a transform
+    is guaranteed to be present.
+    """
+    @functools.wraps(func)
+    def wrapper(self, *args, **kwargs):
+        if kwargs.pop('transform_first', False):
+            if len(args) < 3:
+                # For the fast-path we need X and Y input points
+                raise ValueError("The X and Y arguments must be provided to "
+                                 "use the transform_first=True fast-path.")
+            x, y, z = (np.array(i) for i in args[:3])
+            if not (x.ndim == y.ndim == 2):
+                raise ValueError("The X and Y arguments must be gridded "
+                                 "2-dimensional arrays")
+
+            # Remove the transform from the keyword arguments
+            t = kwargs.pop('transform')
+            # Transform all of the x and y points
+            pts = self.projection.transform_points(t, x, y)
+            x = pts[..., 0].reshape(x.shape)
+            y = pts[..., 1].reshape(y.shape)
+            # The x coordinates could be wrapped, but matplotlib expects
+            # them to be sorted, so we will reorganize the arrays based on x
+            ind = np.argsort(x, axis=1)
+            x = np.take_along_axis(x, ind, axis=1)
+            y = np.take_along_axis(y, ind, axis=1)
+            z = np.take_along_axis(z, ind, axis=1)
+
+            # Use the new points as the input arguments
+            args = (x, y, z) + args[3:]
+        return func(self, *args, **kwargs)
+    return wrapper
+
+
 class GeoAxes(matplotlib.axes.Axes):
     """
     A subclass of :class:`matplotlib.axes.Axes` which represents a
@@ -1593,6 +1637,7 @@ def set_boundary(self, path, transform=None, use_as_clip_path=None):
         self.spines['geo'].set_boundary(path, transform)
 
     @_add_transform
+    @_add_transform_first
     def contour(self, *args, **kwargs):
         """
         Add the "transform" keyword to :func:`~matplotlib.pyplot.contour`.
@@ -1602,6 +1647,16 @@ def contour(self, *args, **kwargs):
         transform
             A :class:`~cartopy.crs.Projection`.
 
+        transform_first : bool, optional
+            If True, this will transform the input arguments into
+            projection-space before computing the contours, which is much
+            faster than computing the contours in data-space and projecting
+            the filled polygons. Using this method does not handle wrapped
+            coordinates as well and can produce misleading contours in the
+            middle of the domain. To use the projection-space method the input
+            arguments X and Y must be provided and be 2-dimensional.
+            The default is False, to compute the contours in data-space.
+
         """
         result = matplotlib.axes.Axes.contour(self, *args, **kwargs)
 
@@ -1621,6 +1676,7 @@ def contour(self, *args, **kwargs):
         return result
 
     @_add_transform
+    @_add_transform_first
     def contourf(self, *args, **kwargs):
         """
         Add the "transform" keyword to :func:`~matplotlib.pyplot.contourf`.
@@ -1630,8 +1686,17 @@ def contourf(self, *args, **kwargs):
         transform
             A :class:`~cartopy.crs.Projection`.
 
-        """
-        t = kwargs['transform']
+        transform_first : bool, optional
+            If True, this will transform the input arguments into
+            projection-space before computing the contours, which is much
+            faster than computing the contours in data-space and projecting
+            the filled polygons. Using this method does not handle wrapped
+            coordinates as well and can produce misleading contours in the
+            middle of the domain. To use the projection-space method the input
+            arguments X and Y must be provided and be 2-dimensional.
+            The default is False, to compute the contours in data-space.
+        """
+        t = kwargs.get('transform')
         if isinstance(t, ccrs.Projection):
             kwargs['transform'] = t = t._as_mpl_transform(self)
         # Set flag to indicate correcting orientation of paths if not ccw

lib/cartopy/tests/mpl/test_contour.py

@@ -8,6 +8,7 @@
 from matplotlib.testing.decorators import cleanup
 import numpy as np
 from numpy.testing import assert_array_almost_equal
+import pytest
 from scipy.interpolate import NearestNDInterpolator
 from scipy.signal import convolve2d
 
@@ -94,3 +95,38 @@ def test_contour_update_bounds():
     # doesn't raise with an Orthographic projection
     # GH issue 1673
     plt.draw()
+
+
+@pytest.mark.parametrize('func', ['contour', 'contourf'])
+@cleanup
+def test_contourf_transform_first(func):
+    """Test the fast-path option for filled contours."""
+    # Gridded data that needs to be wrapped
+    x = np.arange(360)
+    y = np.arange(-25, 26)
+    xx, yy = np.meshgrid(x, y)
+    z = xx + yy**2
+
+    ax = plt.axes(projection=ccrs.PlateCarree())
+    test_func = getattr(ax, func)
+    # Can't handle just Z input with the transform_first
+    with pytest.raises(ValueError,
+                       match="The X and Y arguments must be provided"):
+        test_func(z, transform=ccrs.PlateCarree(),
+                  transform_first=True)
+    # X and Y must also be 2-dimensional
+    with pytest.raises(ValueError,
+                       match="The X and Y arguments must be gridded"):
+        test_func(x, y, z, transform=ccrs.PlateCarree(),
+                  transform_first=True)
+
+    # When calculating the contour in projection-space the extent
+    # will now be the extent of the transformed points (-179, 180, -25, 25)
+    test_func(xx, yy, z, transform=ccrs.PlateCarree(),
+              transform_first=True)
+    assert_array_almost_equal(ax.get_extent(), (-179, 180, -25, 25))
+
+    # The extent without the transform_first should be all the way out to -180
+    test_func(xx, yy, z, transform=ccrs.PlateCarree(),
+              transform_first=False)
+    assert_array_almost_equal(ax.get_extent(), (-180, 180, -25, 25))