From dfe3531738cbda88ba0a3fdf0fa8a2bdeec2f420 Mon Sep 17 00:00:00 2001
From: Ian Sullivan <sullii@uw.edu>
Date: Fri, 8 Dec 2023 16:39:55 -0800
Subject: [PATCH] Respond to review

---
 python/lsst/ip/diffim/detectAndMeasure.py | 22 +++++++++++++---------
 tests/test_detectAndMeasure.py            |  7 +++----
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/python/lsst/ip/diffim/detectAndMeasure.py b/python/lsst/ip/diffim/detectAndMeasure.py
index 4dade02f..92569c48 100644
--- a/python/lsst/ip/diffim/detectAndMeasure.py
+++ b/python/lsst/ip/diffim/detectAndMeasure.py
@@ -143,7 +143,7 @@ class DetectAndMeasureConfig(pipeBase.PipelineTaskConfig,
     )
     badSourceFlags = lsst.pex.config.ListField(
         dtype=str,
-        doc="Do not include sources with any of these flags set in the output catalog.",
+        doc="Sources with any of these flags set are removed before writing the output catalog.",
         default=("base_PixelFlags_flag_offimage",
                  ),
     )
@@ -397,17 +397,21 @@ def removeBadSources(self, diaSources):
             The updated catalog of detected sources, with any source that has a
             flag in ``config.badSourceFlags`` set removed.
         """
-        flags = np.ones(len(diaSources), dtype=bool)
+        nBadTotal = 0
+        selector = np.ones(len(diaSources), dtype=bool)
         for flag in self.config.badSourceFlags:
             try:
-                flags &= ~diaSources[flag]
-            except Exception as e:
+                flags = diaSources[flag]
+            except KeyError as e:
                 self.log.warning("Could not apply source flag: %s", e)
-        nBad = np.count_nonzero(~flags)
-        if nBad > 0:
-            self.log.warning(f"Found and removed {nBad} unphysical sources.")
-        diaSources = diaSources[flags].copy(deep=True)
-        self.metadata.add("nRemovedBadFlaggedSources", nBad)
+                continue
+            nBad = np.count_nonzero(flags)
+            if nBad > 0:
+                self.log.info("Found and removed %d unphysical sources with flag %s.", nBad, flag)
+                selector &= ~flags
+                nBadTotal += nBad
+        diaSources = diaSources[selector].copy(deep=True)
+        self.metadata.add("nRemovedBadFlaggedSources", nBadTotal)
         return diaSources
 
     def addSkySources(self, diaSources, mask, seed):
diff --git a/tests/test_detectAndMeasure.py b/tests/test_detectAndMeasure.py
index be202969..7834b3dd 100644
--- a/tests/test_detectAndMeasure.py
+++ b/tests/test_detectAndMeasure.py
@@ -93,7 +93,7 @@ def _check_values(self, values, minValue=None, maxValue=None):
             self.assertTrue(np.all(values <= maxValue))
 
     def _setup_detection(self, doApCorr=False, doMerge=False,
-                         doSkySources=False, doForcedMeasurement=False, nSkySources=5, badSourceFlags=[]):
+                         doSkySources=False, doForcedMeasurement=False, nSkySources=5, **kwargs):
         """Setup and configure the detection and measurement PipelineTask.
 
         Parameters
@@ -117,9 +117,9 @@ def _setup_detection(self, doApCorr=False, doMerge=False,
         config.doMerge = doMerge
         config.doSkySources = doSkySources
         config.doForcedMeasurement = doForcedMeasurement
-        config.badSourceFlags = badSourceFlags
         if doSkySources:
             config.skySources.nSources = nSkySources
+        config.update(**kwargs)
         return self.detectionTask(config=config)
 
 
@@ -191,9 +191,8 @@ def test_measurements_finite(self):
         self._check_values(output.diaSources.getPsfInstFlux())
 
     def test_remove_unphysical(self):
+        """Check that sources with specified flags are removed from the catalog.
         """
-        """
-
         # Set up the simulated images
         noiseLevel = 1.
         staticSeed = 1