Merge pull request #19 from lsst/tickets/DM-23364

tickets/DM-23364

python/lsst/meas/extensions/scarlet/deblend.py

@@ -22,6 +22,7 @@
 from functools import partial
 
 import numpy as np
+import scarlet
 from scarlet.psf import PSF, gaussian
 from scarlet import PointSource, ExtendedSource, MultiComponentSource
 
@@ -37,7 +38,7 @@
 import lsst.afw.detection as afwDet
 import lsst.afw.table as afwTable
 
-from .source import init_source, modelToHeavy
+from .source import initSource, modelToHeavy
 from .blend import LsstBlend, checkBlendConvergence
 from .observation import LsstFrame, LsstObservation
 
@@ -108,7 +109,7 @@ def _computePsfImage(self, position=None):
     return psfImage
 
 
-def getFootprintMask(footprint, mExposure, config):
+def getFootprintMask(footprint, mExposure):
     """Mask pixels outside the footprint
 
     Parameters
@@ -118,8 +119,6 @@ def getFootprintMask(footprint, mExposure, config):
           mask, and variance data
     footprint : `lsst.detection.Footprint`
         - The footprint of the parent to deblend
-    config : `ScarletDeblendConfig`
-        - Configuration of the deblending task
 
     Returns
     -------
@@ -159,7 +158,7 @@ def deblend(mExposure, footprint, config):
         weights = np.ones_like(images)
 
     # Mask out the pixels outside the footprint
-    mask = getFootprintMask(footprint, mExposure, config)
+    mask = getFootprintMask(footprint, mExposure)
     weights *= ~mask
 
     psfs = _computePsfImage(mExposure, footprint.getCentroid()).array.astype(np.float32)
@@ -190,9 +189,10 @@ def deblend(mExposure, footprint, config):
         else:
             raise ValueError("Unrecognized sourceModel")
 
-        source = init_source(frame=frame, peak=center, observation=observation, bbox=bbox,
-                             symmetric=config.symmetric, monotonic=config.monotonic,
-                             thresh=config.morphThresh, components=components)
+        source = initSource(frame=frame, peak=center, observation=observation, bbox=bbox,
+                            symmetric=config.symmetric, monotonic=config.monotonic,
+                            thresh=config.morphThresh, components=components,
+                            edgeDistance=config.edgeDistance, shifting=False)
         if source is not None:
             sources.append(source)
         else:
@@ -222,36 +222,22 @@ class ScarletDeblendConfig(pexConfig.Config):
                                          "iterations to exit fitter"))
 
     # Blend Configuration options
-    recenterPeriod = pexConfig.Field(dtype=int, default=5,
-                                     doc=("Number of iterations between recentering"))
-    exactLipschitz = pexConfig.Field(dtype=bool, default=True,
-                                     doc=("Calculate exact Lipschitz constant in every step"
-                                          "(True) or only calculate the approximate"
-                                          "Lipschitz constant with significant changes in A,S"
-                                          "(False)"))
+    edgeDistance = pexConfig.Field(dtype=int, default=1,
+                                   doc="All sources with flux within `edgeDistance` from the edge "
+                                       "will be considered edge sources.")
 
     # Constraints
-    sparse = pexConfig.Field(dtype=bool, default=True, doc="Make models compact and sparse")
     morphThresh = pexConfig.Field(dtype=float, default=1,
                                   doc="Fraction of background RMS a pixel must have"
                                       "to be included in the initial morphology")
     monotonic = pexConfig.Field(dtype=bool, default=True, doc="Make models monotonic")
     symmetric = pexConfig.Field(dtype=bool, default=False, doc="Make models symmetric")
-    symmetryThresh = pexConfig.Field(dtype=float, default=1.0,
-                                     doc=("Strictness of symmetry, from"
-                                          "0 (no symmetry enforced) to"
-                                          "1 (perfect symmetry required)."
-                                          "If 'S' is not in `constraints`, this argument is ignored"))
 
     # Other scarlet paremeters
     useWeights = pexConfig.Field(
         dtype=bool, default=True,
         doc=("Whether or not use use inverse variance weighting."
              "If `useWeights` is `False` then flat weights are used"))
-    usePsfConvolution = pexConfig.Field(
-        dtype=bool, default=True,
-        doc=("Whether or not to convolve the morphology with the"
-             "PSF in each band or use the same morphology in all bands"))
     modelPsfSize = pexConfig.Field(
         dtype=int, default=11,
         doc="Model PSF side length in pixels")
@@ -264,8 +250,6 @@ class ScarletDeblendConfig(pexConfig.Config):
     processSingles = pexConfig.Field(
         dtype=bool, default=False,
         doc="Whether or not to process isolated sources in the deblender")
-    storeHistory = pexConfig.Field(dtype=bool, default=False,
-                                   doc="Whether or not to store the history for each source")
     sourceModel = pexConfig.Field(
         dtype=str, default="single",
         doc=("How to determine which model to use for sources, from\n"
@@ -414,6 +398,10 @@ def _addSchemaKeys(self, schema):
         self.modelTypeKey = schema.addField("deblend_modelType", type="String", size=20,
                                             doc="The type of model used, for example "
                                                 "MultiComponentSource, ExtendedSource, PointSource")
+        self.edgeFluxFlagKey = schema.addField("deblend_edgeFluxFlag", type="Flag",
+                                               doc="Source has flux on the edge of the image")
+        self.scarletFluxKey = schema.addField("deblend_scarletFlux", type=np.float32,
+                                              doc="Flux measurement from scarlet")
         # self.log.trace('Added keys to schema: %s', ", ".join(str(x) for x in
         #               (self.nChildKey, self.tooManyPeaksKey, self.tooBigKey))
         #               )
@@ -598,13 +586,14 @@ def deblend(self, mExposure, sources):
                 # morph = source.morphToHeavy(xy0=bbox.getMin())
                 # sed = source.sed / source.sed.sum()
 
-                for f in filters:
+                flux = scarlet.measure.flux(source)
+                for fidx, f in enumerate(filters):
                     if len(models[f].getPeaks()) != 1:
                         err = "Heavy footprint should have a single peak, got {0}"
                         raise ValueError(err.format(len(models[f].peaks)))
                     cat = templateCatalogs[f]
                     child = self._addChild(parentId, cat, models[f], source, converged,
-                                           xy0=bbox.getMin())
+                                           xy0=bbox.getMin(), flux=flux[fidx])
                     if parentId == 0:
                         child.setId(src.getId())
                         child.set(self.runtimeKey, runtime)
@@ -683,7 +672,7 @@ def _skipParent(self, source, masks):
                 mask.addMaskPlane(self.config.notDeblendedMask)
                 fp.spans.setMask(mask, mask.getPlaneBitMask(self.config.notDeblendedMask))
 
-    def _addChild(self, parentId, sources, heavy, scarlet_source, blend_converged, xy0):
+    def _addChild(self, parentId, sources, heavy, scarlet_source, blend_converged, xy0, flux):
         """Add a child to a catalog
 
         This creates a new child in the source catalog,
@@ -719,4 +708,10 @@ def _addChild(self, parentId, sources, heavy, scarlet_source, blend_converged, x
         cx = np.max([np.min([int(np.round(cx)), morph.shape[1]-1]), 0])
         src.set(self.modelCenterFlux, morph[cy, cx])
         src.set(self.modelTypeKey, scarlet_source.__class__.__name__)
+        src.set(self.edgeFluxFlagKey, scarlet_source.isEdge)
+        # Include the source flux in the model space in the catalog.
+        # This uses the narrower model PSF, which ensures that all sources
+        # not located on an edge have all of their flux included in the
+        # measurement.
+        src.set(self.scarletFluxKey, flux)
         return src

python/lsst/meas/extensions/scarlet/source.py

@@ -28,14 +28,56 @@
 import lsst.log
 import lsst.afw.detection as afwDet
 
-__all__ = ["init_source", "morphToHeavy", "modelToHeavy"]
+__all__ = ["initSource", "morphToHeavy", "modelToHeavy"]
 
 logger = lsst.log.Log.getLogger("meas.deblender.deblend")
 
 
-def init_source(frame, peak, observation, bbox,
-                symmetric=False, monotonic=True,
-                thresh=5, components=1):
+def hasEdgeFlux(source, edgeDistance=1):
+    """hasEdgeFlux
+
+    Determine whether or not a source has flux within `edgeDistance`
+    of the edge.
+
+    Parameters
+    ----------
+    source : `scarlet.Component`
+        The source to check for edge flux
+    edgeDistance : int
+        The distance from the edge of the image to consider
+        a source an edge source. For example if `edgeDistance=3`
+        then any source within 3 pixels of the edge will be
+        considered to have edge flux. The minimum value of
+        `edgeDistance` is one, meaning the rows and columns
+        of pixels on the edge.
+
+    Returns
+    -------
+    isEdge: `bool`
+        Whether or not the source has flux on the edge.
+    """
+    assert edgeDistance > 0
+
+    # Use the first band that has a non-zero SED
+    if hasattr(source, "sed"):
+        band = np.min(np.where(source.sed > 0)[0])
+    else:
+        band = np.min(np.where(source.components[0].sed > 0)[0])
+    model = source.get_model()[band]
+    for edge in range(edgeDistance):
+        if (
+            np.any(model[edge-1] > 0) or
+            np.any(model[-edge] > 0) or
+            np.any(model[:, edge-1] > 0) or
+            np.any(model[:, -edge] > 0)
+        ):
+            return True
+    return False
+
+
+def initSource(frame, peak, observation, bbox,
+               symmetric=False, monotonic=True,
+               thresh=5, components=1, edgeDistance=1, shifting=False):
     """Initialize a Source
 
     The user can specify the number of desired components
@@ -85,11 +127,14 @@ def init_source(frame, peak, observation, bbox,
     while components > 1:
         try:
             source = MultiComponentSource(frame, center, observation, symmetric=symmetric,
-                                          monotonic=monotonic, thresh=thresh)
+                                          monotonic=monotonic, thresh=thresh, shifting=shifting)
             if (np.any([np.any(np.isnan(c.sed)) for c in source.components]) or
-                    np.any([np.all(c.sed <= 0) for c in source.components])):
+                    np.any([np.all(c.sed <= 0) for c in source.components]) or
+                    np.any([np.any(~np.isfinite(c.morph)) for c in source.components])):
                 logger.warning("Could not initialize")
                 raise ValueError("Could not initialize source")
+            if hasEdgeFlux(source, edgeDistance):
+                source.shifting = True
             break
         except Exception:
             # If the MultiComponentSource failed to initialize
@@ -99,7 +144,7 @@ def init_source(frame, peak, observation, bbox,
     if components == 1:
         try:
             source = ExtendedSource(frame, center, observation, thresh=thresh,
-                                    symmetric=symmetric, monotonic=monotonic)
+                                    symmetric=symmetric, monotonic=monotonic, shifting=shifting)
             if np.any(np.isnan(source.sed)) or np.all(source.sed <= 0) or np.sum(source.morph) == 0:
                 logger.warning("Could not initialize")
                 raise ValueError("Could not initialize source")
@@ -116,6 +161,26 @@ def init_source(frame, peak, observation, bbox,
             # so skip this source
             return None
 
+    if hasEdgeFlux(source, edgeDistance):
+        # The detection algorithm implemented in meas_algorithms
+        # does not place sources within the edge mask
+        # (roughly 5 pixels from the edge). This results in poor
+        # deblending of the edge source, which for bright sources
+        # may ruin an entire blend.
+        # By turning on shifting we allow exxtended sources to be shifted
+        # by a fraction of a pixel, which is computationally expensive and
+        # not necessary for non-edge sources.
+        # Due to the complexities of scarlet initialization it is easier
+        # to reinitialize edge sources to allow for shifting than it is
+        # to update this parameter on the fly.
+        if not isinstance(source, PointSource) and not shifting:
+            return initSource(frame, peak, observation, bbox,
+                              symmetric, monotonic, thresh, components,
+                              edgeDistance, shifting=True)
+        source.isEdge = True
+    else:
+        source.isEdge = False
+
     source.detectedPeak = peak
     return source
 

tests/test_scarlet.py

@@ -62,7 +62,7 @@ def test_init(self):
         # init stack objects
         foot, peak, bbox = numpyToStack(images, center, (15, 3))
         # init source
-        src = mes.source.init_source(frame=frame, peak=peak, observation=observation, bbox=bbox, thresh=0)
+        src = mes.source.initSource(frame=frame, peak=peak, observation=observation, bbox=bbox, thresh=0)
 
         self.assertFloatsAlmostEqual(src.sed/3, trueSed)
         self.assertFloatsAlmostEqual(src.morph*3, trueMorph, rtol=1e-7)
@@ -81,7 +81,7 @@ def test_to_heavy(self):
         frame = mes.LsstFrame(shape, psfs=targetPsfImage[None])
         observation = mes.LsstObservation(images, psfs=psfImages).match(frame)
         foot, peak, bbox = numpyToStack(images, coords[0], (15, 3))
-        src = mes.source.init_source(frame=frame, peak=peak, observation=observation, bbox=bbox, thresh=0)
+        src = mes.source.initSource(frame=frame, peak=peak, observation=observation, bbox=bbox, thresh=0)
         # Get the HeavyFootprint
         peakSchema = PeakTable.makeMinimalSchema()
         hFoot = mes.morphToHeavy(src, peakSchema=peakSchema)