Added mask planes from template and science fake plane into difference

python/lsst/ip/diffim/subtractImages.py

@@ -195,12 +195,12 @@ class AlardLuptonSubtractBaseConfig(lsst.pex.config.Config):
     )
     badMaskPlanes = lsst.pex.config.ListField(
         dtype=str,
-        default=("NO_DATA", "BAD", "SAT", "EDGE"),
+        default=("NO_DATA", "BAD", "SAT", "EDGE", "FAKE"),
         doc="Mask planes to exclude when selecting sources for PSF matching."
     )
     preserveTemplateMask = lsst.pex.config.ListField(
         dtype=str,
-        default=("NO_DATA", "BAD", "SAT", "INJECTED", "INJECTED_CORE"),
+        default=("NO_DATA", "BAD", "SAT", "FAKE", "INJECTED", "INJECTED_CORE"),
         doc="Mask planes from the template to propagate to the image difference."
     )
 
@@ -461,7 +461,8 @@ def runConvolveTemplate(self, template, science, selectSources):
                                      psfMatchingKernel=kernelResult.psfMatchingKernel)
 
     def runConvolveScience(self, template, science, selectSources):
-        """Convolve the science image with a PSF-matching kernel and subtract the template image.
+        """Convolve the science image with a PSF-matching kernel and subtract
+        the template image.
 
         Parameters
         ----------
@@ -561,20 +562,14 @@ def finalize(self, template, science, difference, kernel,
         """
         # Erase existing detection mask planes.
         #  We don't want the detection mask from the science image
-        mask = difference.mask
-        mask &= ~(mask.getPlaneBitMask("DETECTED") | mask.getPlaneBitMask("DETECTED_NEGATIVE"))
 
-        # We have cleared the template mask plane, so copy the mask plane of
-        # the image difference so that we can calculate correct statistics
-        # during decorrelation. Do this regardless of whether decorrelation is
-        # used for consistency.
-        template[science.getBBox()].mask.array[...] = difference.mask.array[...]
+        self.updateMasks(template, science, difference)
+
         if self.config.doDecorrelation:
             self.log.info("Decorrelating image difference.")
             # We have cleared the template mask plane, so copy the mask plane of
             # the image difference so that we can calculate correct statistics
             # during decorrelation
-            template[science.getBBox()].mask.array[...] = difference.mask.array[...]
             correctedExposure = self.decorrelate.run(science, template[science.getBBox()], difference, kernel,
                                                      templateMatched=templateMatched,
                                                      preConvMode=preConvMode,
@@ -585,6 +580,36 @@ def finalize(self, template, science, difference, kernel,
             correctedExposure = difference
         return correctedExposure
 
+    def updateMasks(self, template, science, difference):
+        """Update the mask planes on images for finalizing."""
+
+        bbox = science.getBBox()
+        mask = difference.mask
+        mask &= ~(mask.getPlaneBitMask("DETECTED") | mask.getPlaneBitMask("DETECTED_NEGATIVE"))
+
+        if "FAKE" in science.mask.getMaskPlaneDict().keys():
+            # propagate the mask plane related to Fake source injection
+            # NOTE: the fake source injection sets FAKE plane, but it should be INJECTED
+            # NOTE: This can be removed in DM-40796
+
+            self.log.info("Adding injected mask planes")
+            mask.addMaskPlane("INJECTED")
+            diffInjectedBitMask = mask.getPlaneBitMask("INJECTED")
+
+            mask.addMaskPlane("INJECTED_TEMPLATE")
+            diffInjTmpltBitMask = mask.getPlaneBitMask("INJECTED_TEMPLATE")
+
+            scienceFakeBitMask = science.mask.getPlaneBitMask('FAKE')
+            tmpltFakeBitMask = template[bbox].mask.getPlaneBitMask('FAKE')
+
+            injScienceMaskArray = ((science.mask.array & scienceFakeBitMask) > 0) * diffInjectedBitMask
+            injTemplateMaskArray = ((template[bbox].mask.array & tmpltFakeBitMask) > 0) * diffInjTmpltBitMask
+
+            mask.array |= injScienceMaskArray
+            mask.array |= injTemplateMaskArray
+
+        template[bbox].mask.array[...] = difference.mask.array[...]
+
     @staticmethod
     def _validateExposures(template, science):
         """Check that the WCS of the two Exposures match, and the template bbox
@@ -722,7 +747,12 @@ def _checkMask(mask, sources, badMaskPlanes):
             kept (True) or rejected (False) based on the value of the
             mask plane at its location.
         """
-        badPixelMask = lsst.afw.image.Mask.getPlaneBitMask(badMaskPlanes)
+        setBadMaskPlanes = [
+            maskPlane for maskPlane in badMaskPlanes if maskPlane in mask.getMaskPlaneDict()
+        ]
+
+        badPixelMask = mask.getPlaneBitMask(setBadMaskPlanes)
+
         xv = np.rint(sources.getX() - mask.getX0())
         yv = np.rint(sources.getY() - mask.getY0())
 
@@ -1057,8 +1087,12 @@ def _interpolateImage(maskedImage, badMaskPlanes, fallbackValue=None):
     result: `float`
         The number of masked pixels that were replaced.
     """
+    imgBadMaskPlanes = [
+        maskPlane for maskPlane in badMaskPlanes if maskPlane in maskedImage.mask.getMaskPlaneDict()
+    ]
+
     image = maskedImage.image.array
-    badPixels = (maskedImage.mask.array & maskedImage.mask.getPlaneBitMask(badMaskPlanes)) > 0
+    badPixels = (maskedImage.mask.array & maskedImage.mask.getPlaneBitMask(imgBadMaskPlanes)) > 0
     image[badPixels] = np.nan
     if fallbackValue is None:
         fallbackValue = np.nanmedian(image)

tests/test_subtractTask.py

@@ -715,6 +715,63 @@ def _compare_apCorrMaps(self, a, b):
             self.assertFloatsAlmostEqual(
                 value.getCoefficients(), value2.getCoefficients(), rtol=0.0)
 
+    def test_fake_mask_plane_propagation(self):
+        """Test that we have the mask planes related to fakes in diffim images.
+        This is testing method called updateMasks
+        """
+        xSize = 200
+        ySize = 200
+        science, sources = makeTestImage(psfSize=2.4, xSize=xSize, ySize=ySize)
+        science_fake_img, science_fake_sources = makeTestImage(
+            psfSize=2.4, xSize=xSize, ySize=ySize, seed=7, nSrc=2, noiseLevel=0.25, fluxRange=1
+        )
+        template, _ = makeTestImage(psfSize=2.4, xSize=xSize, ySize=ySize, doApplyCalibration=True)
+        tmplt_fake_img, tmplt_fake_sources = makeTestImage(
+            psfSize=2.4, xSize=xSize, ySize=ySize, seed=9, nSrc=2, noiseLevel=0.25, fluxRange=1
+        )
+        # created fakes and added them to the images
+        science.image.array += science_fake_img.image.array
+        template.image.array += tmplt_fake_img.image.array
+
+        # TODO: DM-40796 update to INJECTED names when source injection gets refactored
+        # adding mask planes to both science and template images
+        science_mask_planes = science.mask.addMaskPlane("FAKE")
+        template_mask_planes = template.mask.addMaskPlane("FAKE")
+
+        for a_science_source in science_fake_sources:
+            # 3 x 3 masking of the source locations is fine
+            bbox = lsst.geom.Box2I(
+                lsst.geom.Point2I(a_science_source.getX(), a_science_source.getY()), lsst.geom.Extent2I(3, 3)
+            )
+            science[bbox].mask.array |= science_mask_planes
+
+        for a_template_source in tmplt_fake_sources:
+            # 3 x 3 masking of the source locations is fine
+            bbox = lsst.geom.Box2I(
+                lsst.geom.Point2I(a_template_source.getX(), a_template_source.getY()),
+                lsst.geom.Extent2I(3, 3)
+            )
+            template[bbox].mask.array |= template_mask_planes
+
+        science_fake_masked = (science.mask.array & science.mask.getPlaneBitMask("FAKE")) > 0
+        template_fake_masked = (template.mask.array & template.mask.getPlaneBitMask("FAKE")) > 0
+
+        config = subtractImages.AlardLuptonSubtractTask.ConfigClass()
+        task = subtractImages.AlardLuptonSubtractTask(config=config)
+        subtraction = task.run(template, science, sources)
+
+        # check subtraction mask plane is set where we set the previous masks
+        diff_mask = subtraction.difference.mask
+
+        # science mask should be now in INJECTED
+        inj_masked = (diff_mask.array & diff_mask.getPlaneBitMask("INJECTED")) > 0
+
+        # template mask should be now in INJECTED_TEMPLATE
+        injTmplt_masked = (diff_mask.array & diff_mask.getPlaneBitMask("INJECTED_TEMPLATE")) > 0
+
+        self.assertEqual(np.sum(inj_masked.astype(int)-science_fake_masked.astype(int)), 0)
+        self.assertEqual(np.sum(injTmplt_masked.astype(int)-template_fake_masked.astype(int)), 0)
+
 
 class AlardLuptonPreconvolveSubtractTest(lsst.utils.tests.TestCase):