Merge pull request #100 from lsst/tickets/dm-41840

tickets/DM-41840: Implement new detection tasks

doc/conf.py

@@ -1,4 +1,3 @@
-
 """Sphinx configuration file for an LSST stack package.
 This configuration only affects single-package Sphinx documentation builds.
 For more information, see:
@@ -7,7 +6,6 @@
 
 from documenteer.conf.pipelinespkg import *
 
-
 project = "meas_extensions_scarlet"
 html_theme_options["logotext"] = project
 html_title = project

python/lsst/__init__.py

@@ -1,2 +1,3 @@
 import pkgutil
+
 __path__ = pkgutil.extend_path(__path__, __name__)

python/lsst/meas/__init__.py

@@ -1,2 +1,3 @@
 import pkgutil
+
 __path__ = pkgutil.extend_path(__path__, __name__)

python/lsst/meas/extensions/__init__.py

@@ -1,2 +1,3 @@
 import pkgutil
+
 __path__ = pkgutil.extend_path(__path__, __name__)

python/lsst/meas/extensions/scarlet/__init__.py

@@ -19,9 +19,9 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-from .version import *
-
-from .scarletDeblendTask import *
+from . import utils
+from .deconvolveExposureTask import *
 from .io import *
 from .metrics import *
-from . import utils
+from .scarletDeblendTask import *
+from .version import *

python/lsst/meas/extensions/scarlet/deconvolveExposureTask.py

@@ -0,0 +1,210 @@
+# This file is part of meas_extensions_scarlet.
+#
+# Developed for the LSST Data Management System.
+# This product includes software developed by the LSST Project
+# (https://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import logging
+
+import lsst.afw.image as afwImage
+import lsst.pex.config as pexConfig
+import lsst.pipe.base as pipeBase
+import lsst.pipe.base.connectionTypes as cT
+import lsst.scarlet.lite as scl
+import numpy as np
+
+from . import utils
+
+log = logging.getLogger(__name__)
+
+__all__ = [
+    "DeconvolveExposureTask",
+    "DeconvolveExposureConfig",
+    "DeconvolveExposureConnections",
+]
+
+
+class DeconvolveExposureConnections(
+    pipeBase.PipelineTaskConnections,
+    dimensions=("tract", "patch", "skymap", "band"),
+    defaultTemplates={"inputCoaddName": "deep"},
+):
+    """Connections for DeconvolveExposureTask"""
+
+    coadd = cT.Input(
+        doc="Exposure to deconvolve",
+        name="{inputCoaddName}Coadd_calexp",
+        storageClass="ExposureF",
+        dimensions=("tract", "patch", "band", "skymap"),
+    )
+
+    deconvolved = cT.Output(
+        doc="Deconvolved exposure",
+        name="deconvolved_{inputCoaddName}_coadd",
+        storageClass="ExposureF",
+        dimensions=("tract", "patch", "band", "skymap"),
+    )
+
+
+class DeconvolveExposureConfig(
+    pipeBase.PipelineTaskConfig,
+    pipelineConnections=DeconvolveExposureConnections,
+):
+    """Configuration for DeconvolveExposureTask"""
+
+    maxIter = pexConfig.Field[int](
+        doc="Maximum number of iterations",
+        default=100,
+    )
+    minIter = pexConfig.Field[int](
+        doc="Minimum number of iterations",
+        default=10,
+    )
+    eRel = pexConfig.Field[float](
+        doc="Relative error threshold",
+        default=1e-3,
+    )
+    backgroundThreshold = pexConfig.Field[float](
+        default=0,
+        doc="Threshold for background subtraction. "
+        "Pixels in the fit below this threshold will be set to zero",
+    )
+
+
+class DeconvolveExposureTask(pipeBase.PipelineTask):
+    """Deconvolve an Exposure using scarlet lite."""
+
+    ConfigClass = DeconvolveExposureConfig
+    _DefaultName = "deconvolveExposure"
+
+    def __init__(self, initInputs=None, **kwargs):
+        if initInputs is None:
+            initInputs = {}
+        super().__init__(initInputs=initInputs, **kwargs)
+
+    def runQuantum(self, butlerQC, inputRefs, outputRefs):
+        inputs = butlerQC.get(inputRefs)
+        outputs = self.run(**inputs)
+        butlerQC.put(outputs, outputRefs)
+
+    def run(self, coadd: afwImage.Exposure) -> pipeBase.Struct:
+        """Deconvolve an Exposure
+
+        Parameters
+        ----------
+        coadd :
+            Coadd image to deconvolve
+
+        Returns
+        -------
+        deconvolved : `pipeBase.Struct`
+            Deconvolved exposure
+        """
+        # Load the scarlet lite Observation
+        observation = self._buildObservation(coadd)
+
+        # Deconvolve.
+        # Store the loss history for debugging purposes.
+        model, self.loss = self._deconvolve(observation)
+
+        # Store the model in an Exposure
+        exposure = self._modelToExposure(model.data[0], coadd)
+        return pipeBase.Struct(deconvolved=exposure)
+
+    def _buildObservation(self, coadd: afwImage.Exposure) -> scl.Observation:
+        """Build a scarlet lite Observation from an Exposure.
+
+        We don't actually use scarlet, but the optimized convolutions
+        using scarlet data products are still useful.
+
+        Parameters
+        ----------
+        coadd :
+            Coadd image to deconvolve.
+        """
+        bands = ("dummy",)
+        model_psf = scl.utils.integrated_circular_gaussian(sigma=0.8)
+
+        image = coadd.image.array
+        psf = coadd.getPsf().computeKernelImage(coadd.getBBox().getCenter()).array
+        weights = np.ones_like(coadd.image.array)
+        badPixelMasks = utils.defaultBadPixelMasks
+        badPixels = coadd.mask.getPlaneBitMask(badPixelMasks)
+        mask = coadd.mask.array & badPixels
+        weights[mask > 0] = 0
+
+        observation = scl.Observation(
+            images=image.copy()[None],
+            variance=coadd.variance.array.copy()[None],
+            weights=weights[None],
+            psfs=psf[None],
+            model_psf=model_psf[None],
+            convolution_mode="fft",
+            bands=bands,
+            bbox=utils.bboxToScarletBox(coadd.getBBox()),
+        )
+        return observation
+
+    def _deconvolve(self, observation: scl.Observation) -> tuple[scl.Image, list[float]]:
+        """Deconvolve the observed image.
+
+        Parameters
+        ----------
+        observation :
+            Scarlet lite Observation.
+        """
+        model = observation.images.copy()
+        loss = []
+        for n in range(self.config.maxIter):
+            residual = observation.images - observation.convolve(model)
+            loss.append(-0.5 * np.sum(residual.data**2))
+            update = observation.convolve(residual, grad=True)
+            model += update
+            model.data[model.data < 0] = 0
+
+            if n > self.config.minIter and np.abs(loss[-1] - loss[-2]) < self.config.eRel * np.abs(loss[-1]):
+                break
+
+        return model, loss
+
+    def _modelToExposure(self, model: np.ndarray, coadd: afwImage.Exposure) -> afwImage.Exposure:
+        """Convert a scarlet lite Image to an Exposure.
+
+        Parameters
+        ----------
+        image :
+            Scarlet lite Image.
+        """
+        image = afwImage.Image(
+            array=model,
+            xy0=coadd.getBBox().getMin(),
+            deep=False,
+            dtype=coadd.image.array.dtype,
+        )
+        maskedImage = afwImage.MaskedImage(
+            image=image,
+            mask=coadd.mask,
+            variance=coadd.variance,
+            dtype=coadd.image.array.dtype,
+        )
+        exposure = afwImage.Exposure(
+            maskedImage=maskedImage,
+            exposureInfo=coadd.getInfo(),
+            dtype=coadd.image.array.dtype,
+        )
+        return exposure