Update diaCatalogSS and tests to new SourceSelector API

python/lsst/ip/diffim/diaCatalogSourceSelector.py

@@ -22,7 +22,6 @@
 
 import numpy as np
 
-from lsst.afw.table import SourceCatalog
 from lsst.pipe.base import Struct
 import lsst.pex.config as pexConfig
 import lsst.afw.display.ds9 as ds9
@@ -108,7 +107,8 @@ def __call__(self, source):
 ## @}
 
 
-class DiaCatalogSourceSelectorTask(measAlg.BaseStarSelectorTask):
+@pexConfig.registerConfigurable("diaCatalog", measAlg.sourceSelectorRegistry)
+class DiaCatalogSourceSelectorTask(measAlg.BaseSourceSelectorTask):
     """!Select sources for Kernel candidates
 
     @anchor DiaCatalogSourceSelectorTask_
@@ -166,16 +166,27 @@ def DebugInfo(name):
     ConfigClass = DiaCatalogSourceSelectorConfig
     usesMatches = True  # selectStars uses (requires) its matches argument
 
-    def selectStars(self, exposure, sourceCat, matches=None):
-        """Select sources for Kernel candidates
-
-        @param[in] exposure  the exposure containing the sources
-        @param[in] sourceCat  catalog of sources that may be stars (an lsst.afw.table.SourceCatalog)
-        @param[in] matches  a match vector as produced by meas_astrom; required
-                            (defaults to None to match the StarSelector API and improve error handling)
-
-        @return an lsst.pipe.base.Struct containing:
-        - starCat  a list of sources to be used as kernel candidates
+    def selectSources(self, sourceCat, matches=None, exposure=None):
+        """Return a selection of sources for Kernel candidates.
+
+        Parameters:
+        -----------
+        sourceCat : `lsst.afw.table.SourceCatalog`
+            Catalog of sources to select from.
+            This catalog must be contiguous in memory.
+        matches : `list` of `lsst.afw.table.ReferenceMatch`
+             A match vector as produced by meas_astrom.
+        exposure : `lsst.afw.image.Exposure` or None
+            The exposure the catalog was built from; used for debug display.
+
+        Return
+        ------
+        struct : `lsst.pipe.base.Struct`
+            The struct contains the following data:
+
+            - selected : `array` of `bool``
+                Boolean array of sources that were selected, same length as
+                sourceCat.
         """
         import lsstDebug
         display = lsstDebug.Info(__name__).display
@@ -195,10 +206,8 @@ def selectStars(self, exposure, sourceCat, matches=None):
         #
         isGoodSource = CheckSource(sourceCat, self.config.fluxLim, self.config.fluxMax, self.config.badFlags)
 
-        #
         # Go through and find all the acceptable candidates in the catalogue
-        #
-        starCat = SourceCatalog(sourceCat.schema)
+        selected = np.zeros(len(sourceCat), dtype=bool)
 
         if display and displayExposure:
             symbs = []
@@ -209,7 +218,7 @@ def selectStars(self, exposure, sourceCat, matches=None):
         refSchema = matches[0][0].schema
         rRefFluxField = measAlg.getRefFluxField(refSchema, "r")
         gRefFluxField = measAlg.getRefFluxField(refSchema, "g")
-        for ref, source, d in matches:
+        for i, (ref, source, d) in enumerate(matches):
             if not isGoodSource(source):
                 if display and displayExposure:
                     symbs.append("+")
@@ -233,7 +242,7 @@ def selectStars(self, exposure, sourceCat, matches=None):
                 isRightType = (self.config.selectStar and isStar) or (self.config.selectGalaxy and not isStar)
                 isRightVar = (self.config.includeVariable) or (self.config.includeVariable is isVar)
                 if isRightType and isRightVar and isRightColor:
-                    starCat.append(source)
+                    selected[i] = True
                     if display and displayExposure:
                         symbs.append("+")
                         ctypes.append(ds9.GREEN)
@@ -253,9 +262,4 @@ def selectStars(self, exposure, sourceCat, matches=None):
             if pauseAtEnd:
                 input("Continue? y[es] p[db] ")
 
-        return Struct(
-            starCat=starCat,
-        )
-
-
-measAlg.starSelectorRegistry.register("diacatalog", DiaCatalogSourceSelectorTask)
+        return Struct(selected=selected)

tests/test_diaCatalogSourceSelector.py

@@ -37,7 +37,7 @@ def setUp(self):
         schema = afwTable.SourceTable.makeMinimalSchema()
         schema.addField("test_flux", type=float)
         schema.addField("test_fluxSigma", type=float)
-        self.sourceSelector = ipDiffim.DiaCatalogSourceSelectorTask(schema=schema)
+        self.sourceSelector = ipDiffim.DiaCatalogSourceSelectorTask()
         for flag in self.sourceSelector.config.badFlags:
             schema.addField(flag, type="Flag")
         table = afwTable.SourceTable.make(schema)
@@ -89,17 +89,17 @@ def testCuts(self):
         refCat = self.makeRefCatalog()
 
         matches = self.makeMatches(refCat, self.srcCat, nSrc)
-        sources = self.sourceSelector.selectStars(self.exposure, self.srcCat, matches).starCat
+        sources = self.sourceSelector.run(self.srcCat, matches=matches, exposure=self.exposure).sourceCat
         self.assertEqual(len(sources), nSrc)
 
         # Set one of the source flags to be bad
         matches[0].second.set(self.sourceSelector.config.badFlags[0], True)
-        sources = self.sourceSelector.selectStars(self.exposure, self.srcCat, matches).starCat
+        sources = self.sourceSelector.run(self.srcCat, matches=matches, exposure=self.exposure).sourceCat
         self.assertEqual(len(sources), nSrc-1)
 
         # Set one of the ref flags to be bad
         matches[1].first.set("photometric", False)
-        sources = self.sourceSelector.selectStars(self.exposure, self.srcCat, matches).starCat
+        sources = self.sourceSelector.run(self.srcCat, matches=matches, exposure=self.exposure).sourceCat
         self.assertEqual(len(sources), nSrc-2)
 
         # Set one of the colors to be bad
@@ -108,13 +108,13 @@ def testCuts(self):
         gFluxField = getRefFluxField(refCat.schema, "g")
         gFlux = 10**(-0.4 * (grMin - 0.1)) * matches[2].first.get(rFluxField)
         matches[2].first.set(gFluxField, gFlux)
-        sources = self.sourceSelector.selectStars(self.exposure, self.srcCat, matches).starCat
+        sources = self.sourceSelector.run(self.srcCat, matches=matches, exposure=self.exposure).sourceCat
         self.assertEqual(len(sources), nSrc-3)
 
         # Set one of the types to be bad
         if self.sourceSelector.config.selectStar and not self.sourceSelector.config.selectGalaxy:
             matches[3].first.set("resolved", True)
-            sources = self.sourceSelector.selectStars(self.exposure, self.srcCat, matches).starCat
+            sources = self.sourceSelector.run(self.srcCat, matches=matches, exposure=self.exposure).sourceCat
             self.assertEqual(len(sources), nSrc-4)