FlatLED.py

import galsim
import os, sys
import numpy as np
from astropy.io import fits
from scipy.interpolate import griddata
import math
import astropy.constants as cons
from astropy.table import Table
from ObservationSim.MockObject.SpecDisperser import SpecDisperser
import time
from scipy import interpolate

from ObservationSim.MockObject.MockObject import MockObject
# from ObservationSim.Straylight import calculateSkyMap_split_g

# flatDir = '/Volumes/EAGET/LED_FLAT/'
LED_name = ['LED1', 'LED2', 'LED3', 'LED4', 'LED5', 'LED6', 'LED7', 'LED8', 'LED9', 'LED10', 'LED11', 'LED12', 'LED13',
            'LED14']
cwaves_name = {'LED1': '275', 'LED2': '310', 'LED3': '430', 'LED4': '505', 'LED5': '545', 'LED6': '590', 'LED7': '670',
          'LED8': '760', 'LED9': '940', 'LED10': '940', 'LED11': '1050', 'LED12': '1050',
          'LED13': '340', 'LED14': '365'}

cwaves = {'LED1': 2750, 'LED2': 3100, 'LED3': 4300, 'LED4': 5050, 'LED5': 5250, 'LED6': 5900, 'LED7': 6700,
          'LED8': 7600, 'LED9': 8800, 'LED10': 9400, 'LED11': 10500, 'LED12': 15500, 'LED13': 3400, 'LED14': 3650}
cwaves_fwhm = {'LED1': 110, 'LED2': 120, 'LED3': 200, 'LED4': 300, 'LED5': 300, 'LED6': 130, 'LED7': 210,
          'LED8': 260, 'LED9': 400, 'LED10': 370, 'LED11': 500, 'LED12': 1400, 'LED13': 90, 'LED14': 100}
# LED_QE = {'LED1': 0.3, 'LED2': 0.4, 'LED13': 0.5, 'LED14': 0.5, 'LED10': 0.4}
# e-/ms
fluxLED = {'LED1': 0.16478729, 'LED2': 0.084220931, 'LED3': 2.263360617, 'LED4': 2.190623489, 'LED5': 0.703504768,
           'LED6': 0.446117963, 'LED7': 0.647122098, 'LED8': 0.922313442,
           'LED9': 0.987278143, 'LED10': 2.043989167, 'LED11': 0.612571429, 'LED12': 1.228915663, 'LED13': 0.17029384,
           'LED14': 0.27842925}

mirro_eff = {'GU':0.61, 'GV':0.8, 'GI':0.8}

class FlatLED(object):
    def __init__(self, chip,filt, flatDir = '/Users/zhangxin/Work/SlitlessSim/csst_sls_calibration/flat_field_cube/models/', logger=None):
        # self.led_type_list = led_type_list
        self.flatDir = flatDir
        self.filt = filt
        self.chip = chip
        self.logger = logger

    ###
    ### return LED flat, e/s
    ###
    def getLEDImage(self, led_type='LED1'):
        # cwave = cwaves[led_type]
        flat = fits.open(self.flatDir + 'model_' + cwaves_name[led_type] + 'nm.fits')
        xlen = flat[0].header['NAXIS1']
        ylen = 601
        x = np.linspace(0, self.chip.npix_x * 6, xlen)
        y = np.linspace(0, self.chip.npix_y * 5, ylen)
        xx, yy = np.meshgrid(x, y)

        a1 = flat[0].data[0:ylen, 0:xlen]
        # z = np.sin((xx+yy+xx**2+yy**2))
        # fInterp = interp2d(xx, yy, z, kind='linear')

        X_ = np.hstack((xx.flatten()[:, None], yy.flatten()[:, None]))
        Z_ = a1.flatten()

        n_x = np.arange(0, self.chip.npix_x * 6, 1)
        n_y = np.arange(0, self.chip.npix_y * 5, 1)

        M, N = np.meshgrid(n_x, n_y)

        i = self.chip.rowID - 1
        j = self.chip.colID - 1
        U = griddata(X_, Z_, (
            M[self.chip.npix_y * i:self.chip.npix_y * (i + 1), self.chip.npix_x * j:self.chip.npix_x * (j + 1)],
            N[self.chip.npix_y * i:self.chip.npix_y * (i + 1), self.chip.npix_x * j:self.chip.npix_x * (j + 1)]),
                     method='cubic')
        U = U/np.mean(U)
        flatImage = U*fluxLED[led_type]*1000
        return flatImage

    def drawObj_LEDFlat_img(self, led_type_list=['LED1'], exp_t_list=[0.1]):
        if len(led_type_list) > len(exp_t_list):
            return np.ones([self.chip.npix_y,self.chip.npix_x])

        ledFlat = np.zeros([self.chip.npix_y,self.chip.npix_x])
        for i in np.arange(len(led_type_list)):
            led_type = led_type_list[i]
            exp_t = exp_t_list[i]
            unitFlatImg = self.getLEDImage(led_type=led_type)
            led_wave = cwaves[led_type]
            led_fwhm = cwaves_fwhm[led_type]
            led_spec = self.gaussian1d_profile_led(led_wave, led_fwhm)
            speci = interpolate.interp1d(led_spec['WAVELENGTH'], led_spec['FLUX'])
            w_list = np.arange(self.filt.blue_limit, self.filt.red_limit, 0.5) #A

            f_spec = speci(w_list)
            ccd_bp = self.chip._getChipEffCurve(self.chip.filter_type)
            ccd_eff = ccd_bp.__call__(w_list / 10.)
            filt_bp = self.filt.filter_bandpass
            fil_eff = filt_bp.__call__(w_list / 10.)
            t_spec = np.trapz(f_spec*ccd_eff*fil_eff, w_list)
            print(i, np.mean(unitFlatImg), t_spec, exp_t)
            unitFlatImg = unitFlatImg * t_spec

            ledFlat = ledFlat+unitFlatImg*exp_t
        return ledFlat

    def drawObj_LEDFlat_slitless(self, led_type_list=['LED1'], exp_t_list=[0.1]):
        if len(led_type_list) != len(exp_t_list):
            return np.ones([self.chip.npix_y,self.chip.npix_x])

        ledFlat = np.zeros([self.chip.npix_y,self.chip.npix_x])
        for i in np.arange(len(led_type_list)):
            led_type = led_type_list[i]
            exp_t = exp_t_list[i]
            unitFlatImg = self.getLEDImage(led_type=led_type)
            ledFlat_ = unitFlatImg*exp_t
            ledFlat_ = ledFlat_ / mirro_eff[self.filt.filter_type]
            ledFlat_.astype(np.float32)
            led_wave = cwaves[led_type]
            led_fwhm = cwaves_fwhm[led_type]
            led_spec = self.gaussian1d_profile_led(led_wave, led_fwhm)
            ledspec_map = self.calculateLEDSpec(
                skyMap=ledFlat_,
                blueLimit=self.filt.blue_limit,
                redLimit=self.filt.red_limit,
                conf=self.chip.sls_conf,
                pixelSize=self.chip.pix_scale,
                isAlongY=0,
                flat_cube=self.chip.flat_cube, led_spec=led_spec)

            ledFlat = ledFlat + ledspec_map

        return ledFlat

    def drawObj_LEDFlat(self, led_type_list=['LED1'], exp_t_list=[0.1]):
        if self.chip.survey_type == "photometric":
            return self.drawObj_LEDFlat_img(led_type_list=led_type_list, exp_t_list=exp_t_list)
        elif self.chip.survey_type == "spectroscopic":
            return self.drawObj_LEDFlat_slitless(led_type_list=led_type_list, exp_t_list=exp_t_list)


    def gaussian1d_profile_led(self, xc=5050, fwhm=300):
        sigma = fwhm/2.355
        x_radii = int(5*sigma + 1)
        xlist = np.arange(xc-x_radii, xc+x_radii, 0.5)
        xlist_ = np.zeros(len(xlist) + 2)
        xlist_[1:-1] = xlist
        xlist_[0] = 2550
        xlist_[-1] = 10000
        data = np.exp((-(xlist-xc)*(xlist-xc))/(2*sigma*sigma))/(np.sqrt(2*math.pi)*sigma)
        data_ = np.zeros(len(xlist) + 2)
        data_[1:-1] = data
        return Table(np.array([xlist_.astype(np.float32), data_.astype(np.float32)]).T, names=('WAVELENGTH', 'FLUX'))

    def calculateLEDSpec(self, skyMap=None, blueLimit=4200, redLimit=6500,
                                conf=[''], pixelSize=0.074, isAlongY=0,
                                split_pos=3685, flat_cube=None, led_spec=None):

        conf1 = conf[0]
        conf2 = conf[0]
        if np.size(conf) == 2:
            conf2 = conf[1]

        skyImg = galsim.Image(skyMap, xmin=0, ymin=0)

        tbstart = blueLimit
        tbend = redLimit

        fimg = np.zeros_like(skyMap)

        fImg = galsim.Image(fimg)

        spec = led_spec
        if isAlongY == 0:
            directParm = 0
        if isAlongY == 1:
            directParm = 1

        if split_pos >= skyImg.array.shape[directParm]:
            skyImg1 = galsim.Image(skyImg.array)
            origin1 = [0, 0]
            # sdp = specDisperser.specDisperser(orig_img=skyImg1, xcenter=skyImg1.center.x, ycenter=skyImg1.center.y,
            #                                   full_img=fimg, tar_spec=spec, band_start=tbstart, band_end=tbend,
            #                                   origin=origin1,
            #                                   conf=conf1)
            # sdp.compute_spec_orders()

            y_len = skyMap.shape[0]
            x_len = skyMap.shape[1]
            delt_x = 100
            delt_y = 100

            sub_y_start_arr = np.arange(0, y_len, delt_y)
            sub_y_end_arr = sub_y_start_arr + delt_y
            sub_y_end_arr[-1] = min(sub_y_end_arr[-1], y_len)

            sub_x_start_arr = np.arange(0, x_len, delt_x)
            sub_x_end_arr = sub_x_start_arr + delt_x
            sub_x_end_arr[-1] = min(sub_x_end_arr[-1], x_len)

            for i, k1 in enumerate(sub_y_start_arr):
                sub_y_s = k1
                sub_y_e = sub_y_end_arr[i]

                sub_y_center = (sub_y_s + sub_y_e) / 2.

                for j, k2 in enumerate(sub_x_start_arr):
                    sub_x_s = k2
                    sub_x_e = sub_x_end_arr[j]

                    skyImg_sub = galsim.Image(skyImg.array[sub_y_s:sub_y_e, sub_x_s:sub_x_e])
                    origin_sub = [sub_y_s, sub_x_s]
                    sub_x_center = (sub_x_s + sub_x_e) / 2.

                    sdp = SpecDisperser(orig_img=skyImg_sub, xcenter=sub_x_center, ycenter=sub_y_center,
                                        origin=origin_sub,
                                        tar_spec=spec,
                                        band_start=tbstart, band_end=tbend,
                                        conf=conf2,
                                        flat_cube=flat_cube, ignoreBeam=['D', 'E'])

                    spec_orders = sdp.compute_spec_orders()

                    for k, v in spec_orders.items():
                        img_s = v[0]
                        origin_order_x = v[1]
                        origin_order_y = v[2]
                        ssImg = galsim.ImageF(img_s)
                        ssImg.setOrigin(origin_order_x, origin_order_y)
                        bounds = ssImg.bounds & fImg.bounds
                        if bounds.area() == 0:
                            continue
                        fImg[bounds] = fImg[bounds] + ssImg[bounds]



        else:


            # sdp.compute_spec_orders()
            y_len = skyMap.shape[0]
            x_len = skyMap.shape[1]
            delt_x = 500
            delt_y = y_len

            sub_y_start_arr = np.arange(0, y_len, delt_y)
            sub_y_end_arr = sub_y_start_arr + delt_y
            sub_y_end_arr[-1] = min(sub_y_end_arr[-1], y_len)

            delt_x = split_pos - 0
            sub_x_start_arr = np.arange(0, split_pos, delt_x)
            sub_x_end_arr = sub_x_start_arr + delt_x
            sub_x_end_arr[-1] = min(sub_x_end_arr[-1], split_pos)

            for i, k1 in enumerate(sub_y_start_arr):
                sub_y_s = k1
                sub_y_e = sub_y_end_arr[i]

                sub_y_center = (sub_y_s + sub_y_e) / 2.

                for j, k2 in enumerate(sub_x_start_arr):
                    sub_x_s = k2
                    sub_x_e = sub_x_end_arr[j]
                    # print(i,j,sub_y_s, sub_y_e,sub_x_s,sub_x_e)
                    T1 = time.time()
                    skyImg_sub = galsim.Image(skyImg.array[sub_y_s:sub_y_e, sub_x_s:sub_x_e])
                    origin_sub = [sub_y_s, sub_x_s]
                    sub_x_center = (sub_x_s + sub_x_e) / 2.

                    sdp = SpecDisperser(orig_img=skyImg_sub, xcenter=sub_x_center, ycenter=sub_y_center,
                                        origin=origin_sub,
                                        tar_spec=spec,
                                        band_start=tbstart, band_end=tbend,
                                        conf=conf1,
                                        flat_cube=flat_cube)

                    spec_orders = sdp.compute_spec_orders()

                    for k, v in spec_orders.items():
                        img_s = v[0]
                        origin_order_x = v[1]
                        origin_order_y = v[2]
                        ssImg = galsim.ImageF(img_s)
                        ssImg.setOrigin(origin_order_x, origin_order_y)
                        bounds = ssImg.bounds & fImg.bounds
                        if bounds.area() == 0:
                            continue
                        fImg[bounds] = fImg[bounds] + ssImg[bounds]

                    T2 = time.time()

                    print('time: %s ms' % ((T2 - T1) * 1000))

            delt_x = x_len - split_pos
            sub_x_start_arr = np.arange(split_pos, x_len, delt_x)
            sub_x_end_arr = sub_x_start_arr + delt_x
            sub_x_end_arr[-1] = min(sub_x_end_arr[-1], x_len)

            for i, k1 in enumerate(sub_y_start_arr):
                sub_y_s = k1
                sub_y_e = sub_y_end_arr[i]

                sub_y_center = (sub_y_s + sub_y_e) / 2.

                for j, k2 in enumerate(sub_x_start_arr):
                    sub_x_s = k2
                    sub_x_e = sub_x_end_arr[j]
                    # print(i,j,sub_y_s, sub_y_e,sub_x_s,sub_x_e)

                    T1 = time.time()

                    skyImg_sub = galsim.Image(skyImg.array[sub_y_s:sub_y_e, sub_x_s:sub_x_e])
                    origin_sub = [sub_y_s, sub_x_s]
                    sub_x_center = (sub_x_s + sub_x_e) / 2.

                    sdp = SpecDisperser(orig_img=skyImg_sub, xcenter=sub_x_center, ycenter=sub_y_center,
                                        origin=origin_sub,
                                        tar_spec=spec,
                                        band_start=tbstart, band_end=tbend,
                                        conf=conf2,
                                        flat_cube=flat_cube)

                    spec_orders = sdp.compute_spec_orders()

                    for k, v in spec_orders.items():
                        img_s = v[0]
                        origin_order_x = v[1]
                        origin_order_y = v[2]
                        ssImg = galsim.ImageF(img_s)
                        ssImg.setOrigin(origin_order_x, origin_order_y)
                        bounds = ssImg.bounds & fImg.bounds
                        if bounds.area() == 0:
                            continue
                        fImg[bounds] = fImg[bounds] + ssImg[bounds]
                    T2 = time.time()

                    print('time: %s ms' % ((T2 - T1) * 1000))

        if isAlongY == 1:
            fimg, tmx, tmy = rotate90(array_orig=fImg.array, xc=0, yc=0, isClockwise=0)
        else:
            fimg = fImg.array

        fimg = fimg * pixelSize * pixelSize

        return fimg