diff --git a/lib/cartopy/mpl/geoaxes.py b/lib/cartopy/mpl/geoaxes.py index 8c47583bb..d053d6834 100644 --- a/lib/cartopy/mpl/geoaxes.py +++ b/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 diff --git a/lib/cartopy/tests/mpl/test_contour.py b/lib/cartopy/tests/mpl/test_contour.py index 07d03dabb..82ac513e0 100644 --- a/lib/cartopy/tests/mpl/test_contour.py +++ b/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))