Stamp.py

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

from observation_sim.mock_objects.MockObject import MockObject
from observation_sim.mock_objects.SpecDisperser import SpecDisperser
from observation_sim.mock_objects._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