test_field_distortion.py

import unittest

import numpy as np
import galsim
import os
import sys
from astropy.table import Table
from scipy import interpolate
import pickle


class test_field_distortion(unittest.TestCase):
    def __init__(self, methodName="runTest"):
        super(test_field_distortion, self).__init__(methodName)
        self.dataMainPath = os.path.join(
            os.getenv("UNIT_TEST_DATA_ROOT"), "csst_msc_sim/field_distortion"
        )
        self.dataInputPath = os.path.join(self.dataMainPath, "input_catalog")
        self.fdModelName = "FieldDistModel_v2.0_test.pickle"

    def test_fd_model(self):
        cat_dir = self.dataInputPath
        model_dir = self.dataMainPath
        model_date = "2024-05-08"
        model_name = self.fdModelName
        field_distortion_model(
            cat_dir,
            model_dir,
            poly_degree=4,
            model_date=model_date,
            model_name=model_name,
        )

    def test_fd_apply(self):
        model_name = self.fdModelName
        model_dir = self.dataMainPath
        cat_dir = self.dataMainPath
        field_distortion_apply(
            model_name, model_dir, cat_dir, ra_cen=60.0, dec_cen=-40.0, img_rot=0.0
        )


def ccdParam():
    """
    Basic CCD size and noise parameters.
    """
    # CCD size
    xt, yt = 59516, 49752
    x0, y0 = 9216, 9232
    xgap, ygap = (534, 1309), 898
    xnchip, ynchip = 6, 5
    ccdSize = xt, yt, x0, y0, xgap, ygap, xnchip, ynchip

    # other parameters
    readNoise = 5.0  # e/pix
    darkNoise = 0.02  # e/pix/s
    pixel_scale = 0.074  # pixel scale
    gain = 1.0
    ccdBase = readNoise, darkNoise, pixel_scale, gain

    return ccdSize, ccdBase


def chipLim(chip):
    ccdSize, ccdBase = ccdParam()
    xt, yt, x0, y0, gx, gy, xnchip, ynchip = ccdSize
    gx1, gx2 = gx

    rowID = ((chip - 1) % 5) + 1
    colID = 6 - ((chip - 1) // 5)

    xrem = 2 * (colID - 1) - (xnchip - 1)
    xcen = (x0 // 2 + gx1 // 2) * xrem
    if chip <= 5 or chip == 10:
        xcen = (x0 // 2 + gx1 // 2) * xrem + (gx2 - gx1)
    if chip >= 26 or chip == 21:
        xcen = (x0 // 2 + gx1 // 2) * xrem - (gx2 - gx1)
    nx0 = xcen - x0 // 2 + 1
    nx1 = xcen + x0 // 2

    yrem = (rowID - 1) - ynchip // 2
    ycen = (y0 + gy) * yrem
    ny0 = ycen - y0 // 2 + 1
    ny1 = ycen + y0 // 2

    return nx0, nx1, ny0, ny1


def chip_filter(nchip):
    """
    return filter name of a given chip
    """
    filtype = ["nuv", "u", "g", "r", "i", "z", "y"]

    # updated configurations
    # if nchip>24 or nchip<7: raise ValueError("!!! Chip ID: [7,24]")
    if nchip in [6, 15, 16, 25]:
        filter_name = "y"
    if nchip in [11, 20]:
        filter_name = "z"
    if nchip in [7, 24]:
        filter_name = "i"
    if nchip in [14, 17]:
        filter_name = "u"
    if nchip in [9, 22]:
        filter_name = "r"
    if nchip in [12, 13, 18, 19]:
        filter_name = "nuv"
    if nchip in [8, 23]:
        filter_name = "g"

    filter_id = filtype.index(filter_name)

    return filter_id, filter_name


def skyLim(wcs, x0, x1, y0, y1):
    """
    The sky coverage of a single exposure image
    """
    r2d = 180.0 / np.pi
    # xt, yt, x0, y0, gx, gy, xnchip, ynchip = ccdSize()
    s1 = wcs.toWorld(galsim.PositionD(x0, y0))
    s2 = wcs.toWorld(galsim.PositionD(x0, y1))

    s3 = wcs.toWorld(galsim.PositionD(x1, y0))
    s4 = wcs.toWorld(galsim.PositionD(x1, y1))
    ra = [s1.ra.rad * r2d, s2.ra.rad * r2d, s3.ra.rad * r2d, s4.ra.rad * r2d]
    dec = [s1.dec.rad * r2d, s2.dec.rad * r2d, s3.dec.rad * r2d, s4.dec.rad * r2d]

    return min(ra), max(ra), min(dec), max(dec)


def wcsMain(imgRotation=0.0, raCenter=0.0, decCenter=0.0):
    ccdSize, ccdBase = ccdParam()
    xsize, ysize, _, _, _, _, _, _ = ccdSize
    _, _, pixelScale, _ = ccdBase
    xmcen, ymcen = 0.0, 0.0
    imrot = imgRotation * galsim.degrees
    racen = raCenter * galsim.degrees
    deccen = decCenter * galsim.degrees

    # define the wcs
    dudx = -np.cos(imrot.rad) * pixelScale
    dudy = +np.sin(imrot.rad) * pixelScale
    dvdx = -np.sin(imrot.rad) * pixelScale
    dvdy = -np.cos(imrot.rad) * pixelScale
    moscen = galsim.PositionD(x=xmcen, y=ymcen)
    skyCenter = galsim.CelestialCoord(ra=racen, dec=deccen)
    affine = galsim.AffineTransform(dudx, dudy, dvdx, dvdy, origin=moscen)
    wcs = galsim.TanWCS(affine, skyCenter, units=galsim.arcsec)

    return wcs


# FD model
def field_distortion_model(
    cat_dir,
    model_dir,
    poly_degree=4,
    model_date="2024-05-08",
    model_name="FieldDistModel_v2.0_test.pickle",
):
    # default parameter setup
    nccd, nwave, npsf = 30, 4, 30 * 30

    # load a CSST-like wcs
    wcs = wcsMain()
    cd11, cd12 = wcs.cd[0, 0], wcs.cd[0, 1]
    cd21, cd22 = wcs.cd[1, 0], wcs.cd[1, 1]
    xmcen, ymcen = wcs.crpix

    # obtain the interpolation model
    fdFunList = {}
    fdFunList["date"] = model_date
    for iwave in range(1, nwave + 1):
        # if iwave!=1: continue
        iwaveKey = "wave%d" % iwave

        # first construct the global interpolation
        xwList, ywList = [], []
        xdList, ydList = [], []
        fdFunList[iwaveKey] = {}
        for iccd in range(1, nccd + 1):
            # if iccd!=9: continue
            iccdKey = "ccd" + str("0%d" % (iccd))[-2:]

            # load PSF data
            ipsfDatn = cat_dir + "ccd%d_%s.dat" % (iccd, iwaveKey)
            ipsfDat = Table.read(ipsfDatn, format="ascii")
            for ipsf in range(1, npsf + 1):
                # if ipsf!=2: continue
                xField = ipsfDat["field_x"][ipsf - 1]
                yField = ipsfDat["field_y"][ipsf - 1]

                # image coordinate with field distortion
                xImage = 100.0 * (
                    ipsfDat["image_x"][ipsf - 1] + ipsfDat["centroid_x"][ipsf - 1]
                )
                yImage = 100.0 * (
                    ipsfDat["image_y"][ipsf - 1] + ipsfDat["centroid_y"][ipsf - 1]
                )

                # image coordinate only with wcs projection
                xwcs = (cd12 * yField - cd22 * xField) / (
                    cd12 * cd21 - cd11 * cd22
                ) + xmcen
                ywcs = (cd21 * xField - cd11 * yField) / (
                    cd12 * cd21 - cd11 * cd22
                ) + ymcen

                xwList += [xwcs]
                ywList += [ywcs]
                xdList += [xImage]
                ydList += [yImage]

        # global interpolation
        xImageFun = interpolate.SmoothBivariateSpline(
            xwList, ywList, xdList, kx=poly_degree, ky=poly_degree
        )
        yImageFun = interpolate.SmoothBivariateSpline(
            xwList, ywList, ydList, kx=poly_degree, ky=poly_degree
        )
        fdFunList[iwaveKey] = {
            "xImagePos": xImageFun,
            "yImagePos": yImageFun,
            "interpLimit": [
                np.min(xwList),
                np.max(xwList),
                np.min(ywList),
                np.max(ywList),
            ],
        }

        # construct the residual interpolation
        fdFunList[iwaveKey]["residual"] = {}
        for iccd in range(1, nccd + 1):
            # if iccd!=1: continue
            iccdKey = "ccd" + str("0%d" % (iccd))[-2:]
            # open the ditortion data
            ipsfDatn = cat_dir + "ccd%d_%s.dat" % (iccd, iwaveKey)
            ipsfDat = Table.read(ipsfDatn, format="ascii")

            ixwList, iywList = [], []
            idxList, idyList = [], []
            for ipsf in range(1, npsf + 1):
                # if ipsf!=1: continue
                print(
                    "^_^ loading: iccd-{:} iwave-{:} ipsf-{:}".format(iccd, iwave, ipsf)
                )
                xField = ipsfDat["field_x"][ipsf - 1]
                yField = ipsfDat["field_y"][ipsf - 1]
                xImage = 100.0 * (
                    ipsfDat["image_x"][ipsf - 1] + ipsfDat["centroid_x"][ipsf - 1]
                )
                yImage = 100.0 * (
                    ipsfDat["image_y"][ipsf - 1] + ipsfDat["centroid_y"][ipsf - 1]
                )

                # image coordinate only with wcs projection
                xwcs = (cd12 * yField - cd22 * xField) / (
                    cd12 * cd21 - cd11 * cd22
                ) + xmcen
                ywcs = (cd21 * xField - cd11 * yField) / (
                    cd12 * cd21 - cd11 * cd22
                ) + ymcen

                ixPred = xImageFun(xwcs, ywcs)[0][0]
                iyPred = yImageFun(xwcs, ywcs)[0][0]
                idx = xImage - ixPred
                idy = yImage - iyPred
                # print(idx, idy)
                ixwList += [xwcs]
                iywList += [ywcs]
                idxList += [idx]
                idyList += [idy]

            # interpolation
            xResFun = interpolate.SmoothBivariateSpline(
                ixwList, iywList, idxList, kx=poly_degree, ky=poly_degree
            )
            yResFun = interpolate.SmoothBivariateSpline(
                ixwList, iywList, idyList, kx=poly_degree, ky=poly_degree
            )

            fdFunList[iwaveKey]["residual"][iccdKey] = {
                "xResidual": xResFun,
                "yResidual": yResFun,
                "interpLimit": [
                    np.min(ixwList),
                    np.max(ixwList),
                    np.min(iywList),
                    np.max(iywList),
                ],
            }

    # save the interpolation functions
    model_name_full = os.path.join(model_dir, model_name)
    with open(model_name_full, "wb") as out:
        pickle.dump(fdFunList, out, pickle.HIGHEST_PROTOCOL)

    return


def field_distortion_apply(
    model_name, model_dir, cat_dir, ra_cen=60.0, dec_cen=-40.0, img_rot=0.0
):
    # CCD and observation
    ccdSize, ccdBase = ccdParam()
    xsize, ysize, xchip, ychip, xgap, ygap, xnchip, ynchip = ccdSize
    nchip = xnchip * ynchip

    #################################################
    xmcen, ymcen = 0.0, 0.0
    #################################################

    badchip = list(range(1, 6)) + list(range(26, 31)) + [10, 21]

    # define the wcs of the image mosaic
    print(
        "^_^ Construct the wcs of the entire image mosaic using Gnomonic/TAN projection"
    )
    wcs = wcsMain(imgRotation=img_rot, raCenter=ra_cen, decCenter=dec_cen)

    #################################################
    # load the field distortion model
    model_name_full = os.path.join(model_dir, model_name)
    with open(model_name_full, "rb") as f:
        fdModel = pickle.load(f)
    #################################################

    raLow, raUp, decLow, decUp = skyLim(
        wcs, -xsize // 2 + 1, xsize // 2, -ysize // 2 + 1, ysize // 2
    )
    dra = (raUp - raLow) * np.cos(dec_cen * np.pi / 180.0)
    ddec = decUp - decLow
    print(
        "    Image pixel size: %d*%d; center: (Ra, Dec)=(%.3f, %.3f)."
        % (xsize, ysize, ra_cen, dec_cen)
    )
    print("    Field of Veiw: %.2f * %.2f deg^2." % (dra, ddec))

    # filters and corresponding bounds in the image mosaic
    fbound = {}
    print("    Model the filter distributions in the image mosaic ...")
    stats = {}
    for i in range(nchip):
        chip_id = i + 1
        if chip_id in badchip:
            continue
        cx0, cx1, cy0, cy1 = chipLim(chip_id)
        chip_bound = galsim.BoundsD(cx0 - 1, cx1 - 1, cy0 - 1, cy1 - 1)
        chip_filter_id, chip_filt = chip_filter(chip_id)
        # print "^_^ CHIP %d, Filter %s"%(chip_id,chip_filter)
        fbound[chip_id] = [chip_filter_id, chip_filt, chip_bound]
        stats[chip_id] = [0, 0, 0]

    # generate object grid
    ra_input = np.arange(ra_cen - 1.0, ra_cen + 1.0, 0.00125)
    dec_input = np.arange(dec_cen - 1.0, dec_cen + 1.0, 0.00125)
    nobj = len(ra_input) * len(dec_input)
    crdCat = np.zeros((nobj, 2))
    cid = 0
    for id1 in range(len(ra_input)):
        ira = ra_input[id1]
        for id2 in range(len(dec_input)):
            idec = dec_input[id2]
            crdCat[cid, :] = ira, idec
            cid += 1
    print("^_^ Total %d objects are generaged" % nobj)

    # main program
    for i in range(nchip):
        # if i not in [6]: continue
        if i + 1 in badchip:
            continue
        filtidk, filtnmk, boundk = fbound[i + 1]
        idStr = str("0%d" % (i + 1))[-2:]

        ###################################################################
        # 1) Use global field distortion model: FieldDistModelGlobal_v2.0.pickle
        ifdModel = fdModel["wave1"]
        irsModel = fdModel["wave1"]["residual"]["ccd" + idStr]
        xLowI, xUpI, yLowI, yUpI = ifdModel["interpLimit"]
        xlLowI, xlUpI, ylLowI, ylUpI = irsModel["interpLimit"]

        # field distortion model along x/y-axis
        ixfdModel = ifdModel["xImagePos"]
        iyfdModel = ifdModel["yImagePos"]
        ixrsModel = irsModel["xResidual"]
        iyrsModel = irsModel["yResidual"]

        # first-order derivatives of the global field distortion model
        ifx_dx = ixfdModel.partial_derivative(1, 0)
        ifx_dy = ixfdModel.partial_derivative(0, 1)
        ify_dx = iyfdModel.partial_derivative(1, 0)
        ify_dy = iyfdModel.partial_derivative(0, 1)
        # first-order derivatives of the residual field distortion model
        irx_dx = ixrsModel.partial_derivative(1, 0)
        irx_dy = ixrsModel.partial_derivative(0, 1)
        iry_dx = iyrsModel.partial_derivative(1, 0)
        iry_dy = iyrsModel.partial_derivative(0, 1)
        ###################################################################

        # construct the image mosaic firstly
        xorigin, yorigin = xmcen - boundk.xmin, ymcen - boundk.ymin
        print("    Construct the chip mosaic ...")
        fimage = galsim.ImageF(xchip, ychip)
        fimage.setOrigin(boundk.xmin, boundk.ymin)

        fimage.wcs = wcs
        raLow, raUp, decLow, decUp = skyLim(
            wcs, boundk.xmin, boundk.xmax, boundk.ymin, boundk.ymax
        )
        dra = (raUp - raLow) * np.cos(dec_cen * np.pi / 180.0)
        ddec = decUp - decLow
        print("    Image coverage: %.2f * %.2f arcmin^2." % (dra * 60.0, ddec * 60.0))
        # enlarge the sky coverage in order to catch the galaxies at the chip edge
        raLow -= 0.2 / 60.0
        decLow -= 0.2 / 60.0
        raUp += 0.2 / 60.0
        decUp += 0.2 / 60.0
        print(
            "    Range: RA=[%.4f, %.4f]; DEC=[%.4f, %.4f]"
            % (raLow, raUp, decLow, decUp)
        )

        # generate the galaxy and star images
        catxxn = os.path.join(
            cat_dir, "csst_mainfocus_field_distortion_ccd%s_%s.cat" % (idStr, filtnmk)
        )
        hdrxx = "#id_obj id_chip filter ra_true dec_ture x_image_ture y_image_ture x_image y_image g1_fd g2_fd\n"
        fmtxx = "%8d %3d %4s %12.6f %12.6f %13.6f %13.6f %13.6f %13.6f %9.5f %9.5f\n"
        catxx = open(catxxn, "w")
        catxx.write(hdrxx)
        oidxx = 0
        for k in range(nobj):
            # if k != 0: continue
            # input galaxy parameters
            rak = crdCat[k, 0]
            deck = crdCat[k, 1]

            # reject objects out of the image
            if (rak - raLow) * (rak - raUp) > 0.0 or (deck - decLow) * (
                deck - decUp
            ) > 0.0:
                continue

            world_pos = galsim.CelestialCoord(
                ra=rak * galsim.degrees, dec=deck * galsim.degrees
            )
            image_pos = fimage.wcs.toImage(world_pos)
            xk_true = image_pos.x
            yk_true = image_pos.y

            #################################################################
            # field distortion
            if (xLowI - xk_true) * (xUpI - xk_true) > 0 or (yLowI - yk_true) * (
                yUpI - yk_true
            ) > 0:
                continue
            xk = ixfdModel(xk_true, yk_true)[0][0]
            yk = iyfdModel(xk_true, yk_true)[0][0]

            # global offset correction
            if (xlLowI - xk) * (xlUpI - xk) > 0 or (ylLowI - yk) * (ylUpI - yk) > 0:
                continue
            dxk = ixrsModel(xk, yk)[0][0]
            dyk = iyrsModel(xk, yk)[0][0]
            xk = xk + dxk
            yk = yk + dyk

            # field distortion induced ellipticity
            ix_dx = ifx_dx(xk, yk) + irx_dx(xk, yk)
            ix_dy = ifx_dy(xk, yk) + irx_dy(xk, yk)
            iy_dx = ify_dx(xk, yk) + iry_dx(xk, yk)
            iy_dy = ify_dy(xk, yk) + iry_dy(xk, yk)

            g1k_fd = 0.0 + (iy_dy - ix_dx) / (iy_dy + ix_dx)
            g2k_fd = 0.0 - (iy_dx + ix_dy) / (iy_dy + ix_dx)
            #################################################################
            dxk_true, dyk_true = xk_true - xmcen, yk_true - ymcen
            xLock_true, yLock_true = dxk_true + xorigin + 1.0, dyk_true + yorigin + 1.0

            dxk, dyk = xk - xmcen, yk - ymcen
            xLock, yLock = dxk + xorigin + 1.0, dyk + yorigin + 1.0

            if (xLock_true < 0) or (xLock_true > xchip):
                continue
            if (yLock_true < 0) or (yLock_true > ychip):
                continue
            if (xLock < 0) or (xLock > xchip):
                continue
            if (yLock < 0) or (yLock > ychip):
                continue
            linexx = fmtxx % (
                k + 1,
                i + 1,
                filtnmk.lower(),
                rak,
                deck,
                xLock_true,
                yLock_true,
                xLock,
                yLock,
                g1k_fd[0][0],
                g2k_fd[0][0],
            )
            catxx.write(linexx)

        catxx.close()

    return


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