DM-53170: Speed up the diffraction spike mask in crowded fields

python/lsst/pipe/tasks/diffractionSpikeMask.py

@@ -34,6 +34,7 @@
 from lsst.pex.config import Config, ConfigField, Field
 from lsst.pipe.base import Struct, Task
 from lsst.meas.algorithms import getRefFluxField, LoadReferenceObjectsConfig
+from lsst.utils.timer import timeMethod
 import lsst.sphgeom
 from .colorterms import ColortermLibrary
 
@@ -162,6 +163,7 @@ def setRefObjLoader(self, refObjLoader):
         """
         self.refObjLoader = refObjLoader
 
+    @timeMethod
     def run(self, exposure):
         """Load reference objects and mask bright stars on an exposure.
 
@@ -200,7 +202,7 @@ def run(self, exposure):
         if nBright > 0:
             xvals, yvals = exposure.wcs.skyToPixelArray(refCat[bright][self.config.raKey],
                                                         refCat[bright][self.config.decKey])
-            spikeCandidates = self.selectSources(xvals, yvals, mask)
+            spikeCandidates = self.selectSources(xvals, yvals, radii, mask)
             nSpike = len(spikeCandidates)
         else:
             nSpike = 0
@@ -217,15 +219,18 @@ def run(self, exposure):
 
         return refCat[bright][spikeCandidates].copy(deep=True)
 
-    def selectSources(self, xvals, yvals, mask):
+    def selectSources(self, xvals, yvals, spikeRadii, mask):
         """Select saturated sources, and bright sources that are off the image.
 
         Parameters
         ----------
         xvals, yvals : `numpy.ndarray`
             Array of x- and y-values of bright sources to mask.
         mask : `lsst.afw.image.Mask`
-            The mask plane of the image to set the BRIGHT mask plane.
+            The mask plane of the image to set the SPIKE mask plane.
+        spikeRadii : `numpy.ndarray`
+            Predicted lengths in pixels of the diffraction spikes for each
+            bright source.
 
         Returns
         -------
@@ -236,30 +241,34 @@ def selectSources(self, xvals, yvals, mask):
         candidates = np.zeros(len(xvals), dtype=bool)
         saturatedBitMask = mask.getPlaneBitMask(self.config.saturatedMaskPlane)
         bbox = mask.getBBox()
+        projection = np.max([abs(np.cos(np.deg2rad(self.angles))),
+                             abs(np.sin(np.deg2rad(self.angles)))])
 
-        for i, (x, y) in enumerate(zip(xvals, yvals)):
+        for i, (x, y, r) in enumerate(zip(xvals, yvals, spikeRadii)):
             srcBox = lsst.geom.Box2I.makeCenteredBox(lsst.geom.Point2D(x, y), lsst.geom.Extent2I(3, 3))
             if bbox.contains(srcBox):
                 candidates[i] = np.all((mask[srcBox].array & saturatedBitMask) > 0)
             else:
-                candidates[i] = True
+                # If the candidate bright source is off the image, only make a
+                # model for it if the diffraction spikes are long enough.
+                candidates[i] = boxSeparation(bbox, x, y) < r*projection
         return candidates
 
-    def maskSources(self, xvals, yvals, radii, mask):
+    def maskSources(self, xvals, yvals, spikeRadii, mask):
         """Apply the SPIKE mask for a given set of coordinates. The mask plane
         will be modified in place.
 
         Parameters
         ----------
         xvals, yvals : `numpy.ndarray`
             Array of x- and y-values of bright sources to mask.
-        radii : `numpy.ndarray`
+        spikeRadii : `numpy.ndarray`
             Array of radius values for each bright source.
         mask : `lsst.afw.image.Mask`
-            The mask plane of the image to set the BRIGHT mask plane.
+            The mask plane of the image to set the SPIKE mask plane.
         """
         bbox = mask.getBBox()
-        for x, y, r in zip(xvals, yvals, radii):
+        for x, y, r in zip(xvals, yvals, spikeRadii):
             maskSingle = self.makeSingleMask(x, y, r)
             singleBBox = maskSingle.getBBox()
             if bbox.overlaps(singleBBox):
@@ -290,10 +299,10 @@ def makeSingleMask(self, x, y, r):
             # side along the spike and a short base. For efficiency, model the
             # diffraction spike in the opposite direction at the same time,
             # so the overall shape is a narrow diamond with equal length sides.
-            xLong = math.cos(np.deg2rad(angle))*r
-            yLong = math.sin(np.deg2rad(angle))*r
-            xShort = -math.sin(np.deg2rad(angle))*r/self.config.spikeAspectRatio
-            yShort = math.cos(np.deg2rad(angle))*r/self.config.spikeAspectRatio
+            xLong = np.cos(np.deg2rad(angle))*r
+            yLong = np.sin(np.deg2rad(angle))*r
+            xShort = -np.sin(np.deg2rad(angle))*r/self.config.spikeAspectRatio
+            yShort = np.cos(np.deg2rad(angle))*r/self.config.spikeAspectRatio
 
             corners = [lsst.geom.Point2D(x + xLong, y + yLong),
                        lsst.geom.Point2D(x + xShort, y + yShort),
@@ -428,6 +437,9 @@ def getRegion(exposure, margin=None):
     ----------
     exposure : `lsst.afw.image.Exposure`
         Exposure object with calibrated WCS.
+    margin : `lsst.sphgeom.Angle`, optional
+        Grow the surrounding region of the exposure by this amount, in order to
+        mask the diffraction spikes of stars off the image.
 
     Returns
     -------
@@ -453,28 +465,26 @@ def getRegion(exposure, margin=None):
     return region
 
 
-def computePsfWidthFromMoments(psf, angle=0.):
-    """Calculate the width of an elliptical PSF along a given direction.
+def boxSeparation(bbox, x, y):
+    """Return the minimum horizontal or vertical distance from a point to the
+    outside edge of a bounding box.
 
     Parameters
     ----------
-    psf : `lsst.afw.detection.Psf`
-        The point spread function of the image.
-    angle : `float`, optional
-        Rotation CCW from the +x axis to calculate the width of the PSF along.
+    bbox : `lsst.geom.Box2I`
+        The bounding box to check.
+    x, y : `float`
+        Coordinates of the point.
 
     Returns
     -------
-    fwhm : `float`
-        Full width at half maximum of the fitted shape of the PSF along the
-        given `angle`. In pixels.
+    distance : `float`
+        The distance in pixels by which the point is outside the box, or 0 if
+        it is inside.
     """
-    psfShape = psf.computeShape(psf.getAveragePosition())
-    c = np.cos(np.deg2rad(angle))
-    s = np.sin(np.deg2rad(angle))
-    sigma2 = c*c*psfShape.getIxx() + 2*c*s*psfShape.getIxy() + s*s*psfShape.getIyy()
-    sigma = np.sqrt(max(sigma2, 0.0))   # rms width in pixels
-
-    # 5) optional: Gaussian-equivalent FWHM along angle
-    fwhm = 2.0 * np.sqrt(2.0 * np.log(2.0)) * sigma
-    return fwhm
+
+    x0, y0 = bbox.getBegin()
+    x1, y1 = bbox.getEnd()
+    dx = 0 if x0 <= x <= x1 else min(abs(x0 - x), abs(x - x1))
+    dy = 0 if y0 <= y <= y1 else min(abs(y0 - y), abs(y - y1))
+    return max(dx, dy)

tests/test_diffractionSpikeMask.py

@@ -77,6 +77,15 @@ def tearDown(self):
         del self.exposure
         del self.refObjLoader
 
+    def test_raiseWithoutLoader(self):
+        """The task should raise an error if no reference catalog loader is
+        configured.
+        """
+        config = DiffractionSpikeMaskConfig()
+        task = DiffractionSpikeMaskTask(config=config)
+        with self.assertRaises(RuntimeError):
+            task.run(self.exposure)
+
     def test_loadAndMaskStars(self):
         """Run the bright star mask with a selection of reference sources."""
 
@@ -121,7 +130,7 @@ def test_loadAndMaskStars(self):
         self.assertGreater(outside, 0)
 
     def test_noBrightStars(self):
-        """Run the bright star mask with no bright stars"""
+        """Run the bright star mask with no bright stars."""
 
         # Set a very high magnitude limit so that no stars are selected
         config = DiffractionSpikeMaskConfig(magnitudeThreshold=0)
@@ -133,3 +142,73 @@ def test_noBrightStars(self):
         exposure.mask.getPlaneBitMask(config.spikeMask)
         # The images should not be modified
         self.assertImagesEqual(self.exposure.image, exposure.image)
+
+    def test_maskSources(self):
+        """Verify that sources on and off the image are masked correctly."""
+        task = DiffractionSpikeMaskTask(self.refObjLoader, config=DiffractionSpikeMaskConfig())
+        task.set_diffraction_angle(self.exposure)
+        self.exposure.mask.addMaskPlane(task.config.spikeMask)
+
+        xSize, ySize = self.bbox.getDimensions()
+        x0, y0 = self.bbox.getBegin()
+        x1, y1 = self.bbox.getEnd()
+
+        nBright = 50
+        rng = np.random.RandomState(3)
+        xLoc = np.arange(x0 - xSize/4, x1 + xSize/4)
+        rng.shuffle(xLoc)
+        xLoc = xLoc[:nBright]
+        yLoc = np.arange(y0 - ySize/4, y1 + ySize/4)
+        rng.shuffle(yLoc)
+        yLoc = yLoc[:nBright]
+        spikeRadii = np.arange(10, 200)
+        rng.shuffle(spikeRadii)
+        spikeRadii = spikeRadii[:nBright]
+        saturatedBox = geom.Box2I(self.bbox.getBegin(), geom.Extent2I(xSize, ySize//2))
+        baseMask = self.exposure.mask.clone()
+        baseMask[saturatedBox].array |= baseMask.getPlaneBitMask(task.config.saturatedMaskPlane)
+        # There are four classes of sources:
+        # 1. Bright sources on the image with saturated cores - masked
+        # 2. Sources on the image without saturated cores - not masked
+        # 3. Bright sources off the image with predicted diffraction spikes that
+        # overlap the image - masked
+        # 4. Bright sources off the image that are far away enough that any
+        # diffraction spikes would not overlap the image - not masked
+        nClass1 = 0
+        nClass2 = 0
+        nClass3 = 0
+        nClass4 = 0
+        selectedSources = task.selectSources(xLoc, yLoc, spikeRadii, baseMask)
+        for x, y, r, selected in zip(xLoc, yLoc, spikeRadii, selectedSources):
+            mask = baseMask.clone()
+            isInImage = self.bbox.contains(geom.Point2I(x, y))
+            # Bright sources on the image with saturated cores.
+            if isInImage and selected:
+                nClass1 += 1
+                task.maskSources([x], [y], [r], mask)
+                self.assertGreater(np.sum(mask.array & mask.getPlaneBitMask(task.config.spikeMask) > 0), 0)
+
+            # Bright sources on the image without saturated cores.
+            if isInImage and not selected:
+                nClass2 += 1
+                # Do *not* run task.maskSources in this case, since these are
+                # skipped.
+
+            # Bright sources off the image that we predict should overlap the
+            # image, should set the SPIKE mask for some pixels.
+            if not isInImage and selected:
+                nClass3 += 1
+                task.maskSources([x], [y], [r], mask)
+                self.assertGreater(np.sum(mask.array & mask.getPlaneBitMask(task.config.spikeMask) > 0), 0)
+
+            # Sources off the image that are skipped in source selection should
+            # not change the mask even if we do calculate their SPIKE mask.
+            if not isInImage and not selected:
+                nClass4 += 1
+                task.maskSources([x], [y], [r], mask)
+                self.assertMasksEqual(mask, baseMask)
+        # Verify that the test points were sufficient to exercise all classes.
+        self.assertGreater(nClass1, 0)
+        self.assertGreater(nClass2, 0)
+        self.assertGreater(nClass3, 0)
+        self.assertGreater(nClass4, 0)