test_BF_CTE.py

import unittest

import sys,os,math
from itertools import islice
import numpy as np
import copy
import ctypes
import galsim
import yaml
from astropy.io import fits

from ObservationSim.Instrument import Chip, Filter, FilterParam, FocalPlane
from ObservationSim.Instrument.Chip import ChipUtils as chip_utils
#from ObservationSim.sim_steps import add_brighter_fatter_CTE
from ObservationSim.Instrument.Chip.libCTI.CTI_modeling import CTI_sim

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


### test FUNCTION --- START ###
def add_brighter_fatter(img):
    #Inital dynamic lib
    try:
        with pkg_resources.files('ObservationSim.Instrument.Chip.libBF').joinpath("libmoduleBF.so") as lib_path:
            lib_bf = ctypes.CDLL(lib_path)
    except AttributeError:
        with pkg_resources.path('ObservationSim.Instrument.Chip.libBF', "libmoduleBF.so") as lib_path:
            lib_bf = ctypes.CDLL(lib_path)
    lib_bf.addEffects.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_float), ctypes.POINTER(ctypes.c_float), ctypes.c_int]

    # Set bit flag
    bit_flag = 1
    bit_flag = bit_flag | (1 << 2)

    nx, ny = img.array.shape
    nn = nx * ny
    arr_ima= (ctypes.c_float*nn)()
    arr_imc= (ctypes.c_float*nn)()

    arr_ima[:]= img.array.reshape(nn)
    arr_imc[:]= np.zeros(nn)

    lib_bf.addEffects(nx, ny, arr_ima, arr_imc, bit_flag)
    img.array[:, :] = np.reshape(arr_imc, [nx, ny])
    del arr_ima, arr_imc
    return img
### test FUNCTION --- END ###

def defineCCD(iccd, config_file):
    with open(config_file, "r") as stream:
        try:
            config = yaml.safe_load(stream)
            #for key, value in config.items():
            #    print (key + " : " + str(value))
        except yaml.YAMLError as exc:
            print(exc)
    chip = Chip(chipID=iccd, config=config)
    chip.img = galsim.ImageF(400, 200) #galsim.ImageF(chip.npix_x, chip.npix_y)
    focal_plane = FocalPlane(chip_list=[iccd])
    chip.img.wcs= focal_plane.getTanWCS(192.8595, 27.1283, -113.4333*galsim.degrees, chip.pix_scale)
    return chip

def defineFilt(chip):
    filter_param = FilterParam()
    filter_id, filter_type = chip.getChipFilter()
    filt = Filter(
        filter_id=filter_id,
        filter_type=filter_type,
        filter_param=filter_param,
        ccd_bandpass=chip.effCurve)
    bandpass_list = filt.bandpass_sub_list
    return filt


class detModule_coverage(unittest.TestCase):
    def __init__(self, methodName='runTest'):
        super(detModule_coverage, self).__init__(methodName)
        ##self.dataPath = "/public/home/chengliang/CSSOSDataProductsSims/csst-simulation/tests/UNIT_TEST_DATA" ##os.path.join(os.getenv('UNIT_TEST_DATA_ROOT'), 'csst_fz_gc1')
        self.dataPath = os.path.join(os.getenv('UNIT_TEST_DATA_ROOT'), 'csst_msc_sim/csst_fz_msc')
        self.iccd = 1


    def test_add_brighter_fatter(self):
        config_file = os.path.join(self.dataPath, 'config_test.yaml')
        chip = defineCCD(self.iccd, config_file)
        filt = defineFilt(chip)
        print(chip.chipID)
        print(chip.cen_pix_x, chip.cen_pix_y)

        #objA-lowSFB
        obj = galsim.Gaussian(sigma=0.2, flux=1000)
        arr = obj.drawImage(nx=64, ny=64, scale=0.074).array
        chip.img.array[(100-32):(100+32),(200-32):(200+32)] = arr[:,:] 
        img_old = copy.deepcopy(chip.img)
        img_new = add_brighter_fatter(img=chip.img)
        arr1= img_old.array
        arr2= img_new.array
        deltaA_max = np.max(np.abs(arr2-arr1))
        print('deltaA-max:', np.max(np.abs(arr2-arr1))) 
        print('deltaA-min:', np.min(np.abs(arr2-arr1))) 

        #objB-highSFB
        obj = galsim.Gaussian(sigma=0.2, flux=10000)
        arr = obj.drawImage(nx=64, ny=64, scale=0.074).array
        chip.img.array[(100-32):(100+32),(200-32):(200+32)] = arr[:,:] 
        img_old = copy.deepcopy(chip.img)
        img_new = add_brighter_fatter(img=chip.img)
        arr3= img_old.array
        arr4= img_new.array
        deltaB_max = np.max(np.abs(arr4-arr3))
        print('deltaB-max:', np.max(np.abs(arr4-arr3))) 
        print('deltaB-min:', np.min(np.abs(arr4-arr3))) 
        
        self.assertTrue( deltaB_max > deltaA_max )


    def test_apply_CTE(self):
        config_file = os.path.join(self.dataPath, 'config_test.yaml')
        chip = defineCCD(self.iccd, config_file)
        filt = defineFilt(chip)
        print(chip.chipID)
        print(chip.cen_pix_x, chip.cen_pix_y)

        print("  Apply CTE Effect")
        nx,ny,noverscan,nsp,nmax = 4608,4616,84,3,10
        ntotal = 4700
        beta,w,c = 0.478,84700,0
        t = np.array([0.74,7.7,37],dtype=np.float32)
        rho_trap = np.array([0.6,1.6,1.4],dtype=np.float32)
        trap_seeds = np.array([0,100,1000],dtype=np.int32)
        release_seed = 500
        image = fits.getdata(os.path.join(self.dataPath, "testCTE_image_before.fits")).astype(np.int32) 
        #get_trap_map(trap_seeds,nx,ny,nmax,rho_trap,beta,c,".")
        #bin2fits("trap.bin",".",nsp,nx,ny,nmax)
        image_cti = CTI_sim(image,nx,ny,noverscan,nsp,nmax,beta,w,c,t,rho_trap,trap_seeds,release_seed)
        fits.writeto(os.path.join(self.dataPath, "testCTE_image_after.fits"),data=image_cti,overwrite=True)


if __name__ == '__main__':
    unittest.main()