positionGalSimFakes.py

Go to the documentation of this file.

import os
import fcntl
 
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 as lsstConfig
 
from lsst.pipe.tasks.fakes import BaseFakeSourcesConfig, BaseFakeSourcesTask
 
from . import makeFakeGalaxy as makeFake
 
 
class PositionGalSimFakesConfig(BaseFakeSourcesConfig):
    galList = lsstConfig.Field(dtype=str,
                               doc="catalog of galaxies to add")
    maxMargin = lsstConfig.Field(dtype=int, default=600,
                                 optional=True,
                                 doc="Size of margin")
    seed = lsstConfig.Field(dtype=int, default=1,
                            doc="Seed for random number generator")
    addShear = lsstConfig.Field(dtype=bool, default=False,
                                doc='include shear in the galaxies')
    addMask = lsstConfig.Field(dtype=bool, default=False,
                               doc='add FAKE mask plane')
    sersic_prec = lsstConfig.Field(dtype=float, default=0.0,
                                   doc='The desired precision for n')
    cosStr = 'Use Galsim.COSMOSlog()'
    serStr = 'Single sersic galaxies added'
    dSerStr = 'Double sersic galaxies added'
    realStr = 'Real HST galaxy images added'
    galType = lsstConfig.ChoiceField(dtype=str, default='sersic',
                                     allowed={'dsersic': dSerStr,
                                              'sersic': serStr,
                                              'real': realStr,
                                              'cosmos': cosStr},
                                     doc='type of GalSim galaxies to add')
    exclusionLevel = lsstConfig.ChoiceField(dtype=str, default='none',
                                            allowed={'none': "None",
                                                     'marginal': "Marginal",
                                                     'bad_stamp': "Bad Stamp",
                                                     'bad_fits': "Bad Fits"},
                                            doc='Exclusion level')
 
 
class PositionGalSimFakesTask(BaseFakeSourcesTask):
    ConfigClass = PositionGalSimFakesConfig
 
    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 galaxies at real positions")
        PARENT = lsst.afw.image.PARENT
        psf = exposure.getPsf()
        md = exposure.getMetadata()
        calib = exposure.getCalib()
        expBBox = exposure.getBBox(PARENT)
        wcs = exposure.getWcs()
 
        """Deal with the skipped ones."""
        skipLog = 'runAddFake.skipped'
        if not os.path.isfile(skipLog):
            os.system('touch ' + skipLog)
 
        if self.config.galType == 'cosmos':
            import galsim
            exLevel = self.config.exclusionLevel
            cosmosCat = galsim.COSMOSCatalog(exclusion_level=exLevel)
        else:
            cosmosCat = None
 
        for igal, gal in enumerate(self.galData):
            try:
                galident = gal["ID"]
            except KeyError:
                galident = igal + 1
 
            try:
                flux = calib.getFlux(float(gal['mag']))
            except KeyError:
                raise KeyError("No mag column in %s" % self.config.galList)
 
            try:
                galCoord = lsst.afw.geom.SpherePoint(gal['RA'], gal['DEC'], lsst.afw.geom.degrees)
            except KeyError:
                raise KeyError("No RA/DEC column in {} table".format(self.config.galList))
 
            skyToPixelMatrix = wcs.linearizeSkyToPixel(galCoord, lsst.afw.geom.degrees)
            skyToPixelMatrix = skyToPixelMatrix.getLinear().getParameterVector() / 3600.0
 
            galXY = wcs.skyToPixel(galCoord)
            bboxI = exposure.getBBox(PARENT)
            bboxI.grow(self.config.maxMargin)
            if not bboxI.contains(lsst.afw.geom.Point2I(galXY)):
                # Will just skip this object
                continue
 
            # Check the magnitude
            if gal['mag'] <= 0:
                self.log.info("Mag <= 0: Skipping %d" % galident)
                self.log.info("  mag: %7.3d" % gal['mag'])
                with open(skipLog, "a") as slog:
                    try:
                        fcntl.flock(slog, fcntl.LOCK_EX)
                        slog.write("%8d , negMag\n" % galident)
                        fcntl.flock(slog, fcntl.LOCK_UN)
                    except IOError:
                        continue
                continue
 
            # This is extrapolating for the PSF, probably not a good idea
            #  Return an Image of the PSF, in a form suitable for convolution.
            #  The returned image is normalized to sum to unity.
            try:
                psfImage = psf.computeKernelImage(galXY)
            except Exception:
                # There may not be any data at this point, and the object
                # should be skipped
                continue
            try:
                addShear = self.config.addShear
                prec = self.config.sersic_prec
                galType = self.config.galType
                if self.config.galType != 'cosmos':
                    galArray = makeFake.makeGalaxy(flux, gal,
                                                   psfImage.getArray(),
                                                   galType=galType,
                                                   cosmosCat=None,
                                                   calib=None,
                                                   addShear=addShear,
                                                   transform=skyToPixelMatrix)
                else:
                    galArray = makeFake.makeGalaxy(flux, gal,
                                                   psfImage.getArray(),
                                                   cosmosCat=cosmosCat,
                                                   calib=calib,
                                                   galType=galType,
                                                   addShear=addShear,
                                                   sersic_prec=prec,
                                                   transform=skyToPixelMatrix)
            except IndexError as ierr:
                self.log.info("GalSim Index Error: Skipping %d" % galident)
                self.log.info(ierr.message)
                with open(skipLog, "a") as slog:
                    try:
                        fcntl.flock(slog, fcntl.LOCK_EX)
                        slog.write("%8d , galsimI\n" % galident)
                        fcntl.flock(slog, fcntl.LOCK_UN)
                    except IOError:
                        continue
                continue
            except KeyError as kerr:
                self.log.info("GalSim Key Error: Skipping %d" % galident)
                self.log.info(kerr.message)
                with open(skipLog, "a") as slog:
                    try:
                        fcntl.flock(slog, fcntl.LOCK_EX)
                        slog.write("%8d , galsimK\n" % galident)
                        fcntl.flock(slog, fcntl.LOCK_UN)
                    except IOError:
                        continue
                continue
            except ValueError as verr:
                self.log.info("GalSim Value Error: Skipping %d" % galident)
                self.log.info(verr.message)
                with open(skipLog, "a") as slog:
                    try:
                        fcntl.flock(slog, fcntl.LOCK_EX)
                        slog.write("%8d , galsimV\n" % galident)
                        fcntl.flock(slog, fcntl.LOCK_UN)
                    except IOError:
                        continue
                continue
            except RuntimeError as rerr:
                self.log.info("GalSim Runtime Error: Skipping %d" % galident)
                self.log.info(rerr.message)
                with open(skipLog, "a") as slog:
                    try:
                        fcntl.flock(slog, fcntl.LOCK_EX)
                        slog.write("%8d , galsimR\n" % galident)
                        fcntl.flock(slog, fcntl.LOCK_UN)
                    except IOError:
                        continue
                continue
            except Exception as uerr:
                self.log.info("Unexpected Error: Skipping %d" % galident)
                self.log.info(uerr.message)
                with open(skipLog, "a") as slog:
                    try:
                        fcntl.flock(slog, fcntl.LOCK_EX)
                        slog.write("%8d , Unexpected\n" % galident)
                        fcntl.flock(slog, fcntl.LOCK_UN)
                    except IOError:
                        continue
                continue
 
            galImage = lsst.afw.image.ImageF(galArray.astype(np.float32))
            galBBox = galImage.getBBox(PARENT)
            galX0 = (galXY.getX() - galBBox.getWidth()/2.0 + 0.5)
            galY0 = (galXY.getY() - galBBox.getHeight()/2.0 + 0.5)
            galImage = lsst.afw.math.offsetImage(galImage,
                                                 galX0, galY0,
                                                 'lanczos3')
            galBBox = galImage.getBBox(PARENT)
 
            # Check that we're within the larger exposure, otherwise crop
            parentBox = galImage.getBBox(PARENT)
            if expBBox.contains(parentBox) is False:
                newBBox = galImage.getBBox(PARENT)
                newBBox.clip(expBBox)
                if newBBox.getArea() <= 0:
                    self.log.info("BBoxEdge Error: Skipping %d" % galident)
                    with open(skipLog, "a") as slog:
                        try:
                            fcntl.flock(slog, fcntl.LOCK_EX)
                            slog.write("%8d , bboxEdge\n" % galident)
                            fcntl.flock(slog, fcntl.LOCK_UN)
                        except IOError:
                            continue
                    continue
                self.log.info("Cropping FAKE%d from %s to %s" % (galident,
                                                                 str(galBBox), str(newBBox)))
                galImage = galImage.Factory(galImage, newBBox,
                                            PARENT)
                galBBox = newBBox
 
            galMaskedImage = lsst.afw.image.MaskedImageF(galImage)
 
            # Put information of the added fake galaxies into the header
            md.set("FAKE%s" % str(galident), "%.3f, %.3f" % (galXY.getX(),
                                                             galXY.getY()))
            self.log.info("Adding fake %s at: %.1f,%.1f" % (str(galident),
                                                            galXY.getX(),
                                                            galXY.getY()))
 
            galMaskedImage.getMask().set(self.bitmask)
            if not self.config.addMask:
                try:
                    galMaskedImage.getMask().removeAndClearMaskPlane('FAKE',
                                                                     True)
                except Exception:
                    pass
            try:
                galMaskedImage.getMask().removeAndClearMaskPlane('CROSSTALK',
                                                                 True)
            except Exception:
                pass
            try:
                galMaskedImage.getMask().removeAndClearMaskPlane('UNMASKEDNAN',
                                                                 True)
            except Exception:
                pass
 
            maskedImage = exposure.getMaskedImage()
            try:
                maskedImage.getMask().removeAndClearMaskPlane('CROSSTALK',
                                                              True)
            except Exception:
                pass
            try:
                maskedImage.getMask().removeAndClearMaskPlane('UNMASKEDNAN',
                                                              True)
            except Exception:
                pass
            if not self.config.addMask:
                try:
                    maskedImage.getMask().removeAndClearMaskPlane('FAKE', True)
                except Exception:
                    pass
 
            BBox = galMaskedImage.getBBox(PARENT)
            subMaskedImage = maskedImage.Factory(exposure.getMaskedImage(),
                                                 BBox,
                                                 PARENT)
            subMaskedImage += galMaskedImage
 
            """
            #
            galMaskedImage.getMask().set(self.bitmask)
 
            maskedImage = exposure.getMaskedImage()
            try:
                maskedImage.getMask().removeAndClearMaskPlane('CROSSTALK',
                                                              True)
            except Exception:
                pass
            try:
                maskedImage.getMask().removeAndClearMaskPlane('UNMASKEDNAN',
                                                              True)
            except Exception:
                pass
            try:
                maskedImage.getMask().removeAndClearMaskPlane('FAKE', True)
            except Exception:
                pass
            """