fit_coadd_psf.py

Go to the documentation of this file.

# This file is part of pipe_tasks.
#
# 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/>.
 
__all__ = [
    "CatalogExposurePsf", "CoaddPsfFitConfig", "CoaddPsfFitConnections",
    "CoaddPsfFitSubConfig", "CoaddPsfFitSubTask", "CoaddPsfFitTask",
]
 
from .fit_multiband import CatalogExposure, CatalogExposureConfig
from lsst.geom import Point2D
from lsst.meas.base import SkyMapIdGeneratorConfig
import lsst.pex.config as pexConfig
import lsst.pipe.base as pipeBase
import lsst.pipe.base.connectionTypes as cT
 
import dataclasses
from abc import ABC, abstractmethod
from pydantic.dataclasses import dataclass
 
 
@dataclass(frozen=True, kw_only=True, config=CatalogExposureConfig)
class CatalogExposurePsf(CatalogExposure):
    def get_catalog(self):
        return self.catalog
 
    def get_psf_image(self, source):
        """Return the PSF image for this object."""
        center = Point2D(round(source.getX()), round(source.getY()))
        return self.exposure.getPsf().computeKernelImage(center).array
 
 
CoaddPsfFitBaseTemplates = {
    "name_coadd": "deep",
    "name_output_method": "multiprofit",
}
 
 
class CoaddPsfFitConnections(
    pipeBase.PipelineTaskConnections,
    dimensions=("tract", "patch", "band", "skymap"),
    defaultTemplates=CoaddPsfFitBaseTemplates,
):
    coadd = cT.Input(
        doc="Coadd image to fit a PSF model to",
        name="{name_coadd}Coadd_calexp",
        storageClass="ExposureF",
        dimensions=("tract", "patch", "band", "skymap"),
    )
    coadd_cell = cT.Input(
        doc="Cell-coadd image to fit a PSF model to",
        name="{name_coadd}CoaddCell",
        storageClass="MultipleCellCoadd",
        dimensions=("tract", "patch", "band", "skymap"),
    )
    background = cT.Input(
        doc="Background model to subtract from the coadd_cell",
        name="{name_coadd}Coadd_calexp_background",
        storageClass="Background",
        dimensions=("tract", "patch", "band", "skymap"),
    )
    cat_meas = cT.Input(
        doc="Deblended single-band source catalog",
        name="{name_coadd}Coadd_meas",
        storageClass="SourceCatalog",
        dimensions=("tract", "patch", "band", "skymap"),
    )
    cat_output = cT.Output(
        doc="Output PSF fit parameter catalog",
        name="{name_coadd}Coadd_psfs_{name_output_method}",
        storageClass="ArrowTable",
        dimensions=("tract", "patch", "band", "skymap"),
    )
 
    def __init__(self, *, config=None):
        super().__init__(config=config)
        if config is None:
            return
 
        if config.image_type == "future":
            self.coadd = dataclasses.replace(self.coadd, storageClass="CellCoadd")
            del self.coadd_cell
            del self.background
        elif config.use_cell_coadds:
            del self.coadd
        else:
            del self.coadd_cell
            del self.background
 
 
class CoaddPsfFitSubConfig(pexConfig.Config):
    """Base config class for the CoaddPsfFitTask.
 
    Implementing classes may add any necessary attributes.
    """
 
 
class CoaddPsfFitSubTask(pipeBase.Task, ABC):
    """Interface for CoaddPsfFitTask subtasks to fit PSFs.
 
    Parameters
    ----------
    **kwargs
        Additional arguments to be passed to the `lsst.pipe.base.Task`
        constructor.
    """
    ConfigClass = CoaddPsfFitSubConfig
 
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
 
    @abstractmethod
    def run(
        self, catexp: CatalogExposurePsf
    ) -> pipeBase.Struct:
        """Fit PSF images at locations of sources in a single exposure.
 
        Parameters
        ----------
        catexp : `CatalogExposurePsf`
            An exposure to fit a model PSF at the position of all
            sources in the corresponding catalog.
 
        Returns
        -------
        retStruct : `lsst.pipe.base.Struct`
            A struct with a cat_output attribute containing the output
            measurement catalog.
 
        Notes
        -----
        Subclasses may have further requirements on the input parameters,
        including:
            - Passing only one catexp per band;
            - Catalogs containing HeavyFootprints with deblended images;
            - Fitting only a subset of the sources.
        If any requirements are not met, the subtask should fail as soon as
        possible.
        """
        raise NotImplementedError()
 
 
class CoaddPsfFitConfig(
    pipeBase.PipelineTaskConfig,
    pipelineConnections=CoaddPsfFitConnections,
):
    """Configure a CoaddPsfFitTask, including a configurable fitting subtask.
    """
    use_cell_coadds = pexConfig.Field(
        dtype=bool,
        default=False,
        doc="Use cell coadd images for PSF fitting",
    )
    fit_coadd_psf = pexConfig.ConfigurableField(
        target=CoaddPsfFitSubTask,
        doc="Task to fit PSF models for a single coadd",
    )
    idGenerator = SkyMapIdGeneratorConfig.make_field()
    image_type = pexConfig.ChoiceField(
        "Which image type to expect for the input coadd. "
        "This option only directly affects connection storage classes and hence 'runQuantum'; the 'run' "
        "method behavior is determined by which type is actually passed in.",
        allowed={
            "legacy": (
                "Read a lsst.cell_coadds.MultipleCellCoadd via 'coadd_cell` and restore 'background' "
                "(if use_cell_coadd) or lsst.afw.image.Exposure via `coadd` (if not use_cell_coadd)."
            ),
            "future": (
                "Read lsst.images.cells.CellCoadd via the 'coadd' connection.  use_cell_coadd is ignored."
            ),
        },
        dtype=str,
        optional=False,
        default="legacy",
    )
 
 
class CoaddPsfFitTask(pipeBase.PipelineTask):
    """Fit a PSF model at the location of sources in a coadd.
 
    This task is intended to fit only a single PSF model at the
    centroid of all of the sources in a single coadd exposure.
    Subtasks may choose to filter which sources they fit,
    and may output whatever columns they desire in addition to
    the minimum of 'id'.
    """
    ConfigClass = CoaddPsfFitConfig
    _DefaultName = "CoaddPsfFit"
 
    def __init__(self, initInputs, **kwargs):
        super().__init__(initInputs=initInputs, **kwargs)
        self.makeSubtask("fit_coadd_psf")
 
    def runQuantum(self, butlerQC, inputRefs, outputRefs):
        inputs = butlerQC.get(inputRefs)
        id_tp = self.config.idGenerator.apply(butlerQC.quantum.dataId).catalog_id
        dataId = inputRefs.cat_meas.dataId
 
        if self.config.image_type == "future":
            coaddDataRef = inputRefs.coadd
            exposure = inputs.pop('coadd').to_legacy()
        elif self.config.use_cell_coadds:
            coaddDataRef = inputRefs.coadd_cell
            multiple_cell_coadd = inputs.pop('coadd_cell')
            background = inputs.pop('background')
            exposure = multiple_cell_coadd.stitch().asExposure()
            exposure.image -= background.getImage()
        else:
            coaddDataRef = inputRefs.coadd
            exposure = inputs.pop('coadd')
 
        for dataRef in (coaddDataRef,):
            if dataRef.dataId != dataId:
                raise RuntimeError(f'{dataRef=}.dataId != {inputRefs.cat_meas.dataId=}')
 
        catalog = inputs.pop('cat_meas')
        catexp = CatalogExposurePsf(
            catalog=catalog, exposure=exposure, dataId=dataId, id_tract_patch=id_tp,
        )
        assert not inputs, "runQuantum got more inputs than expected"
        outputs = self.run(catexp=catexp)
        butlerQC.put(outputs, outputRefs)
 
    def run(self, catexp: CatalogExposurePsf) -> pipeBase.Struct:
        """Fit a PSF model at the location of sources in a coadd.
 
        Parameters
        ----------
        catexp : `typing.List [CatalogExposurePsf]`
            A list of catalog-exposure pairs in a given band.
 
        Returns
        -------
        retStruct : `lsst.pipe.base.Struct`
            A struct with a cat_output attribute containing the output
            measurement catalog.
 
        Notes
        -----
        Subtasks may have further requirements; see `CoaddPsfFitSubTask.run`.
        """
        cat_output = self.fit_coadd_psf.run(catexp).output
        retStruct = pipeBase.Struct(cat_output=cat_output)
        return retStruct