DM-40798: Exclude sky sources from RB classification

python/lsst/ap/association/__init__.py

@@ -19,11 +19,12 @@
 # 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 .trailedSourceFilter import *
 from .association import *
 from .diaForcedSource import *
+from .filterDiaSourceCatalog import *
 from .loadDiaCatalogs import *
 from .packageAlerts import *
 from .diaPipe import *
+from .trailedSourceFilter import *
 from .transformDiaSourceCatalog import *
+from .version import *

python/lsst/ap/association/filterDiaSourceCatalog.py

@@ -0,0 +1,125 @@
+# 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__ = (
+    "FilterDiaSourceCatalogConfig",
+    "FilterDiaSourceCatalogTask",
+)
+
+import numpy as np
+
+from lsst.afw.table import SourceCatalog
+import lsst.pex.config as pexConfig
+import lsst.pipe.base as pipeBase
+import lsst.pipe.base.connectionTypes as connTypes
+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"),
+    )
+
+    rejectedDiaSources = connTypes.Output(
+        doc="Optional output storing all the rejected DiaSources.",
+        name="{fakesType}{coaddName}Diff_rejectedDiaSrc",
+        storageClass="SourceCatalog",
+        dimensions={"instrument", "visit", "detector"},
+    )
+
+    def __init__(self, *, config=None):
+        super().__init__(config=config)
+        if not self.config.doWriteRejectedSources:
+            self.outputs.remove("rejectedDiaSources")
+
+
+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.",
+    )
+
+    doWriteRejectedSources = pexConfig.Field(
+        dtype=bool,
+        default=True,
+        doc="Store the output DiaSource catalog containing all the rejected "
+        "sky sources."
+    )
+
+
+class FilterDiaSourceCatalogTask(pipeBase.PipelineTask):
+    """Filter out sky sources from a DiaSource catalog."""
+
+    ConfigClass = FilterDiaSourceCatalogConfig
+    _DefaultName = "filterDiaSourceCatalog"
+
+    @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
+        -------
+        filterResults : `lsst.pipe.base.Struct`
+
+            ``filteredDiaSourceCat`` : `lsst.afw.table.SourceCatalog`
+                The catalog of filtered sources.
+            ``rejectedDiaSources`` : `lsst.afw.table.SourceCatalog`
+                The catalog of rejected sources.
+        """
+        rejectedSkySources = None
+        if self.config.doRemoveSkySources:
+            sky_source_column = diaSourceCat["sky_source"]
+            num_sky_sources = np.sum(sky_source_column)
+            rejectedSkySources = diaSourceCat[sky_source_column].copy(deep=True)
+            diaSourceCat = diaSourceCat[~sky_source_column].copy(deep=True)
+            self.log.info(f"Filtered {num_sky_sources} sky sources.")
+        if not rejectedSkySources:
+            rejectedSkySources = SourceCatalog(diaSourceCat.getSchema())
+        filterResults = pipeBase.Struct(filteredDiaSourceCat=diaSourceCat,
+                                        rejectedDiaSources=rejectedSkySources)
+        return filterResults

python/lsst/ap/association/transformDiaSourceCatalog.py

@@ -49,7 +49,7 @@ class TransformDiaSourceCatalogConnections(pipeBase.PipelineTaskConnections,
     )
     diaSourceCat = connTypes.Input(
         doc="Catalog of DiaSources produced during image differencing.",
-        name="{fakesType}{coaddName}Diff_diaSrc",
+        name="{fakesType}{coaddName}Diff_candidateDiaSrc",
         storageClass="SourceCatalog",
         dimensions=("instrument", "visit", "detector"),
     )

tests/test_filterDiaSourceCatalog.py

@@ -0,0 +1,81 @@
+# 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/>.
+
+import unittest
+
+from lsst.ap.association.filterDiaSourceCatalog import (FilterDiaSourceCatalogConfig,
+                                                        FilterDiaSourceCatalogTask)
+import lsst.geom as geom
+import lsst.meas.base.tests as measTests
+import lsst.utils.tests
+
+
+class TestFilterDiaSourceCatalogTask(unittest.TestCase):
+
+    def setUp(self):
+        self.nSources = 10
+        self.nSkySources = 5
+        self.yLoc = 100
+        self.expId = 4321
+        self.bbox = geom.Box2I(geom.Point2I(0, 0),
+                               geom.Extent2I(1024, 1153))
+        dataset = measTests.TestDataset(self.bbox)
+        for srcIdx in range(self.nSources):
+            dataset.addSource(10000.0, geom.Point2D(srcIdx, self.yLoc))
+        schema = dataset.makeMinimalSchema()
+        schema.addField("sky_source", type="Flag", doc="Sky objects.")
+        _, self.diaSourceCat = dataset.realize(10.0, schema, randomSeed=1234)
+        self.diaSourceCat[0:self.nSkySources]["sky_source"] = True
+        self.config = FilterDiaSourceCatalogConfig()
+
+    def test_run_without_filter(self):
+        self.config.doRemoveSkySources = False
+        self.config.doWriteRejectedSources = True
+        filterDiaSourceCatalogTask = FilterDiaSourceCatalogTask(config=self.config)
+        result = filterDiaSourceCatalogTask.run(self.diaSourceCat)
+        self.assertEqual(len(result.filteredDiaSourceCat), len(self.diaSourceCat))
+        self.assertEqual(len(result.rejectedDiaSources), 0)
+        self.assertEqual(len(self.diaSourceCat), self.nSources)
+
+    def test_run_with_filter(self):
+        self.config.doRemoveSkySources = True
+        self.config.doWriteRejectedSources = True
+        filterDiaSourceCatalogTask = FilterDiaSourceCatalogTask(config=self.config)
+        result = filterDiaSourceCatalogTask.run(self.diaSourceCat)
+        nExpectedFilteredSources = self.nSources - self.nSkySources
+        self.assertEqual(len(result.filteredDiaSourceCat),
+                         len(self.diaSourceCat[~self.diaSourceCat['sky_source']]))
+        self.assertEqual(len(result.filteredDiaSourceCat), nExpectedFilteredSources)
+        self.assertEqual(len(result.rejectedDiaSources), self.nSkySources)
+        self.assertEqual(len(self.diaSourceCat), self.nSources)
+
+
+class MemoryTester(lsst.utils.tests.MemoryTestCase):
+    pass
+
+
+def setup_module(module):
+    lsst.utils.tests.init()
+
+
+if __name__ == "__main__":
+    lsst.utils.tests.init()
+    unittest.main()