Slow looping when plotting the aperture over the TPF of saturated stars

[astrobatty]

Problem description

Aperture mask plotting over large TPFs is very slow, because lightkurve loops over each masked pixel and use matplotlib Rectangle to highlight pixels. This happens here.

Example

import lightkurve as lk

tpf = lk.search_targetpixelfile('Beta Gem', sector=20).download()
tpf.plot(aperture_mask='pipeline')

Expected behavior

According to this Stack Overflow issue, lightkurve should use PolyCollection instead of Rectangle.

Environment

• platform independent
• lightkurve version (e.g. 1.0b6): 2.4.2
• installation method (e.g. pip, conda, source): any

[orionlee]

An estimate of the actual impact using a bright star (TIC 233798897, Tmag 4.2) sector 76 SPOC TPF.

• 11x33 pixels in the TPF, with 78 pixels in SPOC aperture

On a recent laptop:

• tpf.plot(aperture_mask=None) - CPU times: total: 62.5 ms, Wall time: 89.5 ms
• tpf.plot(aperture_mask="pipeline") - CPU times: total: 391 ms, Wall time: 720 ms

---

---

The wall time of the typical 11x11 SPOC TPFs - no aperture mask: ~50ms, with aperture mask: ~150ms

[Nschanche]

Thanks @astrobatty for pointing this out. The loop is really slowing things down in this case. Changing the function to

        if aperture_mask is not None:
            aperture_mask = self._parse_aperture_mask(aperture_mask)
            in_aperture = np.where(aperture_mask)
            if hasattr(ax, "wcs"):
                ap_row = in_aperture[0] - 0.5
                ap_col = in_aperture[1] - 0.5
            else:
                ap_row = in_aperture[0] + self.row - 0.5
                ap_col = in_aperture[1] + self.column - 0.5    
            for ii in range(len(ap_row)):
                
                rect=patches.Rectangle((ap_col[ii],ap_row[ii]),1,1, facecolor=none, hatch="//", edgecolor=mask_color)
                ax.add_patch(rect)

speeds things up a lot (1min 24s ± 1.96 s per loop to 1.33 s ± 15.1 ms per loop for the original example). I'm testing this now, and will put in a PR if it looks good.