DM-55432

python/lsst/ap/pipe/createApFakes.py

@@ -390,6 +390,52 @@ class CreateVisitDetectorFakesConfig(
         dtype=float,
         default=26,
     )
+    doAddBlendedFakes = pexConfig.Field(
+        doc="Whether to add blended fakes to the visit detector.",
+        dtype=bool,
+        default=False,
+    )
+    fracBlendedFakes = pexConfig.RangeField(
+        doc="Fraction of blended fakes to add to the visit detector.",
+        dtype=float,
+        default=0.5,
+        min=0,
+        max=1,
+    )
+    fracHostedBlendedFakes = pexConfig.RangeField(
+        doc="Fraction of blended fakes that are hosted by stars.",
+        dtype=float,
+        default=0.5,
+        min=0,
+        max=1,
+    )
+    blendedFakeMagOffset = pexConfig.Field(
+        doc="Standard deviation of magnitude offset for blended fakes.",
+        dtype=float,
+        default=0.5,
+    )
+    blendedFakeMaxOffset = pexConfig.Field(
+        doc="Maximum positional offset for blended fakes in arcseconds.",
+        dtype=float,
+        default=3.0,
+    )
+    blendedFakeMinOffset = pexConfig.Field(
+        doc="Minimum positional offset for blended fakes in arcseconds.",
+        dtype=float,
+        default=0.2,
+    )
+    maxHostedBlendedFakesPerHost = pexConfig.RangeField(
+        doc="Maximum number of hosted blended fakes assigned to the same host in one detector.",
+        dtype=int,
+        default=2,
+        min=1,
+    )
+    maxHostedBlendedFakesTotal = pexConfig.RangeField(
+        doc="Hard cap on number of hosted blended fakes per detector. Set to -1 to disable.",
+        dtype=int,
+        default=-1,
+        min=-1,
+    )
 
 
 class CreateVisitDetectorFakesTask(PipelineTask):
@@ -404,6 +450,39 @@ class CreateVisitDetectorFakesTask(PipelineTask):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         self.log = logging.getLogger(__name__)
+        self._table_dtypes = (
+            [
+                ('x', '<f8'),
+                ('y', '<f8'),
+                ('mag', '<f8'),
+                ('ra', '<f8'),
+                ('dec', '<f8'),
+                ('source_type', '<U4'),
+                ('isVisitSource', '?'),
+                ('isTemplateSource', '?'),
+                ('host_id', '<i8'),
+                ('host_flux', '<f8'),
+                ('host_mag', '<f8'),
+                ('host_ra', '<f8'),
+                ('host_dec', '<f8'),
+                ('delta_ra', '<f8'),
+                ('delta_dec', '<f8'),
+                ('delta_mag', '<f8'),
+                ('host_a', '<f8'),
+                ('host_b', '<f8'),
+                ('host_pa', '<f8'),
+                ('hosted_fake', '?'),
+                ('injection_id', '<i8'),
+                ('isVariable', '?'),
+                ('mag_offset', '<f8'),
+                ('twin_id', '<i8'),
+                ('isBlended', 'bool'),
+                ('visit', '<i8'),
+                ('detector', '<i8'),
+                ('run', '<U62'),
+                ('band', '<U1')
+            ]
+        )
 
     def runQuantum(self, butlerQC, inputRefs, outputRefs):
         inputs = butlerQC.get(inputRefs)
@@ -424,6 +503,27 @@ def _make_unique_injection_ids(self, n_ids, used_ids=None):
 
         return np.asarray(injection_ids, dtype=np.int64)
 
+    def _draw_offset_components_arcsec(self, rng, n_points):
+        """Draw isotropic offsets in arcseconds from an annulus."""
+        if n_points <= 0:
+            return np.zeros(0, dtype=float), np.zeros(0, dtype=float)
+
+        rmin = max(0.0, float(self.config.blendedFakeMinOffset))
+        rmax = max(rmin, float(self.config.blendedFakeMaxOffset))
+        radius = np.sqrt(rng.uniform(rmin**2, rmax**2, size=n_points))
+        theta = rng.uniform(0.0, 2.0 * np.pi, size=n_points)
+        return radius * np.cos(theta), radius * np.sin(theta)
+
+    def _cap_hosted_blended_count(self, requested, n_hosts):
+        """Apply configured limits to hosted blended fake counts."""
+        if requested <= 0 or n_hosts <= 0:
+            return 0
+
+        cap = n_hosts * self.config.maxHostedBlendedFakesPerHost
+        if self.config.maxHostedBlendedFakesTotal > 0:
+            cap = min(cap, self.config.maxHostedBlendedFakesTotal)
+        return min(requested, cap)
+
     def run(self, sourceCat, visit_image):
         """Create a set of visit detector fakes.
 
@@ -439,6 +539,8 @@ def run(self, sourceCat, visit_image):
         outputCat : `astropy.table.Table`
             Catalog of fake sources to draw inputs from.
         """
+        zero_table = Table(dtype=self._table_dtypes)
+
         # Use the visit+detector ids as the random seed.
         visitId = visit_image.getInfo().getVisitInfo().id
         detId = visit_image.detector.getId()
@@ -623,24 +725,151 @@ def run(self, sourceCat, visit_image):
             variable_fakes["isVariable"] = True
 
             catalog = vstack([catalog, variable_fakes])
+            catalog["isBlended"] = False
+
+        if self.config.doAddBlendedFakes:
+            self.log.info("Generating blended fakes.")
+            n_blended_fakes_total = int(len(catalog) * self.config.fracBlendedFakes)
+            n_star_hosted_blended_fakes = int(n_blended_fakes_total * self.config.fracHostedBlendedFakes)
+            idx = rng.choice(
+                len(catalog), size=n_blended_fakes_total - n_star_hosted_blended_fakes, replace=False)
+            # striaghtforward blended fakes
+            blended_fakes = catalog[idx].copy()
+            n_blended_fakes = len(blended_fakes)
+            # blended fakes will be copies of randomly chosen fakes, with a small
+            # magnitude offset and a small positional offset, and will be flagged as blended.
+            # The original fakes will also be flagged as blended.
+            # The idea is to have two fakes in the same location, but with different
+            # magnitudes, to test the deblending capabilities of the pipeline.
+            # We work as if the original fake was a "host"of the blend
+            blended_fakes["delta_mag"] = rng.normal(
+                loc=0.0,
+                scale=self.config.blendedFakeMagOffset,
+                size=n_blended_fakes
+            )
+            delta_ra, delta_dec = self._draw_offset_components_arcsec(rng, n_blended_fakes)
+            blended_fakes["delta_ra"] = delta_ra
+            blended_fakes["delta_dec"] = delta_dec
+
+            blended_fakes["host_mag"] = blended_fakes["mag"]
+            blended_fakes["mag"] += blended_fakes["delta_mag"]
+            blended_fakes["host_ra"] = blended_fakes["ra"]
+            blended_fakes["host_dec"] = blended_fakes["dec"]
+
+            blended_fakes["ra"] += blended_fakes["delta_ra"] / 3600.0
+            blended_fakes["dec"] += blended_fakes["delta_dec"] / 3600
+            blended_fakes["x"], blended_fakes["y"] = wcs.skyToPixelArray(
+                np.deg2rad(blended_fakes["ra"]), np.deg2rad(blended_fakes["dec"])
+            )
+
+            blended_fakes["source_type"] = "Star"
+            blended_fakes["host_id"] = blended_fakes["injection_id"]
+            blended_fakes["twin_id"] = blended_fakes["injection_id"]
+            blended_fakes["isBlended"] = True
+            blended_fakes["hosted_fake"] = True
+
+            if self.config.fracHostedBlendedFakes > 0:
+                hostcatalog = photoCalib.calibrateCatalog(sourceCat).asAstropy()
+                star_hosts = self.select_host_stars(hostcatalog)
+                # if len(star_hosts) is less than the blended fakes, then use replacement
+                if len(star_hosts) == 0:
+                    self.log.warning(
+                        "Hosted blended fake generation requested, but no valid star hosts were selected."
+                    )
+                    n_star_hosted_blended_fakes = 0
+                else:
+                    requested = n_star_hosted_blended_fakes
+                    n_star_hosted_blended_fakes = self._cap_hosted_blended_count(
+                        requested, len(star_hosts)
+                    )
+                    if n_star_hosted_blended_fakes < requested:
+                        self.log.warning(
+                            "Reducing hosted blended fakes from %d to %d to respect host guard rails.",
+                            requested,
+                            n_star_hosted_blended_fakes,
+                        )
+
+                if n_star_hosted_blended_fakes > 0:
+                    # Build a finite host pool so each host appears at most
+                    # maxHostedBlendedFakesPerHost times.
+                    host_pool = np.repeat(
+                        np.arange(len(star_hosts), dtype=int),
+                        self.config.maxHostedBlendedFakesPerHost,
+                    )
+                    idx = rng.choice(host_pool, size=n_star_hosted_blended_fakes, replace=False)
+                    hostcat = star_hosts[idx]
 
-        if len(catalog) > len(np.unique(catalog["injection_id"])):
-            self.log.warning("Duplicate injection IDs detected after catalog assembly; reassigning them.")
-            old_injection_ids = np.asarray(catalog["injection_id"], dtype=np.int64)
-            new_injection_ids = self._make_unique_injection_ids(len(catalog))
-            # re-assign fresh injection ids
-            catalog["injection_id"] = new_injection_ids
-            if "twin_id" in catalog.colnames:
-                id_map = {old_id: new_id for old_id, new_id in zip(old_injection_ids, new_injection_ids)}
-                catalog["twin_id"] = np.asarray(
-                    [id_map.get(int(twin_id), int(twin_id)) for twin_id in catalog["twin_id"]],
-                    dtype=np.int64,
+                    x_hosts = hostcat['slot_Centroid_x']
+                    y_hosts = hostcat['slot_Centroid_y']
+                    ra_hosts = np.rad2deg(hostcat['coord_ra'])
+                    dec_hosts = np.rad2deg(hostcat['coord_dec'])
+                    mag_hosts = hostcat['slot_PsfFlux_mag']
+
+                    delta_ra, delta_dec = self._draw_offset_components_arcsec(
+                        rng, n_star_hosted_blended_fakes
+                    )
+
+                    ra_ssi = ra_hosts + delta_ra / 3600.0
+                    dec_ssi = dec_hosts + delta_dec / 3600.0
+
+                    x_ssi, y_ssi = wcs.skyToPixelArray(ra_ssi, dec_ssi, degrees=True)
+
+                    delta_mag = rng.normal(loc=1, scale=1, size=n_star_hosted_blended_fakes)
+                    mags = mag_hosts + delta_mag
+
+                    #  Create the table of hosted fakes
+                    hosted_fakes = Table()
+                    hosted_fakes["x"] = x_ssi
+                    hosted_fakes["y"] = y_ssi
+                    hosted_fakes["mag"] = mags
+                    hosted_fakes["ra"] = ra_ssi
+                    hosted_fakes["dec"] = dec_ssi
+                    hosted_fakes["host_id"] = hostcat['id']
+                    hosted_fakes["host_flux"] = hostcat['slot_PsfFlux_flux']
+                    hosted_fakes["host_mag"] = hostcat['slot_PsfFlux_mag']
+                    hosted_fakes["host_ra"] = ra_hosts
+                    hosted_fakes["host_dec"] = dec_hosts
+                    hosted_fakes["delta_ra"] = delta_ra
+                    hosted_fakes["delta_dec"] = delta_dec
+                    hosted_fakes["delta_mag"] = delta_mag
+                    hosted_fakes["source_type"] = "Star"
+                    hosted_fakes["hosted_fake"] = True
+                    hosted_fakes["isVisitSource"] = True
+                    hosted_fakes["isTemplateSource"] = False
+                    hosted_fakes["isBlended"] = True
+
+                    blended_fakes = vstack([blended_fakes, hosted_fakes])
+
+                blended_fakes["injection_id"] = self._make_unique_injection_ids(
+                    len(blended_fakes),
+                    used_ids=catalog["injection_id"],
                 )
+                catalog = vstack([catalog, blended_fakes])
+
+            if len(catalog) > len(np.unique(catalog["injection_id"])):
+                self.log.warning("Duplicate injection IDs detected after catalog assembly; reassigning them.")
+                old_injection_ids = np.asarray(catalog["injection_id"], dtype=np.int64)
+                new_injection_ids = self._make_unique_injection_ids(len(catalog))
+                # re-assign fresh injection ids
+                catalog["injection_id"] = new_injection_ids
+                if "twin_id" in catalog.colnames:
+                    id_map = {old_id: new_id for old_id, new_id in zip(old_injection_ids, new_injection_ids)}
+                    catalog["twin_id"] = np.asarray(
+                        [id_map.get(int(twin_id), int(twin_id)) for twin_id in catalog["twin_id"]],
+                        dtype=np.int64,
+                    )
 
         catalog["visit"] = visitId
         catalog["detector"] = detId
+        if catalog["ra"].unit is not None:
+            catalog["ra"] = catalog["ra"].value
+            catalog["dec"] = catalog["dec"].value
+            catalog["delta_ra"] = catalog["delta_ra"].value
+            catalog["delta_dec"] = catalog["delta_dec"].value
+            catalog["host_ra"] = catalog["host_ra"].value
+            catalog["host_dec"] = catalog["host_dec"].value
 
-        return Struct(outputCat=catalog)
+        return Struct(outputCat=vstack([zero_table, catalog]))
 
     def select_hosts(self, sourceCat):
         """
@@ -681,6 +910,45 @@ def select_hosts(self, sourceCat):
             skySourceCut & flagCut & extendednessCut & snrCut].copy()
         return hostCat
 
+    def select_host_stars(self, sourceCat):
+        """
+        Selects host sources from a given source catalog based on a series of classification and flux cuts.
+        The selection criteria are:
+            - The 'base_ClassificationSizeExtendedness_flag' and
+                  'base_ClassificationExtendedness_flag' must both be False.
+            - The 'base_ClassificationSizeExtendedness_value' must be greater than 0.9.
+            - The 'base_ClassificationExtendedness_value' must be equal to 1.
+            - The 'base_PsfFlux_flux' must be greater than 0.
+        Parameters
+        ----------
+        sourceCat : SourceCatalog
+            The source catalog containing the columns required for selection.
+        *args, **kwargs
+            Additional arguments (not used).
+        Returns
+        -------
+        hostCat : ArrowAstropy
+            A deep copy of the subset of the source catalog that passes all selection criteria.
+        """
+
+        # Avoid calibration stars or psf stars; remove flagged sources sky_sources
+        skySourceCut = ~sourceCat['sky_source']
+
+        flagCut = ~sourceCat['base_ClassificationSizeExtendedness_flag']
+        flagCut &= ~sourceCat['base_ClassificationExtendedness_flag']
+        flagCut &= ~sourceCat['slot_Shape_flag']
+        flagCut &= ~sourceCat['slot_Centroid_flag']
+        flagCut &= ~sourceCat['base_PixelFlags_flag']
+
+        extendednessCut = sourceCat['base_ClassificationSizeExtendedness_value'] < 0.9
+        extendednessCut &= sourceCat['base_ClassificationExtendedness_value'] != 1
+
+        snrCut = sourceCat['slot_PsfFlux_flux']/sourceCat['slot_PsfFlux_fluxErr'] > 30
+
+        hostCat = sourceCat[
+            skySourceCut & flagCut & extendednessCut & snrCut].copy()
+        return hostCat
+
     def get_PA_and_axes(self, Ixx, Ixy, Iyy):
         '''
         Calculates the orientation and extent of an object based on its second moments.