Stamp.py

import os, sys
import numpy as np
import galsim
import astropy.constants as cons
from astropy.table import Table
from scipy import interpolate

from ObservationSim.MockObject.MockObject import MockObject
from ObservationSim.MockObject.SpecDisperser import SpecDisperser
from ObservationSim.MockObject._util import eObs, integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getObservedSED, getABMAG,convolveGaussXorders

class Stamp(MockObject):
    def __init__(self, param, logger=None):
        super().__init__(param, logger=logger)

    def unload_SED(self):
        """(Test) free up SED memory
        """
        del self.sed

    def drawObj_multiband(self, tel, pos_img, psf_model, bandpass_list, filt, chip, nphotons_tot=None, g1=0, g2=0, exptime=150., fd_shear=None):
        if nphotons_tot == None:
            nphotons_tot = self.getElectronFluxFilt(filt, tel, exptime)


        try:
            full = integrate_sed_bandpass(sed=self.sed, bandpass=filt.bandpass_full)
        except Exception as e:
            print(e)
            self.logger.error(e)
            return False

        #nphotons_sum = 0
        #photons_list = []
        #xmax, ymax = 0, 0

        if self.getMagFilter(filt) <= 15:
            folding_threshold = 5.e-4
        else:
            folding_threshold = 5.e-3
        gsp = galsim.GSParams(folding_threshold=folding_threshold)

        self.real_pos = self.getRealPos(chip.img, global_x=self.posImg.x, global_y=self.posImg.y,
                                        img_real_wcs=self.chip_wcs)

        x, y = self.real_pos.x + 0.5, self.real_pos.y + 0.5
        x_nominal = int(np.floor(x + 0.5))
        y_nominal = int(np.floor(y + 0.5))
        dx = x - x_nominal
        dy = y - y_nominal
        offset = galsim.PositionD(dx, dy)

        chip_wcs_local = self.chip_wcs.local(self.real_pos)
        is_updated = 0

        if fd_shear:
            g1 += fd_shear.g1
            g2 += fd_shear.g2
        gal_shear = galsim.Shear(g1=g1, g2=g2)

        for i in range(len(bandpass_list)):
            bandpass = bandpass_list[i]
            try:
                sub = integrate_sed_bandpass(sed=self.sed, bandpass=bandpass)
            except Exception as e:
                print(e)
                if self.logger:
                    self.logger.error(e)
                continue
            ratio = sub/full
            if not (ratio == -1 or (ratio != ratio)):
                nphotons = ratio * nphotons_tot
            else:
                continue
            #nphotons_sum += nphotons

            psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass, folding_threshold=folding_threshold)

            _gal  = self.param['image']
            galImg= galsim.ImageF(_gal, scale=self.param['pixScale'])
            gal_temp= galsim.InterpolatedImage(galImg)
            gal_temp= gal_temp.shear(gal_shear)
            gal_temp= gal_temp.withFlux(nphotons)

            gal_temp= galsim.Convolve(psf, gal_temp)

            if i == 0:
                gal = gal_temp
            else:
                gal = gal + gal_temp

        stamp = gal.drawImage(wcs=chip_wcs_local, offset=offset)
        if np.sum(np.isnan(stamp.array)) > 0:
            # ERROR happens
            return 2, pos_shear

        stamp.setCenter(x_nominal, y_nominal)
        bounds = stamp.bounds & galsim.BoundsI(0, chip.npix_x - 1, 0, chip.npix_y - 1)

        if bounds.area() > 0:
            chip.img.setOrigin(0, 0)
            chip.img[bounds] += stamp[bounds]
            is_updated = 1
            chip.img.setOrigin(chip.bound.xmin, chip.bound.ymin)
            del stamp

        if is_updated == 0:
            print("fits obj %s missed"%(self.id))
            if self.logger:
                self.logger.info("fits obj %s missed"%(self.id))
            return 0, pos_shear

        return 1, pos_shear


    def drawObj_slitless(self, tel, pos_img, psf_model, bandpass_list, filt, chip, nphotons_tot=None, g1=0, g2=0,
                         exptime=150., normFilter=None, grating_split_pos=3685, fd_shear=None):
        if normFilter is not None:
            norm_thr_rang_ids = normFilter['SENSITIVITY'] > 0.001
            sedNormFactor = getNormFactorForSpecWithABMAG(ABMag=self.param['mag_use_normal'], spectrum=self.sed,
                                                        norm_thr=normFilter,
                                                        sWave=np.floor(normFilter[norm_thr_rang_ids][0][0]),
                                                        eWave=np.ceil(normFilter[norm_thr_rang_ids][-1][0]))
            if sedNormFactor == 0:
                return 2, None
        else:
            sedNormFactor = 1.
        normalSED = Table(np.array([self.sed['WAVELENGTH'], self.sed['FLUX'] * sedNormFactor]).T,
                          names=('WAVELENGTH', 'FLUX'))

        self.real_pos = self.getRealPos(chip.img, global_x=self.posImg.x, global_y=self.posImg.y,
                                        img_real_wcs=self.chip_wcs)

        x, y = self.real_pos.x + 0.5, self.real_pos.y + 0.5
        x_nominal = int(np.floor(x + 0.5))
        y_nominal = int(np.floor(y + 0.5))
        dx = x - x_nominal
        dy = y - y_nominal
        offset = galsim.PositionD(dx, dy)

        chip_wcs_local = self.chip_wcs.local(self.real_pos)


        if self.getMagFilter(filt) <= 15:
            folding_threshold = 5.e-4
        else:
            folding_threshold = 5.e-3
        gsp = galsim.GSParams(folding_threshold=folding_threshold)
        # nphotons_sum = 0

        flat_cube = chip.flat_cube

        xOrderSigPlus = {'A':1.3909419820029296,'B':1.4760376591236062,'C':4.035447379743442,'D':5.5684364343742825,'E':16.260021029735388}
        grating_split_pos_chip = 0 + grating_split_pos

        branges = np.zeros([len(bandpass_list), 2])

        # print(hasattr(psf_model, 'bandranges'))

        if hasattr(psf_model, 'bandranges'):
            if psf_model.bandranges is None:
                return 2, None
            if len(psf_model.bandranges) != len(bandpass_list):
                return 2, None
            branges = psf_model.bandranges
        else:
            for i in range(len(bandpass_list)):
                branges[i, 0] = bandpass_list[i].blue_limit * 10
                branges[i, 1] = bandpass_list[i].red_limit * 10

        for i in range(len(bandpass_list)):
            # bandpass = bandpass_list[i]
            brange = branges[i]

            # psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass, folding_threshold=folding_threshold)

            _gal  = self.param['image']
            galImg= galsim.ImageF(_gal, scale=self.param['pixScale'])
            gal   = galsim.InterpolatedImage(galImg)

            # (TEST) Random knots
            # knots = galsim.RandomKnots(npoints=100, profile=disk)
            # kfrac = np.random.random()*(1.0 - self.bfrac)
            # gal = self.bfrac * bulge + (1.0 - self.bfrac - kfrac) * disk + kfrac * knots

            gal = gal.withFlux(tel.pupil_area * exptime)
            if fd_shear:
                g1 += fd_shear.g1
                g2 += fd_shear.g2
            gal_shear = galsim.Shear(g1=g1, g2=g2)
            gal = gal.shear(gal_shear)
            # gal = galsim.Convolve(psf, gal)

            # if not big_galaxy: # Not apply PSF for very big galaxy
            #     gal = galsim.Convolve(psf, gal)
            #     # if fd_shear is not None:
            #     #     gal = gal.shear(fd_shear)

            starImg = gal.drawImage(wcs=chip_wcs_local, offset=offset,method = 'real_space')

            origin_star = [y_nominal - (starImg.center.y - starImg.ymin),
                           x_nominal - (starImg.center.x - starImg.xmin)]
            starImg.setOrigin(0, 0)
            gal_origin = [origin_star[0], origin_star[1]]
            gal_end = [origin_star[0] + starImg.array.shape[0] - 1, origin_star[1] + starImg.array.shape[1] - 1]

            if gal_origin[1] < grating_split_pos_chip < gal_end[1]:
                subSlitPos = int(grating_split_pos_chip - gal_origin[1] + 1)
                ## part img disperse

                subImg_p1 = starImg.array[:, 0:subSlitPos]
                star_p1 = galsim.Image(subImg_p1)
                star_p1.setOrigin(0, 0)
                origin_p1 = origin_star
                xcenter_p1 = min(x_nominal,grating_split_pos_chip-1) - 0
                ycenter_p1 = y_nominal-0

                sdp_p1 = SpecDisperser(orig_img=star_p1, xcenter=xcenter_p1,
                                    ycenter=ycenter_p1, origin=origin_p1,
                                    tar_spec=normalSED,
                                    band_start=brange[0], band_end=brange[1],
                                    conf=chip.sls_conf[0],
                                    isAlongY=0,
                                    flat_cube=flat_cube)

                # self.addSLStoChipImage(sdp=sdp_p1, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=chip_wcs_local)
                pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp_p1, chip=chip, pos_img_local=[xcenter_p1, ycenter_p1],
                                                          psf_model=psf_model, bandNo=i + 1,
                                                          grating_split_pos=grating_split_pos,
                                                          local_wcs=chip_wcs_local, pos_img = pos_img)

                subImg_p2 = starImg.array[:, subSlitPos+1:starImg.array.shape[1]]
                star_p2 = galsim.Image(subImg_p2)
                star_p2.setOrigin(0, 0)
                origin_p2 = [origin_star[0], grating_split_pos_chip]
                xcenter_p2 = max(x_nominal, grating_split_pos_chip - 1) - 0
                ycenter_p2 = y_nominal - 0

                sdp_p2 = SpecDisperser(orig_img=star_p2, xcenter=xcenter_p2,
                                       ycenter=ycenter_p2, origin=origin_p2,
                                       tar_spec=normalSED,
                                       band_start=brange[0], band_end=brange[1],
                                       conf=chip.sls_conf[1],
                                       isAlongY=0,
                                       flat_cube=flat_cube)

                # self.addSLStoChipImage(sdp=sdp_p2, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=chip_wcs_local)
                pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp_p2, chip=chip, pos_img_local=[xcenter_p2, ycenter_p2],
                                                          psf_model=psf_model, bandNo=i + 1,
                                                          grating_split_pos=grating_split_pos,
                                                          local_wcs=chip_wcs_local, pos_img = pos_img)

                del sdp_p1
                del sdp_p2
            elif grating_split_pos_chip<=gal_origin[1]:
                sdp = SpecDisperser(orig_img=starImg, xcenter=x_nominal - 0,
                                    ycenter=y_nominal - 0, origin=origin_star,
                                    tar_spec=normalSED,
                                    band_start=brange[0], band_end=brange[1],
                                    conf=chip.sls_conf[1],
                                    isAlongY=0,
                                    flat_cube=flat_cube)
                # self.addSLStoChipImage(sdp=sdp, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=chip_wcs_local)
                pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp, chip=chip, pos_img_local=[x_nominal, y_nominal],
                                                          psf_model=psf_model, bandNo=i + 1,
                                                          grating_split_pos=grating_split_pos,
                                                          local_wcs=chip_wcs_local, pos_img = pos_img)
                del sdp
            elif grating_split_pos_chip>=gal_end[1]:
                sdp = SpecDisperser(orig_img=starImg, xcenter=x_nominal - 0,
                                    ycenter=y_nominal - 0, origin=origin_star,
                                    tar_spec=normalSED,
                                    band_start=brange[0], band_end=brange[1],
                                    conf=chip.sls_conf[0],
                                    isAlongY=0,
                                    flat_cube=flat_cube)
                # self.addSLStoChipImage(sdp=sdp, chip=chip, xOrderSigPlus = xOrderSigPlus, local_wcs=chip_wcs_local)
                pos_shear = self.addSLStoChipImageWithPSF(sdp=sdp, chip=chip, pos_img_local=[x_nominal, y_nominal],
                                                          psf_model=psf_model, bandNo=i + 1,
                                                          grating_split_pos=grating_split_pos,
                                                          local_wcs=chip_wcs_local, pos_img = pos_img)
                del sdp

            # print(self.y_nominal, starImg.center.y, starImg.ymin)
            # del psf
        return 1, pos_shear