Add matching of psf_stars to stars

Propagate psf and aperture correction fields from psf_stars to stars,
without overwriting anything else in stars.
Some of this code was lifted from `CalibrateTask.copyIcSourceFields`.

python/lsst/pipe/tasks/calibrateImage.py

@@ -365,6 +365,17 @@ def __init__(self, initial_stars_schema=None, **kwargs):
         # star measurement subtasks
         if initial_stars_schema is None:
             initial_stars_schema = afwTable.SourceTable.makeMinimalSchema()
+
+        # These fields let us track which sources were used for psf and
+        # aperture correction calculations.
+        self.psf_fields = ("calib_psf_candidate", "calib_psf_used", "calib_psf_reserved",
+                           # TODO DM-39203: these can be removed once apcorr is gone.
+                           "apcorr_slot_CalibFlux_used", "apcorr_base_GaussianFlux_used",
+                           "apcorr_base_PsfFlux_used")
+        for field in self.psf_fields:
+            item = self.psf_schema.find(field)
+            initial_stars_schema.addField(item.getField())
+
         afwTable.CoordKey.addErrorFields(initial_stars_schema)
         self.makeSubtask("star_detection", schema=initial_stars_schema)
         self.makeSubtask("star_mask_streaks")
@@ -638,6 +649,60 @@ def _find_stars(self, exposure, background):
         else:
             return result.sourceCat
 
+    def _match_psf_stars(self, psf_stars, stars):
+        """Match calibration stars to psf stars, to identify which were psf
+        candidates, and which were used or reserved during psf measurement.
+
+        Parameters
+        ----------
+        psf_stars : `lsst.afw.table.SourceCatalog`
+            PSF candidate stars that were sent to the psf determiner. Used to
+            populate psf-related flag fields.
+        stars : `lsst.afw.table.SourceCatalog`
+            Stars that will be used for calibration; psf-related fields will
+            be updated in-place.
+
+        Notes
+        -----
+        This code was adapted from CalibrateTask.copyIcSourceFields().
+        """
+        control = afwTable.MatchControl()
+        # Return all matched objects, to separate blends.
+        control.findOnlyClosest = False
+        matches = afwTable.matchXy(psf_stars, stars, 3.0, control)
+        deblend_key = stars.schema["deblend_nChild"].asKey()
+        matches = [m for m in matches if m[1].get(deblend_key) == 0]
+
+        # Because we had to allow multiple matches to handle parents, we now
+        # need to prune to the best (closest) matches.
+        # Closest matches is a dict of psf_stars source ID to Match record
+        # (psf_stars source, sourceCat source, distance in pixels).
+        best = {}
+        for match0, match1, d in matches:
+            id0 = match0.getId()
+            match = best.get(id0)
+            if match is None or d <= match[2]:
+                best[id0] = (match0, match1, d)
+        matches = list(best.values())
+        ids = np.array([(match0.getId(), match1.getId()) for match0, match1, d in matches]).T
+
+        # Check that no stars sources are listed twice; we already know
+        # that each match has a unique psf_stars id, due to using as the key
+        # in best above.
+        n_matches = len(matches)
+        n_unique = len(set(m[1].getId() for m in matches))
+        if n_unique != n_matches:
+            self.log.warning("%d psf_stars matched only %d stars; ",
+                             n_matches, n_unique)
+
+        # The indices of the IDs, so we can update the flag fields as arrays.
+        idx0 = np.searchsorted(psf_stars["id"], ids[0])
+        idx1 = np.searchsorted(stars["id"], ids[1])
+        for field in self.psf_fields:
+            result = np.zeros(len(stars), dtype=bool)
+            result[idx0] = psf_stars[field][idx1]
+            stars[field] = result
+
     def _fit_astrometry(self, exposure, stars):
         """Fit an astrometric model to the data and return the reference
         matches used in the fit, and the fitted WCS.

tests/test_calibrateImage.py

@@ -297,6 +297,31 @@ def test_photometry(self):
         self.assertFloatsAlmostEqual(stars['slot_PsfFlux_flux'], self.truth_cat['truth_flux'][idx], rtol=0.1)
         self.assertFloatsAlmostEqual(stars['slot_PsfFlux_mag'], self.truth_cat['truth_mag'][idx], rtol=0.01)
 
+    def test_match_psf_stars(self):
+        """Test that _match_psf_stars() flags the correct stars as psf stars
+        and candidates.
+        """
+        calibrate = CalibrateImageTask(config=self.config)
+        psf_stars, background, candidates = calibrate._compute_psf(self.exposure)
+        calibrate._measure_aperture_correction(self.exposure, psf_stars)
+        stars = calibrate._find_stars(self.exposure, background)
+
+        # There should be no psf-related flags set at first.
+        self.assertEqual(stars["calib_psf_candidate"].sum(), 0)
+        self.assertEqual(stars["calib_psf_used"].sum(), 0)
+        self.assertEqual(stars["calib_psf_reserved"].sum(), 0)
+
+        calibrate._match_psf_stars(psf_stars, stars)
+
+        # Sort in order of brightness; the psf stars are the 3 brightest.
+        stars.sort(stars.getPsfFluxSlot().getMeasKey())
+        # sort() above leaves the catalog non-contiguous.
+        stars = stars.copy(deep=True)
+        np.testing.assert_array_equal(stars["calib_psf_candidate"], [False, False, True, True, True])
+        np.testing.assert_array_equal(stars["calib_psf_used"], [False, False, True, True, True])
+        # Too few sources to reserve any in these tests.
+        self.assertEqual(stars["calib_psf_reserved"].sum(), 0)
+
 
 class CalibrateImageTaskRunQuantumTests(lsst.utils.tests.TestCase):
     """Tests of ``CalibrateImageTask.runQuantum``, which need a test butler,