positionStarFakes.py

Go to the documentation of this file.

import numpy as np
from astropy.io import fits

import lsst.afw.geom as afwGeom
import lsst.afw.math as afwMath
import lsst.afw.image as afwImage
import lsst.pex.config as afwConfig

from lsst.pipe.tasks.fakes import BaseFakeSourcesConfig, BaseFakeSourcesTask
from lsst.pex.exceptions import InvalidParameterError


class PositionStarFakesConfig(BaseFakeSourcesConfig):
    starList = afwConfig.Field(dtype=str,
                               doc="Catalog of stars with mags ra/dec")
    seed = afwConfig.Field(dtype=int, default=1,
                           doc="Seed for random number generator")


class PositionStarFakesTask(BaseFakeSourcesTask):
    ConfigClass = PositionStarFakesConfig

    def __init__(self, **kwargs):
        BaseFakeSourcesTask.__init__(self, **kwargs)
        print("RNG seed:", self.config.seed)
        self.rng = afwMath.Random(seed=self.config.seed)
        self.npRand = np.random.RandomState(self.config.seed)
        try:
            self.starData = fits.open(self.config.starList)[1].data
        except Exception:
            raise

    def run(self, exposure, background):

        self.log.info("Adding fake stars at real positions")
        psf = exposure.getPsf()
        psfBBox = psf.computeImage().getBBox()
        margin = max(psfBBox.getWidth(), psfBBox.getHeight())/2 + 1

        PARENT = afwImage.PARENT
        md = exposure.getMetadata()
        expBBox = exposure.getBBox(PARENT)
        wcs = exposure.getWcs()

        for istar, star in enumerate(self.starData):
            try:
                starident = star["ID"]
            except KeyError:
                starident = istar + 1

            try:
                flux = exposure.getCalib().getFlux(float(star['mag']))
            except KeyError:
                raise KeyError("No mag column in %s" % self.config.starList)

            try:
                starCoord = afwGeom.SpherePoint(star['RA'], star['DEC'], afwGeom.degrees)
            except KeyError:
                raise KeyError("No RA/DEC column in table".format(self.config.starList))

            starXY = wcs.skyToPixel(starCoord)
            bboxI = exposure.getBBox(PARENT)
            bboxI.grow(int(margin))
            if not bboxI.contains(afwGeom.Point2I(starXY)):
                continue

            try:
                starImage = psf.computeImage(starXY)
            except InvalidParameterError:
                # This means an image was computed in an area where there was no data
                # continue on to the next star. This is most likely to occur when inserting
                # into coadds
                logmsg = "Skipping fake {} because no input images present at point {}"
                self.log.info(logmsg.format(starident, starXY))
                continue

            starImage *= flux
            starBBox = starImage.getBBox(PARENT)

            # Check that we're within the larger exposure, otherwise crop
            if expBBox.contains(starBBox) is False:
                newBBox = starImage.getBBox(PARENT)
                newBBox.clip(expBBox)
                if newBBox.getArea() <= 0:
                    self.log.info("Skipping fake %d" % starident)
                    continue
                self.log.info("Cropping FAKE%d from %s to %s" % (starident,
                                                                 str(starBBox), str(newBBox)))
                starImage = starImage.Factory(starImage, newBBox, PARENT)
                starBBox = newBBox

            starMaskedImage = afwImage.MaskedImageF(starImage.convertF())

            starMaskedImage.getMask().set(self.bitmask)

            md.set("FAKE%s" % str(starident), "%.3f, %.3f" % (starXY.getX(),
                                                              starXY.getY()))
            self.log.info("Adding fake %s at: %.1f,%.1f" % (str(starident),
                                                            starXY.getX(),
                                                            starXY.getY()))

            maskedImage = exposure.getMaskedImage()
            BBox = starMaskedImage.getBBox(PARENT)
            subMaskedImage = maskedImage.Factory(exposure.getMaskedImage(),
                                                 BBox,
                                                 PARENT)
            subMaskedImage += starMaskedImage