diffimTools.py

Go to the documentation of this file.

#
# LSST Data Management System
# Copyright 2008-2016 LSST Corporation.
#
# This product includes software developed by the
# LSST Project (http://www.lsst.org/).
#
# 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 LSST License Statement and
# the GNU General Public License along with this program.  If not,
# see <http://www.lsstcorp.org/LegalNotices/>.
#

__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 dimensions and shape
        and its expectation value is the value of im at each pixel

    Returns
    -------
    noiseIm : 'lsst.afw.image.Image'
        TODO: DM-17458

    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):
    """TODO: DM-17458

    Parameters
    ----------
    sizeCell : `int`, optional
        TODO: DM-17458
    nCell : `int`, optional
        TODO: DM-17458
    deltaFunctionCounts : `float`, optional
        TODO: DM-17458
    tGaussianWidth : `float`, optional
        TODO: DM-17458
    addNoise : `bool`, optional
        TODO: DM-17458
    bgValue : `float`, optional
        TODO: DM-17458
    display : `bool`, optional
        TODO: DM-17458

    Returns
    -------
    TODO: DM-17458
        TODO: DM-17458
    """
    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.ds9 as ds9
        ds9.mtv(tMi, frame=1)
        ds9.mtv(sMi, frame=2)

    # 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]