gen_calib_catalog.py

from cgitb import reset
from tkinter.tix import INTEGER
import astropy.coordinates as coord
from astropy import units as u
from pylab import *
import numpy as np
import galsim
import math
from astropy.table import Table, vstack
from astropy.io import fits
from astropy.wcs import WCS
import os
import json

import ImageHeader

# from numba import jit
ALL_FILTERS = ["NUV", "u", "g", "r", "i", "z", "y", "GU", "GV", "GI", "FGS"]


class Chip(object):
    def __init__(self, chipID):
        self.chipID = chipID

        self.nchip_x = 6
        self.nchip_y = 5
        self.npix_tot_x = 59516
        self.npix_tot_y = 49752
        self.npix_gap_x = (534, 1309)
        self.npix_gap_y = 898

        self.cen_pix_x = 0
        self.cen_pix_y = 0

        self.npix_x = 9216
        self.npix_y = 9232
        self.pix_scale = 0.074

        chip_dict = json.load(open("chip_definition.json", "r"))[str(self.chipID)]
        for key in chip_dict:
            setattr(self, key, chip_dict[key])
        self.filter_id, self.filter_type = self.getChipFilter()

    def getTanWCS(
        self, ra, dec, img_rot, pix_scale=None, xcen=None, ycen=None, logger=None
    ):
        """Get the WCS of the image mosaic using Gnomonic/TAN projection

        Parameter:
            ra, dec:    float
                        (RA, Dec) of pointing of optical axis
            img_rot:    galsim Angle object
                        Rotation of image
            pix_scale:  float
                        Pixel size in unit of as/pix
        Returns:
            WCS of the focal plane
        """
        if logger is not None:
            logger.info(
                "    Construct the wcs of the entire image mosaic using Gnomonic/TAN projection"
            )
        if (xcen == None) or (ycen == None):
            xcen = self.cen_pix_x
            ycen = self.cen_pix_y
        if pix_scale == None:
            pix_scale = self.pix_scale
        # dudx =  -np.cos(img_rot.rad) * pix_scale
        # dudy =  -np.sin(img_rot.rad) * pix_scale
        # dvdx =  -np.sin(img_rot.rad) * pix_scale
        # dvdy =  +np.cos(img_rot.rad) * pix_scale

        dudx = -np.cos(img_rot.rad) * pix_scale
        dudy = +np.sin(img_rot.rad) * pix_scale
        dvdx = -np.sin(img_rot.rad) * pix_scale
        dvdy = -np.cos(img_rot.rad) * pix_scale
        moscen = galsim.PositionD(x=xcen, y=ycen)
        sky_center = galsim.CelestialCoord(
            ra=ra * galsim.degrees, dec=dec * galsim.degrees
        )
        affine = galsim.AffineTransform(dudx, dudy, dvdx, dvdy, origin=moscen)
        WCS = galsim.TanWCS(affine, sky_center, units=galsim.arcsec)

        return WCS

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

    def getChipCenter(self):
        """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 observation.
        Parameters:
            chipID:         int
                            the index of the chip
        Returns:
            A galsim BoundsD object
        """

        chipID = self.chipID

        rowID, colID = self.getChipRowCol(chipID)
        gx1, gx2 = self.npix_gap_x
        gy = self.npix_gap_y

        # xlim of a given CCD chip
        xrem = 2 * (colID - 1) - (self.nchip_x - 1)
        xcen = (self.npix_x // 2 + gx1 // 2) * xrem
        if chipID >= 26 or chipID == 21:
            xcen = (self.npix_x // 2 + gx1 // 2) * xrem - (gx2 - gx1)
        if chipID <= 5 or chipID == 10:
            xcen = (self.npix_x // 2 + gx1 // 2) * xrem + (gx2 - gx1)
        # nx0 = xcen - self.npix_x//2 + 1
        # nx1 = xcen + self.npix_x//2

        # ylim of a given CCD chip
        yrem = (rowID - 1) - self.nchip_y // 2
        ycen = (self.npix_y + gy) * yrem
        # ny0 = ycen - self.npix_y//2 + 1
        # ny1 = ycen + self.npix_y//2

        return galsim.PositionD(xcen, ycen)

    def getChipFilter(self, chipID=None):
        """Return the filter index and type for a given chip #(chipID)"""
        filter_type_list = ALL_FILTERS
        if chipID == 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 transRaDec2D(ra, dec):
    x1 = np.cos(dec / 57.2957795) * np.cos(ra / 57.2957795)
    y1 = np.cos(dec / 57.2957795) * np.sin(ra / 57.2957795)
    z1 = np.sin(dec / 57.2957795)
    return np.array([x1, y1, z1])


def getobsPA(ra, dec):
    l1 = np.array([0, 0, 1])
    l2 = transRaDec2D(ra, dec)
    polar_ec = coord.SkyCoord(
        0 * u.degree, 90 * u.degree, frame="barycentrictrueecliptic"
    )
    polar_eq = polar_ec.transform_to("icrs")

    # print(polar_eq.ra.value,polar_eq.dec.value)
    polar_d = transRaDec2D(polar_eq.ra.value, polar_eq.dec.value)
    l1l2cross = np.cross(l2, l1)
    pdl2cross = np.cross(l2, polar_d)
    angle = math.acos(
        np.dot(l1l2cross, pdl2cross)
        / (np.linalg.norm(l1l2cross) * np.linalg.norm(pdl2cross))
    )

    angle = (angle) / math.pi * 180

    # if (ra>90 and ra< 270):
    #     angle=-angle
    return angle


def getChipRangeInfo(chipID=8, ra=60.0, dec=-40.0, pa=113.0):
    """_summary_

    Args:
        chipID (int, optional): Chip ID.
        ra (_type_, optional): Chip center ra.
        dec (_type_, optional): Chip center dec.

    Returns:
        _type_: [ra, dec, rotation angle]
    """
    # chip_center = [ra, dec]

    chip = Chip(chipID)

    h_ext = ImageHeader.generateExtensionHeader(
        chip=chip,
        xlen=chip.npix_x,
        ylen=chip.npix_y,
        ra=ra,
        dec=dec,
        pa=pa,
        gain=chip.gain,
        readout=5.0,
        dark=0.02,
        saturation=90000,
        pixel_scale=chip.pix_scale,
        pixel_size=chip.pix_size,
        xcen=chip.x_cen,
        ycen=chip.y_cen,
        extName="SCI",
        timestamp=1621915200,
        exptime=150,
        readoutTime=chip.readout_time,
    )
    chip_wcs = WCS(h_ext)

    return chip_wcs, chip.npix_x, chip.npix_y, chip.chipID, chip.filter_type


def formateSEDfile(inputDir, outputDir):
    fileFns = os.listdir(inputDir)

    for fn in fileFns:
        if fn[0] == ".":
            continue
        if ".fits" not in fn:
            continue
        dfn = os.path.join(inputDir, fn)
        outfn = os.path.join(outputDir, "GP_" + fn[0:-5] + ".fits")
        f = fits.open(dfn)
        data_tab = Table(f[1].data)
        col_names = data_tab.colnames
        if "loglam" in col_names:
            print("1:" + fn)
            w1 = pow(10, f[1].data["loglam"])
            f1 = f[1].data["flux"] * 1e-17
        elif "WAVELENGTH" in col_names:
            print("2:" + fn)
            w1 = np.array(data_tab["WAVELENGTH"].tolist()[0])
            f1 = np.array(data_tab["FLUX"].tolist()[0]) * 1e-17
        w = np.zeros(w1.shape[0] + 4)
        w[0] = 2000
        w[1] = w1[0] - 1
        w[-2] = w1[-1] + 1
        w[-1] = 15000
        w[2:-2] = w1
        fl = np.zeros(f1.shape[0] + 4)
        fl[0] = 0
        fl[1] = 0
        fl[-2] = 0
        fl[-1] = 0
        fl[2:-2] = f1
        t = Table([w, fl], names=["WAVELENGTH", "FLUX"])
        # t = Table([pow(10,f[1].data['loglam']),f[1].data['flux']*1e-17],names=['WAVELENGTH','FLUX'])
        t.write(outfn, overwrite=True)


def genOneChipCat(
    ra=40.0,
    dec=-40.0,
    pa=23.5,
    chipID=1,
    pointSource_flag=True,
    sersic_index=1.0,
    re=0.078,
    specFn="LMC-SMP-58_spectrum_csst.fits",
):

    chip_wcs, chip_xlen, chip_ylen, _, filterT = getChipRangeInfo(
        chipID=chipID, ra=ra, dec=dec, pa=pa
    )

    sppos = 3685

    xSize = 15
    ySize = 30
    y_interval = 20

    X2 = np.linspace(100, sppos - 100, xSize)
    Y2 = np.linspace(100, chip_ylen - 300, ySize)
    y_gap = np.arange(0, y_interval * xSize, y_interval)
    x2, y2 = np.meshgrid(X2, Y2)
    y2 = y2 + y_gap
    x2 = x2.flatten()
    y2 = y2.flatten()

    X1 = np.linspace(sppos + 300, chip_xlen - 100, xSize)
    Y1 = np.linspace(150, chip_ylen - 250, ySize)
    y_gap = np.arange(0, y_interval * xSize, y_interval)
    x1, y1 = np.meshgrid(X1, Y1)
    y1 = y1 + y_gap
    x1 = x1.flatten()
    y1 = y1.flatten()

    tab_size = x1.size + x2.size
    x = np.zeros(tab_size)
    y = np.zeros_like(x)
    x[0 : x1.size] = x1
    x[x1.size :] = x2
    y[0 : x1.size] = y1
    y[x1.size :] = y2

    radecs = chip_wcs.pixel_to_world(x, y)

    # ids = np.arange(0, tab_size, 1)
    mags = np.zeros(tab_size) + -99
    sns = np.zeros(tab_size) + sersic_index
    res = np.zeros(tab_size) + re
    namelist = []
    ids = []
    for i in np.arange(tab_size):
        namelist.append(specFn)
        ids.append(format(chipID, "02.0f") + format(i + 1, "04.0f"))

    if pointSource_flag:
        p_flag = np.ones(tab_size, dtype=bool)  # True
    else:
        p_flag = np.zeros(tab_size, dtype=bool)  # False

    cat_t = Table(
        [ids, radecs.ra.deg, radecs.dec.deg, res, sns, mags, namelist, p_flag],
        names=(
            "ID",
            "RA",
            "DEC",
            "RE",
            "SERSIC_N",
            "MAG_g",
            "SPEC_FN",
            "POINTSOURCE_FLAG",
        ),
    )
    return cat_t


if __name__ == "__main__":
    ra = 244.972773
    dec = 39.895901
    pa = 109.59452000342756
    chipIDs = [1, 2, 3, 4, 5, 10, 21, 26, 27, 28, 29, 30]
    cat_t = Table()
    for chipID in chipIDs:
        cat_one = genOneChipCat(
            ra=ra,
            dec=dec,
            pa=pa,
            chipID=chipID,
            pointSource_flag=True,
            specFn="LMC-SMP-58_spectrum_csst.fits",
        )
        cat_t = vstack([cat_t, cat_one])
    cat_t.write(
        "calibrationPNESPEC_CHIP_ALL_SLS.fits",
        format="fits",
        overwrite=True,
    )