psfConfig.py

"""
CSST image simulation module (in python3): Point Spread Function (PSF)
author:: Wei Chengliang <chengliangwei@pmo.ac.cn>
"""

import sys
from itertools import islice

import numpy as np
import matplotlib.pyplot as plt

import scipy.io
from scipy.io import loadmat
#import xlrd

from scipy import ndimage
from scipy.interpolate import RectBivariateSpline

#from astropy.modeling.models import Ellipse2D
#from astropy.coordinates import Angle
#import matplotlib.patches as mpatches

import ctypes
import galsim



def setupPSFimg(iccd, iwave, psfPath="/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_ciomp"):
    """
    psf model setup for csst-sim
   
    Parameters:
        iccd, iwave (int, int): psf model on iccd & iwave
        psfPath (string, optional): path to psf matrix

    Returns:
        psf_model (psf_class): psf model
  
    Methods:
        psf_model.PSFinplace(self, px, py, interpScheme=1): psf interpolation
        psf_model.PSFspin(self, psf, sigSpin, sigGauss, dx, dy): psf rotation (from Yudong)
    """
    psf_model = PSFimg(iccd, iwave, psfPath)
    return psf_model


##################################################
#       A. psf matrix loading & checking         #
##################################################
def psfPixelLayout(nrows, ncols, cenPosRow, cenPosCol, pixSizeInMicrons=5.0):
    """
    convert psf pixels to physical position
    
    Parameters:
        nrows, ncols (int, int): psf sampling with [nrows, ncols].
        cenPosRow, cenPosCol (float, float): A physical position of the chief ray for a given psf.
        pixSizeInMicrons (float-optional): The pixel size in microns from the psf sampling.
        
    Returns:
        psfPixelPos (numpy.array-float): [posx, posy] in mm for [irow, icol]
 
    Notes:
        1. show positions on ccd, but not position on image only [+/- dy]
    """
    psfPixelPos = np.zeros([2, nrows, ncols])
    if nrows % 2 != 0:
        sys.exit()
    if ncols % 2 != 0:
        sys.exit()
        
    cenPix_row = nrows/2 + 1 #中心主光线对应pixle [由长光定义]
    cenPix_col = ncols/2 + 1

    for irow in range(nrows):
        for icol in range(ncols):
            delta_row = ((irow + 1) - cenPix_row)*pixSizeInMicrons*1e-3
            delta_col = ((icol + 1) - cenPix_col)*pixSizeInMicrons*1e-3
            psfPixelPos[0, irow, icol] = cenPosCol + delta_col
            psfPixelPos[1, irow, icol] = cenPosRow - delta_row  #note-1
            
    return psfPixelPos


def imSigRange(img, fraction=0.80):
    """
    extract the image within x-percent (DISCARD)
    
    Parameters:
        img (numpy.array-float): image
        fraction (float-optional): a percentage

    Returns:
        im1 (numpy.array-float): image
    """
    im1 = img.copy()
    im1size = im1.shape
    im2 = np.sort(im1.reshape(im1size[0]*im1size[1]))
    im2 = im2[::-1]
    im3 = np.cumsum(im2)/np.sum(im2)
    loc = np.where(im3 > fraction)
    #print(im3[loc[0][0]], im2[loc[0][0]])
    im1[np.where(im1 <= im2[loc[0][0]])]=0

    return im1


def imPlotInRange(img):
    """
    plot image within a selected range
    
    Parameters:
        img (numpy.array-float): image

    Returns:
    """
    im1 = img.copy()
    im1size = im1.shape
    X,Y = np.meshgrid(range(im1size[1]),range(im1size[0]))
    Z = im1

    resolution = 25

    f = lambda x,y: Z[int(y),int(x) ]
    g = np.vectorize(f)

    x = np.linspace(0,Z.shape[1], Z.shape[1]*resolution)
    y = np.linspace(0,Z.shape[0], Z.shape[0]*resolution)
    X2, Y2= np.meshgrid(x[:-1],y[:-1])
    Z2 = g(X2,Y2)

    #plt.pcolormesh(X,Y, Z)
    #plt.imshow(img, origin='lower')
    plt.contour(X2-0.5,Y2-0.5,Z2, [0.], colors='w', linestyles='--',  linewidths=[1])

    return 


def findMaxPix(img):
    """
    get the pixel position of the maximum-value
    
    Parameters:
        img (numpy.array-float): image
    
    Returns:
        imgMaxPix_x, imgMaxPix_y (int, int): pixel position in columns & rows
    """
    maxIndx = np.argmax(img)
    maxIndx = np.unravel_index(maxIndx, np.array(img).shape)
    imgMaxPix_x = maxIndx[1]
    imgMaxPix_y = maxIndx[0]

    return imgMaxPix_x, imgMaxPix_y


def psfTailor(img, apSizeInArcsec=0.5, psfSampleSizeInMicrons=5, focalLengthInMeters=28):
    """
    psf tailor within a given aperture size
 
    Parameters:
        img (numpy.array-float): image
        apSizeInArcsec (float-optional): aperture size in arcseconds.
        psfSampleSizeInMicrons (float-optional): psf pixel size in microns.
        focalLengthInMeters (float-optional): csst focal length im meters.
    Returns:
        imgT (numpy.array-float): image
    """
    imgMaxPix_x, imgMaxPix_y = findMaxPix(img)
    apSizeInMicrons = np.deg2rad(apSizeInArcsec/3600.)*focalLengthInMeters*1e6
    apSizeInPix = apSizeInMicrons/psfSampleSizeInMicrons
    apSizeInPix = np.int(np.ceil(apSizeInPix))
    imgT = np.zeros_like(img)
    imgT[imgMaxPix_y-apSizeInPix:imgMaxPix_y+apSizeInPix+1, 
         imgMaxPix_x-apSizeInPix:imgMaxPix_x+apSizeInPix+1] = \
    img[imgMaxPix_y-apSizeInPix:imgMaxPix_y+apSizeInPix+1, 
        imgMaxPix_x-apSizeInPix:imgMaxPix_x+apSizeInPix+1]    
    return imgT


def psfEncircle(img, fraction=0.8, psfSampleSizeInMicrons=5, focalLengthInMeters=28):
    """
    psf tailor within a given percentage.
 
    Parameters:
        img (numpy.array-float): image
        fraction (float-optional): a percentage for psf tailor.
        psfSampleSizeInMicrons (float-optional): psf pixel size in microns.
        focalLengthInMeters (float-optional): csst focal length im meters.
    Returns:
        img*wgt (numpy.array-float): image
        REE80 (float): radius of REE80 in arcseconds.
    """
    imgMaxPix_x, imgMaxPix_y = findMaxPix(img)
    im1 = img.copy()
    im1size = im1.shape
    
    dis = np.zeros_like(img)
    for irow in range(im1size[0]):
        for icol in range(im1size[1]):
            dx = icol - imgMaxPix_x
            dy = irow - imgMaxPix_y
            dis[irow, icol] = np.hypot(dx, dy)
            
    nn = im1size[1]*im1size[0]
    disX = dis.reshape(nn)
    disXsortId = np.argsort(disX)

    imgX = img.reshape(nn)
    imgY = imgX[disXsortId]
    psfFrac = np.cumsum(imgY)/np.sum(imgY)
    ind = np.where(psfFrac > fraction)[0][0]
    
    wgt = np.ones_like(dis)
    wgt[np.where(dis > dis[np.where(img == imgY[ind])])] = 0
    
    REE80 = np.rad2deg(dis[np.where(img == imgY[ind])]*psfSampleSizeInMicrons*1e-6/focalLengthInMeters)*3600
    return img*wgt, REE80


def psfSecondMoments(psfMat, cenX, cenY, pixSize=1):
    """
    estimate the psf ellipticity by the second moment of surface brightness

    Parameters:
        psfMat (numpy.array-float): image
        cenX, cenY (float, float): pixel position of the psf center
        pixSize (float-optional): pixel size
 
    Returns:
        sz (float): psf size
        e1, e2 (float, float): psf ellipticity
    """
    I = psfMat
    ncol = I.shape[1]
    nrow = I.shape[0] 
    w   = 0.0
    w11 = 0.0
    w12 = 0.0
    w22 = 0.0
    for icol in range(ncol):
        for jrow in range(nrow):
            x = icol*pixSize - cenX
            y = jrow*pixSize - cenY
            w   += I[jrow, icol]
            w11 += x*x*I[jrow, icol]
            w12 += x*y*I[jrow, icol]
            w22 += y*y*I[jrow, icol]
    w11 /= w
    w12 /= w
    w22 /= w
    sz = w11 + w22
    e1 = (w11 - w22)/sz
    e2 = 2.0*w12/sz
    
    return sz, e1, e2


def LoadPSF(iccd, iwave, ipsf, psfPath, psfSampleSize=5, CalcPSFsize=True, CalcPSFcenter=True, SigRange=False, TailorScheme=1, InputMaxPixelPos=False):
    '''加载psf信息'''
    """
    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.
        CalcPSFsize (bool-optional): whether calculate psf size & ellipticity. Default: True
        CalcPSFcenter (bool-optional): whether calculate psf center. Default: True 
        SigRange (bool-optional): whether use psf tailor. Default: False
        TailorScheme (int-optional): which method for psf tailor. Default: 1
    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矩阵
    fpathMat = fpath +'/' +'5_psf_array' +'/' +'psf_{:}.mat'.format(ipsf)
    data = scipy.io.loadmat(fpathMat)
    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
        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 = 180  #psf采样大小, in pixels
    #ncols = 180
    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

    #计算psf size & ellipticity
    if CalcPSFsize is True:
        psfMat = psfInfo['psfMat'].copy()
        cenX, cenY, sz, e1, e2, REE80 = psfSizeCalculator(psfMat, psfSampleSize=psfSampleSize, CalcPSFcenter=CalcPSFcenter, SigRange=SigRange, TailorScheme=TailorScheme)
        
        psfInfo['psfCenX_img'] = cenX  #in local pixels, psf质心位置, in pixels
        psfInfo['psfCenY_img'] = cenY  #in local pixels
        psfInfo['psfSize'] = sz
        psfInfo['psf_e1'] = e1
        psfInfo['psf_e2'] = e2
        psfInfo['REE80'] = REE80
    
    return psfInfo


def psfSizeCalculator(psfMat, psfSampleSize=5, CalcPSFcenter=True, SigRange=False, TailorScheme=1):
    """
    calculate psf size & ellipticity
    
    Parameters:
        psfMat (numpy.array): image
        psfSampleSize (float-optional): psf size in microns.
        CalcPSFcenter (bool-optional): whether calculate psf center. Default: True
        SigRange (bool-optional): whether use psf tailor. Default: False
        TailorScheme (int-optional): which method for psf tailor. Default: 1
    Returns:
        cenX, cenY (float, float): the pixel position of the mass center 
        sz (float): psf size
        e1, e2 (float, float): psf ellipticity
        REE80 (float): radius of REE80 in arcseconds
    """
    psfSampleSize = psfSampleSize*1e-3 #mm

    REE80 = -1.0  ##encircling 80% energy
    if SigRange is True:
        if TailorScheme == 1:
            psfMat = imSigRange(psfMat, fraction=0.80)
            psfInfo['psfMat'] = psfMat  #set on/off
        if TailorScheme == 2:
            img = psfTailor(psfMat, apSizeInArcsec=0.5)
            imgX, REE80 = psfEncircle(psfMat)
            psfMat = img
            REE80 = REE80[0]
                
    if CalcPSFcenter is True:
        img = psfMat/np.sum(psfMat)
        y,x = ndimage.center_of_mass(img)  #y-rows, x-cols
        cenX = x
        cenY = y        
    if CalcPSFcenter is False:
        cenPix_X = psfMat.shape[1]/2 #90
        cenPix_Y = psfMat.shape[0]/2 #90
        cenX = cenPix_X + psfInfo['centroid_x']/psfSampleSize
        cenY = cenPix_Y - psfInfo['centroid_y']/psfSampleSize

    pixSize = 1 
    sz, e1, e2 = psfSecondMoments(psfMat, cenX, cenY, pixSize=pixSize)
        
    return cenX, cenY, sz, e1, e2, REE80


def psfStack(*psfMat):
    """
    stacked image from the different psfs

    Parameters:
        *psfMat (numpy.array): the different psfs for stacking
 
    Returns:
        img (numpy.array): image
    """
    nn = len(psfMat)
    img = np.zeros_like(psfMat[0])
    for ii in range(nn):
        img += psfMat[ii]/np.sum(psfMat[ii])
    img /= np.sum(img)
    return img



##################################################
#             B. psf interpolation               #
##################################################
def img2fits(img, fitsName=None):
    """
    saving image to fits file

    Parameters:
        img (numpy.array): image
        fitsName (string): path+filename of fits

    Returns:
    """
    from astropy.io import fits
    grey    = fits.PrimaryHDU(img)
    greyHDU = fits.HDUList([grey])
    if fitsName != None:
        greyHDU.writeto(fitsName)


def psfMatLoad(iccd, iwave, psfPath, psfSampleSize=5, CalcPSFsize=False, CalcPSFcenter=True):
    """
    load psf for interpolation
 
    Parameters:
        iccd, iwave, psfPath: # of ccd/wave and path for psfs
        CalcPSFsize (bool-optional): whether calculate psf size & ellipticity. Default: False
        CalcPSFcenter (bool-optional): whether calculate psf center. Default: True

    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, 101):
        psfInfo = LoadPSF(iccd, iwave, ipsf, psfPath, CalcPSFsize=CalcPSFsize, CalcPSFcenter=CalcPSFcenter, InputMaxPixelPos=False)
        psfSet.append(psfInfo)

    npsf = len(psfSet)
    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):
        PSFMat[ipsf, :, :] = psfSet[ipsf]['psfMat']
        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]['imgMaxPosx_ccd'] #cx
            cen_row[ipsf] = psfSet[ipsf]['imgMaxPosy_ccd'] #cy
    return PSFMat, cen_col, cen_row


def findNeighbors(tx, ty, px, py, dr=0.1, dn=1, OnlyDistance=True):
    """
    find nearest meighbors by 2D-KDTree

    Parameters:
        tx, ty (float, float): a given position
        px, py (numpy.array, numpy.array): position data for tree
        dr (float-optional): distance
        dn (int-optional): nearest-N
        OnlyDistance (bool-optional): only use distance to find neighbors. Default: True

    Returns:
        dataq (numpy.array): index
    """
    import scipy.spatial as spatial

    datax = px
    datay = py
    tree = spatial.KDTree(list(zip(datax.ravel(), datay.ravel())))

    dataq=[]
    rr = dr
    if OnlyDistance == True:
        dataq = tree.query_ball_point([tx, ty], rr)
    if OnlyDistance == False:
        while len(dataq) < dn:
            dataq = tree.query_ball_point([tx, ty], rr)
            rr += dr
        dd = np.hypot(datax[dataq]-tx, datay[dataq]-ty)
        ddSortindx = np.argsort(dd)
        dataq = np.array(dataq)[ddSortindx[0:dn]]
    return dataq


def psfCentering(img, apSizeInArcsec=4., psfSampleSizeInMicrons=5, focalLengthInMeters=28, CenteringMode=1):
    """
    centering psf within an aperture
    
    Parameters:
        img (numpy.array): image
        apSizeInArcsec (float-optional): aperture size in arcseconds.
        psfSampleSizeInMicrons (float-optional): psf pixel size in microns.
        focalLengthInMeters (float-optional): csst focal length im meters.
        CenteringMode (int-optional): how to center psf images

    Returns:
        imgT (numpy.array)
    """
    if CenteringMode == 1:
        imgMaxPix_x, imgMaxPix_y = findMaxPix(img)
    if CenteringMode == 2:
        y,x = ndimage.center_of_mass(img)  #y-rows, x-cols
        imgMaxPix_x = int(x)
        imgMaxPix_y = int(y)
    apSizeInMicrons = np.deg2rad(apSizeInArcsec/3600.)*focalLengthInMeters*1e6
    apSizeInPix = apSizeInMicrons/psfSampleSizeInMicrons
    apSizeInPix = np.int(np.ceil(apSizeInPix))
    imgT = np.zeros_like(img)
    ngy, ngx = img.shape
    cy = int(ngy/2)
    cx = int(ngx/2)
    imgT[cy-apSizeInPix:cy+apSizeInPix+1, 
         cx-apSizeInPix:cx+apSizeInPix+1] = \
    img[imgMaxPix_y-apSizeInPix:imgMaxPix_y+apSizeInPix+1, 
        imgMaxPix_x-apSizeInPix:imgMaxPix_x+apSizeInPix+1]    
    return imgT


def psfMaker_IDW(px, py, PSFMat, cen_col, cen_row, IDWindex=2, OnlyNeighbors=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:
        neigh = findNeighbors(px, py, cen_col, cen_row, dr=5., dn=9, OnlyDistance=False)
        neighFlag = np.zeros(npsf)
        neighFlag[neigh] = 1
        # print("neigh:", neigh)

    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()
        iPSFMat = psfCentering(iPSFMat, CenteringMode=1)
        ipsfWeight = psfWeight[ipsf]
        psfMaker += iPSFMat * ipsfWeight
    psfMaker /= np.nansum(psfMaker)
    
    return psfMaker


def psfMaker_PCA(px, py, PSFMat, cen_col, cen_row, OnlyNeighbors=False, libPCApath='libPCA/libPCA.so'):
    """
    psf interpolation by PCA
    
    Parameters:

    Returns:
    """
    ref_col = px
    ref_row = py
    
    ngy, ngx = PSFMat[0, :, :].shape
    npsf = PSFMat[:, :, :].shape[0]
    
    neigh   = findNeighbors(px, py, cen_col, cen_row, dr=0.3, dn=5, OnlyDistance=False)
    npsfX   = len(neigh)
    psfMatX = np.zeros([npsfX, ngy, ngx])
    cen_colX= np.zeros(npsfX)
    cen_rowX= np.zeros(npsfX)
    for ipsf in range(npsfX):
        psfMatX[ipsf, :, :] = PSFMat[neigh[ipsf], :, :]
        cen_colX[ipsf] = cen_col[neigh[ipsf]]
        cen_rowX[ipsf] = cen_row[neigh[ipsf]]

    psfMaker  = np.zeros((ngy, ngx), dtype='float64')
    if OnlyNeighbors == True:
        PCAbasef, PCAcoeff = psfPCA_generator(psfMatX, npsfX, ngx, libPCApath)
        nPCA  = npsfX
        for iPCA in range(nPCA):
            coeffX = fitPoly(ref_col, ref_row, cen_colX, cen_rowX, PCAcoeff[:, iPCA], order=2)
            psfMaker += coeffX*PCAbasef[iPCA, :, :]          
    
    return psfMaker


def psfPCA_generator(psfMat, npsf, npix, libPCApath):
    """
    generate PCs from psfMat
  
    Parameters:

    Returns:

    """
    libPCA = ctypes.CDLL(libPCApath)  # CDLL加载库

    libPCA.psfPCA.argtypes = [ctypes.POINTER(ctypes.c_float), ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double)]

    Nstar = npsf
    Mp    = npix*npix
    NM    = Nstar*Mp
    NN    = Nstar*Nstar
    arr   = (ctypes.c_float*NM)()
    basef = (ctypes.c_double*NM)()
    coeff = (ctypes.c_double*NN)()

    psfT  = np.zeros(NM)
    for ipsf in range(npsf):
        lp = 0 + ipsf*Mp
        up = Mp+ ipsf*Mp
        ipsfMat = psfMat[ipsf, :, :]
        psfT[lp:up] = ipsfMat.reshape(Mp)

    arr[:] = psfT
    libPCA.psfPCA(arr, Nstar, Mp, basef, coeff)

    PCAbasef = np.zeros([npsf, npix, npix])
    PCAcoeff = np.zeros([npsf, npsf])
    for ipsf in range(npsf):
        lp = 0 + ipsf*Mp
        up = Mp+ ipsf*Mp
        PCAbasef[ipsf, :, :] = np.array(basef[lp:up]).reshape(npix, npix)

        lp = 0   + ipsf*npsf
        up = npsf+ ipsf*npsf
        PCAcoeff[ipsf, :] = np.array(coeff[lp:up])

    return PCAbasef, PCAcoeff

def fitPoly(px, py, datax, datay, dataz, order = 2):
    if order == 1:
        # best-fit linear plane
        A = np.c_[datax, datay, np.ones(datax.shape[0])]
        C,_,_,_ = scipy.linalg.lstsq(A, dataz)    # coefficients
        pz = C[0]*px + C[1]*py + C[2]
    elif order == 2:
        # best-fit quadratic curve
        A = np.c_[np.ones(datax.shape[0]), np.c_[datax, datay], np.prod(np.c_[datax, datay], axis=1), np.c_[datax, datay]**2]
        C,_,_,_ = scipy.linalg.lstsq(A, dataz)
        pz = np.dot(np.c_[1, px, py, px*py, px**2, py**2], C)
    """
    elif order == 3:
       # best-fit cubic curve
        A = np.c_[np.ones(datax.shape[0]), np.c_[datax, datay], np.prod(np.c_[datax, datay], axis=1), np.c_[datax, datay]**2, np.c_[datax, datay]**3]
        C,_,_,_ = scipy.linalg.lstsq(A, dataz)
        pz = np.dot(np.c_[1, px, py, px*py, px**2, py**2, px**3, py**3], C)
    """
    return pz




"""
############################
### not used temporarily ###
############################
def psfSplineMake(px, py, PSFMat, cen_col, cen_row, OnlyNeighbors=False):
    minimum_psf_weight = 1e-8
    ref_col = px
    ref_row = py
    
    cdelt1p = 1
    cdelt2p = 1
    
    ngy, ngx = PSFMat[0, :, :].shape
    psfx = np.linspace(0, ngx-1, ngx)
    psfy = np.linspace(0, ngy-1, ngy)

    npsf = PSFMat[:, :, :].shape[0]
    psfWeight = np.zeros([npsf])
    for ipsf in range(npsf):
        psfWeight[ipsf] = np.sqrt((ref_col - cen_col[ipsf])**2 + (ref_row - cen_row[ipsf])**2)
        psfWeight[ipsf] = max(psfWeight[ipsf], minimum_psf_weight)
        psfWeight[ipsf] = 1./psfWeight[ipsf]
    psfWeight /= np.sum(psfWeight)

    psf  = np.zeros((ngy, ngx), dtype='float64')
    for ipsf in range(npsf):
        iPSFMat = PSFMat[ipsf, :, :]
        ipsfWeight = psfWeight[ipsf]
        psf += iPSFMat * ipsfWeight

    psf /= (np.nansum(psf) * cdelt1p * cdelt2p)
    psfSpline = RectBivariateSpline(psfy, psfx, psf)
    return psf, psfSpline


def psfToImage(psfSpline, cutoff_radius=180):
    ng = 180
    img = np.zeros([ng, ng], dtype='float64')
    for i in range(ng):
        for j in range(ng):
            star_row = 5
            star_column = 5
            if np.sqrt((j-star_column)**2 + (i-star_row)**2) <= cutoff_radius:
                star_flux = 8
                column_cen = j #j - star_column
                row_cen = i #i - star_row
                img[i,j] += star_flux * psfSpline.integral(row_cen-0.5, row_cen+0.5, column_cen-0.5, column_cen+0.5)
    return img
"""



##################################################
#                C. csstPSF class                #
##################################################
class PSFimg(object):
    def __init__(self, iccd, iwave, psfPath):
        self.iccd = iccd
        self.iwave= iwave
        self.psfPath = psfPath

        #loading psfSet >>>
        """
        psfSet = []
        for ipsf in range(1, 901):
            psfInfo = LoadPSF(iccd, iwave, ipsf, psfPath, CalcPSFsize=True, CalcPSFcenter=True, SigRange=False)
            psfSet.append(psfInfo)
        self.psfSet  = psfSet
        """
        a, b, c = psfMatLoad(iccd, iwave, psfPath)
        self.psfMat = a
        self.cenPosx= b
        self.cenPosy= c

    def PSFinplace(self, px, py, interpScheme=1):
        if interpScheme == 1:
            idwIndx = 2
            psf = psfMaker_IDW(px, py, self.psfMat, self.cenPosx, self.cenPosy, IDWindex=idwIndx, OnlyNeighbors=True)
        if interpScheme ==2:
            libPCA = "/Users/chengliangwei/Desktop/csstPSF/libPCA/libPCA.so"
            psf = psfMaker_PCA(px, py, self.psfMat, self.cenPosx, self.cenPosy, OnlyNeighbors=True, libPCApath=libPCA)

        img = galsim.ImageF(psf, scale=0.074/2)
        xpsf = galsim.InterpolatedImage(img)
        
        return xpsf

    """
    def gcPlot(self, psf,pscale=0.074,figout="GC.png"):
        size = np.size(psf,axis=0)
        cxy = 0.5*(size-1)
        width = 0.5*size
        # log scale
        radius = np.arange(np.log10(0.2),np.log10(width),0.01)
        radius = 10.0**radius
        nr = len(radius)
        gc = []
        for i in range(nr): 
                iflux, iferr, xflag = sep.sum_circle(psf,cxy,cxy,radius[i],subpix=0)
                gc += [iflux.tolist()]

        # Estimate the radius for a given flux ratio
        fratio = 0.8
        mid = [i for i in range(nr) if gc[i]<=fratio and gc[i+1]>fratio][0]
        r0, r1 = radius[mid], radius[mid+1]
        gc0, gc1 = gc[mid], gc[mid+1]
        r5 = (fratio-gc0)/(gc1-gc0)*(r1-r0) + r0
        hlf = r5*pscale
        # plot
        pfit = interp1d(radius, gc, kind='cubic')
        fig = pl.figure(figsize=(5.5,4.0))
        ax = fig.add_axes([0.16,0.15,0.80,0.81])
        ax.plot(radius*pscale, pfit(radius), "k", linewidth=2.0)
        ax.plot(radius*pscale, gc, "*r", markersize=5.0,mec="r",alpha=0.3)
        ax.plot([hlf,hlf],[0,fratio],"k",linewidth=2.5)
        ax.plot([0,hlf],[fratio,fratio],"k",linewidth=2.5)
        ax.text(radius[10]*pscale,0.6,"$r_{%.1f}$=%.2f\""%(fratio,hlf))
        ax.set_xlabel("Radius (arcsec)",fontsize=15)
        ax.set_ylabel("Growth of Curve",fontsize=15)
        ax.set_xscale("log")
        ax.set_xlim(radius[0]*pscale,radius[-1]*pscale)
        ax.set_ylim(0.0,1.0)
        for tick in ax.xaxis.get_major_ticks(): tick.label.set_fontsize(15)
        for tick in ax.yaxis.get_major_ticks(): tick.label.set_fontsize(15)
        pl.savefig(figout)
        pl.clf()
        pl.close()

        return
    """

    def PSFspin(self, psf, sigSpin, sigGauss, dx, dy):
        """
        The PSF profile at a given image position relative to the axis center

        Parameters:
        theta : spin angles in a given exposure in unit of [arcsecond]
        dx, dy: relative position to the axis center in unit of [pixels]

        Return:
        Spinned PSF: g1, g2 and axis ratio 'a/b'
        """
        a2Rad = np.pi/(60.0*60.0*180.0)

        ff = sigGauss * 0.107 * (1000.0/10.0) # in unit of [pixels]
        rc = np.sqrt(dx*dx + dy*dy)
        cpix = rc*(sigSpin*a2Rad)

        beta = (np.arctan2(dy,dx) + np.pi/2)
        ell = cpix**2/(2.0*ff**2+cpix**2)
        #ell *= 10.0
        qr = np.sqrt((1.0+ell)/(1.0-ell))

        #psfShape = galsim.Shear(e=ell, beta=beta)
        #g1, g2 = psfShape.g1, psfShape.g2
        #qr = np.sqrt((1.0+ell)/(1.0-ell))

        #return ell, beta, qr
        PSFshear = galsim.Shear(e=ell, beta=beta*galsim.radians)
        return psf.shear(PSFshear), PSFshear(base)



##################################################
#                  D. TEST                       #
##################################################
def psfMaker_IDW_test(tpsf, px, py, PSFMat, cen_col, cen_row, IDWindex=2, OnlyNeighbors=False):
    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:
        neigh = findNeighbors(px, py, cen_col, cen_row, dr=0.1, dn=9, OnlyDistance=False)
        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:
            iy, ix = np.unravel_index(ipsf, (10,10))
            ty, tx = np.unravel_index(tpsf, (10,10))
            if np.abs(iy - ty) > 1 or np.abs(ix - tx) > 1:
                continue
        """
        if OnlyNeighbors == True:
            if neighFlag[ipsf] != 1:
                continue

        if ipsf == tpsf:
            continue

        iPSFMat = PSFMat[ipsf, :, :].copy()
        iPSFMat = psfCentering(iPSFMat, CenteringMode=1)
        ipsfWeight = psfWeight[ipsf]
        psfMaker += iPSFMat * ipsfWeight
    psfMaker /= np.nansum(psfMaker)
    
    return psfMaker


def psfMaker_PCA_test(tpsf, px, py, PSFMat, cen_col, cen_row, OnlyNeighbors=False, libPCApath='libPCA/libPCA.so'):
    """
    psf interpolation by PCA
    
    Parameters:

    Returns:
    """
    ref_col = px
    ref_row = py
    
    ngy, ngx = PSFMat[0, :, :].shape
    npsf = PSFMat[:, :, :].shape[0]
    
    neigh   = findNeighbors(px, py, cen_col, cen_row, dr=0.3, dn=9, OnlyDistance=False)
    npsfX   = len(neigh)

    #去掉tpsf,neigh中排在第一个是最近的psf
    print("CHECK:::", cen_col[neigh[0]], cen_row[neigh[0]], cen_col[tpsf], cen_row[tpsf], cen_col[neigh[0]]-cen_col[tpsf], cen_row[neigh[0]]-cen_row[tpsf])
    psfMatX = np.zeros([npsfX-1, ngy, ngx])
    cen_colX= np.zeros(npsfX-1)
    cen_rowX= np.zeros(npsfX-1)
    for ipsf in range(npsfX):
        if ipsf == 0:
            continue
        psfMatX[ipsf-1, :, :] = PSFMat[neigh[ipsf], :, :]
        cen_colX[ipsf-1] = cen_col[neigh[ipsf]]
        cen_rowX[ipsf-1] = cen_row[neigh[ipsf]]

    psfMaker  = np.zeros((ngy, ngx), dtype='float64')
    if OnlyNeighbors == True:
        PCAbasef, PCAcoeff = psfPCA_generator(psfMatX, npsfX-1, ngx, libPCApath)
        nPCA  = npsfX-1
        for iPCA in range(nPCA):
            coeffX = fitPoly(ref_col, ref_row, cen_colX, cen_rowX, PCAcoeff[:, iPCA], order=2)
            psfMaker += coeffX*PCAbasef[iPCA, :, :]          
    
    return psfMaker


def test_loadPSF():
    iccd = 1  #[1, 30]
    iwave= 1  #[1, 4]
    ipsf = 1  #[1, 100]
    psfPath = '/Users/chengliangwei/csstPSFdata/CSSOS_psf_ciomp'

    psfSet = []
    for ipsf in range(1, 901):
        psfInfo = LoadPSF(iccd, iwave, ipsf, psfPath, CalcPSFsize=True, CalcPSFcenter=True, SigRange=False)
        psfSet.append(psfInfo)

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

def test_psfPCA():
    #load psf
    print('load psf...')
    psfSet = test_loadPSF()

    #set input for psfPCA calc.
    print('PCA calc...')
    npsf = 5
    npix = 180
    psfMat = np.zeros([npsf, npix, npix])
    libPCApath = './libPCA/libPCA.so'
 
    for ipsf in range(5):
        psfMat[ipsf, :, :] = psfSet[ipsf]['psfMat']
    PCAbasef, PCAcoeff = psfPCA_generator(psfMat, npsf, npix, libPCApath)

    #plot check
    print('plot...')
    fig = plt.figure(figsize=(20, 10))
    cc = 90
    dcc= 15