Scaling psf model before fitting

This is controled by a new config parameter `scalePsfModel` added to
`BrightStarCutoutTask` and `BrightStarSubtractTask`.

Author: bazkiaei

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

@@ -180,6 +180,11 @@ class BrightStarCutoutConfig(
         },
     )
 
+    scalePsfModel = Field[bool](
+        doc="If True, uses a scale factor to bring the PSF model data to the same level of the star data.",
+        default=True,
+    )
+
     # PSF Fitting
     useExtendedPsf = Field[bool](
         doc="Use the extended PSF model to normalize bright star cutouts.",
@@ -348,10 +353,18 @@ def run(
             # Fit a scaled PSF and a pedestal to each bright star cutout
             psf = WarpedPsf(inputExposure.getPsf(), pixToPolar, warpingControl)
             constantPsf = KernelPsf(FixedKernel(psf.computeKernelImage(Point2D(0, 0))))
+            # TODO: discuss with Lee whether we should warp the psf here as well?
             if self.config.useExtendedPsf:
                 psfImage = extendedPsf  # Assumed to be warped, center at [0,0]
             else:
                 psfImage = constantPsf.computeKernelImage(constantPsf.getAveragePosition())
+            if self.config.scalePsfModel:
+                psfNeg = psfImage.array < 0
+                self.modelScale = np.nanmean(stampMI.image.array) / np.nanmean(psfImage.array[~psfNeg])
+                psfImage.array *= self.modelScale ######## model scale correction ########
+            else:
+                self.modelScale = 1
+
             fitPsfResults = {}
             if self.config.doFitPsf:
                 fitPsfResults = self._fitPsf(stampMI, psfImage)
@@ -586,7 +599,8 @@ def _fitPsf(self, stampMI: MaskedImageF, psfImage: ImageD | ImageF) -> dict[str,
         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)
+        psfMaskedFluxFrac = np.dot(psfImage.array.flat, psfMaskedPixels.array.flat).astype(np.float64) / psfImage.array.sum()
+        # psfMaskedFluxFrac = np.dot(psfImage.array.astype(bool).flat, psfMaskedPixels.array.flat).astype(np.float64) / psfImage.array.astype(bool).sum()
         if psfMaskedFluxFrac > self.config.psfMaskedFluxFracThreshold:
             return {}  # Handle cases where the PSF image is mostly masked
 
@@ -622,13 +636,14 @@ def _fitPsf(self, stampMI: MaskedImageF, psfImage: ImageD | ImageF) -> dict[str,
             return {}  # Handle singular matrix errors
         if sumSquaredResiduals.size == 0:
             return {}  # Handle cases where sum of the squared residuals are empty
-        scale = solutions[0]
+        # scale = solutions[0]
+        scale = solutions[0] * self.modelScale ######## model scale correction ########
         if scale <= 0:
             return {}  # Handle cases where the PSF scale fit has failed
-        scaleErr = np.sqrt(covarianceMatrix[0, 0])
+        scaleErr = np.sqrt(covarianceMatrix[0, 0]) * self.modelScale ######## model scale correction ########
         pedestal = solutions[1]
         pedestalErr = np.sqrt(covarianceMatrix[1, 1])
-        scalePedestalCov = covarianceMatrix[0, 1]
+        scalePedestalCov = covarianceMatrix[0, 1] * self.modelScale ######## model scale correction ########
         xGradient = solutions[3]
         yGradient = solutions[2]
 
@@ -641,6 +656,7 @@ def _fitPsf(self, stampMI: MaskedImageF, psfImage: ImageD | ImageF) -> dict[str,
         psfBBoxGoodSpans = goodSpans.clippedTo(psfImage.getBBox())
         psfBBoxGoodSpansX, psfBBoxGoodSpansY = psfBBoxGoodSpans.indices()
         psfBBoxData = psfBBoxGoodSpans.flatten(stampMI.image.array, stampMI.getXY0())
+        paddedPsfImage.array /= self.modelScale  ######## model scale correction ########
         psfBBoxModel = (
             psfBBoxGoodSpans.flatten(paddedPsfImage.array, stampMI.getXY0()) * scale
             + pedestal
@@ -652,7 +668,6 @@ def _fitPsf(self, stampMI: MaskedImageF, psfImage: ImageD | ImageF) -> dict[str,
         psfBBoxChiSquared = np.sum(psfBBoxResiduals)
         psfBBoxDegreesOfFreedom = len(psfBBoxData) - 4
         psfBBoxReducedChiSquared = psfBBoxChiSquared / psfBBoxDegreesOfFreedom
-
         return dict(
             scale=scale,
             scaleErr=scaleErr,

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

@@ -24,9 +24,9 @@
 __all__ = ["BrightStarStackConnections", "BrightStarStackConfig", "BrightStarStackTask"]
 
 import numpy as np
-from lsst.afw.image import ImageF, MaskedImageF
+from lsst.afw.image import ImageF
 from lsst.afw.math import StatisticsControl, statisticsStack, stringToStatisticsProperty
-from lsst.geom import Box2I, Extent2I, Point2I
+from lsst.geom import Point2I
 from lsst.meas.algorithms import BrightStarStamps
 from lsst.pex.config import Field, ListField
 from lsst.pipe.base import PipelineTask, PipelineTaskConfig, PipelineTaskConnections, Struct
@@ -66,7 +66,7 @@ class BrightStarStackConfig(
 
     subsetStampNumber = Field[int](
         doc="Number of stamps per subset to generate stacked images for.",
-        default=20,
+        default=2,
     )
     globalReducedChiSquaredThreshold = Field[float](
         doc="Threshold for global reduced chi-squared for bright star stamps.",
@@ -172,9 +172,14 @@ def run(
 
         subsetStampMIs = []
         tempStampMIs = []
+        all_stars = 0
+        used_stars = 0
         for stampsDDH in brightStarStamps:
             stamps = stampsDDH.get()
+            all_stars += len(stamps)
             for stamp in stamps:
+                # print("globalReducedChiSquared: stamp ", stamp.globalReducedChiSquared, "config ", self.config.globalReducedChiSquaredThreshold)
+                # print("psfReducedChiSquared: stamp ", stamp.psfReducedChiSquared, "config ", self.config.psfReducedChiSquaredThreshold)
                 if (
                     stamp.globalReducedChiSquared > self.config.globalReducedChiSquaredThreshold
                     or stamp.psfReducedChiSquared > self.config.psfReducedChiSquaredThreshold
@@ -193,7 +198,11 @@ def run(
                     subsetStampMIs.append(statisticsStack(tempStampMIs, stackTypeProperty, statisticsControl))
                     # TODO: what to do with remaining temp stamps?
                     tempStampMIs = []
+                    used_stars += self.config.subsetStampNumber
 
+        self.metadata["psfStarCount"] = {}
+        self.metadata["psfStarCount"]["all"] = all_stars
+        self.metadata["psfStarCount"]["used"] = used_stars
         # 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"