tickets/DM-48377 draft

Author: bazkiaei

python/lsst/pipe/tasks/brightStarSubtraction/brightStarCutout.py

@@ -60,6 +60,7 @@
 from lsst.pipe.base import PipelineTask, PipelineTaskConfig, PipelineTaskConnections, Struct
 from lsst.pipe.base.connectionTypes import Input, Output, PrerequisiteInput
 from lsst.utils.timer import timeMethod
+from copy import deepcopy
 
 NEIGHBOR_MASK_PLANE = "NEIGHBOR"
 
@@ -293,6 +294,8 @@ def run(
         # Restore original subtracted background
         inputMI = inputExposure.getMaskedImage()
         inputMI += inputBackground.getImage()
+        # Amir: the above addition to inputMI, also adds to the inputExposure.
+        # Amir: but the calibration, three lines later, only is applied to the inputMI.
 
         # Set up NEIGHBOR mask plane; associate footprints with stars
         inputExposure.mask.addMaskPlane(NEIGHBOR_MASK_PLANE)
@@ -317,9 +320,11 @@ def run(
             # Set NEIGHBOR footprints in the mask plane
             if footprintIndex:
                 neighborFootprints = [fp for i, fp in enumerate(allFootprints) if i != footprintIndex]
-                self._setFootprints(inputExposure, neighborFootprints, NEIGHBOR_MASK_PLANE)
+                # self._setFootprints(inputExposure, neighborFootprints, NEIGHBOR_MASK_PLANE)
+                self._setFootprints(inputMI, neighborFootprints, NEIGHBOR_MASK_PLANE)
             else:
-                self._setFootprints(inputExposure, allFootprints, NEIGHBOR_MASK_PLANE)
+                # self._setFootprints(inputExposure, allFootprints, NEIGHBOR_MASK_PLANE)
+                self._setFootprints(inputMI, allFootprints, NEIGHBOR_MASK_PLANE)
 
             # Define linear shifting to recenter stamps
             coordFocalPlaneTan = pixToFocalPlaneTan.applyForward(pixCoord)  # center of warped star
@@ -329,7 +334,8 @@ def run(
             pixToPolar = pixToFocalPlaneTan.then(shift).then(rotation)
 
             # Apply the warp to the star stamp (in-place)
-            warpImage(stampMI, inputExposure.maskedImage, pixToPolar, warpingControl)
+            # warpImage(stampMI, inputExposure.maskedImage, pixToPolar, warpingControl)
+            warpImage(stampMI, inputMI, pixToPolar, warpingControl)
 
             # Trim to the base stamp size, check mask coverage, update metadata
             stampMI = stampMI[self.stampBBox]
@@ -579,6 +585,7 @@ def _fitPsf(self, stampMI: MaskedImageF, psfImage: ImageD | ImageF) -> dict[str,
         psfMaskedPixels = ImageF(psfImage.getBBox())
         psfMaskedPixels.array[:, :] = (stampMI.mask[psfImage.getBBox()].array & badMaskBitMask).astype(bool)
         # TODO: This is np.float64, else FITS metadata serialization fails
+        # Amir: the following tries to find the fraction of the psf flux in the masked area of the psf image.
         psfMaskedFluxFrac = np.dot(psfImage.array.flat, psfMaskedPixels.array.flat).astype(np.float64)
         if psfMaskedFluxFrac > self.config.psfMaskedFluxFracThreshold:
             return {}  # Handle cases where the PSF image is mostly masked
@@ -587,8 +594,10 @@ def _fitPsf(self, stampMI: MaskedImageF, psfImage: ImageD | ImageF) -> dict[str,
         paddedPsfImage = ImageF(stampMI.getBBox())
         paddedPsfImage[psfImage.getBBox()] = psfImage.convertF()
 
+        mask = self.add_psf_mask(paddedPsfImage, stampMI)
         # Create consistently masked data
-        badSpans = SpanSet.fromMask(stampMI.mask, badMaskBitMask)
+        # badSpans = SpanSet.fromMask(stampMI.mask, badMaskBitMask)
+        badSpans = SpanSet.fromMask(mask, badMaskBitMask)
         goodSpans = SpanSet(stampMI.getBBox()).intersectNot(badSpans)
         varianceData = goodSpans.flatten(stampMI.variance.array, stampMI.getXY0())
         if self.config.useMedianVariance:
@@ -598,7 +607,6 @@ def _fitPsf(self, stampMI: MaskedImageF, psfImage: ImageD | ImageF) -> dict[str,
         imageData /= sigmaData
         psfData = goodSpans.flatten(paddedPsfImage.array, paddedPsfImage.getXY0())
         psfData /= sigmaData
-
         # Fit the PSF scale factor and global pedestal
         nData = len(imageData)
         coefficientMatrix = np.ones((nData, 4), dtype=float)  # A
@@ -659,3 +667,10 @@ def _fitPsf(self, stampMI: MaskedImageF, psfImage: ImageD | ImageF) -> dict[str,
             psfDegreesOfFreedom=psfBBoxDegreesOfFreedom,
             psfMaskedFluxFrac=psfMaskedFluxFrac,
         )
+
+    def add_psf_mask(self, psfImage, stampMI):
+        cond = np.isnan(psfImage.array)
+        cond |= psfImage.array < 0
+        mask = deepcopy(stampMI.mask)
+        mask.array[cond] = np.bitwise_or(mask.array[cond], 1)
+        return mask

python/lsst/pipe/tasks/brightStarSubtraction/brightStarStack.py

@@ -24,7 +24,7 @@
 __all__ = ["BrightStarStackConnections", "BrightStarStackConfig", "BrightStarStackTask"]
 
 import numpy as np
-from lsst.afw.image import ImageF
+from lsst.afw.image import ImageF, MaskedImageF
 from lsst.afw.math import StatisticsControl, statisticsStack, stringToStatisticsProperty
 from lsst.geom import Box2I, Extent2I, Point2I
 from lsst.meas.algorithms import BrightStarStamps
@@ -187,19 +187,24 @@ def run(
                 badMaskBitMask = stampMI.mask.getPlaneBitMask(self.config.badMaskPlanes)
                 statisticsControl.setAndMask(badMaskBitMask)
 
+                # Amir: In case the total number of stamps is less than 20, the following will result in an
+                # empty subsetStampMIs list.
                 if len(tempStampMIs) == self.config.subsetStampNumber:
                     subsetStampMIs.append(statisticsStack(tempStampMIs, stackTypeProperty, statisticsControl))
                     # TODO: what to do with remaining temp stamps?
                     tempStampMIs = []
 
         # TODO: which stamp mask plane to use here?
+        # TODO: Amir: there might be cases where subsetStampMIs is an empty list. What do we want to do then?
+        # Currently, we get an "IndexError: list index out of range"
         badMaskBitMask = subsetStampMIs[0].mask.getPlaneBitMask(self.config.badMaskPlanes)
         statisticsControl.setAndMask(badMaskBitMask)
         extendedPsfMI = statisticsStack(subsetStampMIs, stackTypeProperty, statisticsControl)
 
         extendedPsfExtent = extendedPsfMI.getBBox().getDimensions()
         extendedPsfOrigin = Point2I(-1 * (extendedPsfExtent.x // 2), -1 * (extendedPsfExtent.y // 2))
         extendedPsfMI.setXY0(extendedPsfOrigin)
+        # return Struct(extendedPsf=[extendedPsfMI])
 
         return Struct(extendedPsf=extendedPsfMI.getImage())