MockObject.py

import os
import galsim
import numpy as np
import astropy.constants as cons
from astropy import wcs
from astropy.table import Table
import astropy.io.fits as fitsio

from observation_sim.mock_objects._util import magToFlux, VC_A, convolveGaussXorders, convolveImg
from observation_sim.mock_objects._util import integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getObservedSED, \
    getABMAG
from observation_sim.mock_objects.SpecDisperser import SpecDisperser

from observation_sim.instruments.chip import chip_utils


class MockObject(object):
    def __init__(self, param, logger=None):
        self.param = param
        for key in self.param:
            setattr(self, key, self.param[key])

        if self.param["star"] == 0:
            self.type = "galaxy"
        elif self.param["star"] == 1:
            self.type = "star"
        elif self.param["star"] == 2:
            self.type = "quasar"
        # mock_stamp_START
        elif self.param["star"] == 3:
            self.type = "stamp"
        # mock_stamp_END
        # for calibration
        elif self.param["star"] == 4:
            self.type = "calib"
        # END

        self.sed = None
        self.fd_shear = None
        # Place holder for outputs
        self.additional_output_str = ""
        self.logger = logger

    def getMagFilter(self, filt):
        if filt.filter_type in ["GI", "GV", "GU"]:
            return self.param["mag_use_normal"]
        return self.param["mag_%s" % filt.filter_type.lower()]

    def getFluxFilter(self, filt):
        return self.param["flux_%s" % filt.filter_type.lower()]

    def getNumPhotons(self, flux, tel, exptime=150.):
        pupil_area = tel.pupil_area * (100.) ** 2  # m^2 to cm^2
        return flux * pupil_area * exptime

    def getElectronFluxFilt(self, filt, tel, exptime=150.):
        # photonEnergy = filt.getPhotonE()
        # flux = magToFlux(self.getMagFilter(filt))
        # factor = 1.0e4 * flux/photonEnergy * VC_A * (1.0/filt.blue_limit - 1.0/filt.red_limit)
        # return factor * filt.efficiency * tel.pupil_area * exptime
        flux = self.getFluxFilter(filt)
        return flux * tel.pupil_area * exptime

    def getPosWorld(self, ra_offset=0., dec_offset=0.):
        ra = self.param["ra"] + ra_offset
        dec = self.param["dec"] + dec_offset
        return galsim.CelestialCoord(ra=ra * galsim.degrees, dec=dec * galsim.degrees)

    def getPosImg_Offset_WCS(self, img, fdmodel=None, chip=None, verbose=True, chip_wcs=None, img_header=None, ra_offset=0., dec_offset=0.):
        self.posImg = img.wcs.toImage(self.getPosWorld(ra_offset, dec_offset))
        self.localWCS = img.wcs.local(self.posImg)
        # Apply field distortion model
        if (fdmodel is not None) and (chip is not None):
            if verbose:
                print("\n")
                print("Before field distortion:\n")
                print("x = %.2f, y = %.2f\n" %
                      (self.posImg.x, self.posImg.y), flush=True)
            self.posImg, self.fd_shear = fdmodel.get_distorted(
                chip=chip, pos_img=self.posImg)
            if self.posImg is None:
                return None, None, None, None, None
            if verbose:
                print("After field distortion:\n")
                print("x = %.2f, y = %.2f\n" %
                      (self.posImg.x, self.posImg.y), flush=True)

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

        # Deal with chip rotation
        if chip_wcs is not None:
            self.chip_wcs = chip_wcs
        elif img_header is not None:
            self.chip_wcs = galsim.FitsWCS(header=img_header)
        else:
            self.chip_wcs = None

        return self.posImg, self.offset, self.localWCS, self.chip_wcs, self.fd_shear

    def getRealPos(self, img, global_x=0., global_y=0., img_real_wcs=None):
        img_global_pos = galsim.PositionD(global_x, global_y)
        cel_pos = img.wcs.toWorld(img_global_pos)
        realPos = img_real_wcs.toImage(cel_pos)
        return realPos

    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 is None:
            nphotons_tot = self.getElectronFluxFilt(filt, tel, exptime)
        # print("nphotons_tot = ", nphotons_tot)

        try:
            full = integrate_sed_bandpass(
                sed=self.sed, bandpass=filt.bandpass_full)
        except Exception as e:
            print(e)
            if self.logger:
                self.logger.error(e)
            return 2, None
        # Set Galsim Parameters
        if self.getMagFilter(filt) <= 15:
            folding_threshold = 5.e-8
        else:
            folding_threshold = 5.e-3
        gsp = galsim.GSParams(folding_threshold=folding_threshold)

        # Get real image position of object (deal with chip rotation w.r.t its center)
        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)
        # Get real local wcs of object (deal with chip rotation w.r.t its center)
        chip_wcs_local = self.chip_wcs.local(self.real_pos)
        is_updated = 0

        # Loop over all sub-bandpasses
        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
            # print("nphotons_sub-band_%d = %.2f"%(i, nphotons))

            # Get PSF model
            EXTRA = False
            if self.getMagFilter(filt) <= filt.mag_saturation-1.:
                EXTRA = True
            psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass,
                                               folding_threshold=folding_threshold, extrapolate=EXTRA)
            # star = galsim.DeltaFunction(gsparams=gsp)
            # star = star.withFlux(nphotons)
            # star = galsim.Convolve(psf, star)
            star = psf.withFlux(nphotons)

            stamp = star.drawImage(method='no_pixel', wcs=chip_wcs_local, offset=offset)
            if np.sum(np.isnan(stamp.array)) > 0:
                continue
            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:
            # Return code 0: object has missed this detector
            print("obj %s missed" % (self.id))
            if self.logger:
                self.logger.info("obj %s missed" % (self.id))
            return 0, pos_shear
        return 1, pos_shear  # Return code 1: draw sucesss

    def addSLStoChipImage(self, sdp=None, chip=None, xOrderSigPlus=None, local_wcs=None):
        spec_orders = sdp.compute_spec_orders()
        for k, v in spec_orders.items():
            img_s = v[0]
            #########################################################
            # DEBUG
            #########################################################
            # print("before convolveGaussXorders, img_s:", img_s)
            # nan_ids = np.isnan(img_s)
            # if img_s[nan_ids].shape[0] > 0:
            #     # img_s[nan_ids] = 0
            #     print("DEBUG: before convolveGaussXorders specImg nan num is",
            #           img_s[nan_ids].shape[0])
            #########################################################
            # img_s, orig_off = convolveGaussXorders(img_s, xOrderSigPlus[k])
            orig_off = 0
            origin_order_x = v[1] - orig_off
            origin_order_y = v[2] - orig_off
            #########################################################
            # DEBUG
            #########################################################
            # print("DEBUG: orig_off is", orig_off)
            # nan_ids = np.isnan(img_s)
            # if img_s[nan_ids].shape[0] > 0:
            #     img_s[nan_ids] = 0
            #     print("DEBUG: specImg nan num is", img_s[nan_ids].shape[0])
            #########################################################
            stamp = galsim.ImageF(img_s)
            stamp.wcs = local_wcs
            stamp.setOrigin(origin_order_x, origin_order_y)

            bounds = stamp.bounds & galsim.BoundsI(
                0, chip.npix_x - 1, 0, chip.npix_y - 1)
            if bounds.area() == 0:
                continue
            chip.img.setOrigin(0, 0)
            chip.img[bounds] = chip.img[bounds]+stamp[bounds]
            chip.img.setOrigin(chip.bound.xmin, chip.bound.ymin)
            del stamp
        del spec_orders

    def addSLStoChipImageWithPSF(self, sdp=None, chip=None, pos_img_local=[1, 1], psf_model=None, bandNo=1, grating_split_pos=3685, local_wcs=None, pos_img=None):
        spec_orders = sdp.compute_spec_orders()
        pos_shear = galsim.Shear(e=0., beta=(np.pi/2)*galsim.radians)
        if chip.slsPSFOptim:
            for k, v in spec_orders.items():
                img_s = v[0]

                nan_ids = np.isnan(img_s)
                if img_s[nan_ids].shape[0] > 0:
                    img_s[nan_ids] = 0
                    print("DEBUG: specImg nan num is", img_s[nan_ids].shape[0])
                #########################################################
                # img_s, orig_off = convolveImg(img_s, psf_img_m)
                orig_off = [0, 0]
                origin_order_x = v[1] - orig_off[0]
                origin_order_y = v[2] - orig_off[1]

                specImg = galsim.ImageF(img_s)

                specImg.wcs = local_wcs
                specImg.setOrigin(origin_order_x, origin_order_y)

                bounds = specImg.bounds & galsim.BoundsI(
                    0, chip.npix_x - 1, 0, chip.npix_y - 1)
                if bounds.area() == 0:
                    continue

                # orders = {'A': 'order1', 'B': 'order0', 'C': 'order2', 'D': 'order-1', 'E': 'order-2'}
                orders = {'A': 'order1', 'B': 'order0',
                          'C': 'order0', 'D': 'order0', 'E': 'order0'}
                gratingN = chip_utils.getChipSLSGratingID(chip.chipID)[1]
                if pos_img_local[0] < grating_split_pos:
                    gratingN = chip_utils.getChipSLSGratingID(chip.chipID)[0]

                chip.img_stack[gratingN][orders[k]
                                         ]['w' + str(bandNo)].setOrigin(0, 0)
                chip.img_stack[gratingN][orders[k]]['w' + str(
                    bandNo)][bounds] = chip.img_stack[gratingN][orders[k]]['w' + str(bandNo)][bounds] + specImg[bounds]
                chip.img_stack[gratingN][orders[k]]['w' +
                                                    str(bandNo)].setOrigin(chip.bound.xmin, chip.bound.ymin)

        else:
            for k, v in spec_orders.items():
                # img_s = v[0]
                # # print(bandNo,k)
                # try:
                #     psf, pos_shear = psf_model.get_PSF(
                #         chip, pos_img_local=pos_img_local, bandNo=bandNo, galsimGSObject=True, g_order=k, grating_split_pos=grating_split_pos)
                # except:
                #     psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img)

                # psf_img = psf.drawImage(nx=100, ny=100, wcs=local_wcs)

                # psf_img_m = psf_img.array

                # #########################################################
                # # DEBUG
                # #########################################################
                # # ids_p = psf_img_m < 0
                # # psf_img_m[ids_p] = 0

                # # from astropy.io import fits
                # # fits.writeto(str(bandNo) + '_' + str(k) + '_psf.fits', psf_img_m)

                # # print("DEBUG: orig_off is", orig_off)
                # nan_ids = np.isnan(img_s)
                # if img_s[nan_ids].shape[0] > 0:
                #     img_s[nan_ids] = 0
                #     print("DEBUG: specImg nan num is", img_s[nan_ids].shape[0])
                # #########################################################
                # img_s, orig_off = convolveImg(img_s, psf_img_m)
                # origin_order_x = v[1] - orig_off[0]
                # origin_order_y = v[2] - orig_off[1]

                # specImg = galsim.ImageF(img_s)
                # # photons = galsim.PhotonArray.makeFromImage(specImg)
                # # photons.x += origin_order_x
                # # photons.y += origin_order_y

                # # xlen_imf = int(specImg.xmax - specImg.xmin + 1)
                # # ylen_imf = int(specImg.ymax - specImg.ymin + 1)
                # # stamp = galsim.ImageF(xlen_imf, ylen_imf)
                # # stamp.wcs = local_wcs
                # # stamp.setOrigin(origin_order_x, origin_order_y)

                # specImg.wcs = local_wcs
                # specImg.setOrigin(origin_order_x, origin_order_y)

                # print('DEBUG: BEGIN -----------',bandNo,k)

                img_s = v[0]

                nan_ids = np.isnan(img_s)
                if img_s[nan_ids].shape[0] > 0:
                    img_s[nan_ids] = 0
                    print("DEBUG: specImg nan num is", img_s[nan_ids].shape[0])
                #########################################################
                origin_order_x = v[1]
                origin_order_y = v[2]
                specImg = galsim.ImageF(img_s)

                specImg.wcs = local_wcs
                specImg.setOrigin(origin_order_x, origin_order_y)
                try:
                    specImg = psf_model.get_PSF_AND_convolve_withsubImg(
                        chip, cutImg=specImg, pos_img_local=pos_img_local, bandNo=bandNo, g_order=k, grating_split_pos=grating_split_pos)
                except:
                    psf, pos_shear = psf_model.get_PSF(
                        chip=chip, pos_img=pos_img)

                    psf_img = psf.drawImage(nx=100, ny=100, wcs=local_wcs)

                    psf_img_m = psf_img.array

                    img_s, orig_off = convolveImg(img_s, psf_img_m)
                    origin_order_x = v[1] - orig_off[0]
                    origin_order_y = v[2] - orig_off[1]

                    specImg = galsim.ImageF(img_s)

                    specImg.wcs = local_wcs
                    specImg.setOrigin(origin_order_x, origin_order_y)

                bounds = specImg.bounds & galsim.BoundsI(
                    0, chip.npix_x - 1, 0, chip.npix_y - 1)
                if bounds.area() == 0:
                    continue
                chip.img.setOrigin(0, 0)
                chip.img[bounds] = chip.img[bounds] + specImg[bounds]
                # stamp[bounds] = chip.img[bounds]
                # # chip.sensor.accumulate(photons, stamp)
                # chip.img[bounds] = stamp[bounds]
                chip.img.setOrigin(chip.bound.xmin, chip.bound.ymin)
                # del stamp
        del spec_orders
        return 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.

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

        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)

        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])

        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)

            # star = galsim.DeltaFunction(gsparams=gsp)
            # star = star.withFlux(tel.pupil_area * exptime)

            # psf list :["A","B","C","D","E"]
            starImg_List = []
            try:
                pos_img_local = [0, 0]
                x_start = chip.x_cen/chip.pix_size - chip.npix_x / 2.
                y_start = chip.y_cen/chip.pix_size - chip.npix_y / 2.
                pos_img_local[0] = pos_img.x - x_start
                pos_img_local[1] = pos_img.y - y_start
                nnx = 0
                nny = 0
                for order in ["B", "A"]:
                    EXTRA = False
                    if self.getMagFilter(filt) <= filt.mag_saturation-2.:
                        EXTRA = True
                        nnx = 2048
                        nny = 2048
                    psf, pos_shear = psf_model.get_PSF(
                        chip, pos_img_local=pos_img_local, bandNo=i+1, galsimGSObject=True, g_order=order, grating_split_pos=grating_split_pos, extrapolate=EXTRA, ngg=3072)
                    # star_p = galsim.Convolve(psf, star)
                    star_p = psf.withFlux(tel.pupil_area * exptime)
                    if nnx == 0:
                        starImg = star_p.drawImage(
                            wcs=chip_wcs_local, offset=offset, method='no_pixel')
                        nnx = starImg.xmax - starImg.xmin + 1
                        nny = starImg.ymax - starImg.ymin + 1
                    else:
                        starImg = star_p.drawImage(
                            nx=nnx, ny=nny, wcs=chip_wcs_local, offset=offset, method='no_pixel')
                    # n1 = np.sum(np.isinf(starImg.array))
                    # n2 = np.sum(np.isnan(starImg.array))
                    # if n1>0 or n2 > 0:
                    #     print("DEBUG: MockObject, inf:%d, nan:%d"%(n1, n2))
                    starImg.setOrigin(0, 0)
                    starImg_List.append(starImg)
                for order in ["C", "D", "E"]:
                    starImg_List.append(starImg)
            except:
                psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img)
                # star_p = galsim.Convolve(psf, star)
                star_p = psf.withFlux(tel.pupil_area * exptime)
                starImg = star_p.drawImage(wcs=chip_wcs_local, offset=offset, method='no_pixel')
                starImg.setOrigin(0, 0)
                for order in ["A", "B", "C", "D", "E"]:
                    starImg_List.append(starImg)

            # psf_tmp = galsim.Gaussian(sigma=0.002)
            # star = galsim.Convolve(psf_tmp, star)

            # starImg = star.drawImage(
            #     nx=60, ny=60, wcs=chip_wcs_local, offset=offset)

            origin_star = [y_nominal - (starImg.center.y - starImg.ymin),
                           x_nominal - (starImg.center.x - starImg.xmin)]

            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])
                # part img disperse
                star_p1s = []
                for starImg in starImg_List:

                    subImg_p1 = starImg.array[:, 0:subSlitPos]
                    star_p1 = galsim.Image(subImg_p1)
                    star_p1.setOrigin(0, 0)
                    star_p1s.append(star_p1)
                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_p1s, 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)

                star_p2s = []
                for starImg in starImg_List:

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

                sdp_p2 = SpecDisperser(orig_img=star_p2s, 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_List, 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_List, 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
            # del psf
        return 1, pos_shear

    def SNRestimate(self, img_obj, flux, noise_level=0.0, seed=31415):
        img_flux = img_obj.added_flux
        stamp = img_obj.copy() * 0.0
        rng = galsim.BaseDeviate(seed)
        gaussianNoise = galsim.GaussianNoise(rng, sigma=noise_level)
        stamp.addNoise(gaussianNoise)
        sig_obj = np.std(stamp.array)
        snr_obj = img_flux / sig_obj
        return snr_obj

    def drawObj_PSF(self, tel, pos_img, psf_model, bandpass_list, filt, chip, nphotons_tot=None, g1=0, g2=0,
                    exptime=150., fd_shear=None, chip_output=None):
        if nphotons_tot is None:
            nphotons_tot = self.getElectronFluxFilt(filt, tel, exptime)
        # print("nphotons_tot = ", nphotons_tot)

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

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

        # Get real image position of object (deal with chip rotation w.r.t its center)
        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)
        # Get real local wcs of object (deal with chip rotation w.r.t its center)
        chip_wcs_local = self.chip_wcs.local(self.real_pos)
        is_updated = 0

        # Loop over all sub-bandpasses
        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

            # Get PSF model
            psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass,
                                               folding_threshold=folding_threshold)
            star_temp = psf.withFlux(nphotons)

            if i == 0:
                star = star_temp
            else:
                star = star+star_temp

        pixelScale = 0.074
        stamp = star.drawImage(wcs=chip_wcs_local, offset=offset)
        # stamp = star.drawImage(nx=256, ny=256, scale=pixelScale)
        if np.sum(np.isnan(stamp.array)) > 0:
            return None

        fn = chip_output.subdir + "/psfIDW"
        os.makedirs(fn, exist_ok=True)
        fn = fn + "/ccd_{:}".format(chip.chipID) + \
            "_psf_"+str(self.param['id'])+".fits"
        if fn is not None:
            if os.path.exists(fn):
                os.remove(fn)
        hdu = fitsio.PrimaryHDU()
        hdu.data = stamp.array
        hdu.header.set('name',      self.type)
        hdu.header.set('pixScale',  pixelScale)
        hdu.header.set('objID',     self.param['id'])
        hdu.writeto(fn)

        del stamp
        return None