SpecGenerator.py 11.9 KiB
Newer Older
xin's avatar
xin committed

'''
Author: zx
Date: 2021-04-08 13:49:35
xin's avatar
xin committed
LastEditTime: 2022-07-05 14:58:42
xin's avatar
xin committed
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):
xin's avatar
xin committed
    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):
xin's avatar
xin committed
        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
xin's avatar
xin committed
        self.saturation = saturation
xin's avatar
xin committed
    
    '''
    @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,2000,10000,0.5)

        # 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))*self.t*self.expNum*math.pi*(self.aper/2)*(self.aper/2)

        origin_star = [self.ycenter - (stamp.center.y - stamp.ymin),
                                self.xcenter - (stamp.center.x - stamp.xmin)]

        sdp = SpecDisperser.SpecDisperser(orig_img=stamp, xcenter=self.xcenter,
                                            ycenter=self.ycenter, 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]
xin's avatar
xin committed
        Aimg_ = Aimg_orig
xin's avatar
xin committed

xin's avatar
xin committed
        Aimg_ = Aimg_ + (self.skybg + self.dark)*self.t*self.expNum
xin's avatar
xin committed

xin's avatar
xin committed
        Aimg_ = np.random.poisson(Aimg_)
xin's avatar
xin committed
        for i in np.arange(self.expNum):
xin's avatar
xin committed
            Aimg_ = self.addReadoutNois(img = Aimg_, readout = self.readout)
xin's avatar
xin committed

xin's avatar
xin committed
        Aimg = Aimg_ - (self.skybg + self.dark)*self.t*self.expNum
xin's avatar
xin committed


        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
        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])
xin's avatar
xin committed
        specRangeImg = []
xin's avatar
xin committed
        for i in np.arange(1, wave_pix.shape[0] - 1):
            w = wave_pix[i]

            if (bRange[0] <= w <= bRange[1]):
                thp_w = thp_i(w)
xin's avatar
xin committed
                deltW = np.abs(w - wave_pix[i - 1]) / 2 + np.abs(wave_pix[i + 1] - w) / 2
xin's avatar
xin committed
                f = spec_pix[wave_pos[0] - 1 + i]
xin's avatar
xin committed
                f = f / self.t / thp_w / deltW /self.expNum
xin's avatar
xin committed
                err = err2_pix[wave_pos[0] - 1 + i]
                # err = err/ t / deltW
                err = np.sqrt(err)/ self.t / deltW/ thp_w /self.expNum
                # err = err / thp_w 
xin's avatar
xin committed
                specRangeImg.append(wave_pos[0] - 1 + i)
xin's avatar
xin committed
            else:
                f = 0
                err = 0

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

xin's avatar
xin committed
        #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)
xin's avatar
xin committed

        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] = fluxRatio
xin's avatar
xin committed
        saturePix[4] = np.amax(Aimg_cal)
        saturePix[5] = np.amin(Aimg_cal)
xin's avatar
xin committed
        
xin's avatar
xin committed
        idx = (wave_pix >= bRange[0]-100)
        idx1 = (wave_pix[idx] <= bRange[1]+100)

        specTab = Table(np.array([wave_pix[idx][idx1],  wave_flux[idx][idx1], err_flux[idx][idx1]]).T,names=('WAVELENGTH', 'FLUX','ERR'))

        # spec_orig = np.loadtxt(sedFile)

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

xin's avatar
xin committed
        # 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()
xin's avatar
xin committed
        return specTab, Aimg, stamp.array, saturePix
xin's avatar
xin committed


    def generateSpec1dforStar(self,limitfluxratio = 0.8):

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

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

        sed = self.generateSEDfromFiles(self.sedFile,2000,10000,0.5)


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

        origin_star = [self.ycenter - (stamp.center.y - stamp.ymin),
                                self.xcenter - (stamp.center.x - stamp.xmin)]

        sdp = SpecDisperser.SpecDisperser(orig_img=stamp, xcenter=self.xcenter,
                                            ycenter=self.ycenter, 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]
xin's avatar
xin committed
        Aimg_ = Aimg_orig
xin's avatar
xin committed

xin's avatar
xin committed
        Aimg_ = Aimg_ + (self.skybg + self.dark)*self.t*self.expNum
xin's avatar
xin committed

xin's avatar
xin committed
        Aimg_ = np.random.poisson(Aimg_)
xin's avatar
xin committed
        for i in np.arange(self.expNum):
xin's avatar
xin committed
            Aimg_ = self.addReadoutNois(img = Aimg_, readout = self.readout)
xin's avatar
xin committed

xin's avatar
xin committed
        Aimg = Aimg_ - (self.skybg + self.dark)*self.t*self.expNum
xin's avatar
xin committed


        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
        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])
xin's avatar
xin committed

        specRangeImg = []

xin's avatar
xin committed
        for i in np.arange(1, wave_pix.shape[0] - 1):
            w = wave_pix[i]

            if (bRange[0] <= w <= bRange[1]):
                thp_w = thp_i(w)
xin's avatar
xin committed
                deltW = np.abs(w - wave_pix[i - 1]) / 2 + np.abs(wave_pix[i + 1] - w) / 2
xin's avatar
xin committed
                f = spec_pix[wave_pos[0] - 1 + i]
                f = f / self.t / thp_w / deltW /self.expNum
                err = err2_pix[wave_pos[0] - 1 + i]
                # err = err/ t / deltW
                err = np.sqrt(err)/ self.t / deltW/ thp_w /self.expNum
xin's avatar
xin committed
                specRangeImg.append(wave_pos[0] - 1 + i)
xin's avatar
xin committed
                # err = err / thp_w 
            else:
                f = 0
                err = 0

            wave_flux[i] = f
            err_flux[i] = err
xin's avatar
xin committed

        Aimg_cal = Aimg_[y_cent_pos-y_range:y_cent_pos+y_range+1, specRangeImg]
        ids = Aimg_cal > self.saturation
        
xin's avatar
xin committed
        #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)
xin's avatar
xin committed

        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] = fluxRatio
xin's avatar
xin committed
        saturePix[4] = np.amax(Aimg_cal)
        saturePix[5] = np.amin(Aimg_cal)
xin's avatar
xin committed
        
        idx = (wave_pix >= bRange[0]-100)
        idx1 = (wave_pix[idx] <= bRange[1]+100)

        specTab = Table(np.array([wave_pix[idx][idx1],  wave_flux[idx][idx1], err_flux[idx][idx1]]).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()
xin's avatar
xin committed
        return specTab, Aimg, stamp.array, saturePix
xin's avatar
xin committed

    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