Add method to BasePsfDeterminerTask to downsample candidates.

python/lsst/meas/algorithms/psfDeterminer.py

@@ -22,6 +22,7 @@
 __all__ = ["BasePsfDeterminerConfig", "BasePsfDeterminerTask", "psfDeterminerRegistry"]
 
 import abc
+import numpy as np
 
 import lsst.pipe.base as pipeBase
 import lsst.pex.config as pexConfig
@@ -38,6 +39,14 @@ class BasePsfDeterminerConfig(pexConfig.Config):
         optional=True,
         check=lambda x: (x > 0) & (x % 2 == 1),
     )
+    maxCandidates = pexConfig.Field[int](
+        doc="Maximum number of candidates to consider. Will down-sample if given more.",
+        default=300,
+    )
+    downsampleRandomSeed = pexConfig.Field[int](
+        doc="Random seed to use to downsample candidates.",
+        default=98765,
+    )
 
 
 class BasePsfDeterminerTask(pipeBase.Task, metaclass=abc.ABCMeta):
@@ -63,6 +72,31 @@ class BasePsfDeterminerTask(pipeBase.Task, metaclass=abc.ABCMeta):
     def __init__(self, config, schema=None, **kwds):
         pipeBase.Task.__init__(self, config=config, **kwds)
 
+    def downsampleCandidates(self, inputCandidateList):
+        """Down-sample candidates from the input candidate list.
+
+        Parameters
+        ----------
+        inputCandidateList : `list` [`lsst.meas.algorithms.PsfCandidate`]
+            Input candidate list.
+
+        Returns
+        -------
+        outputCandidateList : `list` [`lsst.meas.algorithms.PsfCandidate`]
+            Down-selected candidate list.
+        """
+        if len(inputCandidateList) <= self.config.maxCandidates:
+            return inputCandidateList
+
+        rng = np.random.RandomState(seed=self.config.downsampleRandomSeed)
+
+        selection = rng.choice(len(inputCandidateList), size=self.config.maxCandidates, replace=False)
+        selection = np.sort(selection)
+
+        outputCandidateList = [inputCandidateList[index] for index in selection]
+
+        return outputCandidateList
+
     @abc.abstractmethod
     def determinePsf(self, exposure, psfCandidateList, metadata=None, flagKey=None):
         """Determine a PSF model.

tests/test_psfDetermination.py

@@ -436,6 +436,33 @@ def testShowPsf(self):
         mos = showPsf(psf, display=testDisplay)
         self.assertTrue(len(mos.images) > 0)
 
+    def testDownsampleBase(self):
+        """Test that the downsampleCandidates function works.
+        """
+        self.setupDeterminer()
+
+        # Note that the downsampleCandidates function is designed to work
+        # with a list of psf candidates, it can work with any list.
+        # For these tests we use a list of integers which allows easier
+        # testing that the sort order is maintained.
+
+        # Try with no downsampling.
+        inputList = list(np.arange(100))
+        candidateList = self.psfDeterminer.downsampleCandidates(inputList)
+        np.testing.assert_array_equal(candidateList, inputList)
+
+        # And with downsampling.
+        inputList = list(np.arange(500))
+        candidateList = self.psfDeterminer.downsampleCandidates(inputList)
+        self.assertEqual(len(candidateList), self.psfDeterminer.config.maxCandidates)
+        np.testing.assert_array_equal(np.sort(candidateList), candidateList)
+        self.assertEqual(len(np.unique(candidateList)), len(candidateList))
+
+    def testDownsamplePca(self):
+        """Test PCA determiner with downsampling.
+        """
+        self.setupDeterminer()
+
 
 class PsfCandidateTestCase(lsst.utils.tests.TestCase):
     def testNumToReject(self):