test_BF_CTE.py

import unittest

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

from observation_sim.instruments import Chip, Filter, FilterParam, FocalPlane
from observation_sim.instruments.chip import chip_utils
from observation_sim.instruments.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('observation_sim.instruments.chip.libBF').joinpath("libmoduleBF.so") as lib_path:
            lib_bf = ctypes.CDLL(lib_path)
    except AttributeError:
        with pkg_resources.path('observation_sim.instruments.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)
    # galsim.ImageF(chip.npix_x, chip.npix_y)
    chip.img = galsim.ImageF(400, 200)
    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()