SkybackgroundMap.py

from ObservationSim.MockObject.SpecDisperser import SpecDisperser
from ObservationSim.MockObject.SpecDisperser import rotate90

import galsim
import numpy as np
from astropy.table import Table
from scipy import interpolate

import galsim
import astropy.constants as cons

import os

import time

try:
    import importlib.resources as pkg_resources
except ImportError:
    # Try backported to PY<37 'importlib_resources'
    import importlib_resources as pkg_resources


###calculate sky map by sky SED

def calculateSkyMap_split_g(skyMap=None, blueLimit=4200, redLimit=6500, skyfn='sky_emiss_hubble_50_50_A.dat', conf=[''], pixelSize=0.074, isAlongY=0,
                            split_pos=3685, flat_cube = None, zoldial_spec = None):
    # skyMap = np.ones([yLen, xLen], dtype='float32')
    #
    # if isAlongY == 1:
    #     skyMap = np.ones([xLen, yLen], dtype='float32')

    # for i in range(len(conf)):
    #     conf[i] = os.path.join(SLSSIM_PATH, conf[i])
    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)
    try:
        with pkg_resources.files('ObservationSim.Straylight.data.sky').joinpath(skyfn) as data_path:
            skySpec = np.loadtxt(data_path)
    except AttributeError:
        with pkg_resources.path('ObservationSim.Straylight.data.sky', skyfn) as data_path:
            skySpec = np.loadtxt(data_path)
    # skySpec = np.loadtxt(skyfn)
    spec = Table(np.array([skySpec[:, 0], skySpec[:, 1]]).T, names=('WAVELENGTH', 'FLUX'))

    if zoldial_spec is not None:
        deltL = 0.5
        lamb = np.arange(2000, 11000, deltL)

        speci = interpolate.interp1d(zoldial_spec['WAVELENGTH'], zoldial_spec['FLUX'])

        y = speci(lamb)
        # erg/s/cm2/A --> photo/s/m2/A
        s_flux = y * lamb / (cons.h.value * cons.c.value) * 1e-13
        spec = Table(np.array([lamb, s_flux]).T, names=('WAVELENGTH', 'FLUX'))
    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]

        # sdp = SpecDisperser(orig_img=skyImg1, xcenter=skyImg1.center.x, ycenter=skyImg1.center.y, origin=origin1,
        #                 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:

        # skyImg1 = galsim.Image(skyImg.array[:, 0:split_pos])
        # origin1 = [0, 0]
        # skyImg2 = galsim.Image(skyImg.array[:, split_pos:])
        # origin2 = [0, split_pos]

        # 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 = 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
    
def calculateSkyMap(xLen=9232, yLen=9126, blueLimit=4200, redLimit=6500,
                    skyfn='sky_emiss_hubble_50_50_A.dat', conf='', pixelSize=0.074, isAlongY=0):
    skyMap = np.ones([yLen, xLen], dtype='float32')

    if isAlongY == 1:
        skyMap = np.ones([xLen, yLen], dtype='float32')

    skyImg = galsim.Image(skyMap)

    tbstart = blueLimit
    tbend = redLimit

    fimg = np.zeros_like(skyMap)
    fImg = galsim.Image(fimg)
    try:
        with pkg_resources.files('ObservationSim.Straylight.data.sky').joinpath(skyfn) as data_path:
            skySpec = np.loadtxt(data_path)
    except AttributeError:
        with pkg_resources.path('ObservationSim.Straylight.data.sky', skyfn) as data_path:
            skySpec = np.loadtxt(data_path)
    # skySpec = np.loadtxt(skyfn)
    
    spec = Table(np.array([skySpec[:, 0], skySpec[:, 1]]).T, names=('WAVELENGTH', 'FLUX'))

    sdp = SpecDisperser(orig_img=skyImg, xcenter=skyImg.center.x, ycenter=skyImg.center.y, origin=[1, 1],
                        tar_spec=spec,
                        band_start=tbstart, band_end=tbend,
                        conf=conf)

    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
        fImg[bounds] = fImg[bounds] + ssImg[bounds]

    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