PSFConfig.py

import sys
from itertools import islice

import numpy as np
import matplotlib.pyplot as plt

import scipy.io
#from scipy.io import loadmat

from scipy import ndimage

import ctypes
import galsim

from PSF.PSFInterp.PSFUtil import *


###加载PSF信息###
def LoadPSF(iccd, iwave, ipsf, psfPath, psfSampleSize=5, InputMaxPixelPos=False, PSFCentroidWgt=False):
    """
    load psf informations from psf matrix.

    Parameters:
        iccd (int): ccd number [1,30].
        iwave(int): wave-index [1,4].
        ipsf (int): psf number [1, 100].
        psfPath (int): path to psf matrix
        psfSampleSize (float-optional): psf size in microns.
        InputMaxPixelPos(bool-optional): only True for 30*30 psf-matrix 
    Returns:
        psfInfo (dirctionary)
    """
    if iccd not in np.linspace(1, 30, 30, dtype='int'):
        print('Error - iccd should be in [1, 30].')
        sys.exit()
    if iwave not in np.linspace(1, 4, 4, dtype='int'):
        print('Error - iwave should be in [1, 4].')
        sys.exit()
    if ipsf not in np.linspace(1, 900, 900, dtype='int'):
        print('Error - ipsf should be in [1, 900].')
        sys.exit()

    psfInfo = {}
    fpath = psfPath +'/' +'ccd{:}'.format(iccd) +'/' + 'wave_{:}'.format(iwave)

    #获取ipsf矩阵
    if not PSFCentroidWgt:
        ##读取PSF原数据
        fpathMat = fpath +'/' +'5_psf_array' +'/' +'psf_{:}.mat'.format(ipsf)
        data = scipy.io.loadmat(fpathMat)
        psfInfo['psfMat'] = data['psf']
    if PSFCentroidWgt:
        ##读取PSFCentroidWgt
        ffpath = psfPath +'_proc/' +'ccd{:}'.format(iccd) +'/' + 'wave_{:}'.format(iwave)
        ffpathMat = ffpath +'/' +'5_psf_array' +'/' +'psf_{:}_centroidWgt.mat'.format(ipsf)
        data = scipy.io.loadmat(ffpathMat)
        psfInfo['psfMat'] = data['psf']


    #获取ipsf波长
    fpathWave = fpath +'/' +'1_wavelength.txt'
    f = open(fpathWave, 'r')
    wavelength = np.float(f.readline())
    f.close()
    psfInfo['wavelength'] = wavelength

    #获取ipsf位置
    fpathCoordinate = fpath +'/' +'4_PSF_coordinate.txt'
    f = open(fpathCoordinate, 'r')
    header = f.readline()
    for line in islice(f, ipsf-1, ipsf):
        line = line.strip()
        columns = line.split()
    f.close()
    icol = 0
    psfInfo['field_x'] = float(columns[icol])    #deg, 视场采样位置
    icol+= 1
    psfInfo['field_y'] = float(columns[icol])    #deg
    icol+= 1
    psfInfo['centroid_x'] = float(columns[icol]) #mm, psf质心相对主光线的偏移量
    icol+= 1
    psfInfo['centroid_y'] = float(columns[icol]) #mm
    icol+= 1
    if InputMaxPixelPos == True:
        psfInfo['max_x'] = float(columns[icol])  #mm, max pixel postion
        icol+= 1
        psfInfo['max_y'] = float(columns[icol])  #mm
        icol+= 1
    psfInfo['image_x'] = float(columns[icol])    #mm, 主光线位置
    icol+= 1
    psfInfo['image_y'] = float(columns[icol])    #mm

    nrows, ncols = psfInfo['psfMat'].shape
    psfPos   = psfPixelLayout(nrows, ncols, psfInfo['image_y'], psfInfo['image_x'], pixSizeInMicrons=5.0)
    imgMaxPix_x, imgMaxPix_y = findMaxPix(psfInfo['psfMat'])
    psfInfo['imgMaxPosx_ccd'] = psfPos[0, imgMaxPix_y, imgMaxPix_x] #cx, psf最大值位置, in mm
    psfInfo['imgMaxPosy_ccd'] = psfPos[1, imgMaxPix_y, imgMaxPix_x] #cy

    
    if PSFCentroidWgt:
        #if ipsf == 1:
            #print('Check:', psfInfo['centroid_x'], psfInfo['centroid_y'], data['cx'][0][0], data['cy'][0][0])
        psfInfo['centroid_x'] = data['cx'][0][0] #mm, psfCentroidWgt质心相对主光线的偏移量
        psfInfo['centroid_y'] = data['cy'][0][0] #mm
    return psfInfo


###加载PSF矩阵信息###
def LoadPSFset(iccd, iwave, npsf, psfPath, psfSampleSize=5, InputMaxPixelPos=False, PSFCentroidWgt=False):
    """
    load psfSet for interpolation

    Parameters:
        iccd, iwave, psfPath: # of ccd/wave and path for psfs
        npsf: 100 or 900 for different psf matrixes

    Returns:
        PSFMat (numpy.array): images
        cen_col, cen_row (numpy.array, numpy.array): position of psf center in the view field
    """
    psfSet = []
    for ipsf in range(1, npsf+1):
        psfInfo = LoadPSF(iccd, iwave, ipsf, psfPath, InputMaxPixelPos=InputMaxPixelPos, PSFCentroidWgt=PSFCentroidWgt)
        psfSet.append(psfInfo)

    ngy, ngx = psfSet[0]['psfMat'].shape
    PSFMat = np.zeros([npsf, ngy, ngx])
    cen_col= np.zeros(npsf)
    cen_row= np.zeros(npsf)
    FieldPos = False
    for ipsf in range(npsf):
        iPSFMat = psfSet[ipsf]['psfMat']
        PSFMat[ipsf, :, :] = psfTailor(iPSFMat, apSizeInArcsec=2.5)

        if FieldPos == True:
            cen_col[ipsf] = psfSet[ipsf]['field_x'] #cx
            cen_row[ipsf] = psfSet[ipsf]['field_y'] #cy
        if FieldPos == False:
            cen_col[ipsf] = psfSet[ipsf]['image_x'] + psfSet[ipsf]['centroid_x'] #cx 
            cen_row[ipsf] = psfSet[ipsf]['image_y'] + psfSet[ipsf]['centroid_y'] #cy
            #cen_col[ipsf] = psfSet[ipsf]['imgMaxPosx_ccd'] #cx --- to be updated
            #cen_row[ipsf] = psfSet[ipsf]['imgMaxPosy_ccd'] #cy
            #cen_col[ipsf] = psfSet[ipsf]['image_x']
            #cen_row[ipsf] = psfSet[ipsf]['image_y']
    return psfSet, PSFMat, cen_col, cen_row


###插值PSF图像-IDW方法###
def psfMaker_IDW(px, py, PSFMat, cen_col, cen_row, IDWindex=2, OnlyNeighbors=True, hoc=None, hoclist=None, PSFCentroidWgt=False):
    """
    psf interpolation by IDW

    Parameters:
        px, py (float, float): position of the target
        PSFMat (numpy.array): image
        cen_col, cen_row (numpy.array, numpy.array): potions of the psf centers
        IDWindex (int-optional): the power index of IDW
        OnlyNeighbors (bool-optional): only neighbors are used for psf interpolation

    Returns:
        psfMaker (numpy.array)
    """
    minimum_psf_weight = 1e-8
    ref_col = px
    ref_row = py

    ngy, ngx = PSFMat[0, :, :].shape
    npsf = PSFMat[:, :, :].shape[0]
    psfWeight = np.zeros([npsf])

    if OnlyNeighbors == True:
        if hoc is None:
            neigh = findNeighbors(px, py, cen_col, cen_row, dr=5., dn=4, OnlyDistance=False)
        if hoc is not None:
            #hoc,hoclist = findNeighbors_hoclist(cen_col, cen_row)
            neigh = findNeighbors_hoclist(cen_col, cen_row, tx=px,ty=py, dn=4, hoc=hoc, hoclist=hoclist)

        #print("neigh:", neigh)
        neighFlag = np.zeros(npsf)
        neighFlag[neigh] = 1

    for ipsf in range(npsf):
        if OnlyNeighbors == True:
            if neighFlag[ipsf] != 1:
                continue

        dist = np.sqrt((ref_col - cen_col[ipsf])**2 + (ref_row - cen_row[ipsf])**2)
        if IDWindex == 1:
            psfWeight[ipsf] = dist
        if IDWindex == 2:
            psfWeight[ipsf] = dist**2
        if IDWindex == 3:
            psfWeight[ipsf] = dist**3
        if IDWindex == 4:
            psfWeight[ipsf] = dist**4
        psfWeight[ipsf] = max(psfWeight[ipsf], minimum_psf_weight)
        psfWeight[ipsf] = 1./psfWeight[ipsf]
    psfWeight /= np.sum(psfWeight)

    psfMaker  = np.zeros((ngy, ngx), dtype='float64')
    for ipsf in range(npsf):
        if OnlyNeighbors == True:
            if neighFlag[ipsf] != 1:
                continue

        iPSFMat = PSFMat[ipsf, :, :].copy()
        #if not PSFCentroidWgt:
        #    iPSFMat = psfCentering(iPSFMat, CenteringMode=2)
            #iPSFMat = psfCentering_FFT(iPSFMat)
        #if PSFCentroidWgt:
            
        ipsfWeight = psfWeight[ipsf]
        psfMaker += iPSFMat * ipsfWeight
        #print("ipsf, ipsfWeight:", ipsf, ipsfWeight)
    psfMaker /= np.nansum(psfMaker)

    return psfMaker



###TEST###
if __name__ == '__main__':
    iccd = 1
    iwave= 1
    ipsf = 1 
    npsf = 100
    psfPath = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_ciomp'

    a, b, c, d = LoadPSFset(iccd, iwave, npsf, psfPath, InputMaxPixelPos=False)
    psfSet = a

    print('psfSet has been loaded.')
    print('Usage: psfSet[i][keys]')
    print('psfSet.keys:', psfSet[0].keys())
    print(c[0:10])
    print(d[0:10])