diffimTools.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/>.

__all__ = ["backgroundSubtract", "writeKernelCellSet", "sourceToFootprintList", "NbasisEvaluator"]

# python
import time
import os
from collections import Counter
import numpy as np

# all the c++ level classes and routines
from . import diffimLib

# all the other LSST packages
import lsst.afw.geom as afwGeom
import lsst.afw.image as afwImage
import lsst.afw.table as afwTable
import lsst.afw.detection as afwDetect
import lsst.afw.math.mathLib as afwMath
from lsst.log import Log
import lsst.pex.config as pexConfig
from .makeKernelBasisList import makeKernelBasisList

# Helper functions for ipDiffim; mostly viewing of results and writing
# debugging info to disk.




def makeFlatNoiseImage(mi, seedStat=afwMath.MAX):
    img = mi.getImage()
    seed = int(10.*afwMath.makeStatistics(mi.getImage(), seedStat).getValue() + 1)
    rdm = afwMath.Random(afwMath.Random.MT19937, seed)
    rdmImage = img.Factory(img.getDimensions())
    afwMath.randomGaussianImage(rdmImage, rdm)
    return rdmImage


def makePoissonNoiseImage(im):
    """Return a Poisson noise image based on im

    Parameters
    ----------
    im : `lsst.afw.image.Image`
        image; the output image has the same dtype, dimensions, and shape
        and its expectation value is the value of ``im`` at each pixel

    Returns
    -------
    noiseIm : `lsst.afw.image.Image`
        Newly constructed image instance, same type as ``im``.

    Notes
    -----
    - Warning: This uses an undocumented numpy API (the documented API
        uses a single float expectation value instead of an array).

    - Uses numpy.random; you may wish to call numpy.random.seed first.
    """
    import numpy.random as rand
    imArr = im.getArray()
    noiseIm = im.Factory(im.getBBox())
    noiseArr = noiseIm.getArray()

    with np.errstate(invalid='ignore'):
        intNoiseArr = rand.poisson(imArr)

    noiseArr[:, :] = intNoiseArr.astype(noiseArr.dtype)
    return noiseIm




def fakeCoeffs():
    kCoeffs = ((1.0, 0.0, 0.0),
               (0.005, -0.000001, 0.000001),
               (0.005, 0.000004, 0.000004),
               (-0.001, -0.000030, 0.000030),
               (-0.001, 0.000015, 0.000015),
               (-0.005, -0.000050, 0.000050))
    return kCoeffs


def makeFakeKernelSet(sizeCell=128, nCell=3,
                      deltaFunctionCounts=1.e4, tGaussianWidth=1.0,
                      addNoise=True, bgValue=100., display=False):
    """Generate test template and science images with sources.

    Parameters
    ----------
    sizeCell : `int`, optional
        Size of the square spatial cells in pixels.
    nCell : `int`, optional
        Number of adjacent spatial cells in both direction in both images.
    deltaFunctionCounts : `float`, optional
        Flux value for the template image sources.
    tGaussianWidth : `float`, optional
        Sigma of the generated Gaussian PSF sources in the template image.
    addNoise : `bool`, optional
        If `True`, Poisson noise is added to both the generated template
        and science images.
    bgValue : `float`, optional
        Background level to be added to the generated science image.
    display : `bool`, optional
        If `True` displays the generated template and science images by
        `lsst.afw.display.Display`.

    Notes
    -----
    - The generated images consist of adjacent ``nCell x nCell`` cells, each
      of pixel size ``sizeCell x sizeCell``.
    - The sources in the science image are generated by convolving the
      template by ``sKernel``. ``sKernel`` is a spatial `LinearCombinationKernel`
      of hard wired kernel bases functions. The linear combination has first
      order polynomial spatial dependence with polynomial parameters from ``fakeCoeffs()``.
    - The template image sources are generated in the center of each spatial
      cell from one pixel, set to `deltaFunctionCounts` counts, then convolved
      by a 2D Gaussian with sigma of `tGaussianWidth` along each axis.
    - The sources are also returned in ``kernelCellSet`` each source is "detected"
      exactly at the center of a cell.

    Returns
    -------
    tMi : `lsst.afw.image.MaskedImage`
        Generated template image.
    sMi : `lsst.afw.image.MaskedImage`
        Generated science image.
    sKernel : `lsst.afw.math.LinearCombinationKernel`
        The spatial kernel used to generate the sources in the science image.
    kernelCellSet : `lsst.afw.math.SpatialCellSet`
        Cell grid of `lsst.afw.math.SpatialCell` instances, containing
        `lsst.ip.diffim.KernelCandidate` instances around all the generated sources
        in the science image.
    configFake : `lsst.ip.diffim.ImagePsfMatchConfig`
        Config instance used in the image generation.
    """
    from . import imagePsfMatch
    configFake = imagePsfMatch.ImagePsfMatchConfig()
    configFake.kernel.name = "AL"
    subconfigFake = configFake.kernel.active
    subconfigFake.alardNGauss = 1
    subconfigFake.alardSigGauss = [2.5, ]
    subconfigFake.alardDegGauss = [2, ]
    subconfigFake.sizeCellX = sizeCell
    subconfigFake.sizeCellY = sizeCell
    subconfigFake.spatialKernelOrder = 1
    subconfigFake.spatialModelType = "polynomial"
    subconfigFake.singleKernelClipping = False   # variance is a hack
    subconfigFake.spatialKernelClipping = False  # variance is a hack
    if bgValue > 0.0:
        subconfigFake.fitForBackground = True

    policyFake = pexConfig.makePolicy(subconfigFake)

    basisList = makeKernelBasisList(subconfigFake)
    kSize = subconfigFake.kernelSize

    # This sets the final extent of each convolved delta function
    gaussKernelWidth = sizeCell//2

    # This sets the scale over which pixels are correlated in the
    # spatial convolution; should be at least as big as the kernel you
    # are trying to fit for
    spatialKernelWidth = kSize

    # Number of bad pixels due to convolutions
    border = (gaussKernelWidth + spatialKernelWidth)//2

    # Make a fake image with a matrix of delta functions
    totalSize = nCell*sizeCell + 2*border
    tim = afwImage.ImageF(afwGeom.Extent2I(totalSize, totalSize))
    for x in range(nCell):
        for y in range(nCell):
            tim[x*sizeCell + sizeCell//2 + border - 1,
                y*sizeCell + sizeCell//2 + border - 1,
                afwImage.LOCAL] = deltaFunctionCounts

    # Turn this into stars with a narrow width; conserve counts
    gaussFunction = afwMath.GaussianFunction2D(tGaussianWidth, tGaussianWidth)
    gaussKernel = afwMath.AnalyticKernel(gaussKernelWidth, gaussKernelWidth, gaussFunction)
    cim = afwImage.ImageF(tim.getDimensions())
    afwMath.convolve(cim, tim, gaussKernel, True)
    tim = cim

    # Trim off border pixels
    bbox = gaussKernel.shrinkBBox(tim.getBBox(afwImage.LOCAL))
    tim = afwImage.ImageF(tim, bbox, afwImage.LOCAL)

    # Now make a science image which is this convolved with some
    # spatial function.  Use input basis list.
    polyFunc = afwMath.PolynomialFunction2D(1)
    kCoeffs = fakeCoeffs()
    nToUse = min(len(kCoeffs), len(basisList))

    # Make the full convolved science image
    sKernel = afwMath.LinearCombinationKernel(basisList[:nToUse], polyFunc)
    sKernel.setSpatialParameters(kCoeffs[:nToUse])
    sim = afwImage.ImageF(tim.getDimensions())
    afwMath.convolve(sim, tim, sKernel, True)

    # Get the good subregion
    bbox = sKernel.shrinkBBox(sim.getBBox(afwImage.LOCAL))

    # Add background
    sim += bgValue

    # Watch out for negative values
    tim += 2*np.abs(np.min(tim.getArray()))

    # Add noise?
    if addNoise:
        sim = makePoissonNoiseImage(sim)
        tim = makePoissonNoiseImage(tim)

    # And turn into MaskedImages
    sim = afwImage.ImageF(sim, bbox, afwImage.LOCAL)
    svar = afwImage.ImageF(sim, True)
    smask = afwImage.Mask(sim.getDimensions())
    smask.set(0x0)
    sMi = afwImage.MaskedImageF(sim, smask, svar)

    tim = afwImage.ImageF(tim, bbox, afwImage.LOCAL)
    tvar = afwImage.ImageF(tim, True)
    tmask = afwImage.Mask(tim.getDimensions())
    tmask.set(0x0)
    tMi = afwImage.MaskedImageF(tim, tmask, tvar)

    if display:
        import lsst.afw.display as afwDisplay
        afwDisplay.Display(frame=1).mtv(tMi)
        afwDisplay.Display(frame=2).mtv(sMi)

    # Finally, make a kernelSet from these 2 images
    kernelCellSet = afwMath.SpatialCellSet(afwGeom.Box2I(afwGeom.Point2I(0, 0),
                                                         afwGeom.Extent2I(sizeCell*nCell,
                                                                          sizeCell*nCell)),
                                           sizeCell,
                                           sizeCell)
    stampHalfWidth = 2*kSize
    for x in range(nCell):
        for y in range(nCell):
            xCoord = x*sizeCell + sizeCell//2
            yCoord = y*sizeCell + sizeCell//2
            p0 = afwGeom.Point2I(xCoord - stampHalfWidth,
                                 yCoord - stampHalfWidth)
            p1 = afwGeom.Point2I(xCoord + stampHalfWidth,
                                 yCoord + stampHalfWidth)
            bbox = afwGeom.Box2I(p0, p1)
            tsi = afwImage.MaskedImageF(tMi, bbox, origin=afwImage.LOCAL)
            ssi = afwImage.MaskedImageF(sMi, bbox, origin=afwImage.LOCAL)

            kc = diffimLib.makeKernelCandidate(xCoord, yCoord, tsi, ssi, policyFake)
            kernelCellSet.insertCandidate(kc)

    tMi.setXY0(0, 0)
    sMi.setXY0(0, 0)
    return tMi, sMi, sKernel, kernelCellSet, configFake




def backgroundSubtract(config, maskedImages):
    """Subtract the background from masked images.

    Parameters
    ----------
    config : TODO: DM-17458
        TODO: DM-17458
    maskedImages : `list` of `lsst.afw.image.MaskedImage`
        TODO: DM-17458

    Returns
    -------
    TODO: DM-17458
        TODO: DM-17458
    """
    backgrounds = []
    t0 = time.time()
    algorithm = config.algorithm
    binsize = config.binSize
    undersample = config.undersampleStyle
    bctrl = afwMath.BackgroundControl(algorithm)
    bctrl.setUndersampleStyle(undersample)
    for maskedImage in maskedImages:
        bctrl.setNxSample(maskedImage.getWidth()//binsize + 1)
        bctrl.setNySample(maskedImage.getHeight()//binsize + 1)
        image = maskedImage.getImage()
        backobj = afwMath.makeBackground(image, bctrl)

        image -= backobj.getImageF()
        backgrounds.append(backobj.getImageF())
        del backobj

    t1 = time.time()
    logger = Log.getLogger("ip.diffim.backgroundSubtract")
    logger.debug("Total time for background subtraction : %.2f s", (t1 - t0))
    return backgrounds




def writeKernelCellSet(kernelCellSet, psfMatchingKernel, backgroundModel, outdir):
    """TODO: DM-17458

    Parameters
    ----------
    kernelCellSet : TODO: DM-17458
        TODO: DM-17458
    psfMatchingKernel : TODO: DM-17458
        TODO: DM-17458
    backgroundModel : TODO: DM-17458
        TODO: DM-17458
    outdir : TODO: DM-17458
        TODO: DM-17458
    """
    if not os.path.isdir(outdir):
        os.makedirs(outdir)

    for cell in kernelCellSet.getCellList():
        for cand in cell.begin(False):  # False = include bad candidates
            if cand.getStatus() == afwMath.SpatialCellCandidate.GOOD:
                xCand = int(cand.getXCenter())
                yCand = int(cand.getYCenter())
                idCand = cand.getId()
                diffIm = cand.getDifferenceImage(diffimLib.KernelCandidateF.ORIG)
                kernel = cand.getKernelImage(diffimLib.KernelCandidateF.ORIG)
                diffIm.writeFits(os.path.join(outdir, 'diffim_c%d_x%d_y%d.fits' % (idCand, xCand, yCand)))
                kernel.writeFits(os.path.join(outdir, 'kernel_c%d_x%d_y%d.fits' % (idCand, xCand, yCand)))

                # Diffim from spatial model
                ski = afwImage.ImageD(kernel.getDimensions())
                psfMatchingKernel.computeImage(ski, False, xCand, yCand)
                sk = afwMath.FixedKernel(ski)
                sbg = backgroundModel(xCand, yCand)
                sdmi = cand.getDifferenceImage(sk, sbg)
                sdmi.writeFits(os.path.join(outdir, 'sdiffim_c%d_x%d_y%d.fits' % (idCand, xCand, yCand)))




def sourceToFootprintList(candidateInList, templateExposure, scienceExposure, kernelSize, config, log):
    """Convert a list of sources for the PSF-matching Kernel to Footprints.

    Parameters
    ----------
    candidateInList : TODO: DM-17458
        Input list of Sources
    templateExposure : TODO: DM-17458
        Template image, to be checked for Mask bits in Source Footprint
    scienceExposure : TODO: DM-17458
        Science image, to be checked for Mask bits in Source Footprint
    kernelSize : TODO: DM-17458
        TODO: DM-17458
    config : TODO: DM-17458
        Config that defines the Mask planes that indicate an invalid Source and Bbox grow radius
    log : TODO: DM-17458
        Log for output

    Returns
    -------
    candidateOutList : `list`
        a list of dicts having a "source" and "footprint" field, to be used for Psf-matching

    Raises
    ------
    RuntimeError
        TODO: DM-17458

    Notes
    -----
    Takes an input list of Sources that were selected to constrain
    the Psf-matching Kernel and turns them into a List of Footprints,
    which are used to seed a set of KernelCandidates.  The function
    checks both the template and science image for masked pixels,
    rejecting the Source if certain Mask bits (defined in config) are
    set within the Footprint.
    """

    candidateOutList = []
    fsb = diffimLib.FindSetBitsU()
    badBitMask = 0
    for mp in config.badMaskPlanes:
        badBitMask |= afwImage.Mask.getPlaneBitMask(mp)
    bbox = scienceExposure.getBBox()

    # Size to grow Sources
    if config.scaleByFwhm:
        fpGrowPix = int(config.fpGrowKernelScaling*kernelSize + 0.5)
    else:
        fpGrowPix = config.fpGrowPix
    log.info("Growing %d kernel candidate stars by %d pixels", len(candidateInList), fpGrowPix)

    for kernelCandidate in candidateInList:
        if not type(kernelCandidate) == afwTable.SourceRecord:
            raise RuntimeError("Candiate not of type afwTable.SourceRecord")
        bm1 = 0
        bm2 = 0
        center = afwGeom.Point2I(scienceExposure.getWcs().skyToPixel(kernelCandidate.getCoord()))
        if center[0] < bbox.getMinX() or center[0] > bbox.getMaxX():
            continue
        if center[1] < bbox.getMinY() or center[1] > bbox.getMaxY():
            continue

        xmin = center[0] - fpGrowPix
        xmax = center[0] + fpGrowPix
        ymin = center[1] - fpGrowPix
        ymax = center[1] + fpGrowPix

        # Keep object centered
        if (xmin - bbox.getMinX()) < 0:
            xmax += (xmin - bbox.getMinX())
            xmin -= (xmin - bbox.getMinX())
        if (ymin - bbox.getMinY()) < 0:
            ymax += (ymin - bbox.getMinY())
            ymin -= (ymin - bbox.getMinY())
        if (bbox.getMaxX() - xmax) < 0:
            xmin -= (bbox.getMaxX() - xmax)
            xmax += (bbox.getMaxX() - xmax)
        if (bbox.getMaxY() - ymax) < 0:
            ymin -= (bbox.getMaxY() - ymax)
            ymax += (bbox.getMaxY() - ymax)
        if xmin > xmax or ymin > ymax:
            continue

        kbbox = afwGeom.Box2I(afwGeom.Point2I(xmin, ymin), afwGeom.Point2I(xmax, ymax))
        try:
            fsb.apply(afwImage.MaskedImageF(templateExposure.getMaskedImage(), kbbox, deep=False).getMask())
            bm1 = fsb.getBits()
            fsb.apply(afwImage.MaskedImageF(scienceExposure.getMaskedImage(), kbbox, deep=False).getMask())
            bm2 = fsb.getBits()
        except Exception:
            pass
        else:
            if not((bm1 & badBitMask) or (bm2 & badBitMask)):
                candidateOutList.append({'source': kernelCandidate,
                                         'footprint': afwDetect.Footprint(afwGeom.SpanSet(kbbox))})
    log.info("Selected %d / %d sources for KernelCandidacy", len(candidateOutList), len(candidateInList))
    return candidateOutList


def sourceTableToCandidateList(sourceTable, templateExposure, scienceExposure, kConfig, dConfig, log,
                               basisList, doBuild=False):
    """Convert a list of Sources into KernelCandidates.

    The KernelCandidates are used for fitting the Psf-matching kernel.

    Parameters
    ----------
    sourceTable : TODO: DM-17458
        TODO: DM-17458
    templateExposure : TODO: DM-17458
        TODO: DM-17458
    scienceExposure : TODO: DM-17458
        TODO: DM-17458
    kConfig : TODO: DM-17458
        TODO: DM-17458
    dConfig : TODO: DM-17458
        TODO: DM-17458
    log : TODO: DM-17458
        TODO: DM-17458
    basisList : TODO: DM-17458
        TODO: DM-17458
    doBuild : `bool`, optional
        TODO: DM-17458

    Returns
    -------
    TODO: DM-17458
        TODO: DM-17458
    """
    kernelSize = basisList[0].getWidth()
    footprintList = sourceToFootprintList(list(sourceTable), templateExposure, scienceExposure,
                                          kernelSize, dConfig, log)
    candList = []

    if doBuild and not basisList:
        doBuild = False
    else:
        policy = pexConfig.makePolicy(kConfig)
        visitor = diffimLib.BuildSingleKernelVisitorF(basisList, policy)

    policy = pexConfig.makePolicy(kConfig)
    for cand in footprintList:
        bbox = cand['footprint'].getBBox()
        tmi = afwImage.MaskedImageF(templateExposure.getMaskedImage(), bbox)
        smi = afwImage.MaskedImageF(scienceExposure.getMaskedImage(), bbox)
        kCand = diffimLib.makeKernelCandidate(cand['source'], tmi, smi, policy)
        if doBuild:
            visitor.processCandidate(kCand)
            kCand.setStatus(afwMath.SpatialCellCandidate.UNKNOWN)
        candList.append(kCand)
    return candList





class NbasisEvaluator(object):
    """A functor to evaluate the Bayesian Information Criterion for the number of basis sets
    going into the kernel fitting"""

    def __init__(self, psfMatchConfig, psfFwhmPixTc, psfFwhmPixTnc):
        self.psfMatchConfig = psfMatchConfig
        self.psfFwhmPixTc = psfFwhmPixTc
        self.psfFwhmPixTnc = psfFwhmPixTnc
        if not self.psfMatchConfig.kernelBasisSet == "alard-lupton":
            raise RuntimeError("BIC only implemnted for AL (alard lupton) basis")

    def __call__(self, kernelCellSet, log):
        d1, d2, d3 = self.psfMatchConfig.alardDegGauss
        bicArray = {}
        for d1i in range(1, d1 + 1):
            for d2i in range(1, d2 + 1):
                for d3i in range(1, d3 + 1):
                    dList = [d1i, d2i, d3i]
                    bicConfig = type(self.psfMatchConfig)(self.psfMatchConfig, alardDegGauss=dList)
                    kList = makeKernelBasisList(bicConfig, self.psfFwhmPixTc, self.psfFwhmPixTnc)
                    k = len(kList)
                    visitor = diffimLib.BuildSingleKernelVisitorF(kList, pexConfig.makePolicy(bicConfig))
                    visitor.setSkipBuilt(False)
                    kernelCellSet.visitCandidates(visitor, bicConfig.nStarPerCell)

                    for cell in kernelCellSet.getCellList():
                        for cand in cell.begin(False):  # False = include bad candidates
                            if cand.getStatus() != afwMath.SpatialCellCandidate.GOOD:
                                continue
                            diffIm = cand.getDifferenceImage(diffimLib.KernelCandidateF.RECENT)
                            bbox = cand.getKernel(diffimLib.KernelCandidateF.RECENT).shrinkBBox(
                                diffIm.getBBox(afwImage.LOCAL))
                            diffIm = type(diffIm)(diffIm, bbox, True)
                            chi2 = diffIm.getImage().getArray()**2/diffIm.getVariance().getArray()
                            n = chi2.shape[0]*chi2.shape[1]
                            bic = np.sum(chi2) + k*np.log(n)
                            if cand.getId() not in bicArray:
                                bicArray[cand.getId()] = {}
                            bicArray[cand.getId()][(d1i, d2i, d3i)] = bic

        bestConfigs = []
        for candId in bicArray:
            cconfig, cvals = list(bicArray[candId].keys()), list(bicArray[candId].values())
            idx = np.argsort(cvals)
            bestConfig = cconfig[idx[0]]
            bestConfigs.append(bestConfig)

        counter = Counter(bestConfigs).most_common(3)
        log.info("B.I.C. prefers basis complexity %s %d times; %s %d times; %s %d times",
                 counter[0][0], counter[0][1],
                 counter[1][0], counter[1][1],
                 counter[2][0], counter[2][1])
        return counter[0][0], counter[1][0], counter[2][0]