Merge pull request #452 from lsst/tickets/DM-54428-v30

DM-54428: Add cosmic ray detection and masking on difference images

Author: isullivan

python/lsst/ip/diffim/detectAndMeasure.py

@@ -32,13 +32,13 @@
 from lsst.ip.diffim.utils import (evaluateMaskFraction, computeDifferenceImageMetrics,
                                   populate_sattle_visit_cache)
 from lsst.meas.algorithms import SkyObjectsTask, SourceDetectionTask, SetPrimaryFlagsTask, MaskStreaksTask
-from lsst.meas.algorithms import FindGlintTrailsTask
+from lsst.meas.algorithms import FindGlintTrailsTask, FindCosmicRaysConfig, findCosmicRays
 from lsst.meas.base import ForcedMeasurementTask, ApplyApCorrTask, DetectorVisitIdGeneratorConfig
 import lsst.meas.deblender
 import lsst.meas.extensions.trailedSources  # noqa: F401
 import lsst.meas.extensions.shapeHSM
 import lsst.pex.config as pexConfig
-from lsst.pex.exceptions import InvalidParameterError
+from lsst.pex.exceptions import InvalidParameterError, LengthError
 import lsst.pipe.base as pipeBase
 import lsst.utils
 from lsst.utils.timer import timeMethod
@@ -81,6 +81,33 @@ def metadata(self):
         return {}
 
 
+class TooManyCosmicRays(pipeBase.AlgorithmError):
+    """Raised if the cosmic ray task fails with too many cosmics.
+
+    Parameters
+    ----------
+    maxCosmicRays : `int`
+        Maximum number of cosmic rays allowed.
+    """
+    def __init__(self, maxCosmicRays, **kwargs):
+        msg = f"Cosmic ray task found more than {maxCosmicRays} cosmics."
+        self.msg = msg
+        self._metadata = kwargs
+        super().__init__(msg, **kwargs)
+        self._metadata["maxCosmicRays"] = maxCosmicRays
+
+    def __str__(self):
+        # Exception doesn't handle **kwargs, so we need a custom str.
+        return f"{self.msg}: {self.metadata}"
+
+    @property
+    def metadata(self):
+        for key, value in self._metadata.items():
+            if not (isinstance(value, int) or isinstance(value, float) or isinstance(value, str)):
+                raise TypeError(f"{key} is of type {type(value)}, but only (int, float, str) are allowed.")
+        return self._metadata
+
+
 class DetectAndMeasureConnections(pipeBase.PipelineTaskConnections,
                                   dimensions=("instrument", "visit", "detector"),
                                   defaultTemplates={"coaddName": "deep",
@@ -202,6 +229,16 @@ class DetectAndMeasureConfig(pipeBase.PipelineTaskConfig,
         target=lsst.meas.algorithms.SubtractBackgroundTask,
         doc="Task to perform final background subtraction, after first detection pass.",
     )
+    doFindCosmicRays = pexConfig.Field(
+        dtype=bool,
+        doc="Detect and mask cosmic rays on the difference image?"
+            "CRs can be interpolated over by setting cosmicray.keepCRs=False",
+        default=True,
+    )
+    cosmicray = pexConfig.ConfigField(
+        dtype=FindCosmicRaysConfig,
+        doc="Options for finding and masking cosmic rays",
+    )
     detection = pexConfig.ConfigurableField(
         target=SourceDetectionTask,
         doc="Final source detection for diaSource measurement",
@@ -377,6 +414,8 @@ def setDefaults(self):
                                                          "DETECTED_NEGATIVE",
                                                          "NO_DATA",
                                                          ]
+        # Cosmic ray detection
+        self.cosmicray.keepCRs = True  # do not interpolate over detected CRs
         # DiaSource Detection
         self.detection.thresholdPolarity = "both"
         self.detection.thresholdValue = 5.0
@@ -621,6 +660,11 @@ def run(self, science, matchedTemplate, difference, kernelSources=None,
 
         self._prepareInputs(detectionExposure)
 
+        if self.config.doFindCosmicRays and not self.config.doSubtractBackground:
+            # Detect and interpolate over any remaining cosmic rays
+            # This only needs to run once, right before the final detection.
+            self.findAndMaskCosmicRays(detectionExposure)
+
         # Don't use the idFactory until after deblend+merge, so that we aren't
         # generating ids that just get thrown away (footprint merge doesn't
         # know about past ids).
@@ -639,6 +683,10 @@ def run(self, science, matchedTemplate, difference, kernelSources=None,
             difference.setMask(detectionExposure.mask)
             background = self.subtractFinalBackground.run(difference).background
 
+            if self.config.doFindCosmicRays:
+                # Detect and interpolate over any remaining cosmic rays
+                self.findAndMaskCosmicRays(difference)
+
             # Re-run detection to get final footprints
             table = afwTable.SourceTable.make(self.schema)
             results = self.detection.run(
@@ -1049,6 +1097,39 @@ def measureForcedSources(self, diaSources, image, wcs, template=False):
         for diaSource, forcedSource in zip(diaSources, forcedSources):
             diaSource.assign(forcedSource, mapper)
 
+    def findAndMaskCosmicRays(self, difference):
+        """Detect and mask cosmic rays on the difference image.
+
+        Parameters
+        ----------
+        difference : `lsst.afw.image.Exposure`
+            The background-subtracted difference image.
+            The mask plane will be modified in place.
+        """
+        # This is run on a background-subtracted difference image, so the
+        # remaining background should be ~0.
+        medianBg = 0.
+        try:
+            crs = findCosmicRays(
+                difference.maskedImage,
+                difference.psf,
+                medianBg,
+                pexConfig.makePropertySet(self.config.cosmicray),
+                self.config.cosmicray.keepCRs,
+            )
+        except LengthError:
+            raise TooManyCosmicRays(self.config.cosmicray.nCrPixelMax) from None
+        num = 0
+        if crs is not None:
+            mask = difference.maskedImage.mask
+            crBit = mask.getPlaneBitMask("CR")
+            afwDetection.setMaskFromFootprintList(mask, crs, crBit)
+            num = len(crs)
+
+        text = "masked" if self.config.cosmicray.keepCRs else "interpolated over"
+        self.log.info("Identified and %s %s cosmic rays.", text, num)
+        self.metadata["cosmic_ray_count"] = num
+
     def calculateMetrics(self, science, difference, diaSources, kernelSources):
         """Add difference image QA metrics to the Task metadata.
 

tests/test_detectAndMeasure.py

@@ -414,6 +414,46 @@ def _detection_wrapper(positive=True):
         _detection_wrapper(positive=True)
         _detection_wrapper(positive=False)
 
+    def test_mask_cosmic_rays(self):
+        """Run detection on a difference image containing a cosmic ray.
+        """
+        # Set up the simulated images
+        noiseLevel = 1.
+        staticSeed = 1
+        fluxLevel = 500
+        xSize = 400
+        ySize = 400
+        kwargs = {"seed": staticSeed, "psfSize": 2.4, "fluxLevel": fluxLevel, "xSize": xSize, "ySize": ySize}
+        science, sources = makeTestImage(noiseLevel=noiseLevel, noiseSeed=6, **kwargs)
+        matchedTemplate, _ = makeTestImage(noiseLevel=noiseLevel/4, noiseSeed=7, **kwargs)
+        crMask = science.mask.getPlaneBitMask("CR")
+
+        # Configure the detection Task
+        detectionTask = self._setup_detection(doMerge=False, doSkySources=True)
+
+        # Test that no CRs are detected
+        transients, transientSources = makeTestImage(noiseLevel=noiseLevel, noiseSeed=8, nSrc=10, **kwargs)
+        science.maskedImage += transients.maskedImage
+        difference = science.clone()
+        difference.maskedImage -= matchedTemplate.maskedImage
+        output = detectionTask.run(science, matchedTemplate, difference, sources)
+        crMaskSet = (output.subtractedMeasuredExposure.mask.array & crMask) > 0
+        self.assertTrue(np.all(crMaskSet == 0))
+
+        crX0 = round(sources.getX()[0] - science.getX0())
+        crY0 = round(sources.getY()[0] - science.getY0())
+        crX1 = round(sources.getX()[1] - science.getX0())
+        crY1 = round(sources.getY()[1] - science.getY0())
+        # Add CR-like shape and check that CR is detected
+        # Pick two locations on top of sources, since that is what is likely to
+        # be missed in the first stage of CR rejection.
+        difference.image.array[crX0:crX0+1, crY0:crY0+5] += 1234  # Arbitrary but large flux for the CRs
+        difference.image.array[crX1:crX1+5, crY1:crY1+1] += 1234
+        output = detectionTask.run(science, matchedTemplate, difference)
+        crMaskSet = (output.subtractedMeasuredExposure.mask.array & crMask) > 0
+        self.assertTrue(np.all(crMaskSet[crX0:crX0+1, crY0:crY0+5]))
+        self.assertTrue(np.all(crMaskSet[crX1:crX1+5, crY1:crY1+1]))
+
     def test_missing_mask_planes(self):
         """Check that detection runs with missing mask planes.
         """