SkybackgroundMap.py

from observation_sim.mock_objects.SpecDisperser import SpecDisperser
from observation_sim.mock_objects.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('observation_sim.sky_background.data.sky').joinpath(skyfn) as data_path:
            skySpec = np.loadtxt(data_path)
    except AttributeError:
        with pkg_resources.path('observation_sim.sky_background.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('observation_sim.sky_background.data.sky').joinpath(skyfn) as data_path:
            skySpec = np.loadtxt(data_path)
    except AttributeError:
        with pkg_resources.path('observation_sim.sky_background.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