Improve contour performance in WCSAxes by 10-1000x - #7568

Background

WCSAxes has a mechanism for overplotting contours from an image in a different coordinate system from the original image. This essentially boils down to:

ax.contour(contour_data, transform=ax.get_transform(contour_wcs))

However, the way Matplotlib deals with this internally is to compute the contours on contour_data, then to represent each line in the contour map as a Matplotlib Path (note - each line, not each level - there might be multiple lines per level), then to transform each path individually using the specified transform. In the case where coordinate conversions are needed, the transform includes calls to SkyCoord, which we know is a little slow, especially if called once for each line in the contour map. For some contour maps I've made, there are sometimes thousands of individual lines.

New approach

The approach implemented here is to overload contour and contourf, and internally call them without transform, but to then modify the resulting ContourSet manually. We do this by extracting all the vertices of the lines, stacking them into a single array, transforming them, and splitting them again. This is far more efficient because we then only effectively call SkyCoord once.

Performance

The following functions test the new approach, the old approach, and also test the performance without contours:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.axes import Axes

from astropy.io import fits
from astropy.wcs import WCS
from astropy.utils.data import get_pkg_data_filename

hdu_2mass = fits.open(get_pkg_data_filename('galactic_center/gc_2mass_h.fits'))[0]
hdu_msx = fits.open(get_pkg_data_filename('galactic_center/gc_msx_e.fits'))[0]


def new():
    ax = plt.subplot(1, 1, 1, projection=WCS(hdu_2mass.header))
    ax.imshow(hdu_2mass.data)
    ax.contour(hdu_msx.data, levels=np.logspace(-6, -3, 10),
               transform=ax.get_transform(WCS(hdu_msx.header)), colors='white')
    plt.savefig('new.png')


def old():
    ax = plt.subplot(1, 1, 1, projection=WCS(hdu_2mass.header))
    ax.imshow(hdu_2mass.data)
    Axes.contour(ax, hdu_msx.data, levels=np.logspace(-6, -3, 10),
                 transform=ax.get_transform(WCS(hdu_msx.header)), colors='white')
    plt.savefig('old.png')


def nocontour():
    ax = plt.subplot(1, 1, 1, projection=WCS(hdu_2mass.header))
    ax.imshow(hdu_2mass.data)
    plt.savefig('nocontour.png')

The results are as follows:

In [10]: %timeit old()
14.5 s ± 710 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [11]: %timeit new()
198 ms ± 8.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [13]: %timeit nocontour()
125 ms ± 1.39 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

If we subtract the base time for the plot excluding the contour, the old approach took 14+ seconds, whereas now it takes ~70ms, a speedup of around 200x.

The exact speedup will depend a lot on the number of lines in the contour map. The worst case scenario (for the new implementation) would be if there was just a single contour line. If I take the example above and set levels=[5.8e-4], I get a single contour line. In that case, the timings look like:

In [5]: %timeit old()
156 ms ± 3.65 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [4]: %timeit new()
145 ms ± 1.85 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [6]: %timeit nocontour()
153 ms ± 7.71 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Basically the differences are in the noise. So there should not really be any cases where this results in worse performance.

I've added a couple of benchmark to astropy-benchmarks: astropy/astropy-benchmarks#58