computeSpatiallySampledMetrics.py

Go to the documentation of this file.

# This file is part of ip_diffim.
#
# Developed for the LSST Data Management System.
# This product includes software developed by the LSST Project
# (https://www.lsst.org).
# See the COPYRIGHT file at the top-level directory of this distribution
# for details of code ownership.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
import dataclasses
 
import numpy as np
import scipy.signal
 
import lsst.geom
 
import lsst.afw.image as afwImage
import lsst.afw.table as afwTable
import lsst.pipe.base as pipeBase
import lsst.pex.config as pexConfig
 
from lsst.ip.diffim.utils import getPsfFwhm, angleMean, getKernelCenterDisplacement
from lsst.meas.algorithms import SkyObjectsTask
from lsst.pex.exceptions import InvalidParameterError, RangeError
from lsst.utils.timer import timeMethod
 
import lsst.utils
 
__all__ = ["SpatiallySampledMetricsConfig", "SpatiallySampledMetricsTask"]
 
 
class SpatiallySampledMetricsConnections(pipeBase.PipelineTaskConnections,
                                         dimensions=("instrument", "visit", "detector"),
                                         defaultTemplates={"coaddName": "deep",
                                                           "warpTypeSuffix": "",
                                                           "fakesType": ""}):
    science = pipeBase.connectionTypes.Input(
        doc="Input science exposure.",
        dimensions=("instrument", "visit", "detector"),
        storageClass="ExposureF",
        name="{fakesType}calexp"
    )
    template = pipeBase.connectionTypes.Input(
        doc="Warped and not PSF-matched template used to create the difference image.",
        dimensions=("instrument", "visit", "detector"),
        storageClass="ExposureF",
        name="{fakesType}{coaddName}Diff_templateExp",
    )
    difference = pipeBase.connectionTypes.Input(
        doc="Difference image with detection mask plane filled in.",
        dimensions=("instrument", "visit", "detector"),
        storageClass="ExposureF",
        name="{fakesType}{coaddName}Diff_differenceExp",
    )
    diaSources = pipeBase.connectionTypes.Input(
        doc="Filtered diaSources on the difference image.",
        dimensions=("instrument", "visit", "detector"),
        storageClass="ArrowAstropy",
        name="{fakesType}dia_source_detector",
    )
    psfMatchingKernel = pipeBase.connectionTypes.Input(
        doc="Kernel used to PSF match the science and template images.",
        dimensions=("instrument", "visit", "detector"),
        storageClass="MatchingKernel",
        name="{fakesType}{coaddName}Diff_psfMatchKernel",
    )
    spatiallySampledMetrics = pipeBase.connectionTypes.Output(
        doc="Summary metrics computed at randomized locations.",
        dimensions=("instrument", "visit", "detector"),
        storageClass="ArrowAstropy",
        name="{fakesType}{coaddName}Diff_spatiallySampledMetrics",
    )
 
    def __init__(self, *, config=None):
        super().__init__(config=config)
        if config.image_type == "future":
            self.science = dataclasses.replace(self.science, storageClass="VisitImage")
            self.template = dataclasses.replace(self.template, storageClass="DifferenceImage")
            self.difference = dataclasses.replace(self.difference, storageClass="DifferenceImage")
            self.psfMatchingKernel = dataclasses.replace(self.psfMatchingKernel,
                                                         storageClass="ConvolutionKernel")
 
 
class SpatiallySampledMetricsConfig(pipeBase.PipelineTaskConfig,
                                    pipelineConnections=SpatiallySampledMetricsConnections):
    """Config for SpatiallySampledMetricsTask
    """
    metricsMaskPlanes = lsst.pex.config.ListField(
        dtype=str,
        doc="List of mask planes to include in metrics. "
            "A mask plane that is not registered reports 0 pixels set.",
        default=('BAD', 'CLIPPED', 'CR', 'DETECTED', 'DETECTED_NEGATIVE', 'EDGE',
                 'INEXACT_PSF', 'INJECTED', 'INJECTED_TEMPLATE', 'INTRP', 'NOT_DEBLENDED',
                 'NO_DATA', 'REJECTED', 'SAT', 'SAT_TEMPLATE', 'SENSOR_EDGE', 'STREAK', 'SUSPECT',
                 'UNMASKEDNAN',
                 ),
    )
    metricSources = pexConfig.ConfigurableField(
        target=SkyObjectsTask,
        doc="Generate QA metric sources",
    )
    image_type = pexConfig.ChoiceField[str](
        doc="Type of the input images and PSF matching kernel. run() takes the"
        " legacy types either way.",
        allowed={
            "legacy": "Read lsst.afw.image.ExposureF and lsst.afw.math.Kernel.",
            "future": "Read lsst.images.VisitImage for the science image,"
                      " lsst.images.DifferenceImage for the template and difference"
                      " image, and lsst.images.convolution_kernels.ConvolutionKernel.",
        },
        optional=False,
        default="legacy",
    )
 
    def setDefaults(self):
        self.metricSources.avoidMask = ["NO_DATA", "EDGE"]
 
 
class SpatiallySampledMetricsTask(lsst.pipe.base.PipelineTask):
    """Detect and measure sources on a difference image.
    """
    ConfigClass = SpatiallySampledMetricsConfig
    _DefaultName = "spatiallySampledMetrics"
 
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
 
        self.makeSubtask("metricSources")
        self.schema = afwTable.SourceTable.makeMinimalSchema()
        self.schema.addField(
            "x", "F",
            "X location of the metric evaluation.",
            units="pixel")
        self.schema.addField(
            "y", "F",
            "Y location of the metric evaluation.",
            units="pixel")
        self.metricSources.skySourceKey = self.schema.addField("sky_source", type="Flag",
                                                               doc="Metric evaluation objects.")
        self.schema.addField(
            "source_density", "F",
            "Density of diaSources at location.",
            units="count/degree^2")
        self.schema.addField(
            "dipole_density", "F",
            "Density of dipoles at location.",
            units="count/degree^2")
        self.schema.addField(
            "dipole_direction", "F",
            "Mean dipole orientation.",
            units="radian")
        self.schema.addField(
            "dipole_separation", "F",
            "Mean dipole separation.",
            units="pixel")
        self.schema.addField(
            "template_value", "F",
            "Median of template at location.",
            units="nJy")
        self.schema.addField(
            "template_variance", "F",
            "Median of template variance at location.",
            units="nJy^2")
        self.schema.addField(
            "science_value", "F",
            "Median of science at location.",
            units="nJy")
        self.schema.addField(
            "science_variance", "F",
            "Median of science variance at location.",
            units="nJy^2")
        self.schema.addField(
            "diffim_value", "F",
            "Median of diffim at location.",
            units="nJy")
        self.schema.addField(
            "diffim_variance", "F",
            "Median of diffim variance at location.",
            units="nJy^2")
        self.schema.addField(
            "diffim_chi2PerPix", "F",
            "Robust normalized noise of diffim at location:"
            " (1.4826*MAD(image))^2 / median(variance), evaluated on background"
            " pixels (DETECTED, DETECTED_NEGATIVE, BAD, SAT, EDGE, NO_DATA excluded)."
            " Expected ~1.0 for a well-decorrelated diffim; values >>1 indicate"
            " residual structure, <<1 indicates over-estimated variance.")
        self.schema.addField(
            "science_psfSize", "F",
            "Width of the science image PSF at location.",
            units="pixel")
        self.schema.addField(
            "template_psfSize", "F",
            "Width of the template image PSF at location.",
            units="pixel")
        for maskPlane in self.config.metricsMaskPlanes:
            self.schema.addField(
                "%s_mask_fraction"%maskPlane.lower(), "F",
                "Fraction of pixels with %s mask"%maskPlane
            )
        self.schema.addField(
            "psfMatchingKernel_sum", "F",
            "PSF matching kernel sum at location.")
        self.schema.addField(
            "psfMatchingKernel_dx", "F",
            "PSF matching kernel centroid offset in x at location.",
            units="pixel")
        self.schema.addField(
            "psfMatchingKernel_dy", "F",
            "PSF matching kernel centroid offset in y at location.",
            units="pixel")
        self.schema.addField(
            "psfMatchingKernel_length", "F",
            "PSF matching kernel centroid offset module.",
            units="arcsecond")
        self.schema.addField(
            "psfMatchingKernel_position_angle", "F",
            "PSF matching kernel centroid offset position angle.",
            units="radian")
        self.schema.addField(
            "psfMatchingKernel_direction", "F",
            "PSF matching kernel centroid offset direction in detector plane.",
            units="radian")
        self.schema.addField(
            "psfMatchingKernel_residualNorm", "F",
            "Shape-only PSF match residual at location:"
            "L2 norm of (K-convolved template PSF - science PSF),"
            " relative to the science PSF L2 norm. Larger values indicate worse"
            " PSF matching. Assumes the kernel was solved to convolve the template.")
 
    def runQuantum(self, butlerQC, inputRefs, outputRefs):
        inputs = butlerQC.get(inputRefs)
        if self.config.image_type == "future":
            inputs["science"] = inputs["science"].to_legacy()
            inputs["template"] = inputs["template"].to_legacy()
            inputs["difference"] = inputs["difference"].to_legacy()
            inputs["psfMatchingKernel"] = inputs["psfMatchingKernel"].to_legacy()
        outputs = self.run(**inputs)
        butlerQC.put(outputs, outputRefs)
 
    @timeMethod
    def run(self, science, template, difference, diaSources, psfMatchingKernel):
        """Calculate difference image metrics on specific locations across the images
 
        Parameters
        ----------
        science : `lsst.afw.image.ExposureF`
            Science exposure that the template was subtracted from.
        template : `lsst.afw.image.ExposureF`
            Warped and non PSF-matched template that was used produce
            the difference image.
        difference : `lsst.afw.image.ExposureF`
            Result of subtracting template from the science image.
        diaSources : `lsst.afw.table.SourceCatalog`
                The catalog of detected sources.
        psfMatchingKernel : `~lsst.afw.math.LinearCombinationKernel`
            The PSF matching kernel of the subtraction to evaluate.
 
        Returns
        -------
        results : `lsst.pipe.base.Struct`
            ``spatiallySampledMetrics`` :  `astropy.table.Table`
                Image quality metrics spatially sampled locations.
        """
 
        idFactory = lsst.meas.base.IdGenerator().make_table_id_factory()
 
        spatiallySampledMetrics = afwTable.SourceCatalog(self.schema)
        spatiallySampledMetrics.getTable().setIdFactory(idFactory)
 
        self.metricSources.run(mask=science.mask, seed=difference.info.id, catalog=spatiallySampledMetrics)
 
        # Mask plane bits are looked up in the global afw dict, not the
        # image's own. A plane that is not registered in this process has no
        # pixels set in any image.
        metricsMaskPlanes = {}
        for maskPlane in self.config.metricsMaskPlanes:
            try:
                metricsMaskPlanes[maskPlane] = afwImage.Mask.getPlaneBitMask(maskPlane)
            except InvalidParameterError:
                self.log.info("Mask plane %s is not registered; its metrics will be zero.", maskPlane)
                metricsMaskPlanes[maskPlane] = 0
 
        for src in spatiallySampledMetrics:
            self._evaluateLocalMetric(src, science, template, difference, diaSources,
                                      metricsMaskPlanes=metricsMaskPlanes,
                                      psfMatchingKernel=psfMatchingKernel)
        spatiallySampledMetrics = spatiallySampledMetrics.copy(deep=True).asAstropy()
        return pipeBase.Struct(spatiallySampledMetrics=spatiallySampledMetrics)
 
    def _evaluateLocalMetric(self, src, science, template, difference, diaSources,
                             metricsMaskPlanes, psfMatchingKernel):
        """Calculate image quality metrics at spatially sampled locations.
 
        Parameters
        ----------
        src : `lsst.afw.table.SourceRecord`
            The source record to be updated with metric calculations.
        diaSources : `lsst.afw.table.SourceCatalog`
            The catalog of detected sources.
        science : `lsst.afw.image.Exposure`
            The science image.
        difference : `lsst.afw.image.Exposure`
            Result of subtracting template from the science image.
        metricsMaskPlanes : `dict` [`str`, `int`]
            Mask planes to calculate metrics from, mapped to their bitmask.
            The bitmask is 0 for a plane that is not registered.
        psfMatchingKernel : `~lsst.afw.math.LinearCombinationKernel`
            The PSF matching kernel of the subtraction to evaluate.
        """
        bbox = src.getFootprint().getBBox()
        pix = bbox.getCenter()
        src.set('science_psfSize', getPsfFwhm(science.psf, position=pix))
        try:
            src.set('template_psfSize', getPsfFwhm(template.psf, position=pix))
        except (InvalidParameterError, RangeError):
            src.set('template_psfSize', np.nan)
 
        metricRegionSize = 100
        bbox.grow(metricRegionSize)
        bbox = bbox.clippedTo(science.getBBox())
        nPix = bbox.getArea()
        pixScale = science.wcs.getPixelScale(bbox.getCenter())
        area = nPix*pixScale.asDegrees()**2
        peak = src.getFootprint().getPeaks()[0]
        src.set('x', peak['i_x'])
        src.set('y', peak['i_y'])
        src.setCoord(science.wcs.pixelToSky(peak['i_x'], peak['i_y']))
        selectSources = diaSources[bbox.contains(diaSources['x'], diaSources['y'])]
        sourceDensity = len(selectSources)/area
        dipoleSources = selectSources[selectSources["isDipole"]]
        dipoleDensity = len(dipoleSources)/area
 
        if dipoleSources:
            meanDipoleOrientation = angleMean(dipoleSources["dipoleAngle"])
            src.set('dipole_direction', meanDipoleOrientation)
            meanDipoleSeparation = np.mean(dipoleSources["dipoleLength"])
            src.set('dipole_separation', meanDipoleSeparation)
 
        templateVal = np.median(template[bbox].image.array)
        templateVar = np.median(template[bbox].variance.array)
        scienceVal = np.median(science[bbox].image.array)
        scienceVar = np.median(science[bbox].variance.array)
        diffimVal = np.median(difference[bbox].image.array)
        diffimVar = np.median(difference[bbox].variance.array)
        src.set('source_density', sourceDensity)
        src.set('dipole_density', dipoleDensity)
        src.set('template_value', templateVal)
        src.set('template_variance', templateVar)
        src.set('science_value', scienceVal)
        src.set('science_variance', scienceVar)
        src.set('diffim_value', diffimVal)
        src.set('diffim_variance', diffimVar)
        src.set('diffim_chi2PerPix', self._diffimChi2PerPix(difference[bbox]))
        differenceMask = difference.mask[bbox].array
        for maskPlane, bitMask in metricsMaskPlanes.items():
            src.set("%s_mask_fraction"%maskPlane.lower(),
                    np.count_nonzero(differenceMask & bitMask)/differenceMask.size
                    )
 
        krnlSum, dx, dy, direction, length = getKernelCenterDisplacement(
            psfMatchingKernel, src.get('x'), src.get('y'))
 
        point1 = lsst.geom.SpherePoint(
            src.get('coord_ra'), src.get('coord_dec'),
            lsst.geom.radians)
        point2 = science.wcs.pixelToSky(src.get('x') + dx, src.get('y') + dy)
        bearing = point1.bearingTo(point2)
        pa_ref_angle = lsst.geom.Angle(np.pi/2, lsst.geom.radians)
        pa = pa_ref_angle - bearing
        # Wrap around to get Delta_RA from -pi to +pi
        pa = pa.wrapCtr()
        position_angle = pa.asRadians()
 
        src.set('psfMatchingKernel_sum', krnlSum)
        src.set('psfMatchingKernel_dx', dx)
        src.set('psfMatchingKernel_dy', dy)
        src.set('psfMatchingKernel_length', length*pixScale.asArcseconds())
        src.set('psfMatchingKernel_position_angle', position_angle)  # in E of N position angle
        src.set('psfMatchingKernel_direction', direction)  # direction offset in detector
 
        src.set('psfMatchingKernel_residualNorm',
                self._psfMatchResidualNorm(psfMatchingKernel, science.psf, template.psf,
                                           src.get('x'), src.get('y')))
 
    def _diffimChi2PerPix(self, difference):
        """Robust normalized noise of the difference image.
 
        Computes ``(1.4826 * MAD(image))^2 / median(variance)`` on
        background-only pixels. The 1.4826 factor rescales the
        Median Absolute Deviation to a Gaussian-sigma estimate, so the
        ratio has expectation 1.0 on pure noise.
 
        Returns NaN if no usable pixels remain or the median variance
        is non-positive.
        """
        image = difference.image.array
        variance = difference.variance.array
        mask = difference.mask
        excludePlanes = [p for p in ("DETECTED", "DETECTED_NEGATIVE", "BAD",
                                     "SAT", "EDGE", "NO_DATA")
                         if p in mask.getMaskPlaneDict()]
        if excludePlanes:
            badBits = mask.getPlaneBitMask(excludePlanes)
            good = (mask.array & badBits) == 0
        else:
            good = np.ones(image.shape, dtype=bool)
        good &= np.isfinite(image) & np.isfinite(variance) & (variance > 0)
        if not np.any(good):
            return np.nan
        varMed = np.median(variance[good])
        if not np.isfinite(varMed) or varMed <= 0:
            return np.nan
        imgGood = image[good]
        mad = np.median(np.abs(imgGood - np.median(imgGood)))
        return float((1.4826*mad)**2/varMed)
 
    def _psfMatchResidualNorm(self, kernel, sciencePsf, templatePsf, x, y):
        """Relative L2 norm of the PSF-match residual at (x, y).
 
        Convolves the template PSF with the matching kernel evaluated at
        (x, y) and compares to the science PSF at the same position, both
        renormalized to unit sum so the result captures shape mismatch only.
        The kernel is assumed to convolve the template.
 
        Returns NaN if either PSF cannot be evaluated at the position or
        the resulting images cannot be normalized.
        """
        point = lsst.geom.Point2D(x, y)
        try:
            psfSci = sciencePsf.computeKernelImage(point).array
            psfTmp = templatePsf.computeKernelImage(point).array
        except (InvalidParameterError, RangeError):
            return np.nan
 
        kImage = afwImage.ImageD(kernel.getDimensions())
        kernel.computeImage(kImage, doNormalize=True, x=x, y=y)
        matched = scipy.signal.fftconvolve(psfTmp, kImage.array, mode='same')
 
        matchedSum = matched.sum()
        psfSciSum = psfSci.sum()
        if matchedSum <= 0 or psfSciSum <= 0:
            return np.nan
        matched = matched/matchedSum
        psfSci = psfSci/psfSciSum
 
        # PSF stamps may differ in size between the two exposures; crop
        # both to a common centered region before differencing.
        h = min(matched.shape[0], psfSci.shape[0])
        w = min(matched.shape[1], psfSci.shape[1])
 
        def _crop(a):
            sy = (a.shape[0] - h)//2
            sx = (a.shape[1] - w)//2
            return a[sy:sy + h, sx:sx + w]
        matched = _crop(matched)
        psfSci = _crop(psfSci)
 
        sciNorm = np.sqrt(np.sum(psfSci**2))
        if sciNorm == 0:
            return np.nan
        return float(np.sqrt(np.sum((matched - psfSci)**2))/sciNorm)