DM-54275: Filter diaSources based on bbox when trailed sources unavailable

python/lsst/ap/association/filterDiaSourceCatalog.py

@@ -160,7 +160,9 @@ class FilterDiaSourceCatalogConfig(
 
 
 class FilterDiaSourceCatalogTask(pipeBase.PipelineTask):
-    """Filter sources from a DiaSource catalog."""
+    """Filter sources from a DiaSource catalog based on their trail length,
+    sky sources, and bad flags.
+    """
 
     ConfigClass = FilterDiaSourceCatalogConfig
     _DefaultName = "filterDiaSourceCatalog"
@@ -169,6 +171,12 @@ class FilterDiaSourceCatalogTask(pipeBase.PipelineTask):
     def run(self, diaSourceCat, diffImVisitInfo):
         """Filter sources from the supplied DiaSource catalog.
 
+        DiaSources are filtered based on criteria including trail length,
+        whether they are a sky source, and whether they have one or more bad
+        flags. Any source which meets a filtering criteria is removed from the
+        main output catalog and optionally saved in one or more separate output
+        catalogs.
+
         Parameters
         ----------
         diaSourceCat : `lsst.afw.table.SourceCatalog`
@@ -210,14 +218,30 @@ def run(self, diaSourceCat, diffImVisitInfo):
 
         if self.config.doTrailedSourceFilter:
             trail_mask = self._check_dia_source_trail(diaSourceCat, exposure_time)
+            bbox_mask = self._check_dia_source_trail_bbox(diaSourceCat, exposure_time)
+            num_trail_filtered = np.sum(trail_mask & ~bbox_mask)
+            num_bbox_filtered = np.sum(bbox_mask)
+            trail_mask |= bbox_mask
             longTrailedDiaSources = diaSourceCat[trail_mask].copy(deep=True)
             rejectedSources = diaSourceCat[rejected_mask].copy(deep=True)
             rejected_mask |= trail_mask
             diaSourceCat = diaSourceCat[~rejected_mask].copy(deep=True)
 
-            self.log.info("%i DiaSources exceed max_trail_length %f arcseconds per second, "
-                          "dropping from source catalog."
-                          % (len(longTrailedDiaSources), self.config.max_trail_length))
+            if num_trail_filtered > 0:
+
+                self.log.info("%i DiaSources exceed max_trail_length "
+                              "(%f arcseconds per second), dropping from "
+                              "source catalog. "
+                              % (num_trail_filtered, self.config.max_trail_length,))
+
+            if num_bbox_filtered > 0:
+                self.log.info(
+                    " %i DiaSources had no trail calculation and their bounding"
+                    " box exceeded max_trail_length (%f arcseconds per second) "
+                    "in either direction or across the diagonal, dropping from "
+                    "source catalog."
+                    % (num_bbox_filtered, self.config.max_trail_length))
+
             self.metadata.add("num_filtered", len(longTrailedDiaSources))
 
             if self.config.doWriteTrailedSources:
@@ -271,6 +295,42 @@ def _check_dia_source_trail(self, dia_sources, exposure_time):
 
         return trail_mask
 
+    def _check_dia_source_trail_bbox(self, dia_sources, exposure_time):
+        """Check bounding boxes of DiaSources with failed trail measurements.
+
+        For sources where the trail measurement flag is set (indicating the
+        trail length is unable to be measured), fall back to checking the footprint
+        bounding box dimensions against the max trail length threshold.
+
+        Parameters
+        ----------
+        dia_sources : `lsst.afw.table.SourceCatalog`
+            Input diaSources to check.
+        exposure_time : `float`
+            Exposure time from difference image.
+
+        Returns
+        -------
+        bbox_mask : `numpy.ndarray`
+            Boolean mask for diaSources whose footprint bounding box width,
+            height, or diagonal exceeds the max trail length threshold.
+        """
+        pixelScale = self._estimate_pixel_scale(dia_sources)
+        max_length_pixels = self.config.max_trail_length * exposure_time / pixelScale
+
+        trail_flag = dia_sources["ext_trailedSources_Naive_flag"]
+        bbox_mask = np.zeros(len(dia_sources), dtype=bool)
+
+        for i in np.where(trail_flag)[0]:
+            bbox = dia_sources[i].getFootprint().getBBox()
+            width = bbox.getWidth()
+            height = bbox.getHeight()
+            diagonal = np.sqrt(width ** 2 + height ** 2)
+            if diagonal >= max_length_pixels:
+                bbox_mask[i] = True
+
+        return bbox_mask
+
     def _estimate_pixel_scale(self, catalog):
         """Quickly calculate the pixel scale from catalog values
 

tests/test_filterDiaSourceCatalog.py

@@ -27,6 +27,8 @@
                                                         FilterDiaSourceReliabilityConfig,
                                                         FilterDiaSourceReliabilityTask)
 import lsst.geom as geom
+import lsst.afw.detection as afwDetect
+import lsst.afw.geom as afwGeom
 import lsst.meas.base.tests as measTests
 import lsst.utils.tests
 import lsst.afw.image as afwImage
@@ -69,6 +71,8 @@ def setUp(self):
                         "centered at DiaSource position.")
         schema.addField("ip_diffim_forced_PsfFlux_instFluxErr", type="F",
                         doc="Estimated uncertainty of ip_diffim_forced_PsfFlux_instFlux.")
+        schema.addField('ext_trailedSources_Naive_flag', type="Flag",
+                        doc="General trail measurement failure flag")
         schema.addField('ext_trailedSources_Naive_flag_off_image', type="Flag",
                         doc="Trail extends off image")
         schema.addField('ext_trailedSources_Naive_flag_suspect_long_trail',
@@ -342,6 +346,162 @@ def test_pixelScale_calculation(self):
         scale = filterDiaSourceCatalogTask._estimate_pixel_scale(self.diaSourceCat)
         self.assertEqual(self.config.estimatedPixelScale, scale)
 
+    def test_run_with_filter_centroid_flag(self):
+        """Test that sources with slot_Centroid_flag set are filtered with the
+        default badFlagList configuration.
+        """
+        self.config.doRemoveSkySources = False
+        self.config.doRemoveNegativeDirectImageSources = False
+        self.config.doWriteRejectedSkySources = False
+        self.config.doTrailedSourceFilter = False
+        # Default badFlagList includes slot_Centroid_flag
+        self.assertIn("slot_Centroid_flag", self.config.badFlagList)
+
+        # Set slot_Centroid_flag on a few sources that aren't already flagged
+        nCentroidFlagged = 3
+        # Pick sources at the end that have no other badFlagList flags
+        trail_offset = (self.nSkySources + self.nCrCenterSources
+                        + self.nFakeFlagSources + self.nNegativeSources)
+        for i in range(nCentroidFlagged):
+            self.diaSourceCat[trail_offset + i]["slot_Centroid_flag"] = True
+
+        filterDiaSourceCatalogTask = FilterDiaSourceCatalogTask(config=self.config)
+        result = filterDiaSourceCatalogTask.run(self.diaSourceCat, self.visitInfo)
+        # Default badFlagList filters crCenter + slot_Centroid_flag + high_varianceCenterAll.
+        # crCenter flags were already set for nCrCenterSources.
+        # Our added centroid flags are on previously unflagged sources.
+        nExpectedFiltered = self.nSources - self.nCrCenterSources - nCentroidFlagged
+        self.assertEqual(len(result.filteredDiaSourceCat), nExpectedFiltered)
+
+    def test_check_dia_source_trail_bbox_no_flag(self):
+        """Test that sources without ext_trailedSources_Naive_flag are not
+        flagged by the bbox check, regardless of bbox size.
+        """
+        self.config.doTrailedSourceFilter = True
+        task = FilterDiaSourceCatalogTask(config=self.config)
+        exposure_time = self.visitInfo.getExposureTime()
+        # None of the sources have the trail flag set by default
+        bbox_mask = task._check_dia_source_trail_bbox(self.diaSourceCat, exposure_time)
+        self.assertEqual(np.sum(bbox_mask), 0)
+
+    def test_check_dia_source_trail_bbox_no_flag_large_bbox(self):
+        """Test that sources without ext_trailedSources_Naive_flag are not
+        flagged by the bbox check even when the bounding box is large enough
+        that it would otherwise trigger the bbox filter.
+        """
+        self.config.doTrailedSourceFilter = True
+        task = FilterDiaSourceCatalogTask(config=self.config)
+        exposure_time = self.visitInfo.getExposureTime()
+        pixelScale = task._estimate_pixel_scale(self.diaSourceCat)
+        max_length_pixels = self.config.max_trail_length * exposure_time / pixelScale
+
+        # Give source 0 a large footprint exceeding the threshold, but do NOT
+        # set ext_trailedSources_Naive_flag — the bbox check should be skipped.
+        srcIdx = 0
+        largeSpan = afwGeom.SpanSet(
+            geom.Box2I(geom.Point2I(0, 0),
+                       geom.Extent2I(int(max_length_pixels) + 10, 5))
+        )
+        footprint = afwDetect.Footprint(largeSpan)
+        self.diaSourceCat[srcIdx].setFootprint(footprint)
+        self.assertFalse(self.diaSourceCat[srcIdx]["ext_trailedSources_Naive_flag"])
+
+        bbox_mask = task._check_dia_source_trail_bbox(self.diaSourceCat, exposure_time)
+        self.assertFalse(bbox_mask[srcIdx])
+        self.assertEqual(np.sum(bbox_mask), 0)
+
+    def test_check_dia_source_trail_bbox_flag_small_bbox(self):
+        """Test that sources with ext_trailedSources_Naive_flag set but a small bounding
+        box are not flagged.
+        """
+        self.config.doTrailedSourceFilter = True
+        task = FilterDiaSourceCatalogTask(config=self.config)
+        exposure_time = self.visitInfo.getExposureTime()
+        # Set the trail flag on a source with a small footprint (PSF-sized)
+        self.diaSourceCat[0]["ext_trailedSources_Naive_flag"] = True
+        bbox_mask = task._check_dia_source_trail_bbox(self.diaSourceCat, exposure_time)
+        # The default PSF footprint is small (~7 pixels), well below the
+        # threshold for max_trail_length * exposure_time / pixelScale
+        self.assertFalse(bbox_mask[0])
+
+    def test_check_dia_source_trail_bbox_flag_large_bbox(self):
+        """Test that sources with ext_trailedSources_Naive_flag set and a
+        large bounding box are flagged.
+        """
+        self.config.doTrailedSourceFilter = True
+        task = FilterDiaSourceCatalogTask(config=self.config)
+        exposure_time = self.visitInfo.getExposureTime()
+        pixelScale = task._estimate_pixel_scale(self.diaSourceCat)
+        max_length_pixels = self.config.max_trail_length * exposure_time / pixelScale
+
+        # Set the trail flag on a source and give it a large footprint
+        srcIdx = 0
+        self.diaSourceCat[srcIdx]["ext_trailedSources_Naive_flag"] = True
+        largeSpan = afwGeom.SpanSet(
+            geom.Box2I(geom.Point2I(0, 0),
+                       geom.Extent2I(int(max_length_pixels) + 10, 5))
+        )
+        footprint = afwDetect.Footprint(largeSpan)
+        self.diaSourceCat[srcIdx].setFootprint(footprint)
+
+        bbox_mask = task._check_dia_source_trail_bbox(self.diaSourceCat, exposure_time)
+        self.assertTrue(bbox_mask[srcIdx])
+
+    def test_check_dia_source_trail_bbox_diagonal(self):
+        """Test a source with a large diagonal but small width/height should still
+        be flagged.
+        """
+        self.config.doTrailedSourceFilter = True
+        task = FilterDiaSourceCatalogTask(config=self.config)
+        exposure_time = self.visitInfo.getExposureTime()
+        pixelScale = task._estimate_pixel_scale(self.diaSourceCat)
+        max_length_pixels = self.config.max_trail_length * exposure_time / pixelScale
+
+        srcIdx = 1
+        self.diaSourceCat[srcIdx]["ext_trailedSources_Naive_flag"] = True
+        # Create a bbox whose individual sides are below max_length_pixels
+        # but whose diagonal exceeds it.
+        side = int(max_length_pixels * 0.8)
+        largeSpan = afwGeom.SpanSet(
+            geom.Box2I(geom.Point2I(0, 0), geom.Extent2I(side, side))
+        )
+        footprint = afwDetect.Footprint(largeSpan)
+        self.diaSourceCat[srcIdx].setFootprint(footprint)
+
+        bbox_mask = task._check_dia_source_trail_bbox(self.diaSourceCat, exposure_time)
+        # diagonal = side * sqrt(2) ~ max_length_pixels * 1.13 > max_length_pixels
+        self.assertTrue(bbox_mask[srcIdx])
+
+    def test_run_with_trail_and_bbox_filter(self):
+        """Test doTrailedSourceFilter=True filters sources
+        via both trail length and bbox fallback.
+        """
+        self.config.doRemoveSkySources = False
+        self.config.badFlagList = []
+        self.config.doRemoveNegativeDirectImageSources = False
+        self.config.doWriteRejectedSkySources = False
+        self.config.doTrailedSourceFilter = True
+        task = FilterDiaSourceCatalogTask(config=self.config)
+        exposure_time = self.visitInfo.getExposureTime()
+        pixelScale = task._estimate_pixel_scale(self.diaSourceCat)
+        max_length_pixels = self.config.max_trail_length * exposure_time / pixelScale
+
+        # Set the trail flag and give a large footprint to a source that
+        # wouldn't be caught by the trail length check (first source).
+        srcIdx = 0
+        self.diaSourceCat[srcIdx]["ext_trailedSources_Naive_flag"] = True
+        largeSpan = afwGeom.SpanSet(
+            geom.Box2I(geom.Point2I(0, 0),
+                       geom.Extent2I(int(max_length_pixels) + 10, 5))
+        )
+        footprint = afwDetect.Footprint(largeSpan)
+        self.diaSourceCat[srcIdx].setFootprint(footprint)
+
+        result = task.run(self.diaSourceCat, self.visitInfo)
+        # The bbox-flagged source should have been removed
+        nExpectedFiltered = self.nSources - self.nFilteredTrailedSources - 1
+        self.assertEqual(len(result.filteredDiaSourceCat), nExpectedFiltered)
+
 
 class MemoryTester(lsst.utils.tests.MemoryTestCase):
     pass