MockObject.py 30.2 KiB
Newer Older
Fang Yuedong's avatar
Fang Yuedong committed
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

Fang Yuedong's avatar
Fang Yuedong committed

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
Fang Yuedong's avatar
Fang Yuedong committed
        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))
Fang Yuedong's avatar
Fang Yuedong committed
        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 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 == 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

            # nphotons_sum += nphotons
            # print("nphotons_sub-band_%d = %.2f"%(i, nphotons))

            # Get PSF model
            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(nphotons)
            # star = galsim.Convolve(psf, star)
            star = psf.withFlux(nphotons)

            stamp = star.drawImage(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])
Fang Yuedong's avatar
Fang Yuedong committed
            #########################################################
            # img_s, orig_off = convolveGaussXorders(img_s, xOrderSigPlus[k])
            orig_off = 0
Fang Yuedong's avatar
Fang Yuedong committed
            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])
Fang Yuedong's avatar
Fang Yuedong committed
            #########################################################
            stamp = galsim.ImageF(img_s)
Fang Yuedong's avatar
Fang Yuedong committed
            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]
Fang Yuedong's avatar
Fang Yuedong committed
            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]
Zhang Xin's avatar
Zhang Xin committed

                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)
Zhang Xin's avatar
Zhang Xin committed
                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'}
Zhang Xin's avatar
Zhang Xin committed
                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]

Zhang Xin's avatar
Zhang Xin committed
                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)
Zhang Xin's avatar
Zhang Xin committed

                # print('DEBUG: BEGIN -----------',bandNo,k)
Zhang Xin's avatar
Zhang Xin committed

Zhang Xin's avatar
Zhang Xin committed

                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:
Zhang Xin's avatar
Zhang Xin committed
                    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)
Zhang Xin's avatar
Zhang Xin committed
                    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
Fang Yuedong's avatar
Fang Yuedong committed
        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-4
        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)

Zhang Xin's avatar
Zhang Xin committed
            # psf list :["A","B","C","D","E"]
Zhang Xin's avatar
Zhang Xin committed
                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
Zhang Xin's avatar
Zhang Xin committed
                for order in ["A", "B"]:
                    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)
                    # star_p = galsim.Convolve(psf, star)
Zhang Xin's avatar
Zhang Xin committed
                    star_p = psf.withFlux(tel.pupil_area * exptime)
Zhang Xin's avatar
Zhang Xin committed
                        starImg = star_p.drawImage(
                            wcs=chip_wcs_local, offset=offset)
                        nnx = starImg.xmax - starImg.xmin + 1
                        nny = starImg.ymax - starImg.ymin + 1
                    else:
Zhang Xin's avatar
Zhang Xin committed
                        starImg = star_p.drawImage(
                            nx=nnx, ny=nny, wcs=chip_wcs_local, offset=offset)
                    # 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)
Zhang Xin's avatar
Zhang Xin committed
                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)
Zhang Xin's avatar
Zhang Xin committed
                star_p = psf.withFlux(tel.pupil_area * exptime)
                starImg = star_p.drawImage(wcs=chip_wcs_local, offset=offset)
                starImg.setOrigin(0, 0)
Zhang Xin's avatar
Zhang Xin committed
                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)
Fang Yuedong's avatar
Fang Yuedong committed

            origin_star = [y_nominal - (starImg.center.y - starImg.ymin),
                           x_nominal - (starImg.center.x - starImg.xmin)]
Zhang Xin's avatar
Zhang Xin committed

Fang Yuedong's avatar
Fang Yuedong committed
            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])
Fang Yuedong's avatar
Fang Yuedong committed
                # part img disperse
Zhang Xin's avatar
Zhang Xin committed
                star_p1s = []
                for starImg in starImg_List:
Fang Yuedong's avatar
Fang Yuedong committed

                    subImg_p1 = starImg.array[:, 0:subSlitPos]
                    star_p1 = galsim.Image(subImg_p1)
                    star_p1.setOrigin(0, 0)
                    star_p1s.append(star_p1)
Fang Yuedong's avatar
Fang Yuedong committed
                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,
Fang Yuedong's avatar
Fang Yuedong committed
                                       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)

Zhang Xin's avatar
Zhang Xin committed
                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)

Zhang Xin's avatar
Zhang Xin committed
                star_p2s = []
                for starImg in starImg_List:
Fang Yuedong's avatar
Fang Yuedong committed

Zhang Xin's avatar
Zhang Xin committed
                                              subSlitPos:starImg.array.shape[1]]
                    star_p2 = galsim.Image(subImg_p2)
                    star_p2.setOrigin(0, 0)
                    star_p2s.append(star_p2)
Fang Yuedong's avatar
Fang Yuedong committed
                origin_p2 = [origin_star[0], grating_split_pos_chip]
                xcenter_p2 = max(x_nominal, grating_split_pos_chip) - 0
Fang Yuedong's avatar
Fang Yuedong committed
                ycenter_p2 = y_nominal - 0

                sdp_p2 = SpecDisperser(orig_img=star_p2s, xcenter=xcenter_p2,
Fang Yuedong's avatar
Fang Yuedong committed
                                       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)

Zhang Xin's avatar
Zhang Xin committed
                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)
Fang Yuedong's avatar
Fang Yuedong committed

                del sdp_p1
                del sdp_p2
            elif grating_split_pos_chip <= gal_origin[1]:
                sdp = SpecDisperser(orig_img=starImg_List, xcenter=x_nominal - 0,
Fang Yuedong's avatar
Fang Yuedong committed
                                    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)
Zhang Xin's avatar
Zhang Xin committed
                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)
Fang Yuedong's avatar
Fang Yuedong committed
                del sdp
            elif grating_split_pos_chip >= gal_end[1]:
                sdp = SpecDisperser(orig_img=starImg_List, xcenter=x_nominal - 0,
Fang Yuedong's avatar
Fang Yuedong committed
                                    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)
Zhang Xin's avatar
Zhang Xin committed
                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)
Fang Yuedong's avatar
Fang Yuedong committed
                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 == 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 != 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