Create pipelineTask to exclude sky sources

python/lsst/ap/association/filterDiaSourceCatalog.py

@@ -0,0 +1,102 @@
+# This file is part of ap_association
+#
+# 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__ = (
+    "FilterDiaSourceCatalogConnections",
+    "FilterDiaSourceCatalogConfig",
+    "FilterDiaSourceCatalogTask",
+)
+
+import lsst.pex.config as pexConfig
+import lsst.pipe.base as pipeBase
+import lsst.pipe.base.connectionTypes as connTypes
+from lsst.meas.base import DetectorVisitIdGeneratorConfig
+from lsst.utils.timer import timeMethod
+
+
+class FilterDiaSourceCatalogConnections(
+    pipeBase.PipelineTaskConnections,
+    dimensions=("instrument", "visit", "detector"),
+    defaultTemplates={"coaddName": "deep", "fakesType": ""},
+):
+    """Connections class for FilterDiaSourceCatalogTask."""
+
+    diaSourceCat = connTypes.Input(
+        doc="Catalog of DiaSources produced during image differencing.",
+        name="{fakesType}{coaddName}Diff_diaSrc",
+        storageClass="SourceCatalog",
+        dimensions=("instrument", "visit", "detector"),
+    )
+
+    filteredDiaSourceCat = connTypes.Output(
+        doc="Output catalog of DiaSources after filtering.",
+        name="{fakesType}{coaddName}Diff_candidateDiaSrc",
+        storageClass="SourceCatalog",
+        dimensions=("instrument", "visit", "detector"),
+    )
+
+
+class FilterDiaSourceCatalogConfig(
+    pipeBase.PipelineTaskConfig, pipelineConnections=FilterDiaSourceCatalogConnections
+):
+    """Config class for FilterDiaSourceCatalogTask."""
+
+    doRemoveSkySources = pexConfig.Field(
+        dtype=bool,
+        default=False,
+        doc="Input DiaSource catalog contains SkySources that should be "
+        "removed before storing the output DiaSource catalog.",
+    )
+
+    idGenerator = DetectorVisitIdGeneratorConfig.make_field()
+
+
+class FilterDiaSourceCatalogTask(pipeBase.PipelineTask):
+    """Filter out sky sources from a DiaSource catalog."""
+
+    ConfigClass = FilterDiaSourceCatalogConfig
+    _DefaultName = "filterDiaSourceCatalog"
+
+    def runQuantum(self, butlerQC, inputRefs, outputRefs):
+        inputs = butlerQC.get(inputRefs)
+        idGenerator = self.config.idGenerator.apply(butlerQC.quantum.dataId)
+        inputs["ccdVisitId"] = idGenerator.catalog_id
+
+        outputs = self.run(**inputs)
+
+        butlerQC.put(outputs, outputRefs)
+
+    @timeMethod
+    def run(self, diaSourceCat):
+        """Filter sky sources from the supplied DiaSource catalog.
+
+        Parameters
+        ----------
+        diaSourceCat : `lsst.afw.table.SourceCatalog`
+            Catalog of sources measured on the difference image.
+
+        Returns
+        -------
+        results : `lsst.pipe.base.Struct`
+            Results struct with components.
+        """
+        filteredDiaSourceCat = diaSourceCat[~diaSourceCat["sky_source"]]
+        return pipeBase.Struct(filteredDiaSourceCat=filteredDiaSourceCat)