DM-38557: Automate mask sizing for masked deblending

python/lsst/ts/wep/cwfs/algorithm.py

@@ -117,6 +117,7 @@ def __init__(self, algoDir):
 
         # True if at least one of images has a blended object
         self.blend_exists = False
+        self.mask_growth_iter = 4
 
     def reset(self):
         """Reset the calculation for the new input images with the same
@@ -800,8 +801,8 @@ def _singleItr(self, I1, I2, model, tol=1e-3):
                     )
                     sys.exit()
 
-                # Check for blends
-                if (np.sum(np.isnan(I1.blendOffsetX)) == 0) and (
+                # Check for blends in either image
+                if (np.sum(np.isnan(I1.blendOffsetX)) == 0) or (
                     np.sum(np.isnan(I2.blendOffsetX)) == 0
                 ):
                     self.blend_exists = True
@@ -810,10 +811,45 @@ def _singleItr(self, I1, I2, model, tol=1e-3):
                 # These will be applied to compensated images, so we will
                 # create the masks with compensated=True
                 boundaryT = self.getBoundaryThickness()
+
+                # Create shifted mask from non-blended mask
+                for compIm in [I1, I2]:
+                    if np.sum(np.isnan(compIm.blendOffsetX)) > 0:
+                        continue
+                    compIm.makeMask(self._inst, model, boundaryT, 1)
+                    dilatedMask, numPaddingIter = compIm.autoDilateBlendMask(
+                        compIm.mask_pupil
+                    )
+
+                    finalMask, shiftedMask = compIm.createBlendedCoadd(
+                        compIm.mask_pupil,
+                        blendPadding=numPaddingIter,
+                        returnShiftedMask=True,
+                    )
+                    # Mask only blended areas in final stamp
+                    compIm.updateImage(
+                        compIm.getImg() * np.invert(np.array(shiftedMask, dtype=bool))
+                    )
+
                 # If the compensable image has no blended centroids
                 # this function will just create a single masked donut
-                I1.makeBlendedMask(self._inst, model, boundaryT, 1, compensated=True)
-                I2.makeBlendedMask(self._inst, model, boundaryT, 1, compensated=True)
+                I1.makeBlendedMask(
+                    self._inst,
+                    model,
+                    boundaryT,
+                    1,
+                    compensated=True,
+                    blendPadding=self.mask_growth_iter,
+                )
+                I2.makeBlendedMask(
+                    self._inst,
+                    model,
+                    boundaryT,
+                    1,
+                    compensated=True,
+                    blendPadding=self.mask_growth_iter,
+                )
+
                 self._makeMasterMask(I1, I2, self.getPoissonSolverName())
 
                 # Load the offAxis correction coefficients

python/lsst/ts/wep/cwfs/compensableImage.py

@@ -33,7 +33,7 @@
     shift,
 )
 from scipy.interpolate import RectBivariateSpline
-from scipy.signal import correlate
+from scipy.signal import correlate, argrelmax
 
 from lsst.ts.wep.paramReader import ParamReader
 from lsst.ts.wep.cwfs.tool import (
@@ -43,7 +43,7 @@
     ZernikeAnnularJacobian,
 )
 from lsst.ts.wep.cwfs.image import Image
-from lsst.ts.wep.utility import DefocalType, CentroidFindType, FilterType, rotMatrix
+from lsst.ts.wep.utility import DefocalType, CentroidFindType, FilterType
 from galsim.utilities import horner2d
 
 
@@ -1503,7 +1503,7 @@ def makeBlendedMask(
         model,
         boundaryT,
         maskScalingFactorLocal,
-        blendPadding=8,
+        blendPadding=None,
         compensated=True,
     ):
         """Create a binary mask of a central source with the area overlapping
@@ -1524,9 +1524,12 @@ def makeBlendedMask(
             zero.
         maskScalingFactorLocal : float
             Mask scaling factor (for fast beam) for local correction.
-        blendPadding : int, optional
+        blendPadding : int or None, optional
             Number of pixels to increase the radius and expand the
-            footprint of the blended source. (the default is 8.)
+            footprint of the blended source. If None then the
+            amount of expansion to cover the blended source will be
+            calculated automatically with `autoDilateBlendMask`. If an integer
+            it must be >= 0. (the default is None.)
         compensated: bool, optional
             Whether to calculate masks for the compensated or original image.
             If True, mask will be in orientation of the compensated image.
@@ -1541,30 +1544,55 @@ def makeBlendedMask(
 
         # Add blended donuts if they exist
         if np.sum(np.isnan(self.blendOffsetX)) == 0:
-            newMaskPupil = self.createBlendedCoadd(self.mask_pupil, blendPadding)
-            newMaskComp = self.createBlendedCoadd(self.mask_comp, blendPadding)
+            newMaskPupil, iterNum = self.autoDilateBlendMask(self.mask_pupil)
+            if blendPadding is None:
+                newMaskComp = self.createBlendedCoadd(
+                    self.mask_comp, blendPadding=iterNum
+                )
+            else:
+                iterNum = 0
+                newMaskPupil = self.createBlendedCoadd(
+                    self.mask_pupil, blendPadding=blendPadding + iterNum
+                )
+                newMaskComp = self.createBlendedCoadd(
+                    self.mask_comp, blendPadding=blendPadding + iterNum
+                )
             self.mask_pupil = newMaskPupil
             self.mask_comp = newMaskComp
 
-    def createBlendedCoadd(self, maskArray, blendPadding):
+    def createBlendedCoadd(self, maskArray, blendPadding=None, returnShiftedMask=False):
         """Cut out regions of the input mask where blended donuts
         overlap with the original donut.
 
         Parameters
         ----------
         maskArray : numpy.ndarray
             Original input binary mask with just a single unblended donut.
-        blendPadding : int
+        blendPadding : int or None, optional
             Number of pixels to increase the radius and expand the
-            footprint of the blended source.
+            footprint of the blended source. If set to None then
+            the blendPadding will automatically be generated by
+            `autoDilateBlendMask`. If set to an integer value it must
+            be >= 0. (the default is None.)
 
         Returns
         -------
         numpy.ndarray [int]
             New maskArray modified to remove the footprint of any
             overlapping donuts from the masked area of the original donut.
+        numpy.ndarray [int]
+            New maskArray that is the union of shifted masks covering
+            all blended objects.
+
+        Raises
+        ------
+        AssertionError
+            blendPadding must be None or an integer > 0.
         """
 
+        if blendPadding is not None:
+            assert blendPadding > 0, "blendPadding must be None or an integer > 0."
+
         # Switch x,y order because x,y locations in the catalogs and the
         # LSST Science Pipeline objects are consistent with each other but
         # numpy arrays are [row, column] so when using the catalog x,y values
@@ -1575,12 +1603,9 @@ def createBlendedCoadd(self, maskArray, blendPadding):
         maskBlends = list()
         for xShift, yShift in pixelShift:
             shiftVector = np.array([xShift, yShift], dtype=int)
-            # Rotate extrafocal donut position 180 degrees
-            if self.defocalType == DefocalType.Extra:
-                theta = np.degrees(np.pi)
-                shiftVector = np.dot(rotMatrix(theta), shiftVector)
             shiftedMask = shift(copy(maskArray), shiftVector)
-            if blendPadding > 0:
+            # Only pad if blendPadding is an int greater than 0
+            if blendPadding is not None:
                 shiftedMask = binary_dilation(shiftedMask, iterations=blendPadding)
             maskBlends.append(shiftedMask)
 
@@ -1597,4 +1622,78 @@ def createBlendedCoadd(self, maskArray, blendPadding):
         maskFinal = maskArray - maskCoadd
         maskFinal[maskFinal < 0] = 0
 
-        return maskFinal
+        if returnShiftedMask is False:
+            return maskFinal
+        else:
+            allShifted = np.zeros(np.shape(maskArray))
+            for maskBlend in maskBlends:
+                allShifted += maskBlend
+            allShifted[allShifted > 1] = 1
+
+            return maskFinal, allShifted
+
+    def autoDilateBlendMask(self, shiftedMask):
+        """Automatically expand the mask footprint
+        of a blended set of donuts by removing a second
+        non-zero peak from a histogram of the pixel values.
+
+        Parameters
+        ----------
+        shiftedMask: numpy.ndarray
+            Model mask array centered on the donut that is
+            blended with the central source.
+
+        Returns
+        -------
+        numpy.ndarray
+            Model mask array with an expanded footprint.
+        """
+
+        def calcNumPeaks(imageArray, blendMask):
+            """Helper function to find number of peaks in masked area."""
+            maskedPixelVals = blendMask * imageArray
+            # Find 95th percentile to set range for histogram to avoid a couple
+            # bad pixels skewing histogram (eg, bad cosmic ray or ISR).
+            p99PixelVal = np.percentile(imageArray, 99)
+            maskedPixelHist = np.histogram(
+                maskedPixelVals, bins=12, range=(0, p99PixelVal)
+            )
+            medianPixel = np.median(
+                maskedPixelVals[np.where(maskedPixelVals > maskedPixelHist[1][1])]
+            )
+            binEdges = np.linspace(0.5 * medianPixel, 1.5 * medianPixel, num=11)
+            binEdges[-1] = 2.5 * medianPixel
+            binEdges = np.append(binEdges, 3.0 * medianPixel)
+            maskedPixelHist = np.histogram(
+                maskedPixelVals,
+                bins=binEdges,  # range=(0.5 * medianPixel, 2.5 * medianPixel)
+            )
+            print(maskedPixelHist)
+            # Find the highest bins
+            binRanking = np.argsort(maskedPixelHist[0])
+            # The highest bin will be the one around 0,
+            # the second highest should be the main donut
+            maxNonZeroBin = binRanking[-1]
+            # Find the peaks
+            histPeaks = argrelmax(maskedPixelHist[0])  # [0]
+            # Count the number of peaks beyond the main donut peak
+            peaksBeyondMax = np.where(histPeaks > maxNonZeroBin)[0]
+            # Any peaks beyond the main donut peak are
+            # the ones we want to remove
+            numHistPeaksBeyond = len(peaksBeyondMax)
+            return numHistPeaksBeyond
+
+        blendMask = self.createBlendedCoadd(copy(shiftedMask), blendPadding=None)
+        numMaskedPeaks = calcNumPeaks(self.getImg(), blendMask)
+
+        # Return if no need to pad mask
+        if numMaskedPeaks == 0:
+            return blendMask, None
+
+        paddingIter = 0
+        while numMaskedPeaks > 0:
+            paddingIter += 1
+            blendMask = self.createBlendedCoadd(shiftedMask, blendPadding=paddingIter)
+            numMaskedPeaks = calcNumPeaks(self.getImg(), blendMask)
+
+        return blendMask, paddingIter

python/lsst/ts/wep/task/cutOutDonutsBase.py

@@ -44,7 +44,6 @@
 from lsst.ts.wep.task.donutStamps import DonutStamps
 from lsst.ts.wep.cwfs.donutTemplateFactory import DonutTemplateFactory
 from scipy.signal import correlate
-from scipy.ndimage import shift, binary_dilation
 
 
 class CutOutDonutsBaseTaskConnections(
@@ -144,9 +143,11 @@ class CutOutDonutsBaseTaskConfig(
         default=False,
     )
     maskGrowthIter = pexConfig.Field(
-        doc="How many iterations of binary dilation to run on the mask model.",
+        doc="How many iterations of binary dilation to run on the mask model. \
+        If None then will run cwfs/compensable/autoDilateBlendMask to grow mask.",
         dtype=int,
-        default=6,
+        default=None,
+        optional=True,
     )
 
 
@@ -499,7 +500,11 @@ def cutOutStamps(self, exposure, donutCatalog, defocalType, cameraName):
             boundaryT = 1
             maskScalingFactorLocal = 1
             donutStamp.makeMasks(
-                inst, self.opticalModel, boundaryT, maskScalingFactorLocal
+                inst,
+                self.opticalModel,
+                boundaryT,
+                maskScalingFactorLocal,
+                blendPadding=self.maskGrowthIter,
             )
             donutStamp.stamp_im.setMask(donutStamp.mask_comp)
 
@@ -508,21 +513,24 @@ def cutOutStamps(self, exposure, donutCatalog, defocalType, cameraName):
                 donutStamp.comp_im.makeMask(
                     inst, self.opticalModel, boundaryT, maskScalingFactorLocal
                 )
-                shiftedMask = shift(
+                if self.maskGrowthIter is None:
+                    (
+                        dilatedMask,
+                        numPaddingIter,
+                    ) = donutStamp.comp_im.autoDilateBlendMask(
+                        donutStamp.comp_im.mask_pupil
+                    )
+                else:
+                    numPaddingIter = self.maskGrowthIter
+                finalMask, shiftedMask = donutStamp.comp_im.createBlendedCoadd(
                     donutStamp.comp_im.mask_pupil,
-                    np.array(
-                        [
-                            donutStamp.comp_im.blendOffsetY,
-                            donutStamp.comp_im.blendOffsetX,
-                        ]
-                    ),
+                    blendPadding=numPaddingIter,
+                    returnShiftedMask=True,
+                )
+                # Mask only blended areas in final stamp
+                donutStamp.stamp_im.image.array *= np.invert(
+                    np.array(shiftedMask, dtype=bool)
                 )
-                shiftedMask = binary_dilation(
-                    shiftedMask, iterations=self.maskGrowthIter
-                ).astype(int)
-                shiftedMask[shiftedMask == 0] += 2
-                shiftedMask -= 1
-                donutStamp.stamp_im.image.array *= shiftedMask
 
             finalStamps.append(donutStamp)
 

python/lsst/ts/wep/task/donutStamp.py

@@ -244,7 +244,9 @@ def _setCompensableImage(self):
         self.mask_pupil = afwImage.Mask()
         self.mask_comp = afwImage.Mask()
 
-    def makeMasks(self, inst, model, boundaryT, maskScalingFactorLocal):
+    def makeMasks(
+        self, inst, model, boundaryT, maskScalingFactorLocal, blendPadding=None
+    ):
         """Get the binary mask which considers the obscuration and off-axis
         correction.
 
@@ -261,6 +263,12 @@ def makeMasks(self, inst, model, boundaryT, maskScalingFactorLocal):
             zero.
         maskScalingFactorLocal : `float`
             Mask scaling factor (for fast beam) for local correction.
+        blendPadding : int or None, optional
+            Number of pixels to increase the radius and expand the
+            footprint of the blended source. If None then the
+            amount of expansion to cover the blended source will be
+            calculated automatically with `autoDilateBlendMask`. If an integer
+            it must be >= 0. (the default is None.)
 
         Returns
         -------
@@ -270,7 +278,9 @@ def makeMasks(self, inst, model, boundaryT, maskScalingFactorLocal):
             Padded mask for use at the offset planes.
         """
 
-        self.comp_im.makeBlendedMask(inst, model, boundaryT, maskScalingFactorLocal)
+        self.comp_im.makeBlendedMask(
+            inst, model, boundaryT, maskScalingFactorLocal, blendPadding=blendPadding
+        )
 
         # 0 flag in mask is part of image that is not donut
         # 1 flag in mask means it is part of the model donut