DM-32034: Create MatchTractCatalog(Probabilistic) Tasks

python/lsst/pipe/tasks/match_tract_catalog.py

@@ -0,0 +1,222 @@
+# 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__ = [
+    'MatchTractCatalogSubConfig', 'MatchTractCatalogSubTask',
+    'MatchTractCatalogConfig', 'MatchTractCatalogTask'
+]
+
+import lsst.afw.geom as afwGeom
+import lsst.pex.config as pexConfig
+import lsst.pipe.base as pipeBase
+import lsst.pipe.base.connectionTypes as cT
+from lsst.skymap import BaseSkyMap
+
+from abc import ABC, abstractmethod
+
+import pandas as pd
+from typing import Tuple, Set
+
+
+MatchTractCatalogBaseTemplates = {
+    "name_input_cat_ref": "truth_summary",
+    "name_input_cat_target": "objectTable_tract",
+}
+
+
+class MatchTractCatalogConnections(
+    pipeBase.PipelineTaskConnections,
+    dimensions=("tract", "skymap"),
+    defaultTemplates=MatchTractCatalogBaseTemplates,
+):
+    cat_ref = cT.Input(
+        doc="Reference object catalog to match from",
+        name="{name_input_cat_ref}",
+        storageClass="DataFrame",
+        dimensions=("tract", "skymap"),
+        deferLoad=True,
+    )
+    cat_target = cT.Input(
+        doc="Target object catalog to match",
+        name="{name_input_cat_target}",
+        storageClass="DataFrame",
+        dimensions=("tract", "skymap"),
+        deferLoad=True,
+    )
+    skymap = cT.Input(
+        doc="Input definition of geometry/bbox and projection/wcs for coadded exposures",
+        name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME,
+        storageClass="SkyMap",
+        dimensions=("skymap",),
+    )
+    cat_output_ref = cT.Output(
+        doc="Reference matched catalog with indices of target matches",
+        name="match_ref_{name_input_cat_ref}_{name_input_cat_target}",
+        storageClass="DataFrame",
+        dimensions=("tract", "skymap"),
+    )
+    cat_output_target = cT.Output(
+        doc="Target matched catalog with indices of reference matches",
+        name="match_target_{name_input_cat_ref}_{name_input_cat_target}",
+        storageClass="DataFrame",
+        dimensions=("tract", "skymap"),
+    )
+
+
+class MatchTractCatalogSubConfig(pexConfig.Config):
+    """Config class for the MatchTractCatalogSubTask to define methods returning
+    values that depend on multiple config settings.
+    """
+    @property
+    @abstractmethod
+    def columns_in_ref(self) -> Set[str]:
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def columns_in_target(self) -> Set[str]:
+        raise NotImplementedError()
+
+
+class MatchTractCatalogSubTask(pipeBase.Task, ABC):
+    """An abstract interface for subtasks of MatchTractCatalogTask to match
+    two tract object catalogs.
+
+    Parameters
+    ----------
+    **kwargs
+        Additional arguments to be passed to the `lsst.pipe.base.Task`
+        constructor.
+    """
+    ConfigClass = MatchTractCatalogSubConfig
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    @abstractmethod
+    def run(
+        self,
+        catalog_ref: pd.DataFrame,
+        catalog_target: pd.DataFrame,
+        wcs: afwGeom.SkyWcs = None,
+    ) -> pipeBase.Struct:
+        """Match sources in a reference tract catalog with a target catalog.
+
+        Parameters
+        ----------
+        catalog_ref : `pandas.DataFrame`
+            A reference catalog to match objects/sources from.
+        catalog_target : `pandas.DataFrame`
+            A target catalog to match reference objects/sources to.
+        wcs : `lsst.afw.image.SkyWcs`
+            A coordinate system to convert catalog positions to sky coordinates.
+
+        Returns
+        -------
+        retStruct : `lsst.pipe.base.Struct`
+            A struct with output_ref and output_target attribute containing the
+            output matched catalogs.
+        """
+        raise NotImplementedError()
+
+
+class MatchTractCatalogConfig(
+    pipeBase.PipelineTaskConfig,
+    pipelineConnections=MatchTractCatalogConnections,
+):
+    """Configure a MatchTractCatalogTask, including a configurable matching subtask.
+    """
+    match_tract_catalog = pexConfig.ConfigurableField(
+        target=MatchTractCatalogSubTask,
+        doc="Task to match sources in a reference tract catalog with a target catalog",
+    )
+
+    def get_columns_in(self) -> Tuple[Set, Set]:
+        """Get the set of input columns required for matching.
+
+        Returns
+        -------
+        columns_ref : `set` [`str`]
+            The set of required input catalog column names.
+        columns_target : `set` [`str`]
+            The set of required target catalog column names.
+        """
+        try:
+            columns_ref, columns_target = (self.match_tract_catalog.columns_in_ref,
+                                           self.match_tract_catalog.columns_in_target)
+        except AttributeError as err:
+            raise RuntimeError(f'{__class__}.match_tract_catalog must have columns_in_ref and'
+                               f' columns_in_target attributes: {err}') from None
+        return set(columns_ref), set(columns_target)
+
+
+class MatchTractCatalogTask(pipeBase.PipelineTask):
+    """Match sources in a reference tract catalog with those in a target catalog.
+    """
+    ConfigClass = MatchTractCatalogConfig
+    _DefaultName = "MatchTractCatalog"
+
+    def __init__(self, initInputs, **kwargs):
+        super().__init__(initInputs=initInputs, **kwargs)
+        self.makeSubtask("match_tract_catalog")
+
+    def runQuantum(self, butlerQC, inputRefs, outputRefs):
+        inputs = butlerQC.get(inputRefs)
+        columns_ref, columns_target = self.config.get_columns_in()
+        skymap = inputs.pop("skymap")
+
+        outputs = self.run(
+            catalog_ref=inputs['cat_ref'].get(parameters={'columns': columns_ref}),
+            catalog_target=inputs['cat_target'].get(parameters={'columns': columns_target}),
+            wcs=skymap[butlerQC.quantum.dataId["tract"]].wcs,
+        )
+        butlerQC.put(outputs, outputRefs)
+
+    def run(
+        self,
+        catalog_ref: pd.DataFrame,
+        catalog_target: pd.DataFrame,
+        wcs: afwGeom.SkyWcs = None,
+    ) -> pipeBase.Struct:
+        """Match sources in a reference tract catalog with a target catalog.
+
+        Parameters
+        ----------
+        catalog_ref : `pandas.DataFrame`
+            A reference catalog to match objects/sources from.
+        catalog_target : `pandas.DataFrame`
+            A target catalog to match reference objects/sources to.
+        wcs : `lsst.afw.image.SkyWcs`
+            A coordinate system to convert catalog positions to sky coordinates,
+            if necessary.
+
+        Returns
+        -------
+        retStruct : `lsst.pipe.base.Struct`
+            A struct with output_ref and output_target attribute containing the
+            output matched catalogs.
+        """
+        output = self.match_tract_catalog.run(catalog_ref, catalog_target, wcs=wcs)
+        if output.exceptions:
+            self.log.warn('Exceptions: %s', output.exceptions)
+        retStruct = pipeBase.Struct(cat_output_ref=output.cat_output_ref,
+                                    cat_output_target=output.cat_output_target)
+        return retStruct

python/lsst/pipe/tasks/match_tract_catalog_probabilistic.py

@@ -0,0 +1,95 @@
+# 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__ = [
+    'MatchTractCatalogProbabilisticConfig', 'MatchTractCatalogProbabilisticTask',
+]
+
+import lsst.afw.geom as afwGeom
+from lsst.meas.astrom.match_probabilistic_task import MatchProbabilisticTask
+from lsst.meas.astrom.matcher_probabilistic import MatchProbabilisticConfig
+import lsst.pipe.base as pipeBase
+
+from .match_tract_catalog import MatchTractCatalogSubConfig, MatchTractCatalogSubTask
+
+import pandas as pd
+from typing import Set
+
+
+class MatchTractCatalogProbabilisticConfig(MatchProbabilisticConfig, MatchTractCatalogSubConfig):
+    """Config class for the MatchTractCatalogSubTask to define methods returning
+    values that depend on multiple config settings.
+    """
+
+    @property
+    def columns_in_ref(self) -> Set[str]:
+        return super().columns_in_ref
+
+    @property
+    def columns_in_target(self) -> Set[str]:
+        return super().columns_in_target
+
+
+class MatchTractCatalogProbabilisticTask(MatchProbabilisticTask, MatchTractCatalogSubTask):
+    """An abstract interface for subtasks of MatchTractCatalogTask to match
+    two tract object catalogs.
+
+    Parameters
+    ----------
+    **kwargs
+        Additional arguments to be passed to the `lsst.pipe.base.Task`
+        constructor.
+    """
+    ConfigClass = MatchTractCatalogProbabilisticConfig
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def run(
+        self,
+        catalog_ref: pd.DataFrame,
+        catalog_target: pd.DataFrame,
+        wcs: afwGeom.SkyWcs = None,
+    ) -> pipeBase.Struct:
+        """Match sources in a reference tract catalog with a target catalog.
+
+        Parameters
+        ----------
+        catalog_ref : `pandas.DataFrame`
+            A reference catalog to match objects/sources from.
+        catalog_target : `pandas.DataFrame`
+            A target catalog to match reference objects/sources to.
+        wcs : `lsst.afw.image.SkyWcs`
+            A coordinate system to convert catalog positions to sky coordinates.
+            Only needed if `config.coords_ref_to_convert` is used to convert
+            reference catalog sky coordinates to pixel positions.
+
+        Returns
+        -------
+        retStruct : `lsst.pipe.base.Struct`
+            A struct with output_ref and output_target attribute containing the
+            output matched catalogs.
+        """
+        return super().run(
+            catalog_ref=catalog_ref,
+            catalog_target=catalog_target,
+            wcs=wcs,
+        )