tickets/DM-47251: Added maxSelect to DonutStampsSelector

doc/versionHistory.rst

@@ -6,6 +6,14 @@
 Version History
 ##################
 
+.. _lsst.ts.wep-12.6.0:
+
+-------------
+12.6.0
+-------------
+
+* Added maxSelect config to DonutStampsSelector
+
 .. _lsst.ts.wep-12.5.0:
 
 -------------

pipelines/production/comCamRapidAnalysisPipeline.yaml

@@ -5,7 +5,7 @@ tasks:
     class: lsst.ts.wep.task.generateDonutDirectDetectTask.GenerateDonutDirectDetectTask
     config:
       donutSelector.useCustomMagLimit: True
-      donutSelector.sourceLimit: 5
+      donutSelector.sourceLimit: 20
   cutOutDonutsScienceSensorGroupTask:
     class: lsst.ts.wep.task.cutOutDonutsScienceSensorTask.CutOutDonutsScienceSensorTask
     config:
@@ -20,6 +20,7 @@ tasks:
       estimateZernikes.maxNollIndex: 28
       estimateZernikes.saveHistory: False
       estimateZernikes.maskKwargs: { "doMaskBlends": False }
+      donutStampSelector.maxSelect: 5
   isr:
     class: lsst.ip.isr.IsrTaskLSST
     config:

python/lsst/ts/wep/task/donutStampSelectorTask.py

@@ -31,6 +31,11 @@
 
 
 class DonutStampSelectorTaskConfig(pexConfig.Config):
+    maxSelect = pexConfig.Field(
+        dtype=int,
+        default=5,
+        doc="Maximum number of donut stamps to select. If -1, all are selected.",
+    )
     selectWithEntropy = pexConfig.Field(
         dtype=bool,
         default=False,
@@ -225,6 +230,10 @@ def selectStamps(self, donutStamps):
         # choose only donuts that satisfy all selected conditions
         selected = entropySelect * snSelect * fracBadPixSelect
 
+        # make sure we don't select more than maxSelect
+        if self.config.maxSelect != -1:
+            selected[np.cumsum(selected) > self.config.maxSelect] = False
+
         # store information about which donuts were selected
         # use QTable even though no units at the moment in
         # case we end up adding more later we don't have to change

tests/task/test_donutStampSelectorTask.py

@@ -100,6 +100,7 @@ def testValidateConfigs(self):
         self.assertEqual(self.OrigTask.config.useCustomSnLimit, False)
 
         # Test changing configs
+        self.config.maxSelect = 10
         self.config.selectWithEntropy = True
         self.config.selectWithSignalToNoise = False
         self.config.selectWithFracBadPixels = False
@@ -108,6 +109,7 @@ def testValidateConfigs(self):
         self.config.maxFracBadPixels = 0.2
         self.ModifiedTask = DonutStampSelectorTask(config=self.config, name="Mod Task")
 
+        self.assertEqual(self.ModifiedTask.config.maxSelect, 10)
         self.assertEqual(self.ModifiedTask.config.selectWithEntropy, True)
         self.assertEqual(self.ModifiedTask.config.selectWithSignalToNoise, False)
         self.assertEqual(self.ModifiedTask.config.selectWithFracBadPixels, False)
@@ -173,7 +175,7 @@ def testSelectStamps(self):
         for v in donutsQuality["SN"][donutsQuality["SN_SELECT"]]:
             self.assertLess(minSignalToNoise, v)
 
-        # finally turn all selections off and make sure everything is selected
+        # turn all selections off and make sure everything is selected
         self.config.selectWithEntropy = False
         self.config.selectWithSignalToNoise = False
         self.config.selectWithFracBadPixels = False
@@ -184,6 +186,15 @@ def testSelectStamps(self):
         self.assertEqual(np.sum(selection.donutsQuality["FRAC_BAD_PIX_SELECT"]), 3)
         self.assertEqual(np.sum(selection.donutsQuality["FINAL_SELECT"]), 3)
 
+        # set maxSelect = 1 and make sure the final selection is only 1
+        self.config.maxSelect = 1
+        task = DonutStampSelectorTask(config=self.config, name="maxSelect=1")
+        selection = task.selectStamps(donutStampsIntra)
+        self.assertEqual(np.sum(selection.donutsQuality["ENTROPY_SELECT"]), 3)
+        self.assertEqual(np.sum(selection.donutsQuality["SN_SELECT"]), 3)
+        self.assertEqual(np.sum(selection.donutsQuality["FRAC_BAD_PIX_SELECT"]), 3)
+        self.assertEqual(np.sum(selection.donutsQuality["FINAL_SELECT"]), 1)
+
     def testTaskRun(self):
         donutStampsIntra = self.butler.get(
             "donutStampsIntra", dataId=self.dataIdExtra, collections=[self.runName]