Add ability to undo bgModel1 in SkyCorrectionTask

python/lsst/pipe/tasks/skyCorrection.py

@@ -178,6 +178,11 @@ class SkyCorrectionConfig(PipelineTaskConfig, pipelineConnections=SkyCorrectionC
         dtype=FocalPlaneBackgroundConfig,
         doc="Initial background model, prior to sky frame subtraction",
     )
+    undoBgModel1 = Field(
+        dtype=bool,
+        default=False,
+        doc="If True, adds back initial background model after sky and removes bgModel1 from the list",
+    )
     sky = ConfigurableField(
         target=SkyMeasurementTask,
         doc="Sky measurement",
@@ -210,6 +215,11 @@ def setDefaults(self):
         self.bgModel2.ySize = 256
         self.bgModel2.smoothScale = 1.0
 
+    def validate(self):
+        super().validate()
+        if self.undoBgModel1 and not self.doSky and not self.doBgModel2:
+            raise ValueError("If undoBgModel1 is True, task requires at least one of doSky or doBgModel2.")
+
 
 class SkyCorrectionTask(PipelineTask):
     """Perform a full focal plane sky correction."""
@@ -278,11 +288,11 @@ def run(self, calExps, calBkgs, skyFrames, camera):
 
         Parameters
         ----------
-        calExps : `list` [`lsst.afw.image.exposure.ExposureF`]
+        calExps : `list` [`lsst.afw.image.ExposureF`]
             Detector calibrated exposure images for the visit.
         calBkgs : `list` [`lsst.afw.math.BackgroundList`]
             Detector background lists matching the calibrated exposures.
-        skyFrames : `list` [`lsst.afw.image.exposure.ExposureF`]
+        skyFrames : `list` [`lsst.afw.image.ExposureF`]
             Sky frame calibration data for the input detectors.
         camera : `lsst.afw.cameraGeom.Camera`
             Camera matching the input data to process.
@@ -292,24 +302,29 @@ def run(self, calExps, calBkgs, skyFrames, camera):
         results : `Struct` containing:
             skyCorr : `list` [`lsst.afw.math.BackgroundList`]
                 Detector-level sky correction background lists.
-            calExpMosaic : `lsst.afw.image.exposure.ExposureF`
+            calExpMosaic : `lsst.afw.image.ExposureF`
                 Visit-level mosaic of the sky corrected data, binned.
                 Analogous to `calexp - skyCorr`.
-            calBkgMosaic : `lsst.afw.image.exposure.ExposureF`
+            calBkgMosaic : `lsst.afw.image.ExposureF`
                 Visit-level mosaic of the sky correction background, binned.
                 Analogous to `calexpBackground + skyCorr`.
         """
         # Restore original backgrounds in-place; optionally refine mask maps
         numOrigBkgElements = [len(calBkg) for calBkg in calBkgs]
-        _ = self._restoreBackgroundRefineMask(calExps, calBkgs)
+        _ = self._restoreOriginalBackgroundRefineMask(calExps, calBkgs)
 
         # Bin exposures, generate full-fp bg, map to CCDs and subtract in-place
         _ = self._subtractVisitBackground(calExps, calBkgs, camera, self.config.bgModel1)
+        initialBackgroundIndex = len(calBkgs[0]._backgrounds) - 1
 
         # Subtract a scaled sky frame from all input exposures
         if self.config.doSky:
             self._subtractSkyFrame(calExps, skyFrames, calBkgs)
 
+        # Adds full-fp bg back onto exposures, removes it from list
+        if self.config.undoBgModel1:
+            _ = self._undoInitialBackground(calExps, calBkgs, initialBackgroundIndex)
+
         # Bin exposures, generate full-fp bg, map to CCDs and subtract in-place
         if self.config.doBgModel2:
             _ = self._subtractVisitBackground(calExps, calBkgs, camera, self.config.bgModel2)
@@ -328,7 +343,7 @@ def run(self, calExps, calBkgs, skyFrames, camera):
 
         return Struct(skyCorr=calBkgs, calExpMosaic=calExpMosaic, calBkgMosaic=calBkgMosaic)
 
-    def _restoreBackgroundRefineMask(self, calExps, calBkgs):
+    def _restoreOriginalBackgroundRefineMask(self, calExps, calBkgs):
         """Restore original background to each calexp and invert the related
         background model; optionally refine the mask plane.
 
@@ -350,17 +365,17 @@ def _restoreBackgroundRefineMask(self, calExps, calBkgs):
 
         Parameters
         ----------
-        calExps : `lsst.afw.image.exposure.ExposureF`
+        calExps : `lsst.afw.image.ExposureF`
             Detector level calexp images to process.
-        calBkgs : `lsst.afw.math._backgroundList.BackgroundList`
+        calBkgs : `lsst.afw.math.BackgroundList`
             Detector level background lists associated with the calexps.
 
         Returns
         -------
-        calExps : `lsst.afw.image.exposure.ExposureF`
-            The calexps with the initially subtracted background restored.
-        skyCorrBases : `lsst.afw.math._backgroundList.BackgroundList`
-            The inverted initial background models; the genesis for skyCorrs.
+        calExps : `lsst.afw.image.ExposureF`
+            The calexps with the originally subtracted background restored.
+        skyCorrBases : `lsst.afw.math.BackgroundList`
+            The inverted original background models; the genesis for skyCorrs.
         """
         skyCorrBases = []
         for calExp, calBkg in zip(calExps, calBkgs):
@@ -379,14 +394,45 @@ def _restoreBackgroundRefineMask(self, calExps, calBkgs):
 
             stats = np.nanpercentile(skyCorrBase.array, [50, 75, 25])
             self.log.info(
-                "Detector %d: Initial background restored; BG median = %.1f counts, BG IQR = %.1f counts",
+                "Detector %d: Original background restored; BG median = %.1f counts, BG IQR = %.1f counts",
                 calExp.getDetector().getId(),
                 -stats[0],
                 np.subtract(*stats[1:]),
             )
             skyCorrBases.append(skyCorrBase)
         return calExps, skyCorrBases
 
+    def _undoInitialBackground(self, calExps, calBkgs, initialBackgroundIndex):
+        """Undo the initial background subtraction (bgModel1) after sky frame
+        subtraction.
+
+        Parameters
+        ----------
+        calExps : `list` [`lsst.afw.image.ExposureF`]
+            Calibrated exposures to be background subtracted.
+        calBkgs : `list` [`lsst.afw.math.BackgroundList`]
+            Background lists associated with the input calibrated exposures.
+        initialBackgroundIndex : `int`
+            Index of the initial background (bgModel1) in the background list.
+
+        Notes
+        -----
+        Inputs are modified in-place.
+        """
+        for calExp, calBkg in zip(calExps, calBkgs):
+            image = calExp.getMaskedImage()
+
+            # Remove bgModel1 from the background list; restore in the image
+            initialBackground = calBkg[initialBackgroundIndex][0].getImageF()
+            image += initialBackground
+            calBkg._backgrounds.pop(initialBackgroundIndex)
+
+            self.log.info(
+                "Detector %d: The initial background model prior to sky frame subtraction (bgModel1) has "
+                "been removed from the background list",
+                calExp.getDetector().getId(),
+            )
+
     def _subtractVisitBackground(self, calExps, calBkgs, camera, config):
         """Perform a full focal-plane background subtraction for a visit.
 
@@ -400,9 +446,9 @@ def _subtractVisitBackground(self, calExps, calBkgs, camera, config):
 
         Parameters
         ----------
-        calExps : `list` [`lsst.afw.image.exposure.ExposureF`]
+        calExps : `list` [`lsst.afw.image.ExposureF`]
             Calibrated exposures to be background subtracted.
-        calBkgs : `list` [`lsst.afw.math._backgroundList.BackgroundList`]
+        calBkgs : `list` [`lsst.afw.math.BackgroundList`]
             Background lists associated with the input calibrated exposures.
         camera : `lsst.afw.cameraGeom.Camera`
             Camera description.
@@ -413,7 +459,7 @@ def _subtractVisitBackground(self, calExps, calBkgs, camera, config):
         -------
         calExps : `list` [`lsst.afw.image.maskedImage.MaskedImageF`]
             Background subtracted exposures for creating a focal plane image.
-        calBkgs : `list` [`lsst.afw.math._backgroundList.BackgroundList`]
+        calBkgs : `list` [`lsst.afw.math.BackgroundList`]
             Updated background lists with a visit-level model appended.
         """
         # Set up empty full focal plane background model object
@@ -478,16 +524,16 @@ def _subtractDetectorBackground(self, calExp, bgModel):
 
         Parameters
         ----------
-        calExp : `lsst.afw.image.exposure.ExposureF`
+        calExp : `lsst.afw.image.ExposureF`
             Exposure to subtract the background model from.
         bgModel : `lsst.pipe.tasks.background.FocalPlaneBackground`
             Full focal plane camera-level background model.
 
         Returns
         -------
-        calExp : `lsst.afw.image.exposure.ExposureF`
+        calExp : `lsst.afw.image.ExposureF`
             Background subtracted input exposure.
-        calBkgElement : `lsst.afw.math._backgroundList.BackgroundList`
+        calBkgElement : `lsst.afw.math.BackgroundList`
             Detector level realization of the full focal plane bg model.
         """
         image = calExp.getMaskedImage()
@@ -510,11 +556,11 @@ def _subtractSkyFrame(self, calExps, skyFrames, calBkgs):
 
         Parameters
         ----------
-        calExps : `list` [`lsst.afw.image.exposure.ExposureF`]
+        calExps : `list` [`lsst.afw.image.ExposureF`]
             Calibrated exposures to be background subtracted.
-        skyFrames : `list` [`lsst.afw.image.exposure.ExposureF`]
+        skyFrames : `list` [`lsst.afw.image.ExposureF`]
             Sky frame calibration data for the input detectors.
-        calBkgs : `list` [`lsst.afw.math._backgroundList.BackgroundList`]
+        calBkgs : `list` [`lsst.afw.math.BackgroundList`]
             Background lists associated with the input calibrated exposures.
         """
         skyFrameBgModels = []
@@ -549,11 +595,11 @@ def _binAndMosaic(self, exposures, camera, binning, ids=None, refExps=None):
             Binning size to be applied to input images.
         ids : `list` [`int`], optional
             List of detector ids to iterate over.
-        refExps : `list` [`lsst.afw.image.exposure.ExposureF`], optional
+        refExps : `list` [`lsst.afw.image.ExposureF`], optional
             If supplied, mask planes from these reference images will be used.
         Returns
         -------
-        mosaicImage : `lsst.afw.image.exposure.ExposureF`
+        mosaicImage : `lsst.afw.image.ExposureF`
             Mosaicked full focal plane image.
         """
         refExps = np.resize(refExps, len(exposures))  # type: ignore