Commit 50fd7c8d authored by JX's avatar JX 😵
Browse files

Merge remote-tracking branch 'origin/develop'

parents f4e3478d 50f8dc5d
Pipeline #4413 passed with stage
in 0 seconds
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 = os.path.join(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 = os.path.join(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()
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment