Chip.py

import galsim
import os
import numpy as np
import pickle
import json
from astropy.table import Table
from numpy.random import Generator, PCG64
from astropy.io import fits
from datetime import datetime

import observation_sim.instruments._util as _util
from observation_sim.instruments.chip import effects
from observation_sim.instruments.FocalPlane import FocalPlane
from observation_sim.config.header import generatePrimaryHeader, generateExtensionHeader
from observation_sim.instruments._util import rotate_conterclockwise
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


class Chip(FocalPlane):
    def __init__(self, chipID, ccdEffCurve_dir=None, CRdata_dir=None, sls_dir=None, config=None, treering_func=None, logger=None):
        # Get focal plane (instance of paraent class) info
        super().__init__()
        self.nsecy = 2
        self.nsecx = 8
        self.gain_channel = np.ones(self.nsecy * self.nsecx)
        self._set_attributes_from_config(config)

        self.logger = logger

        # A chip ID must be assigned
        self.chipID = int(chipID)
        self.chip_name = str(chipID).rjust(2, '0')

        # Get corresponding filter info
        self.filter_id, self.filter_type = self.getChipFilter()
        self.survey_type = self._getSurveyType()

        if self.filter_type != "FGS":
            self._getChipRowCol()

        # Set the relavent specs for detectors
        try:
            with pkg_resources.files('observation_sim.instruments.data.ccd').joinpath("chip_definition.json") as chip_definition:
                with open(chip_definition, "r") as f:
                    chip_dict = json.load(f)[str(self.chipID)]
        except AttributeError:
            with pkg_resources.path('observation_sim.instruments.data.ccd', "chip_definition.json") as chip_definition:
                with open(chip_definition, "r") as f:
                    chip_dict = json.load(f)[str(self.chipID)]
        for key in chip_dict:
            setattr(self, key, chip_dict[key])

        self.fdModel = None
        if self.filter_type == "FGS":
            fgs_name = self.chip_name[0:4]
            try:
                with pkg_resources.files('observation_sim.instruments.data.field_distortion').joinpath("FieldDistModelGlobal_pr4_%s.pickle" % (fgs_name.lower())) as field_distortion:
                    with open(field_distortion, "rb") as f:
                        self.fdModel = pickle.load(f)
            except AttributeError:
                with pkg_resources.path('observation_sim.instruments.data.field_distortion', "FieldDistModelGlobal_pr4_%s.pickle" % (fgs_name.lower())) as field_distortion:
                    with open(field_distortion, "rb") as f:
                        self.fdModel = pickle.load(f)
        else:
            # Get the corresponding field distortion model
            try:
                with pkg_resources.files('observation_sim.instruments.data.field_distortion').joinpath("FieldDistModel_v2.0.pickle") as field_distortion:
                    with open(field_distortion, "rb") as f:
                        self.fdModel = pickle.load(f)
            except AttributeError:
                with pkg_resources.path('observation_sim.instruments.data.field_distortion', "FieldDistModelGlobal_mainFP_v1.0.pickle") as field_distortion:
                    with open(field_distortion, "rb") as f:
                        self.fdModel = pickle.load(f)

        # Get boundary (in pix)
        self.bound = self.getChipLim()

        self.ccdEffCurve_dir = ccdEffCurve_dir
        self.CRdata_dir = CRdata_dir

        slsconfs = chip_utils.getChipSLSConf(chipID=self.chipID)
        if np.size(slsconfs) == 1:
            try:
                with pkg_resources.files('observation_sim.instruments.data.sls_conf').joinpath(slsconfs) as conf_path:
                    self.sls_conf = str(conf_path)
            except AttributeError:
                with pkg_resources.path('observation_sim.instruments.data.sls_conf', slsconfs) as conf_path:
                    self.sls_conf = str(conf_path)
        else:
            # self.sls_conf = [os.path.join(self.sls_dir, slsconfs[0]), os.path.join(self.sls_dir, slsconfs[1])]
            self.sls_conf = []
            try:
                with pkg_resources.files('observation_sim.instruments.data.sls_conf').joinpath(slsconfs[0]) as conf_path:
                    self.sls_conf.append(str(conf_path))
            except AttributeError:
                with pkg_resources.path('observation_sim.instruments.data.sls_conf', slsconfs[0]) as conf_path:
                    self.sls_conf.append(str(conf_path))
            try:
                with pkg_resources.files('observation_sim.instruments.data.sls_conf').joinpath(slsconfs[1]) as conf_path:
                    self.sls_conf.append(str(conf_path))
            except AttributeError:
                with pkg_resources.path('observation_sim.instruments.data.sls_conf', slsconfs[1]) as conf_path:
                    self.sls_conf.append(str(conf_path))

        self.effCurve = self._getChipEffCurve(self.filter_type)
        self._getCRdata()

        # # Define the sensor model
        self.sensor = galsim.Sensor()

        self.flat_cube = None  # for spectroscopic flat field cube simulation

    def _set_attributes_from_config(self, config):
        # Default setting
        self.read_noise = 5.0   # e/pix
        self.dark_noise = 0.02  # e/pix/s
        self.rotate_angle = 0.
        self.overscan = 1000
        # Override default values
        # for key in ["gain", "bias_level, dark_exptime", "flat_exptime", "readout_time", "full_well", "read_noise", "dark_noise", "overscan"]:
        #     if key in config["ins_effects"]:
        #         setattr(self, key, config["ins_effects"][key])

    def _getChipRowCol(self):
        self.rowID, self.colID = self.getChipRowCol(self.chipID)

    def getChipRowCol(self, chipID):
        rowID = ((chipID - 1) % 5) + 1
        colID = 6 - ((chipID - 1) // 5)
        return rowID, colID

    def _getSurveyType(self):
        if self.filter_type in _util.SPEC_FILTERS:
            return "spectroscopic"
        elif self.filter_type in _util.PHOT_FILTERS:
            return "photometric"
        # elif self.filter_type in ["FGS"]:
        #     return "FGS"

    def _getChipEffCurve(self, filter_type):
        # CCD efficiency curves
        if filter_type in ['NUV', 'u', 'GU']:
            filename = 'UV0.txt'
        if filter_type in ['g', 'r', 'GV', 'FGS']:
            # TODO, need to switch to the right efficiency curvey for FGS CMOS
            filename = 'Astro_MB.txt'
        if filter_type in ['i', 'z', 'y', 'GI']:
            filename = 'Basic_NIR.txt'
        try:
            with pkg_resources.files('observation_sim.instruments.data.ccd').joinpath(filename) as ccd_path:
                table = Table.read(ccd_path, format='ascii')
        except AttributeError:
            with pkg_resources.path('observation_sim.instruments.data.ccd', filename) as ccd_path:
                table = Table.read(ccd_path, format='ascii')
        throughput = galsim.LookupTable(
            x=table['col1'], f=table['col2'], interpolant='linear')
        bandpass = galsim.Bandpass(throughput, wave_type='nm')
        return bandpass

    def _getCRdata(self):
        try:
            with pkg_resources.files('observation_sim.instruments.data').joinpath("wfc-cr-attachpixel.dat") as cr_path:
                self.attachedSizes = np.loadtxt(cr_path)
        except AttributeError:
            with pkg_resources.path('observation_sim.instruments.data', "wfc-cr-attachpixel.dat") as cr_path:
                self.attachedSizes = np.loadtxt(cr_path)

    # def loadSLSFLATCUBE(self, flat_fn='flat_cube.fits'):
    #     try:
    #         with pkg_resources.files('observation_sim.instruments.data').joinpath(flat_fn) as data_path:
    #             flat_fits = fits.open(data_path, ignore_missing_simple=True)
    #     except AttributeError:
    #         with pkg_resources.path('observation_sim.instruments.data', flat_fn) as data_path:
    #             flat_fits = fits.open(data_path, ignore_missing_simple=True)

    #     fl = len(flat_fits)
    #     fl_sh = flat_fits[0].data.shape
    #     assert fl == 4, 'FLAT Field Cube is Not 4 layess!!!!!!!'
    #     self.flat_cube = np.zeros([fl, fl_sh[0], fl_sh[1]])
    #     for i in np.arange(0, fl, 1):
    #         self.flat_cube[i] = flat_fits[i].data

    def getChipFilter(self, chipID=None):
        """Return the filter index and type for a given chip #(chipID)
        """
        filter_type_list = _util.ALL_FILTERS
        if chipID is None:
            chipID = self.chipID

        # updated configurations
        if chipID > 42 or chipID < 1:
            raise ValueError("!!! Chip ID: [1,42]")
        if chipID in [6, 15, 16, 25]:
            filter_type = "y"
        if chipID in [11, 20]:
            filter_type = "z"
        if chipID in [7, 24]:
            filter_type = "i"
        if chipID in [14, 17]:
            filter_type = "u"
        if chipID in [9, 22]:
            filter_type = "r"
        if chipID in [12, 13, 18, 19]:
            filter_type = "NUV"
        if chipID in [8, 23]:
            filter_type = "g"
        if chipID in [1, 10, 21, 30]:
            filter_type = "GI"
        if chipID in [2, 5, 26, 29]:
            filter_type = "GV"
        if chipID in [3, 4, 27, 28]:
            filter_type = "GU"
        if chipID in range(31, 43):
            filter_type = 'FGS'
        filter_id = filter_type_list.index(filter_type)

        return filter_id, filter_type

    def getChipLim(self, chipID=None):
        """Calculate the edges in pixel for a given CCD chip on the focal plane
        NOTE: There are 5*4 CCD chips in the focus plane for photometric / spectroscopic observation.
        Parameters:
            chipID:         int
                            the index of the chip
        Returns:
            A galsim BoundsD object
        """
        xmin, xmax, ymin, ymax = 1e10, -1e10, 1e10, -1e10
        xcen = self.x_cen / self.pix_size
        ycen = self.y_cen / self.pix_size
        sign_x = [-1., 1., -1., 1.]
        sign_y = [-1., -1., 1., 1.]
        for i in range(4):
            x = xcen + sign_x[i] * self.npix_x / 2.
            y = ycen + sign_y[i] * self.npix_y / 2.
            x, y = _util.rotate_conterclockwise(
                x0=xcen, y0=ycen, x=x, y=y, angle=self.rotate_angle)
            xmin, xmax = min(xmin, x), max(xmax, x)
            ymin, ymax = min(ymin, y), max(ymax, y)
        return galsim.BoundsD(xmin, xmax, ymin, ymax)

    def getSkyCoverage(self, wcs):
        # print("In getSkyCoverage: xmin = %.3f, xmax = %.3f, ymin = %.3f, ymax = %.3f"%(self.bound.xmin, self.bound.xmax, self.bound.ymin, self.bound.ymax))
        return super().getSkyCoverage(wcs, self.bound.xmin, self.bound.xmax, self.bound.ymin, self.bound.ymax)

    def getSkyCoverageEnlarged(self, wcs, margin=0.5):
        """The enlarged sky coverage of the chip
        """
        margin /= 60.0
        bound = self.getSkyCoverage(wcs)
        return galsim.BoundsD(bound.xmin - margin, bound.xmax + margin, bound.ymin - margin, bound.ymax + margin)

    def isContainObj(self, ra_obj=None, dec_obj=None, x_image=None, y_image=None, wcs=None, margin=1):
        # magin in number of pix
        if (ra_obj is not None) and (dec_obj is not None):
            if wcs is None:
                wcs = self.img.wcs
            pos_obj = wcs.toImage(galsim.CelestialCoord(
                ra=ra_obj*galsim.degrees, dec=dec_obj*galsim.degrees))
            x_image, y_image = pos_obj.x, pos_obj.y
        elif (x_image is None) or (y_image is None):
            raise ValueError(
                "Either (ra_obj, dec_obj) or (x_image, y_image) should be given")

        xmin, xmax = self.bound.xmin - margin, self.bound.xmax + margin
        ymin, ymax = self.bound.ymin - margin, self.bound.ymax + margin
        if (x_image - xmin) * (x_image - xmax) > 0.0 or (y_image - ymin) * (y_image - ymax) > 0.0:
            return False
        return True

    def getChipNoise(self, exptime=150.0):
        noise = self.dark_noise * exptime + self.read_noise**2
        return noise