randomGalSimFakes.py

Go to the documentation of this file.

import numpy as np
from astropy.io import fits

import lsst.afw.image
import lsst.afw.geom
import lsst.afw.math
import lsst.afw.cameraGeom
import lsst.pex.config
from lsst.pipe.tasks.fakes import BaseFakeSourcesConfig, BaseFakeSourcesTask
import lsst.synpipe.makeFakeGalaxy as makeFake


class RandomGalSimFakesConfig(BaseFakeSourcesConfig):
    galList = lsst.pex.config.Field(dtype=str, doc="catalog of galaxies to add")
    margin = lsst.pex.config.Field(dtype=int, default=None, optional=True,
                                   doc="Size of margin at edge that should not be added")
    seed = lsst.pex.config.Field(dtype=int, default=1,
                                 doc="Seed for random number generator")
    galType = lsst.pex.config.ChoiceField(dtype=str, default='sersic',
                                          allowed={'dsersic': 'double sersic galaxies added',
                                                   'sersic': 'single sersic galaxies added',
                                                   'real': 'real HST galaxy images added'},
                                          doc='type of GalSim galaxies to add')
    nGal = lsst.pex.config.Field(dtype=int, doc="""number of galaxies to add, if 0, then everything in catalog,
 otherwise a random subset of nGal from the catalog""", default=0)


class RandomGalSimFakesTask(BaseFakeSourcesTask):
    ConfigClass = RandomGalSimFakesConfig

    def __init__(self, **kwargs):
        BaseFakeSourcesTask.__init__(self, **kwargs)
        print("RNG seed:", self.config.seed)
        self.rng = lsst.afw.math.Random(seed=self.config.seed)
        self.npRand = np.random.RandomState(self.config.seed)
        self.galData = fits.open(self.config.galList)[1].data

    def run(self, exposure, background):

        self.log.info("Adding fake random galaxies")
        psf = exposure.getPsf()
        psfBBox = psf.computeImage().getBBox()
        minMargin = int(np.floor(max(psfBBox.getWidth(), psfBBox.getHeight())/2)) + 1
        md = exposure.getMetadata()
        expBBox = exposure.getBBox()
        scalingMatrix = np.array([[0.0, 1.0], [1.0, 0.0]]) / exposure.getWcs().pixelScale().asArcseconds()

        if self.config.nGal == 0:
            doGal = enumerate(self.galData)
        else:
            inds = self.npRand.choice(list(range(len(self.galData))), size=self.config.nGal, replace=False)
            doGal = list(zip(inds, self.galData[inds]))

        for igal, gal in doGal:
            try:
                galident = gal["ID"]
            except KeyError:
                galident = igal + 1

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

            # don't put the galaxy within one PSF box of the edge
            # or within the given pixel margin
            if self.config.margin is not None:
                margin = self.config.margin
            else:
                margin = minMargin
            bboxI = (exposure.getBBox(lsst.afw.image.PARENT))
            bboxI.grow(-margin)
            bboxD = lsst.afw.geom.BoxD(bboxI)
            x = self.rng.flat(bboxD.getMinX(), bboxD.getMaxX())
            y = self.rng.flat(bboxD.getMinY(), bboxD.getMaxY())
            # TODO: check for overlaps here and regenerate x,y if necessary

            psfImage = psf.computeKernelImage(lsst.afw.geom.Point2D(x, y))
            galArray = makeFake.makeGalaxy(flux, gal, psfImage.getArray(), self.config.galType,
                                           transform=scalingMatrix)
            galImage = lsst.afw.image.ImageF(galArray.astype(np.float32))
            galBBox = galImage.getBBox(lsst.afw.image.PARENT)
            galImage = lsst.afw.math.offsetImage(galImage,
                                                 x - galBBox.getWidth()/2.0 + 0.5,
                                                 y - galBBox.getHeight()/2.0 + 0.5,
                                                 'lanczos3')
            galBBox = galImage.getBBox(lsst.afw.image.PARENT)

            # check that we're within the larger exposure, otherwise crop
            if expBBox.contains(galImage.getBBox(lsst.afw.image.PARENT)) is False:
                newBBox = galImage.getBBox(lsst.afw.image.PARENT)
                newBBox.clip(expBBox)
                self.log.info("Cropping FAKE%d from %s to %s"%(galident, str(galBBox), str(newBBox)))
                galImage = galImage.Factory(galImage, newBBox, lsst.afw.image.PARENT)
                galBBox = newBBox

            galMaskedImage = lsst.afw.image.MaskedImageF(galImage)

            mask = galMaskedImage.getMask()
            mask.set(self.bitmask)

            md.set("FAKE%d" % gal['ID'], "%.3f, %.3f" % (x, y))
            self.log.info("Adding fake at: %.1f,%.1f" % (x, y))

            # TODO: set the mask
            galMaskedImage.getMask().set(self.bitmask)
            subMaskedImage = exposure.getMaskedImage().Factory(exposure.getMaskedImage(),
                                                               galMaskedImage.getBBox(lsst.afw.image.PARENT),
                                                               lsst.afw.image.PARENT)
            subMaskedImage += galMaskedImage