SpecGenerator.py

'''
Author: zx
Date: 2021-04-08 13:49:35
LastEditTime: 2023-02-24 00:57:20
LastEditors: xin zhangxinbjfu@gmail.com
Description: In User Settings Edit
FilePath: /undefined/Users/zhangxin/Work/SlitlessSim/sls_lit_demo/simDemo.py
'''

import galsim
import SpecDisperser
# from numpy import *
import numpy as np
from scipy import interpolate
import astropy.constants as acon
from astropy.table import Table
import math
from astropy.io import fits
import random

from astropy.table import Table
import matplotlib.pyplot as plt

import mpi4py.MPI as MPI

import os,sys

from . import Config


class SpecGenerator(object):
    def __init__(self,sedFn = 'a.txt', grating = 'GI', beam = 'A', aper = 2.0, xcenter = 5000,ycenter = 5000, p_size = 0.074, psf = None, skybg = 0.3, dark = 0.02, readout = 5, t = 150, expNum = 1, config = None, saturation = 90000):
        self.sedFile = sedFn
        self.grating = grating
        self.beam = beam
        self.aper = aper
        self.xcenter = xcenter
        self.ycenter = ycenter
        self.p_size = p_size
        self.psf = psf
        self.skybg = skybg
        self.dark = dark
        self.readout = readout
        self.t = t
        self.expNum = expNum
        self.config = config
        self.saturation = saturation
    
    '''
    @description: 
    @param {*} fn: file name, include 2 column, wavelength(A)  and flux(erg/s/cm2/A) 
    @param {*} s: band start , unit A
    @param {*} e; end, unit A
    @param {*} deltL: sample interval for SED
    @return {*} sed, unit photo/s/m2/A
    '''
    def generateSEDfromFiles(self, fn, s, e, deltL):
        """
        s: lambda start, unit A
        e: lambda end, unit A

        return:
        SEDs is array, 2-dim, (gal_num+1)*(wavelength size), last row is wavelength
        """
        lamb = np.arange(s, e + deltL, deltL)
        spec_orig = np.loadtxt(fn)

        speci = interpolate.interp1d(spec_orig[:, 0], spec_orig[:, 1])
        y = speci(lamb)
        # erg/s/cm2/A --> photo/s/m2/A
        flux = y * lamb / (acon.h.value * acon.c.value) * 1e-13

        SED = Table(np.array([lamb, flux]).T,names=('WAVELENGTH', 'FLUX'))

        return SED


    def generateSpec1dforGal(self, s_n = 1.0, re = 1, pa = 90,q_ell = 0.6,limitfluxratio=0.9):

        specConfile = self.config.conFiles[self.grating]

        throughput_f = self.config.senFisle[self.grating] + self.config.orderIDs[self.beam] + '.fits'

        sed = self.generateSEDfromFiles(self.sedFile,2500,10000,0.01)

        x_nominal = int(np.floor(self.xcenter + 0.5))
        y_nominal = int(np.floor(self.ycenter + 0.5))
        dx = self.xcenter - x_nominal+0.5
        dy = self.ycenter - y_nominal+0.5
        offset = galsim.PositionD(dx, dy)

        # print(skybg)
        # print(specConfile)
        # print(throughput_f)

        # plt.figure()
        # plt.plot(sed['WAVELENGTH'], sed['FLUX'])

        gal = galsim.Sersic(s_n, half_light_radius=re)

        gal_pa = pa * galsim.degrees
        gal_ell = gal.shear(q=q_ell, beta=gal_pa)

        conv_gal = galsim.Convolve([gal_ell,self.psf])


        stamp = conv_gal.drawImage(wcs=galsim.PixelScale(self.p_size), offset=offset)*self.t*self.expNum*math.pi*(self.aper/2)*(self.aper/2)
        stamp.setOrigin(0,0)
        origin_star = [y_nominal - (stamp.center.y - stamp.ymin),
                                x_nominal - (stamp.center.x - stamp.xmin)]


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

        sdp = SpecDisperser.SpecDisperser(orig_img=stamp, xcenter=x_nominal,
                                            ycenter=y_nominal, origin=origin_star,
                                            tar_spec=sed,
                                            conf=specConfile,
                                            isAlongY=0)

        spec_orders = sdp.compute_spec_orders()
        
        thp = Table.read(throughput_f)
        thp_i = interpolate.interp1d(thp['WAVELENGTH'], thp['SENSITIVITY'])

        Aimg_orig = spec_orders[self.beam][0]
        Aimg_ = Aimg_orig

        Aimg_ = Aimg_ + (self.skybg + self.dark)*self.t*self.expNum

        Aimg_ = np.random.poisson(Aimg_)
        for i in np.arange(self.expNum):
            Aimg_ = self.addReadoutNois(img = Aimg_, readout = self.readout)

        Aimg = Aimg_ - (self.skybg + self.dark)*self.t*self.expNum


        wave_pix = spec_orders[self.beam][5]
        wave_pos = spec_orders[self.beam][3]

        wave_pos_y=spec_orders[self.beam][4]

        sh = Aimg.shape
        spec_pix = np.zeros(sh[1])
        err2_pix = np.zeros(sh[1])

        # print(spec_orders[beamOrder][4])
        # print(sh)
        # plt.figure()
        # plt.imshow(Aimg)
        y_cent_pos = int(np.round(np.mean(wave_pos_y)))

        tFlux = np.sum(spec_orders[self.beam][0])
        # print(tFlux)
        fluxRatio = 0
        for i in range(int(sh[0]/2)):
            pFlux = np.sum(spec_orders[self.beam][0][y_cent_pos-i:y_cent_pos+i+1])
            
            fluxRatio = pFlux/tFlux
            if fluxRatio>limitfluxratio:
                break
        
        f1 = spec_orders[self.beam][0][y_cent_pos-i:y_cent_pos+i+1].sum(0)
        f2 = spec_orders[self.beam][0].sum(0)
        ratio_vec = np.zeros_like(f1)
        nozero_flag = f2 != 0
        
        ratio_vec[nozero_flag] = f1[nozero_flag]/f2[nozero_flag]
        # ratio_vec = spec_orders[self.beam][0][y_cent_pos-i:y_cent_pos+i+1].sum(0)/spec_orders[self.beam][0].sum(0)
        y_range = i
        # print(y_range, fluxRatio)
        y_len_pix = 2 * y_range + 1
        for i in range(sh[1]):
            spec_pix[i] = sum(Aimg[y_cent_pos-y_range:y_cent_pos+y_range+1, i])
            err2_pix[i] = sum(Aimg_orig[y_cent_pos-y_range:y_cent_pos+y_range+1, i]) + (self.skybg + self.dark)*self.t * y_len_pix * self.expNum + self.readout*self.readout * y_len_pix * self.expNum

        bRange = self.config.bandRanges[self.grating]
        wave_flux = np.zeros(wave_pix.shape[0])
        err_flux = np.zeros(wave_pix.shape[0])
        specRangeImg = []

        true_center = stamp.center + galsim.PositionD(self.xcenter-x_nominal, self.ycenter-y_nominal)
        wavePos_x = true_center.x + wave_pos - wave_pos[0]

        wavePos_x_interp = np.arange(int(wavePos_x[0]), int(wavePos_x[-1]))
        lam_trace = np.interp(wavePos_x_interp,wavePos_x,wave_pix)

        wave_flux = np.zeros(lam_trace.shape[0])
        err_flux = np.zeros(lam_trace.shape[0])



        for i in np.arange(1, lam_trace.shape[0] - 1):
            w = lam_trace[i]
            wave2pix_pos=wavePos_x_interp[i]

            if (bRange[0] <= w <= bRange[1]):
                thp_w = thp_i(w)
                deltW = np.abs(w - lam_trace[i - 1]) / 2 + np.abs(lam_trace[i + 1] - w) / 2
                f = spec_pix[wave2pix_pos]
                f_ratio = ratio_vec[wave2pix_pos]
                if f_ratio==0:
                    f_ratio=1
                f = f / self.t / thp_w / deltW /self.expNum/f_ratio
                err = err2_pix[wave2pix_pos]
                # err = err/ t / deltW
                err = np.sqrt(err)/ self.t / deltW/ thp_w /self.expNum/f_ratio
                specRangeImg.append(wave2pix_pos)
                # err = err / thp_w 
            else:
                f = 0
                err = 0

            wave_flux[i] = f
            err_flux[i] = err
        
        Aimg_cal = Aimg_[y_cent_pos-y_range:y_cent_pos+y_range+1, specRangeImg]
        ids = Aimg_cal > self.saturation

        #1. saturation pixel number, 2. total pixel number, 3 saturation ratio, 4.flux ratio in photo aperture,5.max value,6.min value
        saturePix = np.zeros(6)

        saturePix[0] = Aimg_cal[ids].shape[0]
        saturePix[1] = Aimg_cal.shape[0]*Aimg_cal.shape[1]
        saturePix[2] = saturePix[0]/saturePix[1]
        saturePix[3] = 1
        saturePix[4] = np.amax(Aimg_cal)
        saturePix[5] = np.amin(Aimg_cal)
        

        idx = (lam_trace >= bRange[0]-100)
        idx1 = (lam_trace[idx] <= bRange[1]+100)

        w_select = lam_trace[idx][idx1]
        f_select = wave_flux[idx][idx1]
        e_select =  err_flux[idx][idx1]
        lam_index = np.argsort(w_select)

        specTab = Table(np.array([w_select[lam_index], f_select[lam_index], e_select[lam_index]]).T,names=('WAVELENGTH', 'FLUX','ERR'))

        # spec_orig = np.loadtxt(sedFile)

        # plt.figure()
        # plt.plot(spec_orig[:,0], spec_orig[:,1])

        # plt.figure()
        # plt.errorbar(wave_pix[idx][idx1], wave_flux[idx][idx1],err_flux[idx][idx1])
        # plt.legend([self.sedFile])
        # # plt.plot(wave_pix[idx][idx1], wave_flux[idx][idx1])
        # plt.show()
        return specTab, Aimg, stamp.array, saturePix


    def generateSpec1dforStar(self,limitfluxratio = 0.8):
        import matplotlib.pyplot as plt
        specConfile = self.config.conFiles[self.grating]

        throughput_f = self.config.senFisle[self.grating] + self.config.orderIDs[self.beam] + '.fits'

        sed = self.generateSEDfromFiles(self.sedFile,2500,10000,0.01)

        x_nominal = int(np.floor(self.xcenter + 0.5))
        y_nominal = int(np.floor(self.ycenter + 0.5))
        dx = self.xcenter - x_nominal+0.5
        dy = self.ycenter - y_nominal+0.5
        offset = galsim.PositionD(dx, dy)

        star = galsim.DeltaFunction()
            # star = star.withFlux(tel.pupil_area * exptime)
        star = galsim.Convolve(self.psf, star)
        stamp = star.drawImage(wcs=galsim.PixelScale(self.p_size), offset=offset,nx=100, ny=100)*self.t*self.expNum*math.pi*(self.aper/2)*(self.aper/2)
        stamp.setOrigin(0,0)
        # print(stamp.center)
        # kk = np.where(stamp.array == np.amax(stamp.array))
        # print(kk[0][0], kk[1][0])
        # plt.figure()
        # plt.imshow(stamp.array)

        # ttt = np.sum(stamp.array)
        # stamp.array[:,:]=0
        # stamp.array[stamp.center.y, stamp.center.x] = ttt
        


        # stamp = self.psf.drawImage(wcs=galsim.PixelScale(self.p_size))*self.t*self.expNum*math.pi*(self.aper/2)*(self.aper/2)

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

        sdp = SpecDisperser.SpecDisperser(orig_img=stamp, xcenter=x_nominal,
                                            ycenter=y_nominal, origin=origin_star,
                                            tar_spec=sed,
                                            conf=specConfile,
                                            isAlongY=0)

        spec_orders = sdp.compute_spec_orders()
        
        thp = Table.read(throughput_f)
        thp_i = interpolate.interp1d(thp['WAVELENGTH'], thp['SENSITIVITY'])

        Aimg_orig = spec_orders[self.beam][0]
        Aimg_ = Aimg_orig

        Aimg_ = Aimg_ + (self.skybg + self.dark)*self.t*self.expNum

        Aimg_ = np.random.poisson(Aimg_)
        for i in np.arange(self.expNum):
            Aimg_ = self.addReadoutNois(img = Aimg_, readout = self.readout)

        Aimg = Aimg_ - (self.skybg + self.dark)*self.t*self.expNum


        wave_pix = spec_orders[self.beam][5]
        wave_pos = spec_orders[self.beam][3]

        wave_pos_y=spec_orders[self.beam][4]

        sh = Aimg.shape
        spec_pix = np.zeros(sh[1])
        err2_pix = np.zeros(sh[1])

        # print(spec_orders[beamOrder][4])
        # print(sh)
        # plt.figure()
        # plt.imshow(Aimg)
        y_cent_pos = int(np.round(np.mean(wave_pos_y)))

        tFlux = np.sum(spec_orders[self.beam][0])
        # print(tFlux)
        fluxRatio = 0
        for i in range(int(sh[0]/2)):
            pFlux = np.sum(spec_orders[self.beam][0][y_cent_pos-i:y_cent_pos+i+1])
            
            fluxRatio = pFlux/tFlux
            if fluxRatio>limitfluxratio:
                break
        f1 = spec_orders[self.beam][0][y_cent_pos-i:y_cent_pos+i+1].sum(0)
        f2 = spec_orders[self.beam][0].sum(0)
        ratio_vec = np.zeros_like(f1)
        nozero_flag = f2 != 0
        
        ratio_vec[nozero_flag] = f1[nozero_flag]/f2[nozero_flag]
        # ratio_vec = spec_orders[self.beam][0][y_cent_pos-i:y_cent_pos+i+1].sum(0)/spec_orders[self.beam][0].sum(0)
        y_range = i
        # print(y_range, fluxRatio)
        y_len_pix = 2 * y_range + 1
        for i in range(sh[1]):
            spec_pix[i] = sum(Aimg[y_cent_pos-y_range:y_cent_pos+y_range+1, i])
            err2_pix[i] = sum(Aimg_orig[y_cent_pos-y_range:y_cent_pos+y_range+1, i]) + (self.skybg + self.dark)*self.t * y_len_pix * self.expNum + self.readout*self.readout * y_len_pix * self.expNum

        bRange = self.config.bandRanges[self.grating]
        

        specRangeImg = []

        true_center = stamp.center + galsim.PositionD(self.xcenter-x_nominal, self.ycenter-y_nominal)
        wavePos_x = true_center.x + wave_pos - wave_pos[0]

        wavePos_x_interp = np.arange(int(wavePos_x[0]), int(wavePos_x[-1]))
        lam_trace = np.interp(wavePos_x_interp,wavePos_x,wave_pix)

        wave_flux = np.zeros(lam_trace.shape[0])
        err_flux = np.zeros(lam_trace.shape[0])

        for i in np.arange(1, lam_trace.shape[0] - 1):
            w = lam_trace[i]
            wave2pix_pos=wavePos_x_interp[i]

            if (bRange[0] <= w <= bRange[1]):
                thp_w = thp_i(w)
                deltW = np.abs(w - lam_trace[i - 1]) / 2 + np.abs(lam_trace[i + 1] - w) / 2
                f = spec_pix[wave2pix_pos]
                f_ratio = ratio_vec[wave2pix_pos]
                if f_ratio==0:
                    f_ratio=1
                f = f / self.t / thp_w / deltW /self.expNum/f_ratio
                err = err2_pix[wave2pix_pos]
                # err = err/ t / deltW
                err = np.sqrt(err)/ self.t / deltW/ thp_w /self.expNum/f_ratio
                specRangeImg.append(wave2pix_pos)
                # err = err / thp_w 
            else:
                f = 0
                err = 0

            wave_flux[i] = f
            err_flux[i] = err

        Aimg_cal = Aimg_[y_cent_pos-y_range:y_cent_pos+y_range+1, specRangeImg]
        ids = Aimg_cal > self.saturation
        
        #1. saturation pixel number, 2. total pixel number, 3 saturation ratio, 4.flux ratio in photo aperture,5.max value,6.min value
        saturePix = np.zeros(6)

        saturePix[0] = Aimg_cal[ids].shape[0]
        saturePix[1] = Aimg_cal.shape[0]*Aimg_cal.shape[1]
        saturePix[2] = saturePix[0]/saturePix[1]
        saturePix[3] = 1
        saturePix[4] = np.amax(Aimg_cal)
        saturePix[5] = np.amin(Aimg_cal)
        
        idx = (lam_trace >= bRange[0]-100)
        idx1 = (lam_trace[idx] <= bRange[1]+100)

        w_select = lam_trace[idx][idx1]
        f_select = wave_flux[idx][idx1]
        e_select =  err_flux[idx][idx1]
        lam_index = np.argsort(w_select)

        specTab = Table(np.array([w_select[lam_index], f_select[lam_index], e_select[lam_index]]).T,names=('WAVELENGTH', 'FLUX','ERR'))

        # spec_orig = np.loadtxt(sedFile)

        # plt.figure()
        # plt.plot(spec_orig[:,0], spec_orig[:,1])

        # plt.figure()
        # plt.errorbar(wave_pix[idx][idx1], wave_flux[idx][idx1],err_flux[idx][idx1])
        # plt.legend([self.sedFile])
        # # plt.plot(wave_pix[idx][idx1], wave_flux[idx][idx1])
        # plt.show()
        return specTab, Aimg, stamp.array, saturePix

    def addReadoutNois(self, img = None, readout = 5):
        for i in range(img.shape[0]):
            for j in range(img.shape[1]):
                img[i,j] += round(random.gauss(mu = 0, sigma = readout))

        return img